diff --git a/pipeline/workflow/ingestion-helper/app_test.py b/pipeline/workflow/ingestion-helper/app_test.py index 8ae6ca7a..f3dfb50b 100644 --- a/pipeline/workflow/ingestion-helper/app_test.py +++ b/pipeline/workflow/ingestion-helper/app_test.py @@ -275,7 +275,8 @@ def test_update_ingestion_status_v1(self, mock_get_ingestion_metrics): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } mock_get_ingestion_metrics.return_value = mock_metrics @@ -377,7 +378,8 @@ def test_update_ingestion_history_success(self, mock_get_ingestion_metrics): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } mock_get_ingestion_metrics.return_value = mock_metrics @@ -415,7 +417,8 @@ def test_update_ingestion_history_failure(self, mock_get_ingestion_metrics): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } mock_get_ingestion_metrics.return_value = mock_metrics @@ -451,7 +454,8 @@ def test_update_ingestion_history_retry(self, mock_get_ingestion_metrics): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } mock_get_ingestion_metrics.return_value = mock_metrics diff --git a/pipeline/workflow/ingestion-helper/clients/schema.sql b/pipeline/workflow/ingestion-helper/clients/schema.sql index eba6443f..aaf4bc9d 100644 --- a/pipeline/workflow/ingestion-helper/clients/schema.sql +++ b/pipeline/workflow/ingestion-helper/clients/schema.sql @@ -126,6 +126,7 @@ CREATE TABLE IngestionHistory ( NodeCount INT64, EdgeCount INT64, ObservationCount INT64, + TimeSeriesCount INT64, ) PRIMARY KEY(WorkflowExecutionID); CREATE TABLE ImportVersionHistory ( @@ -138,6 +139,7 @@ CREATE TABLE ImportVersionHistory ( NodeCount INT64, EdgeCount INT64, ObservationCount INT64, + TimeSeriesCount INT64, Comment STRING(MAX), ) PRIMARY KEY (ImportName, UpdateTimestamp DESC); diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index 172f9dca..1ac42c9d 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -317,7 +317,8 @@ def _insert(transaction: Transaction): columns = [ "CompletionTimestamp", "IngestionFailure", "WorkflowExecutionID", "DataflowJobID", "IngestedImports", - "ExecutionTime", "NodeCount", "EdgeCount", "ObservationCount" + "ExecutionTime", "NodeCount", "EdgeCount", "ObservationCount", + "TimeSeriesCount" ] m = metrics if metrics else {} values = [[ @@ -329,7 +330,8 @@ def _insert(transaction: Transaction): m.get('execution_time'), m.get('node_count'), m.get('edge_count'), - m.get('obs_count') + m.get('obs_count'), + m.get('ts_count') ]] transaction.insert_or_update(table="IngestionHistory", columns=columns, @@ -401,7 +403,8 @@ def _update(transaction: Transaction): 'execution_time': 'ExecutionTime', 'node_count': 'NodeCount', 'edge_count': 'EdgeCount', - 'obs_count': 'ObservationCount' + 'obs_count': 'ObservationCount', + 'ts_count': 'TimeSeriesCount', } for key, column in metric_map.items(): if key in metrics: @@ -446,7 +449,7 @@ def _insert(transaction: Transaction): version_history_columns = [ "ImportName", "Version", "UpdateTimestamp", "WorkflowExecutionID", "Status", "ExecutionTime", "NodeCount", - "EdgeCount", "ObservationCount", "Comment" + "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ] m = metrics if metrics else {} version_history_values = [] @@ -459,6 +462,7 @@ def _insert(transaction: Transaction): m.get('node_count'), m.get('edge_count'), m.get('obs_count'), + m.get('ts_count'), "ingestion-workflow:" + workflow_id ]) @@ -558,7 +562,7 @@ def _record(transaction: Transaction): columns = [ "ImportName", "Version", "UpdateTimestamp", "WorkflowExecutionID", "Status", "ExecutionTime", "NodeCount", - "EdgeCount", "ObservationCount", "Comment" + "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ] m = metrics if metrics else {} values = [[ @@ -568,6 +572,7 @@ def _record(transaction: Transaction): m.get('node_count'), m.get('edge_count'), m.get('obs_count'), + m.get('ts_count'), comment ]] transaction.insert(table="ImportVersionHistory", diff --git a/pipeline/workflow/ingestion-helper/clients/spanner_test.py b/pipeline/workflow/ingestion-helper/clients/spanner_test.py index 12a04965..4cb6834b 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner_test.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner_test.py @@ -444,7 +444,8 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } ) @@ -454,7 +455,8 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): self.assertEqual(kwargs['columns'], [ "CompletionTimestamp", "IngestionFailure", "WorkflowExecutionID", "DataflowJobID", "IngestedImports", - "ExecutionTime", "NodeCount", "EdgeCount", "ObservationCount" + "ExecutionTime", "NodeCount", "EdgeCount", "ObservationCount", + "TimeSeriesCount" ]) self.assertEqual(kwargs['values'], [[ spanner.COMMIT_TIMESTAMP, @@ -465,7 +467,8 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): 120, 1000, 2000, - 500 + 500, + 300 ]]) @patch('google.cloud.spanner.Client') @@ -542,7 +545,8 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } ) @@ -559,6 +563,7 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): self.assertIn("DataflowJobID", kwargs['columns']) self.assertIn("ExecutionTime", kwargs['columns']) self.assertIn("NodeCount", kwargs['columns']) + self.assertIn("TimeSeriesCount", kwargs['columns']) # Verify the specific values wf_idx = kwargs['columns'].index("WorkflowExecutionID") @@ -566,6 +571,7 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): failure_idx = kwargs['columns'].index("IngestionFailure") job_idx = kwargs['columns'].index("DataflowJobID") node_idx = kwargs['columns'].index("NodeCount") + ts_idx = kwargs['columns'].index("TimeSeriesCount") comp_time_idx = kwargs['columns'].index("CompletionTimestamp") values = kwargs['values'][0] @@ -574,6 +580,7 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): self.assertEqual(values[failure_idx], False) self.assertEqual(values[job_idx], "job-456") self.assertEqual(values[node_idx], 1000) + self.assertEqual(values[ts_idx], 300) self.assertEqual(values[comp_time_idx], spanner.COMMIT_TIMESTAMP) @patch('google.cloud.spanner.Client') @@ -598,7 +605,8 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } ) @@ -663,7 +671,8 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): 'execution_time': 120, 'node_count': 1000, 'edge_count': 2000, - 'obs_count': 500 + 'obs_count': 500, + 'ts_count': 300, } ) @@ -702,11 +711,11 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): self.assertEqual(kwargs['columns'], [ "ImportName", "Version", "UpdateTimestamp", "WorkflowExecutionID", "Status", "ExecutionTime", "NodeCount", - "EdgeCount", "ObservationCount", "Comment" + "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ]) self.assertEqual(kwargs['values'], [[ "test_import", "v1.2.3", spanner.COMMIT_TIMESTAMP, "wf-123", - None, None, None, None, None, "ingestion-workflow:wf-123" + None, None, None, None, None, None, "ingestion-workflow:wf-123" ]]) @patch('google.cloud.spanner.Client') @@ -730,11 +739,11 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): self.assertEqual(kwargs['columns'], [ "ImportName", "Version", "UpdateTimestamp", "WorkflowExecutionID", "Status", "ExecutionTime", "NodeCount", - "EdgeCount", "ObservationCount", "Comment" + "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ]) self.assertEqual(kwargs['values'], [[ "test_import", "v1.2.3", spanner.COMMIT_TIMESTAMP, "wf-789", - None, None, None, None, None, "ingestion-workflow:wf-789" + None, None, None, None, None, None, "ingestion-workflow:wf-789" ]]) if __name__ == '__main__': diff --git a/pipeline/workflow/ingestion-helper/utils/imports.py b/pipeline/workflow/ingestion-helper/utils/imports.py index 33f9d1fa..caed3ec8 100644 --- a/pipeline/workflow/ingestion-helper/utils/imports.py +++ b/pipeline/workflow/ingestion-helper/utils/imports.py @@ -127,13 +127,14 @@ def get_ingestion_metrics(project_id, location, job_id): job_id: The Dataflow job ID. Returns: - A dictionary containing 'obs_count', 'node_count', 'edge_count', and 'execution_time'. + A dictionary containing 'obs_count', 'node_count', 'edge_count', 'ts_count', and 'execution_time'. """ dataflow = build('dataflow', 'v1b3', cache_discovery=False) # Fetch Dataflow metrics node_count = 0 edge_count = 0 obs_count = 0 + ts_count = 0 execution_time = 0 if project_id and job_id: try: @@ -164,6 +165,8 @@ def get_ingestion_metrics(project_id, location, job_id): edge_count += int(metric['scalar']) elif name == 'graph_observation_count': obs_count += int(metric['scalar']) + elif name == 'graph_timeseries_count': + ts_count += int(metric['scalar']) except HttpError as e: logging.error( f"Error fetching dataflow metrics for job {job_id}: {e}") @@ -171,5 +174,6 @@ def get_ingestion_metrics(project_id, location, job_id): 'obs_count': obs_count, 'node_count': node_count, 'edge_count': edge_count, + 'ts_count': ts_count, 'execution_time': execution_time } diff --git a/pipeline/workflow/spanner-ingestion-workflow.yaml b/pipeline/workflow/spanner-ingestion-workflow.yaml index c53b2205..444d3414 100644 --- a/pipeline/workflow/spanner-ingestion-workflow.yaml +++ b/pipeline/workflow/spanner-ingestion-workflow.yaml @@ -15,6 +15,7 @@ main: - helper_url: ${"https://ingestion-helper-service-" + sys.get_env("PROJECT_NUMBER") + "." + location + ".run.app"} - import_list: ${default(map.get(args, "importList"), [])} - skip_aggregation: ${default(map.get(args, "skipAggregation"), true)} + - import_info: null - dataflow_job_id: null - execution_error: null - acquire_ingestion_lock: @@ -39,14 +40,28 @@ main: try: steps: - get_import_info: - call: http.post - args: - url: ${helper_url + "/imports/info"} - auth: - type: OIDC - body: - importList: ${import_list} - result: import_info + steps: + - fetch_import_info: + call: http.post + args: + url: ${helper_url + "/imports/info"} + auth: + type: OIDC + body: + importList: ${import_list} + result: import_info + - update_ingestion_history_pending: + call: http.post + args: + url: ${helper_url + "/imports/ingestion-history"} + auth: + type: OIDC + body: + workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' + status: 'PENDING' + stage: 'dataflow' + importList: '${import_info.body}' + result: pending_history_response - run_ingestion_job: call: run_dataflow_job args: @@ -73,32 +88,62 @@ main: project_id: ${project_id} location: ${location} - update_ingestion_status: - call: http.post - args: - url: ${helper_url + "/imports/ingestion-status"} - auth: - type: OIDC - body: - workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' - jobId: '${dataflow_job_id}' - importList: '${import_info.body}' - status: 'SUCCESS' - result: function_response + steps: + - update_status_success: + call: http.post + args: + url: ${helper_url + "/imports/ingestion-status"} + auth: + type: OIDC + body: + workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' + jobId: '${dataflow_job_id}' + importList: '${import_info.body}' + status: 'SUCCESS' + result: function_response + - update_ingestion_history_success: + call: http.post + args: + url: ${helper_url + "/imports/ingestion-history"} + auth: + type: OIDC + body: + workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' + jobId: '${dataflow_job_id}' + importList: '${import_info.body}' + status: 'SUCCESS' + stage: 'dataflow' + result: success_history_response except: as: e steps: - record_failure: - call: http.post - args: - url: ${helper_url + "/imports/ingestion-status"} - auth: - type: OIDC - body: - workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' - jobId: ${default(dataflow_job_id, default(map.get(e, "job_id"), "N/A"))} - importList: '${import_info.body}' - status: 'RETRY' - result: retry_response + steps: + - update_status_retry: + call: http.post + args: + url: ${helper_url + "/imports/ingestion-status"} + auth: + type: OIDC + body: + workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' + jobId: ${default(dataflow_job_id, default(map.get(e, "job_id"), "N/A"))} + importList: '${default(import_info, {"body": []}).body}' + status: 'RETRY' + result: retry_response + - update_ingestion_history_retry: + call: http.post + args: + url: ${helper_url + "/imports/ingestion-history"} + auth: + type: OIDC + body: + workflowId: '${sys.get_env("GOOGLE_CLOUD_WORKFLOW_EXECUTION_ID")}' + jobId: ${default(dataflow_job_id, default(map.get(e, "job_id"), "N/A"))} + importList: '${default(import_info, {"body": []}).body}' + status: 'RETRY' + stage: 'dataflow' + result: retry_history_response - capture_error: assign: - execution_error: ${e} @@ -159,26 +204,40 @@ run_dataflow_job: - condition: ${import_list == "[]"} return: '' - launch_dataflow_job: - call: googleapis.dataflow.v1b3.projects.locations.flexTemplates.launch - args: - projectId: '${project_id}' - location: '${location}' - body: - launchParameter: - containerSpecGcsPath: '${template_gcs_path}' - jobName: '${jobName}' - parameters: - importList: '${import_list}' - projectId: '${spanner_project_id}' - spannerInstanceId: '${spanner_instance_id}' - spannerDatabaseId: '${spanner_database_id}' - # Revert this once switched to the new DB. - spannerObservationTableName: 'Observation_final_v3' - spannerTimeSeriesTableName: 'TimeSeries_final_v3' - environment: - numWorkers: ${numWorkers} - machineType: ${machineType} - result: launch_result + steps: + - launch_job: + call: googleapis.dataflow.v1b3.projects.locations.flexTemplates.launch + args: + projectId: '${project_id}' + location: '${location}' + body: + launchParameter: + containerSpecGcsPath: '${template_gcs_path}' + jobName: '${jobName}' + parameters: + importList: '${import_list}' + projectId: '${spanner_project_id}' + spannerInstanceId: '${spanner_instance_id}' + spannerDatabaseId: '${spanner_database_id}' + # Revert this once switched to the new DB. + spannerObservationTableName: 'Observation_final_v3' + spannerTimeSeriesTableName: 'TimeSeries_final_v3' + environment: + numWorkers: ${numWorkers} + machineType: ${machineType} + result: launch_result + - update_ingestion_history_running: + call: http.post + args: + url: ${helper_url + "/imports/ingestion-history"} + auth: + type: OIDC + body: + workflowId: '${workflow_id}' + status: 'RUNNING' + stage: 'dataflow' + jobId: '${launch_result.job.id}' + result: running_history_response - wait_for_job_completion: call: sys.sleep args: