Update data files in Google Cloud Buckets on publish (#249)

* Update data files in Google Cloud Buckets on publish
Fixes #248

* Remove testing code

* Add Google auth

* Adjust job permissions

* Update permissions

* Remove auth step from PR action

* Update data files in Google Cloud Buckets on publish
Fixes #248

Author: nikhilwoodruff

.github/workflows/code_changes.yaml

@@ -21,6 +21,10 @@ jobs:
           with:
               args: ". -l 79 --check"
   Test:
+      permissions:
+        contents: "read"
+        # Required to auth against gcp
+        id-token: "write"
       runs-on: larger-runner
       steps:
           - name: Checkout repo
@@ -32,6 +36,10 @@ jobs:
             uses: actions/setup-python@v2
             with:
                 python-version: '3.11'
+          - uses: "google-github-actions/auth@v2"
+            with:
+              workload_identity_provider: "projects/322898545428/locations/global/workloadIdentityPools/policyengine-research-id-pool/providers/prod-github-provider"
+              service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
 
           - name: Install package
             run: uv pip install -e .[dev] --system

changelog_entry.yaml

@@ -1,4 +1,4 @@
 - bump: patch
   changes:
-    changed:
-    - Methodology to directly impute auto loan interest instead of assuming a 2% interest rate on auto loan balance.
+    fixed:
+    - Upload to GCP on dataset build.

policyengine_us_data/storage/upload_completed_datasets.py

@@ -5,10 +5,20 @@
 )
 from policyengine_us_data.storage import STORAGE_FOLDER
 from policyengine_us_data.utils.huggingface import upload
+from google.cloud import storage
 
 
 def upload_datasets():
-    for dataset in [EnhancedCPS_2024, Pooled_3_Year_CPS_2023, CPS_2023]:
+    storage_client = storage.Client()
+    bucket = storage_client.bucket("policyengine-us-data")
+
+    datasets_to_upload = [
+        EnhancedCPS_2024,
+        Pooled_3_Year_CPS_2023,
+        CPS_2023,
+    ]
+
+    for dataset in datasets_to_upload:
         dataset = dataset()
         if not dataset.exists:
             raise ValueError(
@@ -21,6 +31,13 @@ def upload_datasets():
             dataset.file_path.name,
         )
 
+        blob = dataset.file_path.name
+        blob = bucket.blob(blob)
+        blob.upload_from_filename(dataset.file_path)
+        print(
+            f"Uploaded {dataset.file_path.name} to GCS bucket policyengine-us-data."
+        )
+
 
 if __name__ == "__main__":
     upload_datasets()

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "microdf_python>=0.4.3",
     "microimpute",
     "pip-system-certs",
+    "google-cloud-storage",
 ]
 
 [project.optional-dependencies]