Merge branch 'datacommonsorg:master' into master

Author: kartik-s21

import-automation/executor/Dockerfile

@@ -36,7 +36,10 @@ fonts-liberation \
 xdg-utils \
 chromium \
 chromium-driver \
-p7zip-full
+p7zip-full \
+libeccodes-dev \
+libeccodes-tools \
+&& rm -rf /var/lib/apt/lists/*
 # Install the Google Cloud CLI
 RUN apt-get update && \
     curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \

import-automation/executor/app/configs.py

@@ -153,6 +153,8 @@ class ExecutorConfig:
     validation_config_file: str = 'tools/import_validation/validation_config.json'
     # Latest import version (overwrite)
     import_version_override: str = ''
+    # Override cron schedule in import spec.
+    cron_schedule_override: str = ''
     # Relative path to version folder for graph files.
     graph_data_path: str = '/*/*/*.mcf'
     # Maximum time venv creation can take in seconds.
@@ -169,6 +171,8 @@ class ExecutorConfig:
     disable_email_notifications: bool = True
     # Skip uploading the data to GCS (for local testing).
     skip_gcs_upload: bool = False
+    # Skip uploading input files to GCS.
+    skip_input_upload: bool = False
     # Maximum time a blocking call to the importer to
     # perform an import can take in seconds.
     importer_import_timeout: float = 20 * 60

import-automation/executor/app/executor/import_executor.py

@@ -82,12 +82,8 @@ def upload_file(self, src: str, dest: str) -> None:
         """
         _strings_not_empty(src, dest)
         dest = self._fix_path(dest)
-        logging.info('GCSFileUploader.upload_file: Uploading %s to %s', src,
-                     dest)
         blob = self.bucket.blob(dest)
         blob.upload_from_filename(src)
-        logging.info('GCSFileUploader.upload_file: Uploaded %s to %s', src,
-                     dest)
 
     def upload_string(self, string: str, dest: str) -> None:
         """Uploads a string to a file in the bucket, overwriting it.
@@ -102,10 +98,8 @@ def upload_string(self, string: str, dest: str) -> None:
         """
         _strings_not_empty(dest)
         dest = self._fix_path(dest)
-        logging.info('GCSFileUploader.upload_string: Uploading to %s', dest)
         blob = self.bucket.blob(dest)
         blob.upload_from_string(string)
-        logging.info('GCSFileUploader.upload_string: Uploaded to %s', dest)
 
     def _fix_path(self, path):
         """Returns {self.path_prefix}/{path}."""

import-automation/executor/app/service/file_uploader.py

@@ -46,7 +46,6 @@ steps:
         python import_test.py
     env:
       - 'PROJECT_ID=$PROJECT_ID'
-      - 'LOCATION=$LOCATION'
       - 'GCS_BUCKET=${_GCS_BUCKET}'
       - 'IMAGE_URI=${_DOCKER_IMAGE}:${COMMIT_SHA}'
     dir: 'import-automation/executor'

import-automation/executor/cloudbuild.yaml

@@ -6,6 +6,7 @@ beautifulsoup4
 chardet
 chromedriver_py
 croniter
+curl_cffi
 dataclasses
 datacommons
 datacommons_client
@@ -41,6 +42,7 @@ omegaconf
 prettytable
 protobuf
 psutil
+pygrib
 pylint
 pyspellchecker
 pytest

import-automation/executor/requirements.txt

@@ -254,11 +254,11 @@ function build_docker {
 # Get the latest import output from GCS
 function get_latest_gcs_import_output {
   echo_log "Looking for import files on GCS at $GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME..."
-  LATEST_VERSION=$(gsutil cat gs://$GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME/latest_version.txt)
-  if [[ -n "$LATEST_VERSION" ]]; then
-    echo_log "latest_version.txt: $LATEST_VERSION"
-    run_cmd gsutil ls -lR gs://$GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME/$LATEST_VERSION
-    echo_log "View latest import files at: https://pantheon.corp.google.com/storage/browser/$GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME/$LATEST_VERSION"
+  STAGING_VERSION=$(gsutil cat gs://$GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME/staging_version.txt)
+  if [[ -n "$STAGING_VERSION" ]]; then
+    echo_log "staging_version.txt: $STAGING_VERSION"
+    run_cmd gsutil ls -lR gs://$GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME/$STAGING_VERSION
+    echo_log "View latest import files at: https://pantheon.corp.google.com/storage/browser/$GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME/$STAGING_VERSION"
   else
     echo_log "No files on GCS at $GCS_BUCKET/$IMPORT_DIR/$IMPORT_NAME"
   fi
@@ -317,7 +317,7 @@ function get_import_config {
     IMPORT_VERSION=$(cut -d: -f2 <<< "$ver_override")
     get_latest_gcs_import_output
     # Add config to update version.
-    add_import_version_notes "$IMPORT_VERSION" "Updating latest $IMPORT_NAME from: $LATEST_VERSION to: $IMPORT_VERSION, $NOTE"
+    add_import_version_notes "$IMPORT_VERSION" "Updating latest $IMPORT_NAME from: $STAGING_VERSION to: $IMPORT_VERSION, $NOTE"
   fi
   for c_v in $config_vals; do
     param=$(cut -d: -f1 <<< "$c_v")

import-automation/executor/run_import.sh

@@ -21,6 +21,7 @@ substitutions:
   _SPANNER_PROJECT_ID: 'datcom-store'
   _SPANNER_INSTANCE_ID: 'dc-kg-test'
   _SPANNER_DATABASE_ID: 'dc_graph_import'
+  _SPANNER_GRAPH_DATABASE_ID: 'dc_graph_2025_11_07'
   _GCS_BUCKET_ID: 'datcom-prod-imports'
   _LOCATION: 'us-central1'
   _GCS_MOUNT_BUCKET: 'datcom-volume-mount'
@@ -33,15 +34,15 @@ steps:
 
 - id: 'spanner-ingestion-workflow'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['workflows', 'deploy', 'spanner-ingestion-workflow', '--project', '${_PROJECT_ID}', '--location', '${_LOCATION}', '--source', 'spanner-ingestion-workflow.yaml', '--set-env-vars', 'LOCATION=${_LOCATION},PROJECT_ID=${_PROJECT_ID},SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID},SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID},SPANNER_DATABASE_ID=${_SPANNER_DATABASE_ID}']
+  args: ['workflows', 'deploy', 'spanner-ingestion-workflow', '--project', '${_PROJECT_ID}', '--location', '${_LOCATION}', '--source', 'spanner-ingestion-workflow.yaml', '--set-env-vars', 'LOCATION=${_LOCATION},PROJECT_ID=${_PROJECT_ID},SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID},SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID},SPANNER_DATABASE_ID=${_SPANNER_GRAPH_DATABASE_ID}']
 
 - id: 'spanner-ingestion-helper'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['functions', 'deploy', 'spanner-ingestion-helper', '--gen2', '--project', '${_PROJECT_ID}', '--region', '${_LOCATION}', '--runtime', 'python312', '--source', 'ingestion-helper', '--no-allow-unauthenticated', '--trigger-http', '--entry-point', 'ingestion_helper', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID},SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID},SPANNER_DATABASE_ID=${_SPANNER_DATABASE_ID},GCS_BUCKET_ID=${_GCS_BUCKET_ID},LOCATION=${_LOCATION}']
+  args: ['functions', 'deploy', 'spanner-ingestion-helper', '--gen2', '--project', '${_PROJECT_ID}', '--region', '${_LOCATION}', '--runtime', 'python312', '--source', 'ingestion-helper', '--no-allow-unauthenticated', '--trigger-http', '--entry-point', 'ingestion_helper', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID},SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID},SPANNER_DATABASE_ID=${_SPANNER_DATABASE_ID},SPANNER_GRAPH_DATABASE_ID=${_SPANNER_GRAPH_DATABASE_ID},GCS_BUCKET_ID=${_GCS_BUCKET_ID},LOCATION=${_LOCATION}']
 
 - id: 'import-aggregation-helper'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['functions', 'deploy', 'import-aggregation-helper', '--runtime', 'python312', '--source', 'aggregation-helper', '--no-allow-unauthenticated', '--trigger-http', '--entry-point', 'aggregation_helper', '--project', '${_PROJECT_ID}', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID},SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID},SPANNER_DATABASE_ID=${_SPANNER_DATABASE_ID},GCS_BUCKET_ID=${_GCS_BUCKET_ID},LOCATION=${_LOCATION},BQ_DATASET_ID=${_BQ_DATASET_ID}']
+  args: ['functions', 'deploy', 'import-aggregation-helper', '--runtime', 'python312', '--source', 'aggregation-helper', '--no-allow-unauthenticated', '--trigger-http', '--entry-point', 'aggregation_helper', '--project', '${_PROJECT_ID}', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID},SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID},SPANNER_DATABASE_ID=${_SPANNER_GRAPH_DATABASE_ID},GCS_BUCKET_ID=${_GCS_BUCKET_ID},LOCATION=${_LOCATION},BQ_DATASET_ID=${_BQ_DATASET_ID}']
 
 - id: 'import-automation-helper'
   name: 'gcr.io/cloud-builders/gcloud'

import-automation/workflow/cloudbuild.yaml

@@ -21,6 +21,7 @@ substitutions:
   _SPANNER_PROJECT_ID: 'datcom-ci'
   _SPANNER_INSTANCE_ID: 'datcom-spanner-test'
   _SPANNER_DATABASE_ID: 'dc-test-db'
+  _SPANNER_GRAPH_DATABASE_ID: 'dc-test-db'
   _GCS_BUCKET_ID: 'datcom-ci-test'
   _GCS_MOUNT_BUCKET: 'datcom-ci-test'
   _BQ_DATASET_ID: 'datacommons'

import-automation/workflow/cloudbuild_main.yaml

@@ -21,6 +21,10 @@ main:
               memory: 32768
               disk: 100
           - resources: ${default(map.get(args, "resources"), defaultResources)}
+          - runIngestion: ${default(map.get(args, "runIngestion"), false)}
+          - ingestionArgs:
+              importList:
+                - ${text.split(importName, ":")[1]}
     - runImportJob:
         try:
           call: googleapis.batch.v1.projects.locations.jobs.create
@@ -99,6 +103,18 @@ main:
             override: false
             comment: '${"import-workflow:" + sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}'
         result: functionResponse
+    - runIngestion:
+        switch:
+          - condition: ${runIngestion}
+            steps:
+              - runSpannerIngestion:
+                  call: googleapis.workflowexecutions.v1.projects.locations.workflows.executions.create
+                  args:
+                    parent: ${"projects/" + projectId + "/locations/" + region + "/workflows/spanner-ingestion-workflow"}
+                    body:
+                      argument: ${json.encode_to_string(ingestionArgs)}
+                    connector_params:
+                      skip_polling: true
     - returnResult:
         return:
           jobId: ${jobId}