Add per-provenance counters to dataflow and spanner

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<String> 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<McfOptimizedGraph> optGraph =
         PipelineUtils.readOptimizedMcfGraph(importName, graphPath, pipeline);
     PCollection<TimeSeries> uniqueSeries =

pipeline/workflow/ingestion-helper/clients/spanner.py

@@ -401,24 +401,34 @@ 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",
                 "WorkflowExecutionID", "Status", "ExecutionTime",
                 "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,22 +522,27 @@ 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",
                 "WorkflowExecutionID", "Status", "ExecutionTime",
                 "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,

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()
 
+

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)

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,
     }