Revert "Add per-provenance counters to dataflow and spanner (#658)" (#694)

Reverting commit af33c21 to allow DCP
rollout. Will resubmit afterwards.

Author: vish-cs

pipeline/ingestion/src/main/java/org/datacommons/ingestion/pipeline/GraphIngestionPipeline.java

@@ -33,6 +33,14 @@ 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");
@@ -129,18 +137,6 @@ 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
@@ -166,8 +162,7 @@ private static void processImport(
     PipelineUtils.InputFormat format = PipelineUtils.resolveFormat(graphPath);
     if (format == PipelineUtils.InputFormat.TFRECORD) {
       TfRecordProcessingResult result =
-          processTfRecordImport(
-              pipeline, importName, graphPath, isBaseDc, obsCounter, timeSeriesCounter);
+          processTfRecordImport(pipeline, importName, graphPath, isBaseDc);
       writeToSpanner(
           pipeline,
           spannerClient,
@@ -331,12 +326,7 @@ private static class TfRecordProcessingResult {
   }
 
   private static TfRecordProcessingResult processTfRecordImport(
-      Pipeline pipeline,
-      String importName,
-      String graphPath,
-      boolean isBaseDc,
-      Counter obsCounter,
-      Counter timeSeriesCounter) {
+      Pipeline pipeline, String importName, String graphPath, boolean isBaseDc) {
     PCollection<McfOptimizedGraph> optGraph =
         PipelineUtils.readOptimizedMcfGraph(importName, graphPath, pipeline);
     PCollection<TimeSeries> uniqueSeries =

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

@@ -401,34 +401,24 @@ 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_name, import_json.get('latestVersion'),
+                    import_json['importName'], import_json['latestVersion'],
                     spanner.COMMIT_TIMESTAMP, workflow_id, status,
                     m.get('execution_time'),
-                    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
+                    m.get('node_count'),
+                    m.get('edge_count'),
+                    m.get('obs_count'),
+                    m.get('ts_count'), "ingestion-workflow:" + workflow_id
                 ])
 
             if version_history_values:
@@ -522,27 +512,22 @@ 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'),
-                node_count,
-                edge_count,
-                obs_count,
-                ts_count, comment
+                m.get('node_count'),
+                m.get('edge_count'),
+                m.get('obs_count'),
+                m.get('ts_count'), comment
             ]]
             transaction.insert(table="ImportVersionHistory",
                                columns=columns,

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

@@ -702,45 +702,7 @@ 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=503,
+                status_code=400,
                 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,8 +127,7 @@ 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', 'execution_time',
-        and 'import_metrics'.
+        A dictionary containing 'obs_count', 'node_count', 'edge_count', 'ts_count', and 'execution_time'.
     """
     dataflow = build('dataflow', 'v1b3', cache_discovery=False)
     # Fetch Dataflow metrics
@@ -137,7 +136,6 @@ 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
@@ -161,35 +159,14 @@ def get_ingestion_metrics(project_id, location, job_id):
                 jobId=job_id).execute()
             for metric in metrics.get('metrics', []):
                 name = metric['name']['name']
-                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
+                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'])
         except HttpError as e:
             logging.error(
                 f"Error fetching dataflow metrics for job {job_id}: {e}")
@@ -198,6 +175,5 @@ 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,
-        'import_metrics': import_metrics,
+        'execution_time': execution_time
     }