Address reviewer comments

vish-cs · vish-cs · commit 1a86bc155568 · 2025-12-29T08:37:11.000Z
diff --git a/import-automation/executor/app/executor/import_executor.py b/import-automation/executor/app/executor/import_executor.py
@@ -822,9 +822,10 @@ def _import_one_helper(
                                     self.config.file_download_timeout)
 
         output_dir = f'{relative_import_dir}/{import_name}'
-        version = _clean_time(utils.pacific_time())
         if version == 'DATE_VERSION_PLACEHOLDER':
             version = datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%d")
+        else:
+            version = _clean_time(utils.pacific_time())
         import_summary.latest_version = 'gs://' + os.path.join(
             self.config.storage_prod_bucket_name, output_dir, version, '*', '*',
             '*.mcf')
diff --git a/import-automation/executor/update_import_version.sh b/import-automation/executor/update_import_version.sh
@@ -12,7 +12,8 @@ if [ "$#" -ne 3 ]; then
     exit 1
 fi
 
-FUNCTION_URL="https://spanner-ingestion-helper-965988403328.us-central1.run.app"
+# Deployed using import-automation/workflow/cloudbuild.yaml
+FUNCTION_URL="https://us-central1-datcom-import-automation-prod.cloudfunctions.net/spanner-ingestion-helper"
 IMPORT_NAME=$1
 VERSION=$2
 REASON=$3
diff --git a/import-automation/workflow/cloudbuild.yaml b/import-automation/workflow/cloudbuild.yaml
@@ -1,17 +1,26 @@
 # gcloud builds submit --config=cloudbuild.yaml . --project=datcom-import-automation-prod 
 
+substitutions:
+  _PROJECT_ID: 'datcom-import-automation-prod'
+  _SPANNER_PROJECT_ID: 'datcom-store'
+  _SPANNER_INSTANCE_ID: 'dc-kg-test'
+  _SPANNER_DATABASE_ID: 'dc_graph_import'
+  _GCS_BUCKET_ID: 'datcom-prod-imports'
+  _LOCATION: 'us-central1'
+  _GCS_MOUNT_BUCKET: 'datcom-volume-mount'
+
 steps:
 - id: 'import-automation-workflow'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['workflows', 'deploy', 'import-automation-workflow', '--source', 'import-automation-workflow.yaml']
+  args: ['workflows', 'deploy', 'import-automation-workflow', '--source', 'import-automation-workflow.yaml', '--set-env-vars', 'LOCATION=${_LOCATION},GCS_BUCKET_ID=${_GCS_BUCKET_ID},GCS_MOUNT_BUCKET=${_GCS_MOUNT_BUCKET}']
 
 - id: 'spanner-ingestion-workflow'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['workflows', 'deploy', 'spanner-ingestion-workflow', '--source', 'spanner-ingestion-workflow.yaml']
+  args: ['workflows', 'deploy', 'spanner-ingestion-workflow', '--source', 'spanner-ingestion-workflow.yaml', '--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}']
 
 - id: 'spanner-ingestion-helper'
   name: 'gcr.io/cloud-builders/gcloud'
-  args: ['functions', 'deploy', 'spanner-ingestion-helper', '--runtime', 'python312', '--source', 'ingestion-helper', '--no-allow-unauthenticated', '--trigger-http', '--entry-point', 'ingestion_helper']
+  args: ['functions', 'deploy', 'spanner-ingestion-helper', '--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}']
 
 - id: 'import-automation-helper'
   name: 'gcr.io/cloud-builders/gcloud'
diff --git a/import-automation/workflow/import-automation-workflow.yaml b/import-automation/workflow/import-automation-workflow.yaml
@@ -4,17 +4,18 @@ main:
     - init:
         assign:
           - projectId: ${sys.get_env("GOOGLE_CLOUD_PROJECT_ID")}
-          - region: "us-central1"
+          - region: ${sys.get_env("LOCATION")}
           - imageUri: "gcr.io/datcom-ci/dc-import-executor:stable"
           - jobId: ${text.substring(args.jobName, 0, 50) + "-" + string(int(sys.now()))}
           - importName: ${args.importName}
           - importConfig: ${args.importConfig}
-          - gcsMountBucket: "datcom-volume-mount"
-          - gcsImportBucket: "datcom-prod-imports"
+          - gcsMountBucket: ${sys.get_env("GCS_MOUNT_BUCKET")}
+          - gcsImportBucket: ${sys.get_env("GCS_BUCKET_ID")}
           - gcsMountPath: "/tmp/gcs"
-          - function_url: 'https://spanner-ingestion-helper-965988403328.us-central1.run.app'
+          - ingestionHelper: "spanner-ingestion-helper" 
+          - functionUrl: ${"https://" + region + "-" + projectId + ".cloudfunctions.net/" + ingestionHelper}
           - startTime: ${sys.now()}
-    - createAndRunBatchJob:
+    - runImportJob:
         try:
           call: googleapis.batch.v1.projects.locations.jobs.create
           args:
@@ -59,39 +60,39 @@ main:
                 initial_delay: 60
                 multiplier: 2
                 max_delay: 600
-          result: createAndRunBatchJobResponse
+          result: importJobResponse
         except:
           as: e
           steps:
-            - update_import_status:
+            - updateImportStatus:
                 call: http.post
                 args:
-                  url: ${function_url}
+                  url: ${functionUrl}
                   auth:
                     type: OIDC
                   body:
-                    actionType: update_import_status
+                    actionType: 'update_import_status'
                     jobId: ${jobId}
                     importName: ${text.split(args.importName, ":")[1]}
                     status: 'FAILED'
                     duration: ${int(sys.now() - startTime)}
                     version: ${"gs://" + gcsImportBucket + "/" + text.replace_all(args.importName, ":", "/")}
                     schedule: ${default(map.get(args, "schedule"), "")}
-                result: function_response
-            - fail_workflow:
+                result: functionResponse
+            - failWorkflow:
                 raise: ${e}
-    - update_import_version:
+    - updateImportVersion:
         call: http.post
         args:
-          url: ${function_url}
+          url: ${functionUrl}
           auth:
             type: OIDC
           body:
-            actionType: update_import_version
+            actionType: 'update_import_version'
             importName: ${args.importName}
             version: 'staging'
             reason: 'import-automation'
-        result: function_response
+        result: functionResponse
     - returnResult:
         return:
           jobId: ${jobId}
diff --git a/import-automation/workflow/ingestion-helper/import_utils.py b/import-automation/workflow/ingestion-helper/import_utils.py
@@ -1,3 +1,17 @@
+# 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.
+
 """Utility functions for the ingestion helper."""
 
 import logging
@@ -22,15 +36,14 @@ def get_caller_identity(request):
             unverified_claims = {}
             try:
                 unverified_claims = jwt.decode(token, verify=False)
-                #logging.info(f"Token claims (unverified): iss={unverified_claims.get('iss')}, aud={unverified_claims.get('aud')}, email={unverified_claims.get('email')}")
-                logging.warning(
-                    f"Could not decode unverified token for debugging: {debug_e}"
-                )
                 id_info = id_token.verify_oauth2_token(token,
                                                        requests.Request())
                 return id_info.get('email', 'unknown_email')
             except Exception as e:
                 if unverified_claims:
+                    logging.warning(
+                        f"Could not decode unverified token for debugging: {e}"
+                    )
                     email = unverified_claims.get('email', 'unknown_email')
                     return f"{email} (unverified)"
                 return 'decode_error'
@@ -111,23 +124,39 @@ def create_import_params(summary) -> dict:
 
 
 def get_ingestion_metrics(project_id, location, job_id):
-    """Fetches graph metrics (nodes, edges, observations) from a Dataflow job.
+    """Fetches graph metrics (nodes, edges, observations) and execution time from a Dataflow job.
 
     Args:
         project_id: The GCP project ID.
         location: The location of the Dataflow job.
         job_id: The Dataflow job ID.
 
     Returns:
-        A dictionary containing 'obs_count', 'node_count', and 'edge_count'.
+        A dictionary containing 'obs_count', 'node_count', 'edge_count', and 'execution_time'.
     """
     dataflow = build('dataflow', 'v1b3', cache_discovery=False)
     # Fetch Dataflow metrics
     node_count = 0
     edge_count = 0
     obs_count = 0
+    execution_time = 0
     if project_id and job_id:
         try:
+            # Fetch Job details for execution time
+            job = dataflow.projects().locations().jobs().get(
+                projectId=project_id, location=location, jobId=job_id).execute()
+
+            start_time_str = job.get('startTime')
+            current_state_time_str = job.get('currentStateTime')
+
+            if start_time_str and current_state_time_str:
+                # Handle potential 'Z' suffix by replacing it with '+00:00' for compatibility
+                start_time = datetime.fromisoformat(
+                    start_time_str.replace('Z', '+00:00'))
+                end_time = datetime.fromisoformat(
+                    current_state_time_str.replace('Z', '+00:00'))
+                execution_time = (end_time - start_time).total_seconds()
+
             metrics = dataflow.projects().locations().jobs().getMetrics(
                 projectId=project_id, location=location,
                 jobId=job_id).execute()
@@ -145,5 +174,6 @@ def get_ingestion_metrics(project_id, location, job_id):
     return {
         'obs_count': obs_count,
         'node_count': node_count,
-        'edge_count': edge_count
+        'edge_count': edge_count,
+        'execution_time': execution_time
     }
diff --git a/import-automation/workflow/ingestion-helper/main.py b/import-automation/workflow/ingestion-helper/main.py
@@ -1,18 +1,48 @@
+# 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.
+
 import functions_framework
 from spanner_client import SpannerClient
 from storage_client import StorageClient
 from flask import jsonify
 import logging
 import import_utils
+import os
+import sys
+from absl import flags
 
 logging.getLogger().setLevel(logging.INFO)
 
-_PROJECT_ID = 'datcom-import-automation-prod'
-_SPANNER_PROJECT_ID = 'datcom-store'
-_SPANNER_INSTANCE_ID = 'dc-kg-test'
-_SPANNER_DATABASE_ID = 'dc_graph_import'
-_GCS_BUCKET_ID = 'datcom-prod-imports'
-_LOCATION = 'us-central1'
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    'project_id', os.environ.get('PROJECT_ID', 'datcom-import-automation-prod'),
+    'GCP Project ID')
+flags.DEFINE_string('spanner_project_id',
+                    os.environ.get('SPANNER_PROJECT_ID', 'datcom-store'),
+                    'Spanner Project ID')
+flags.DEFINE_string('spanner_instance_id',
+                    os.environ.get('SPANNER_INSTANCE_ID', 'dc-kg-test'),
+                    'Spanner Instance ID')
+flags.DEFINE_string('spanner_database_id',
+                    os.environ.get('SPANNER_DATABASE_ID', 'dc_graph_import'),
+                    'Spanner Database ID')
+flags.DEFINE_string('gcs_bucket_id',
+                    os.environ.get('GCS_BUCKET_ID', 'datcom-prod-imports'),
+                    'GCS Bucket ID')
+flags.DEFINE_string('location', os.environ.get('LOCATION', 'us-central1'),
+                    'GCP Location')
 
 
 def _validate_params(request_json, required_params):
@@ -36,9 +66,9 @@ def ingestion_helper(request):
         return (validation_error, 400)
 
     actionType = request_json['actionType']
-    spanner = SpannerClient(_SPANNER_PROJECT_ID, _SPANNER_INSTANCE_ID,
-                            _SPANNER_DATABASE_ID)
-    storage = StorageClient(_GCS_BUCKET_ID)
+    spanner = SpannerClient(FLAGS.spanner_project_id, FLAGS.spanner_instance_id,
+                            FLAGS.spanner_database_id)
+    storage = StorageClient(FLAGS.gcs_bucket_id)
 
     if actionType == 'get_import_list':
         # Gets the list of imports that are ready for ingestion.
@@ -76,8 +106,8 @@ def ingestion_helper(request):
         import_list = request_json['importList']
         workflow_id = request_json['workflowId']
         job_id = request_json['jobId']
-        metrics = import_utils.get_ingestion_metrics(_PROJECT_ID, _LOCATION,
-                                                     job_id)
+        metrics = import_utils.get_ingestion_metrics(FLAGS.project_id,
+                                                     FLAGS.location, job_id)
         spanner.update_ingestion_status(import_list, workflow_id, job_id,
                                         metrics)
         return ('Updated ingestion status', 200)
diff --git a/import-automation/workflow/ingestion-helper/requirements.txt b/import-automation/workflow/ingestion-helper/requirements.txt
@@ -3,4 +3,5 @@ google-cloud-spanner
 croniter
 google-api-python-client
 google-cloud-storage
-google-auth
+google-auth
+absl-py
diff --git a/import-automation/workflow/ingestion-helper/spanner_client.py b/import-automation/workflow/ingestion-helper/spanner_client.py
@@ -1,3 +1,17 @@
+# 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.
+
 import logging
 from google.cloud import spanner
 from google.cloud.spanner_v1 import Transaction
@@ -158,11 +172,13 @@ def _record(transaction: Transaction):
             # 2. Insert into the IngestionHistory table
             columns = [
                 "CompletionTimestamp", "WorkflowExecutionID", "DataflowJobId",
-                "IngestedImports", "NodeCount", "EdgeCount", "ObservationCount"
+                "IngestedImports", "ExecutionTime", "NodeCount", "EdgeCount",
+                "ObservationCount"
             ]
             values = [[
                 spanner.COMMIT_TIMESTAMP, workflow_id, job_id,
-                succeeded_imports, metrics['node_count'], metrics['edge_count'],
+                succeeded_imports, metrics['execution_time'],
+                metrics['node_count'], metrics['edge_count'],
                 metrics['obs_count']
             ]]
             transaction.insert_or_update(table="IngestionHistory",
diff --git a/import-automation/workflow/ingestion-helper/storage_client.py b/import-automation/workflow/ingestion-helper/storage_client.py
@@ -1,3 +1,17 @@
+# 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.
+
 """Storage client for the ingestion helper."""
 
 import logging
