Merge branch 'master' into implmented_golden_checks_annual_pop

Author: niveditasing

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

@@ -77,6 +77,7 @@
 AUTO_IMPORT_JOB_STATUS = "auto-import-job-status"
 IMPORT_SUMMARY_FILE = "import_summary.json"
 STAGING_VERSION_FILE = "staging_version.txt"
+MAX_LOG_CHUNK_SIZE = 50000
 
 
 class ImportStatus(Enum):
@@ -480,7 +481,8 @@ def _invoke_import_tool(self, absolute_import_dir: str,
                              "latency": timer.time(),
                              "input_index": input_index,
                              "import_input": import_prefix,
-                         })
+                         },
+                         skip_stream_logging=True)
             process.check_returncode()
             logging.info(
                 f'Generated resolved mcf for {import_prefix} in {output_path}.')
@@ -809,7 +811,8 @@ def _invoke_import_job(self, absolute_import_dir: str,
                                  "latency_secs": timer.time(),
                                  "script_index": script_index,
                                  "script_path": path,
-                             })
+                             },
+                             skip_stream_logging=True)
                 process.check_returncode()
 
         import_summary.import_stats['script_execution_time'] = start_timer.time(
@@ -1118,6 +1121,33 @@ def run_and_handle_exception(
         return ExecutionResult(ImportStatus.FAILURE, [], message)
 
 
+def _stream_payload_in_chunks(label: str, payload) -> None:
+    """Helper function to split text and log in safe chunks to avoid
+
+    "Log entry too large" (InvalidArgument) errors.
+    """
+    if not payload:
+        return
+
+    if isinstance(payload, bytes):
+        payload_str = payload.decode('utf-8', errors='replace')
+    else:
+        payload_str = str(payload)
+
+    chunk_size = MAX_LOG_CHUNK_SIZE
+    total_len = len(payload_str)
+
+    if total_len <= chunk_size:
+        logging.info(f'{label}:\n{payload_str}')
+        return
+
+    logging.info(f'--- Start of {label} (Total length: {total_len} chars) ---')
+    for i in range(0, total_len, chunk_size):
+        chunk = payload_str[i:i + chunk_size]
+        logging.info(f'[{label} Part {i//chunk_size + 1}]:\n{chunk}')
+    logging.info(f'--- End of {label} ---')
+
+
 @log_function_call
 def _run_with_timeout_async(args: List[str],
                             timeout: float,
@@ -1217,8 +1247,8 @@ def _run_with_timeout(args: List[str],
                                  cwd=cwd,
                                  env=env)
         logging.info(
-            f'Completed command: {args}, retcode: {process.returncode}, stdout:'
-            f' {process.stdout}, stderr: {process.stderr}')
+            f'Completed command: {args}, retcode: {process.returncode}')
+
         return process
     except Exception as e:
         message = traceback.format_exc()
@@ -1378,28 +1408,37 @@ def _construct_process_message(message: str,
     message = (f'{message}\n'
                f'[Subprocess command]: {command}\n'
                f'[Subprocess return code]: {process.returncode}')
-    if process.stdout:
-        message += f'\n[Subprocess stdout]:\n{process.stdout}'
-    if process.stderr:
-        message += f'\n[Subprocess stderr]:\n{process.stderr}'
     return message
 
 
 def _log_process(process: subprocess.CompletedProcess,
                  import_name: str = '',
-                 metrics: dict = {}) -> None:
+                 metrics: dict = None,
+                 skip_stream_logging: bool = False) -> None:
     """Logs the result of a subprocess.
 
   Args:
       process: subprocess.CompletedProcess object whose arguments, return code,
         stdout, and stderr are to be logged.
+      import_name: Name of the import for labeling logs.
+      metrics: Dictionary to store execution metrics.
+      skip_stream_logging: Whether to skip chunked logging of stdout and stderr.
   """
+    if metrics is None:
+        metrics = {}
     process_message = 'Subprocess succeeded'
     if process.returncode:
         process_message = 'Subprocess failed'
+
     message = _construct_process_message(process_message, process)
     logging.info(message)
 
+    if not skip_stream_logging:
+        _stream_payload_in_chunks(f'[{import_name}] Subprocess stdout',
+                                  process.stdout)
+        _stream_payload_in_chunks(f'[{import_name}] Subprocess stderr',
+                                  process.stderr)
+
     status = ImportStatus.FAILURE if process.returncode else ImportStatus.SUCCESS
     if import_name:
         metrics["import_name"] = import_name

import-automation/executor/test/import_executor_test.py

@@ -83,11 +83,7 @@ def test_construct_process_message(self):
         expected = (
             'message\n'
             '[Subprocess command]: printf "out" & >&2 printf "err" & exit 1\n'
-            '[Subprocess return code]: 1\n'
-            '[Subprocess stdout]:\n'
-            'out\n'
-            '[Subprocess stderr]:\n'
-            'err')
+            '[Subprocess return code]: 1')
         self.assertEqual(expected, message)
 
     def test_construct_process_message_no_output(self):

import-automation/workflow/cloudbuild.yaml

@@ -12,40 +12,88 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-# Child Cloud Build configuration for deploying to a specific environment.
-# Defaults are set to PRODUCTION. Staging builds must override these values.
-#
+# Parent Cloud Build configuration that orchestrates Staging and Production deployments.
+# Usage: gcloud builds submit . --config=cloudbuild.yaml --project=datcom-ci
+
 substitutions:
-  # Production config.
-  _PROJECT_ID: 'datcom-import-automation-prod'
-  _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'
+  # Staging Configuration (Overrides defaults in child build)
+  _PROJECT_ID: 'datcom-ci'
+  _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'
-  _PROJECT_NUMBER: '965988403328'
+  _LOCATION: 'us-central1'
+  _PROJECT_NUMBER: '879489846695'
   _AR_REPO_URL: 'us-docker.pkg.dev/datcom-ci/gcr.io'
-  _BQ_SPANNER_CONN_ID: 'projects/datcom-import-automation-prod/locations/us/connections/bq_spanner_conn'
+  _BQ_SPANNER_CONN_ID: 'projects/datcom-ci/locations/us-central1/connections/bq_spanner_conn_test'
+  _VERSION: '${SHORT_SHA}'
+  _DATAFLOW_TEMPLATE_PATH: 'gs://datcom-templates/templates/flex/ingestion.json'
+  _PROD_TAG: 'latest' # TODO: rename to prod
 
 steps:
-- id: 'ingestion-helper-service'
+
+# 1. Build and push helper images
+- id: 'build-ingestion-helper'
+  name: 'gcr.io/cloud-builders/gcloud'
+  args: ['builds', 'submit', 'ingestion-helper', '--config', 'ingestion-helper/cloudbuild.yaml', '--substitutions', '_AR_REPO_URL=${_AR_REPO_URL},_VERSION=${_VERSION}']
+  dir: 'import-automation/workflow'
+
+- id: 'build-import-helper'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['run', 'deploy', 'ingestion-helper-service', '--image', '${_AR_REPO_URL}/datacommons-ingestion-helper:latest', '--region', '${_LOCATION}', '--project', '${_PROJECT_ID}', '--no-allow-unauthenticated', '--timeout', '60m', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},LOCATION=${_LOCATION},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},BQ_DATASET_ID=${_BQ_DATASET_ID},BQ_SPANNER_CONN_ID=${_BQ_SPANNER_CONN_ID}']
+  args: ['builds', 'submit', 'import-helper', '--config', 'import-helper/cloudbuild.yaml', '--substitutions', '_AR_REPO_URL=${_AR_REPO_URL},_VERSION=${_VERSION}']
+  dir: 'import-automation/workflow'
 
-- id: 'import-helper-service'
+# 2. Trigger Staging Build (Child)
+# Overrides default (Production) values with Staging values.
+- id: 'deploy-staging'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['run', 'deploy', 'import-helper-service', '--image', '${_AR_REPO_URL}/datacommons-import-helper:latest', '--region', '${_LOCATION}', '--project', '${_PROJECT_ID}', '--no-allow-unauthenticated', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},LOCATION=${_LOCATION},PROJECT_NUMBER=${_PROJECT_NUMBER},GCS_BUCKET_ID=${_GCS_BUCKET_ID}']
+  args:
+    - 'builds'
+    - 'submit'
+    - '.'
+    - '--config=deploy-services.yaml'
+    - '--project=${_PROJECT_ID}'
+    - '--substitutions=_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},_GCS_MOUNT_BUCKET=${_GCS_MOUNT_BUCKET},_BQ_DATASET_ID=${_BQ_DATASET_ID},_PROJECT_NUMBER=${_PROJECT_NUMBER},_BQ_SPANNER_CONN_ID=${_BQ_SPANNER_CONN_ID},_VERSION=${_VERSION},_DATAFLOW_TEMPLATE_PATH=${_DATAFLOW_TEMPLATE_PATH}'
+  dir: 'import-automation/workflow'
+
+# 2. Run E2E Tests on Staging
+- id: 'e2e-test-staging'
+  name: 'python:3.11'
+  entrypoint: 'bash'
+  args:
+    - '-c'
+    - |
+      pip install google-cloud-spanner google-cloud-workflows absl-py
+      python spanner_ingestion_test.py
+  env:
+    - 'PROJECT_ID=${_PROJECT_ID}'
+    - 'LOCATION=${_LOCATION}'
+    - 'SPANNER_PROJECT_ID=${_SPANNER_PROJECT_ID}'
+    - 'SPANNER_INSTANCE_ID=${_SPANNER_INSTANCE_ID}'
+    - 'SPANNER_DATABASE_ID=${_SPANNER_DATABASE_ID}'
+    - 'GCS_BUCKET_ID=${_GCS_BUCKET_ID}'
+  dir: 'import-automation/workflow'
 
-- id: 'import-automation-workflow'
+- id: 'tag-prod'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['workflows', 'deploy', 'import-automation-workflow', '--project', '${_PROJECT_ID}', '--location', '${_LOCATION}', '--source', 'import-automation-workflow.yaml', '--set-env-vars', 'LOCATION=${_LOCATION},GCS_BUCKET_ID=${_GCS_BUCKET_ID},GCS_MOUNT_BUCKET=${_GCS_MOUNT_BUCKET},PROJECT_NUMBER=${_PROJECT_NUMBER}']
+  args: ['builds', 'submit', '.', '--config', 'tag-prod.yaml', '--substitutions', '_AR_REPO_URL=${_AR_REPO_URL},_PROD_TAG=${_PROD_TAG},_VERSION=${_VERSION}']
+  dir: 'import-automation/workflow'
 
-- id: 'spanner-ingestion-workflow'
+# 3. Trigger Production Build (Child)
+# Uses default (Production) values defined in cloudbuild.yaml.
+- id: 'deploy-prod'
   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_GRAPH_DATABASE_ID},PROJECT_NUMBER=${_PROJECT_NUMBER}']
+  args:
+    - 'builds'
+    - 'submit'
+    - '.'
+    - '--config=deploy-services.yaml'
+    - '--project=${_PROJECT_ID}'  # Build runs in CI project, deploys to Prod
+    - '--substitutions=_VERSION=${_PROD_TAG}'
+  dir: 'import-automation/workflow'
 
 options:
   logging: CLOUD_LOGGING_ONLY

import-automation/workflow/cloudbuild_main.yaml

@@ -0,0 +1,53 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Child Cloud Build configuration for deploying to a specific environment.
+# Defaults are set to PRODUCTION. Staging builds must override these values.
+#
+substitutions:
+  # Production config.
+  _PROJECT_ID: 'datcom-import-automation-prod'
+  _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'
+  _BQ_DATASET_ID: 'datacommons'
+  _PROJECT_NUMBER: '965988403328'
+  _AR_REPO_URL: 'us-docker.pkg.dev/datcom-ci/gcr.io'
+  _BQ_SPANNER_CONN_ID: 'projects/datcom-import-automation-prod/locations/us/connections/bq_spanner_conn'
+  _VERSION: '${SHORT_SHA}'
+  _DATAFLOW_TEMPLATE_PATH: 'gs://datcom-templates/templates/flex/ingestion-${_VERSION}.json'
+
+steps:
+- id: 'ingestion-helper-service'
+  name: 'gcr.io/cloud-builders/gcloud'
+  args: ['run', 'deploy', 'ingestion-helper-service', '--image', '${_AR_REPO_URL}/datacommons-ingestion-helper:${_VERSION}', '--region', '${_LOCATION}', '--project', '${_PROJECT_ID}', '--no-allow-unauthenticated', '--timeout', '60m', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},LOCATION=${_LOCATION},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},BQ_DATASET_ID=${_BQ_DATASET_ID},BQ_SPANNER_CONN_ID=${_BQ_SPANNER_CONN_ID}']
+
+- id: 'import-helper-service'
+  name: 'gcr.io/cloud-builders/gcloud'
+  args: ['run', 'deploy', 'import-helper-service', '--image', '${_AR_REPO_URL}/datacommons-import-helper:${_VERSION}', '--region', '${_LOCATION}', '--project', '${_PROJECT_ID}', '--no-allow-unauthenticated', '--set-env-vars', 'PROJECT_ID=${_PROJECT_ID},LOCATION=${_LOCATION},PROJECT_NUMBER=${_PROJECT_NUMBER},GCS_BUCKET_ID=${_GCS_BUCKET_ID}']
+
+- id: 'import-automation-workflow'
+  name: 'gcr.io/cloud-builders/gcloud'
+  args: ['workflows', 'deploy', 'import-automation-workflow', '--project', '${_PROJECT_ID}', '--location', '${_LOCATION}', '--source', 'import-automation-workflow.yaml', '--set-env-vars', 'LOCATION=${_LOCATION},GCS_BUCKET_ID=${_GCS_BUCKET_ID},GCS_MOUNT_BUCKET=${_GCS_MOUNT_BUCKET},PROJECT_NUMBER=${_PROJECT_NUMBER}']
+
+- 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_GRAPH_DATABASE_ID},PROJECT_NUMBER=${_PROJECT_NUMBER},DATAFLOW_TEMPLATE_PATH=${_DATAFLOW_TEMPLATE_PATH}']
+
+options:
+  logging: CLOUD_LOGGING_ONLY