diff --git a/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java b/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java index 1efbe88b..6c74942a 100644 --- a/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java +++ b/pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java @@ -33,14 +33,6 @@ public class GraphIngestionPipeline { private static final Logger LOGGER = LoggerFactory.getLogger(GraphIngestionPipeline.class); private static final Counter nodeInvalidTypeCounter = Metrics.counter(GraphIngestionPipeline.class, "mcf_nodes_without_type"); - private static final Counter nodeCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_node_count"); - private static final Counter edgeCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_edge_count"); - private static final Counter obsCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_observation_count"); - private static final Counter timeSeriesCounter = - Metrics.counter(GraphIngestionPipeline.class, "graph_timeseries_count"); // List of imports that require node combination. private static final List IMPORTS_TO_COMBINE = List.of("Schema", "Place", "Provenance"); @@ -137,6 +129,18 @@ private static void processImport( boolean isBaseDc = options.getIsBaseDc(); String provenance = ProvenanceUtils.getProvenanceDcid(importName, isBaseDc); + String shortName = + importName.contains(":") + ? importName.substring(importName.lastIndexOf(':') + 1) + : importName; + + Counter nodeCounter = Metrics.counter(GraphIngestionPipeline.class, "node_count:" + shortName); + Counter edgeCounter = Metrics.counter(GraphIngestionPipeline.class, "edge_count:" + shortName); + Counter obsCounter = + Metrics.counter(GraphIngestionPipeline.class, "observation_count:" + shortName); + Counter timeSeriesCounter = + Metrics.counter(GraphIngestionPipeline.class, "timeseries_count:" + shortName); + // 1. Prepare Deletes: // Generate mutations to delete existing data for this import/provenance. // Create a dummy signal if deletes are skipped, so downstream dependencies are satisfied @@ -162,7 +166,8 @@ private static void processImport( PipelineUtils.InputFormat format = PipelineUtils.resolveFormat(graphPath); if (format == PipelineUtils.InputFormat.TFRECORD) { TfRecordProcessingResult result = - processTfRecordImport(pipeline, importName, graphPath, isBaseDc); + processTfRecordImport( + pipeline, importName, graphPath, isBaseDc, obsCounter, timeSeriesCounter); writeToSpanner( pipeline, spannerClient, @@ -326,7 +331,12 @@ private static class TfRecordProcessingResult { } private static TfRecordProcessingResult processTfRecordImport( - Pipeline pipeline, String importName, String graphPath, boolean isBaseDc) { + Pipeline pipeline, + String importName, + String graphPath, + boolean isBaseDc, + Counter obsCounter, + Counter timeSeriesCounter) { PCollection optGraph = PipelineUtils.readOptimizedMcfGraph(importName, graphPath, pipeline); PCollection uniqueSeries = diff --git a/pipeline/workflow/ingestion-helper/clients/spanner.py b/pipeline/workflow/ingestion-helper/clients/spanner.py index 593b581d..5c1a874c 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner.py @@ -401,6 +401,9 @@ def update_import_version_history(self, logging.info( f"Updating ImportVersionHistory table for workflow {workflow_id}") + m = metrics if metrics else {} + import_metrics = m.get('import_metrics', {}) + def _insert(transaction: Transaction): columns = [ "ImportName", "Version", "UpdateTimestamp", @@ -408,17 +411,24 @@ def _insert(transaction: Transaction): "NodeCount", "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ] - m = metrics if metrics else {} version_history_values = [] for import_json in import_list_json: + import_name = import_json.get('importName') + if not import_name: + continue + + short_name = import_name.split(':')[-1] + counts = import_metrics.get(short_name) or import_json + version_history_values.append([ - import_json['importName'], import_json['latestVersion'], + import_name, import_json.get('latestVersion'), spanner.COMMIT_TIMESTAMP, workflow_id, status, m.get('execution_time'), - m.get('node_count'), - m.get('edge_count'), - m.get('obs_count'), - m.get('ts_count'), "ingestion-workflow:" + workflow_id + counts.get('node_count') or counts.get('nodeCount'), + counts.get('edge_count') or counts.get('edgeCount'), + counts.get('obs_count') or counts.get('obsCount'), + counts.get('ts_count') or counts.get('tsCount'), + "ingestion-workflow:" + workflow_id ]) if version_history_values: @@ -512,6 +522,12 @@ def update_version_history(self, import_name = import_name.split(':')[-1] logging.info(f"Updating version history for {import_name} to {version}") + m = metrics if metrics else {} + node_count = m.get('node_count') + edge_count = m.get('edge_count') + obs_count = m.get('obs_count') + ts_count = m.get('ts_count') + def _record(transaction: Transaction): columns = [ "ImportName", "Version", "UpdateTimestamp", @@ -519,15 +535,14 @@ def _record(transaction: Transaction): "NodeCount", "EdgeCount", "ObservationCount", "TimeSeriesCount", "Comment" ] - m = metrics if metrics else {} values = [[ import_name, version, spanner.COMMIT_TIMESTAMP, workflow_id, status, m.get('execution_time'), - m.get('node_count'), - m.get('edge_count'), - m.get('obs_count'), - m.get('ts_count'), comment + node_count, + edge_count, + obs_count, + ts_count, comment ]] transaction.insert(table="ImportVersionHistory", columns=columns, diff --git a/pipeline/workflow/ingestion-helper/clients/spanner_test.py b/pipeline/workflow/ingestion-helper/clients/spanner_test.py index bba47135..fb547d23 100644 --- a/pipeline/workflow/ingestion-helper/clients/spanner_test.py +++ b/pipeline/workflow/ingestion-helper/clients/spanner_test.py @@ -702,7 +702,45 @@ def run_in_transaction_side_effect(callback, *args, **kwargs): None, None, None, None, None, None, "ingestion-workflow:wf-789" ]]) + @patch('google.cloud.spanner.Client') + def test_update_import_version_history_per_import_counts(self, mock_spanner_client): + mock_instance = MagicMock() + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value = mock_instance + mock_instance.database.return_value = mock_db + + mock_transaction = MagicMock() + def run_in_transaction_side_effect(callback, *args, **kwargs): + return callback(mock_transaction, *args, **kwargs) + mock_db.run_in_transaction.side_effect = run_in_transaction_side_effect + + client = SpannerClient("project", "instance", "database") + + import_list = [ + {"importName": "import1", "latestVersion": "v1.0"}, + {"importName": "import2", "latestVersion": "v2.0"} + ] + metrics = { + 'execution_time': 120, + 'import_metrics': { + "import1": {'node_count': 10, 'edge_count': 20, 'ts_count': 5, 'obs_count': 15}, + "import2": {'node_count': 30, 'edge_count': 40, 'ts_count': 25, 'obs_count': 35}, + } + } + client.update_import_version_history( + import_list, workflow_id="wf-123", status="SUCCESS", metrics=metrics + ) + + mock_transaction.insert.assert_called_once() + _, kwargs = mock_transaction.insert.call_args + self.assertEqual(kwargs['table'], 'ImportVersionHistory') + self.assertEqual(kwargs['values'], [ + ["import1", "v1.0", spanner.COMMIT_TIMESTAMP, "wf-123", "SUCCESS", 120, 10, 20, 15, 5, "ingestion-workflow:wf-123"], + ["import2", "v2.0", spanner.COMMIT_TIMESTAMP, "wf-123", "SUCCESS", 120, 30, 40, 35, 25, "ingestion-workflow:wf-123"] + ]) + if __name__ == '__main__': unittest.main() + diff --git a/pipeline/workflow/ingestion-helper/routes/database.py b/pipeline/workflow/ingestion-helper/routes/database.py index d5a89f12..da439462 100644 --- a/pipeline/workflow/ingestion-helper/routes/database.py +++ b/pipeline/workflow/ingestion-helper/routes/database.py @@ -58,7 +58,7 @@ def acquire_ingestion_lock(req: LockAcquireRequest, spanner: SpannerClient = Dep status_ok = spanner.acquire_lock(req.workflowId, req.timeout) if not status_ok: raise HTTPException( - status_code=400, + status_code=503, detail=f"Failed to acquire lock: Lock already held or acquisition timed out for workflow {req.workflowId}" ) return BaseResponse(status=ResponseStatus.OK) diff --git a/pipeline/workflow/ingestion-helper/utils/imports.py b/pipeline/workflow/ingestion-helper/utils/imports.py index caed3ec8..f9364847 100644 --- a/pipeline/workflow/ingestion-helper/utils/imports.py +++ b/pipeline/workflow/ingestion-helper/utils/imports.py @@ -127,7 +127,8 @@ 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', 'ts_count', and 'execution_time'. + A dictionary containing 'obs_count', 'node_count', 'edge_count', 'ts_count', 'execution_time', + and 'import_metrics'. """ dataflow = build('dataflow', 'v1b3', cache_discovery=False) # Fetch Dataflow metrics @@ -136,6 +137,7 @@ def get_ingestion_metrics(project_id, location, job_id): obs_count = 0 ts_count = 0 execution_time = 0 + import_metrics = {} if project_id and job_id: try: # Fetch Job details for execution time @@ -159,14 +161,35 @@ def get_ingestion_metrics(project_id, location, job_id): jobId=job_id).execute() for metric in metrics.get('metrics', []): name = metric['name']['name'] - if name == 'graph_node_count': - node_count += int(metric['scalar']) - elif name == 'graph_edge_count': - 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']) + scalar = int(metric.get('scalar', 0)) + match = re.match(r"^([^:]+)(?::(.+))?$", name) + metric_type = match.group(1) if match else name + imp_name = match.group(2).split(':')[-1] if match and match.group(2) else None + + if imp_name and imp_name not in import_metrics: + import_metrics[imp_name] = { + 'node_count': 0, + 'edge_count': 0, + 'obs_count': 0, + 'ts_count': 0, + } + + if metric_type == 'node_count': + node_count += scalar + if imp_name: + import_metrics[imp_name]['node_count'] = scalar + elif metric_type == 'edge_count': + edge_count += scalar + if imp_name: + import_metrics[imp_name]['edge_count'] = scalar + elif metric_type == 'observation_count': + obs_count += scalar + if imp_name: + import_metrics[imp_name]['obs_count'] = scalar + elif metric_type == 'timeseries_count': + ts_count += scalar + if imp_name: + import_metrics[imp_name]['ts_count'] = scalar except HttpError as e: logging.error( f"Error fetching dataflow metrics for job {job_id}: {e}") @@ -175,5 +198,6 @@ def get_ingestion_metrics(project_id, location, job_id): 'node_count': node_count, 'edge_count': edge_count, 'ts_count': ts_count, - 'execution_time': execution_time + 'execution_time': execution_time, + 'import_metrics': import_metrics, }