diff --git a/pipeline/workflow/aggregation-helper/.dockerignore b/pipeline/workflow/aggregation-helper/.dockerignore new file mode 100644 index 00000000..f7e2d9fb --- /dev/null +++ b/pipeline/workflow/aggregation-helper/.dockerignore @@ -0,0 +1,7 @@ +.venv +__pycache__ +*.pyc +.pytest_cache +.coverage +.git +.gitignore diff --git a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py index 8e1789dd..dccb1dd6 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/aggregation_test.py @@ -236,20 +236,71 @@ def test_run_all(self): mock_job = MagicMock() self.mock_executor.execute.return_value = mock_job - jobs = generator.run_all(LinkedEdgeConfig(import_names=["import1", "import2"])) + jobs = generator.run_all(LinkedEdgeConfig( + import_names=["import1", "import2"] + )) - self.assertEqual(len(jobs), 3) # Should run 3 queries + self.assertEqual(len(jobs), 3) # Runs the 3 scoped linked edge queries self.assertEqual(self.mock_executor.execute.call_count, 3) - # Verify queries contain import names and connection id + # Verify queries contain connection id and spanner destination uri calls = self.mock_executor.execute.call_args_list for call in calls: query = call[0][0] self.assertIn("test-conn", query) - self.assertIn("import1", query) - self.assertIn("import2", query) self.assertIn("spanner-uri", query) - self.assertIn("dc/base/import1", query) # Since is_base_dc=True + + def test_run_topic_list_edges(self): + generator = LinkedEdgeGenerator(self.mock_executor, is_base_dc=True) + + mock_job = MagicMock() + self.mock_executor.execute.return_value = mock_job + + job = generator.run_topic_list_edges() + self.assertEqual(job, mock_job) + self.mock_executor.execute.assert_called_once() + + query = self.mock_executor.execute.call_args[0][0] + self.assertIn("temp_raw_topic_edges", query) + self.assertIn("temp_topic_types", query) + self.assertIn("temp_topic_nodes", query) + self.assertIn("temp_svpg_nodes", query) + self.assertIn("relevantVariableList", query) + self.assertIn("memberList", query) + self.assertIn("STRING_AGG(DISTINCT e.object_id, ',' ORDER BY e.object_id)", query) + self.assertIn("CONCAT(SUBSTR(TRIM(list_value), 1, 16), ':', TO_HEX(SHA256(TRIM(list_value))))", query) + self.assertIn('spanner_options = \'{"table": "Node"}\'', query) + self.assertIn('spanner_options = \'{"table": "Edge"}\'', query) + self.assertIn("dc/base/generated/TopicLists", query) + + def test_run_topic_list_edges_not_base_dc(self): + generator = LinkedEdgeGenerator(self.mock_executor, is_base_dc=False) + + mock_job = MagicMock() + self.mock_executor.execute.return_value = mock_job + + job = generator.run_topic_list_edges() + self.assertEqual(job, mock_job) + + query = self.mock_executor.execute.call_args[0][0] + self.assertIn("generated/TopicLists", query) + self.assertNotIn("dc/base/generated/TopicLists", query) + + def test_run_linked_member(self): + generator = LinkedEdgeGenerator(self.mock_executor, is_base_dc=True) + + mock_job = MagicMock() + self.mock_executor.execute.return_value = mock_job + + job = generator.run_linked_member(import_names=["import1"]) + self.assertEqual(job, mock_job) + + query = self.mock_executor.execute.call_args[0][0] + self.assertIn("temp_topic_types", query) + self.assertIn("temp_topic_nodes", query) + self.assertIn("temp_svpg_nodes", query) + self.assertIn("relevantVariable", query) + self.assertIn("linkedMember", query) class TestProvenanceSummaryGenerator(unittest.TestCase): diff --git a/pipeline/workflow/aggregation-helper/aggregation/deleter.py b/pipeline/workflow/aggregation-helper/aggregation/deleter.py index 61b2d851..d5505e97 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/deleter.py +++ b/pipeline/workflow/aggregation-helper/aggregation/deleter.py @@ -14,7 +14,6 @@ """Deletes aggregated data in Spanner using Partitioned DML.""" -import concurrent.futures import logging from typing import List from google.cloud import spanner @@ -67,24 +66,15 @@ def delete_aggregated_data(self, imports_to_delete: List[str]) -> None: ("KeyValueStore", "DELETE FROM KeyValueStore WHERE type = 'ProvenanceSummary' AND provenance IN UNNEST(@provenances)", "") ] - def _execute_delete(table_name: str, sql: str, extra_desc: str) -> int: - rows = db.execute_partitioned_dml( - sql, params=params, param_types=param_types - ) - logging.info(f"Deleted {rows} rows from {table_name} table{extra_desc}.") - return rows - - try: - with concurrent.futures.ThreadPoolExecutor(max_workers=len(delete_queries)) as executor: - futures = [ - executor.submit(_execute_delete, table, sql, desc) - for table, sql, desc in delete_queries - ] - for future in concurrent.futures.as_completed(futures): - future.result() # Propagate any worker thread exceptions to main thread - except Exception as e: - logging.error(f"Failed to execute partitioned DML for deletions: {e}") - raise + for table_name, sql, extra_desc in delete_queries: + try: + rows = db.execute_partitioned_dml( + sql, params=params, param_types=param_types + ) + logging.info(f"Deleted {rows} rows from {table_name} table{extra_desc}.") + except Exception as e: + logging.error(f"Failed to execute partitioned DML for deletions on {table_name}: {e}") + raise def delete_stat_var_group_edges(self) -> int: """Deletes all generated StatVarGroup edges across all provenances in Spanner.""" @@ -115,3 +105,17 @@ def delete_linked_edges(self, imports_to_delete: List[str]) -> int: rows = self.spanner_database.execute_partitioned_dml(sql, params=params, param_types=param_types) logging.info(f"Deleted {rows} linked relationship edges for imports: {imports_to_delete}") return rows + + def delete_topic_list_edges(self) -> int: + """Deletes consolidated topic and peer group list edges from Spanner.""" + provenance_name = get_provenance_name("generated/TopicLists", self.is_base_dc) + sql = ( + "DELETE FROM Edge " + "WHERE provenance = @provenance " + "AND predicate IN ('relevantVariableList', 'memberList')" + ) + params = {"provenance": provenance_name} + param_types = {"provenance": spanner.param_types.STRING} + rows = self.spanner_database.execute_partitioned_dml(sql, params=params, param_types=param_types) + logging.info(f"Deleted {rows} topic and peer group list edges for provenance: {provenance_name}") + return rows diff --git a/pipeline/workflow/aggregation-helper/aggregation/deleter_test.py b/pipeline/workflow/aggregation-helper/aggregation/deleter_test.py index 98e3ac16..146a6c57 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/deleter_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/deleter_test.py @@ -87,12 +87,14 @@ def test_delete_aggregated_data_not_base_dc(self, mock_spanner_client): @patch('aggregation.deleter.spanner.Client') def test_delete_aggregated_data_exception_propagates(self, mock_spanner_client): - """Verifies that an exception raised in a worker thread is re-raised by delete_aggregated_data.""" + """Verifies that an exception raised during partitioned DML is re-raised by delete_aggregated_data.""" mock_db = MagicMock() mock_db.execute_partitioned_dml.side_effect = RuntimeError("Spanner deletion error") mock_spanner_client.return_value.instance.return_value.database.return_value = mock_db deleter = AggregationDeleter("proj", "inst", "db") + with self.assertRaises(RuntimeError): + deleter.delete_aggregated_data(["ImportA"]) @patch('aggregation.deleter.spanner.Client') def test_delete_stat_var_group_edges(self, mock_spanner_client): @@ -127,8 +129,48 @@ def test_delete_linked_edges(self, mock_spanner_client): self.assertIn("DELETE FROM Edge", sql) self.assertIn("provenance IN UNNEST(@provenances)", sql) self.assertIn("'linkedContainedInPlace'", sql) + self.assertIn("'linkedMemberOf'", sql) + self.assertIn("'linkedMember'", sql) + self.assertNotIn("'relevantVariableList'", sql) + self.assertNotIn("'memberList'", sql) self.assertEqual(params, {"provenances": ["dc/base/generated/ImportA"]}) + @patch('aggregation.deleter.spanner.Client') + def test_delete_topic_list_edges_base_dc(self, mock_spanner_client): + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value.database.return_value = mock_db + + deleter = AggregationDeleter("proj", "inst", "db", is_base_dc=True) + deleter.delete_topic_list_edges() + + mock_db.execute_partitioned_dml.assert_called_once() + call_args = mock_db.execute_partitioned_dml.call_args + sql = call_args[0][0] + params = call_args[1]["params"] + self.assertIn("DELETE FROM Edge", sql) + self.assertIn("provenance = @provenance", sql) + self.assertIn("'relevantVariableList'", sql) + self.assertIn("'memberList'", sql) + self.assertEqual(params, {"provenance": "dc/base/generated/TopicLists"}) + + @patch('aggregation.deleter.spanner.Client') + def test_delete_topic_list_edges_custom_dc(self, mock_spanner_client): + mock_db = MagicMock() + mock_spanner_client.return_value.instance.return_value.database.return_value = mock_db + + deleter = AggregationDeleter("proj", "inst", "db", is_base_dc=False) + deleter.delete_topic_list_edges() + + mock_db.execute_partitioned_dml.assert_called_once() + call_args = mock_db.execute_partitioned_dml.call_args + sql = call_args[0][0] + params = call_args[1]["params"] + self.assertIn("DELETE FROM Edge", sql) + self.assertIn("provenance = @provenance", sql) + self.assertIn("'relevantVariableList'", sql) + self.assertIn("'memberList'", sql) + self.assertEqual(params, {"provenance": "generated/TopicLists"}) + if __name__ == '__main__': unittest.main() diff --git a/pipeline/workflow/aggregation-helper/aggregation/e2e_tests/linked_edge_generator_test.py b/pipeline/workflow/aggregation-helper/aggregation/e2e_tests/linked_edge_generator_test.py index e6dcd183..bec5cf61 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/e2e_tests/linked_edge_generator_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/e2e_tests/linked_edge_generator_test.py @@ -302,5 +302,98 @@ def test_linked_edges_multiple_imports(self): self.assertEqual(len(res_b), 3, "ImportB should have 3 scoped linked edges.") + def test_topic_and_svpg_list_edges(self): + """Tests materialization of relevantVariableList and memberList edges and literal nodes.""" + import_name = 'TopicTest_Import' + + # 1. Setup mock topics and SVPGs + self.add_node('dc/topic/Environment', 'Environment', types=['Topic']) + self.add_node('custom/topic/AirQuality', 'Air Quality') # Type provided via typeOf edge + self.add_node('dc/svpg/AgeGroups', 'Age Groups', types=['StatVarPeerGroup']) + + self.add_edge('custom/topic/AirQuality', 'typeOf', 'Topic', import_name) + self.add_edge('dc/svpg/AgeGroups', 'typeOf', 'StatVarPeerGroup', import_name) + + # Add relevantVariable and member arcs + self.add_edge('dc/topic/Environment', 'relevantVariable', 'Count_Person', import_name) + self.add_edge('dc/topic/Environment', 'relevantVariable', 'dc/topic/Water', import_name) + + self.add_edge('custom/topic/AirQuality', 'relevantVariable', 'AirPollution_PM25', import_name) + self.add_edge('custom/topic/AirQuality', 'relevantVariable', 'AirPollution_O3', import_name) + + self.add_edge('dc/svpg/AgeGroups', 'member', 'Count_Person_18To64', import_name) + self.add_edge('dc/svpg/AgeGroups', 'member', 'Count_Person_0To17', import_name) + + self.flush_to_spanner() + + calculations = [ + { + "name": "Linked Edges With Topic Lists", + "type": "LINKED_EDGES", + "stage": 1, + "input_imports": [import_name], + "generate_topic_list_edges": True + } + ] + res = self.run_orchestrator(calculations=calculations, active_imports=[import_name]) + self.assertTrue(res.success) + + expected_provenance = f'dc/base/generated/TopicLists' if self.is_base_dc else 'generated/TopicLists' + + with self.database.snapshot(multi_use=True) as snapshot: + # Verify Edge records + edge_query = """ + SELECT subject_id, predicate, provenance + FROM Edge + WHERE predicate IN ('relevantVariableList', 'memberList') + ORDER BY subject_id + """ + edges = list(snapshot.execute_sql(edge_query)) + self.assertEqual(len(edges), 3) + self.assertEqual(tuple(edges[0]), ('custom/topic/AirQuality', 'relevantVariableList', expected_provenance)) + self.assertEqual(tuple(edges[1]), ('dc/svpg/AgeGroups', 'memberList', expected_provenance)) + self.assertEqual(tuple(edges[2]), ('dc/topic/Environment', 'relevantVariableList', expected_provenance)) + + # Verify Node records contain the aggregated CSV string values + node_query = """ + SELECT n.value + FROM Node n + JOIN Edge e ON n.subject_id = e.object_id + WHERE e.predicate IN ('relevantVariableList', 'memberList') + ORDER BY e.subject_id + """ + nodes = [r[0] for r in snapshot.execute_sql(node_query)] + self.assertEqual(nodes, [ + 'AirPollution_O3,AirPollution_PM25', + 'Count_Person_0To17,Count_Person_18To64', + 'Count_Person,dc/topic/Water' + ]) + + def test_topic_list_edges_disabled(self): + """Tests that disabling generate_topic_list_edges skips topic/SVPG list edge materialization.""" + import_name = 'TopicDisabledTest_Import' + + self.add_node('dc/topic/Environment', 'Environment', types=['Topic']) + self.add_edge('dc/topic/Environment', 'relevantVariable', 'Count_Person', import_name) + self.flush_to_spanner() + + calculations = [ + { + "name": "Linked Edges Without Topic Lists", + "type": "LINKED_EDGES", + "stage": 1, + "input_imports": [import_name], + "generate_topic_list_edges": False + } + ] + res = self.run_orchestrator(calculations=calculations, active_imports=[import_name]) + self.assertTrue(res.success) + + with self.database.snapshot() as snapshot: + query = "SELECT count(*) FROM Edge WHERE predicate = 'relevantVariableList'" + count = list(snapshot.execute_sql(query))[0][0] + self.assertEqual(count, 0) + + class LinkedEdgeGeneratorCustomDcTest(LinkedEdgeGeneratorIntegrationTest): is_base_dc = False diff --git a/pipeline/workflow/aggregation-helper/aggregation/linked_edge_generator.py b/pipeline/workflow/aggregation-helper/aggregation/linked_edge_generator.py index fcd42fd8..86ba594a 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/linked_edge_generator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/linked_edge_generator.py @@ -26,6 +26,7 @@ class LinkedEdgeConfig: """Configuration for linked edge generation.""" import_names: Optional[List[str]] = None + generate_topic_list_edges: bool = False class LinkedEdgeGenerator: @@ -219,16 +220,34 @@ def run_linked_member( SELECT * FROM EXTERNAL_QUERY("{self.executor.connection_id}", "SELECT subject_id, predicate, object_id, provenance FROM Edge WHERE predicate IN ('relevantVariable', 'member'){provenance_filter}"); + CREATE OR REPLACE TEMPORARY TABLE `temp_topic_types` AS + SELECT * FROM EXTERNAL_QUERY("{self.executor.connection_id}", + "SELECT subject_id, object_id FROM Edge WHERE predicate = 'typeOf' AND object_id IN ('Topic', 'StatVarPeerGroup')"); + + CREATE OR REPLACE TEMPORARY TABLE `temp_topic_nodes` AS + SELECT DISTINCT subject_id FROM `temp_topic_types` WHERE object_id = 'Topic' + UNION DISTINCT + SELECT DISTINCT subject_id FROM `temp_base_member` WHERE predicate = 'relevantVariable'; + + CREATE OR REPLACE TEMPORARY TABLE `temp_svpg_nodes` AS + SELECT DISTINCT subject_id FROM `temp_topic_types` WHERE object_id = 'StatVarPeerGroup'; + + CREATE OR REPLACE TEMPORARY TABLE `temp_all_topic_svpg_nodes` AS + SELECT subject_id FROM `temp_topic_nodes` + UNION DISTINCT + SELECT subject_id FROM `temp_svpg_nodes`; + CREATE OR REPLACE TEMPORARY TABLE `temp_topic_hierarchy` AS - SELECT DISTINCT subject_id, object_id, provenance - FROM `temp_base_member` - WHERE (subject_id LIKE 'dc/topic%' OR subject_id LIKE 'dc/svpg%'); + SELECT DISTINCT b.subject_id, b.object_id, b.provenance + FROM `temp_base_member` b + JOIN `temp_all_topic_svpg_nodes` n ON b.subject_id = n.subject_id; EXPORT DATA OPTIONS( uri="{dest}", format='CLOUD_SPANNER', spanner_options = '{{"table": "Edge"}}' ) AS - WITH RECURSIVE Descendants AS ( + WITH RECURSIVE + Descendants AS ( SELECT subject_id, object_id AS descendant, @@ -259,9 +278,8 @@ def run_linked_member( {prov_expr} as provenance FROM Descendants - WHERE subject_id LIKE 'dc/topic%' - AND descendant NOT LIKE 'dc/topic%' - AND descendant NOT LIKE 'dc/svpg%' + WHERE subject_id IN (SELECT subject_id FROM temp_topic_nodes) + AND descendant NOT IN (SELECT subject_id FROM temp_all_topic_svpg_nodes) ) SELECT subject_id, @@ -272,3 +290,98 @@ def run_linked_member( NewEdges """ return self.executor.execute(query) + + def run_topic_list_edges(self) -> Optional[bigquery.job.QueryJob]: + """Materializes relevantVariableList on Topics and memberList on SVPGs.""" + dest = self.executor.get_spanner_destination_uri() + prefix = BASE_PROVENANCE_PREFIX if self.is_base_dc else "" + output_provenance = f"{prefix}generated/TopicLists" + + query = f""" # nosec + -- Pull raw relevantVariable and member arcs across all active topics + CREATE OR REPLACE TEMPORARY TABLE `temp_raw_topic_edges` AS + SELECT * FROM EXTERNAL_QUERY("{self.executor.connection_id}", + "SELECT subject_id, predicate, object_id FROM Edge WHERE predicate IN ('relevantVariable', 'member')"); + + -- Pull global topic & SVPG type definitions to support cross-import schemas + CREATE OR REPLACE TEMPORARY TABLE `temp_topic_types` AS + SELECT * FROM EXTERNAL_QUERY("{self.executor.connection_id}", + "SELECT subject_id, object_id FROM Edge WHERE predicate = 'typeOf' AND object_id IN ('Topic', 'StatVarPeerGroup')"); + + CREATE OR REPLACE TEMPORARY TABLE `temp_topic_nodes` AS + SELECT DISTINCT subject_id FROM `temp_topic_types` WHERE object_id = 'Topic' + UNION DISTINCT + SELECT DISTINCT subject_id FROM `temp_raw_topic_edges` WHERE predicate = 'relevantVariable'; + + CREATE OR REPLACE TEMPORARY TABLE `temp_svpg_nodes` AS + SELECT DISTINCT subject_id FROM `temp_topic_types` WHERE object_id = 'StatVarPeerGroup'; + + -- Aggregate relevantVariable -> relevantVariableList for Topic nodes + CREATE OR REPLACE TEMPORARY TABLE `temp_aggregated_relevant_variable_list` AS + SELECT + e.subject_id, + 'relevantVariableList' AS predicate, + STRING_AGG(DISTINCT e.object_id, ',' ORDER BY e.object_id) AS list_value, + '{output_provenance}' AS provenance + FROM `temp_raw_topic_edges` e + JOIN `temp_topic_nodes` t ON e.subject_id = t.subject_id + WHERE e.predicate = 'relevantVariable' + GROUP BY e.subject_id; + + -- Aggregate member -> memberList for StatVarPeerGroup nodes + CREATE OR REPLACE TEMPORARY TABLE `temp_aggregated_member_list` AS + SELECT + e.subject_id, + 'memberList' AS predicate, + STRING_AGG(DISTINCT e.object_id, ',' ORDER BY e.object_id) AS list_value, + '{output_provenance}' AS provenance + FROM `temp_raw_topic_edges` e + JOIN `temp_svpg_nodes` s ON e.subject_id = s.subject_id + WHERE e.predicate = 'member' + GROUP BY e.subject_id; + + -- Generate DCGraph hashed keys for terminal literal nodes + CREATE OR REPLACE TEMPORARY TABLE `temp_all_list_edges` AS + WITH combined AS ( + SELECT subject_id, predicate, list_value, provenance FROM `temp_aggregated_relevant_variable_list` + UNION ALL + SELECT subject_id, predicate, list_value, provenance FROM `temp_aggregated_member_list` + ) + SELECT + subject_id, + predicate, + list_value, + CONCAT(SUBSTR(TRIM(list_value), 1, 16), ':', TO_HEX(SHA256(TRIM(list_value)))) AS object_id, + provenance + FROM combined; + + -- Export Node records (stores raw CSV string in Node.value) + EXPORT DATA + OPTIONS( + uri="{dest}", + format='CLOUD_SPANNER', + spanner_options = '{{"table": "Node"}}' + ) AS + SELECT DISTINCT + object_id AS subject_id, + list_value AS value, + CAST(NULL AS BYTES) AS bytes, + '' AS name, + CAST([] AS ARRAY) AS types + FROM `temp_all_list_edges`; + + -- Export Edge records (links Topic/SVPG to Node key) + EXPORT DATA + OPTIONS( + uri="{dest}", + format='CLOUD_SPANNER', + spanner_options = '{{"table": "Edge"}}' + ) AS + SELECT DISTINCT + subject_id, + predicate, + object_id, + provenance + FROM `temp_all_list_edges`; + """ + return self.executor.execute(query) diff --git a/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py b/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py index f348523b..12a7e661 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/orchestrator.py @@ -107,6 +107,7 @@ class OrchestratorConfig: enable_embeddings: bool = False bq_dataset_id: str = "datacommons" generate_stat_var_groups: bool = True + generate_topic_list_edges: bool = False max_parallel_imports: int = 10 @@ -278,15 +279,20 @@ def _process_single_import(self, single_import: str, dry_run: bool) -> ImportExe ) def _run_global_calculations(self, dry_run: bool = True) -> Optional[ImportExecutionResult]: - """Runs global, import-independent calculation steps (e.g., EMBEDDING_GENERATION).""" + """Runs global, import-independent calculation steps (e.g., EMBEDDING_GENERATION, TOPIC_LIST_EDGES).""" global_calcs = [ calc for calc in self.calculations if calc.get("type") in GLOBAL_CALCULATION_TYPES and not calc.get("disabled", False) ] - if not global_calcs: + should_generate_topic_lists = getattr(self.config, "generate_topic_list_edges", False) or any( + calc.get("generate_topic_list_edges", False) + for calc in self.calculations + if not calc.get("disabled", False) + ) + if not global_calcs and not should_generate_topic_lists: return None - logging.info(f"=== Starting Global Import-Independent Calculations ({len(global_calcs)} step(s)) ===") + logging.info(f"=== Starting Global Import-Independent Calculations ({len(global_calcs)} config step(s)) ===") for calc in global_calcs: step_type = calc.get("type") if dry_run: @@ -313,6 +319,31 @@ def _run_global_calculations(self, dry_run: bool = True) -> Optional[ImportExecu error_message=str(e) ) + if should_generate_topic_lists: + logging.info("Triggering global step: 'TOPIC_LIST_EDGES' (Consolidated relevantVariableList & memberList)...") + if dry_run: + logging.info("[DRY RUN] Would execute global step: Topic & SVPG List Edges") + else: + try: + generator = LinkedEdgeGenerator(self.executor, self.is_base_dc) + topic_job = generator.run_topic_list_edges() + if topic_job and hasattr(topic_job, "job_id"): + logging.info(f"Submitted global topic list edge job: {topic_job.job_id}") + self._wait_for_jobs( + job_ids=[topic_job.job_id], + poll_interval=self.poll_interval, + step_name="Topic & SVPG List Edges", + single_import="GLOBAL" + ) + except Exception as e: + logging.error(f"Global topic list edge generation failed: {e}") + return ImportExecutionResult( + import_name="GLOBAL", + success=False, + stages_executed=[], + error_message=str(e) + ) + return ImportExecutionResult( import_name="GLOBAL", success=True, @@ -334,6 +365,11 @@ def _delete_previous_aggregations( to_delete = set() linked_to_delete = set() delete_stat_var_groups = False + should_delete_topic_lists = getattr(self.config, "generate_topic_list_edges", False) or any( + calc.get("generate_topic_list_edges", False) + for calc in self.calculations + if not calc.get("disabled", False) + ) for single_import in imports: for calc in self.calculations: if self._calc_applies_to_import(calc, single_import): @@ -345,7 +381,7 @@ def _delete_previous_aggregations( if calc.get("type") == CalculationType.STAT_VAR_GROUPS: delete_stat_var_groups = True - if not to_delete and not linked_to_delete and not delete_stat_var_groups: + if not to_delete and not linked_to_delete and not delete_stat_var_groups and not should_delete_topic_lists: logging.info("No existing aggregated data resolved for deletion.") return @@ -358,6 +394,8 @@ def _delete_previous_aggregations( logging.info(f"[Dry Run] Would delete linked relationship edges for imports: {linked_to_delete_list}") if delete_stat_var_groups: logging.info("[Dry Run] Would delete StatVarGroup edges across all provenances.") + if should_delete_topic_lists: + logging.info("[Dry Run] Would delete topic and peer group list edges across all provenances.") else: if to_delete_list: self.deleter.delete_aggregated_data(to_delete_list) @@ -365,6 +403,8 @@ def _delete_previous_aggregations( self.deleter.delete_linked_edges(linked_to_delete_list) if delete_stat_var_groups: self.deleter.delete_stat_var_group_edges() + if should_delete_topic_lists: + self.deleter.delete_topic_list_edges() def _get_active_stages_for_import(self, single_import: str) -> List[int]: """Returns a sorted list of unique active stage numbers for a single import. diff --git a/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py b/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py index 5f7573ca..078c6fc6 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py +++ b/pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py @@ -531,6 +531,32 @@ def test_generate_stat_var_groups_disabled(self, mock_executor): )) self.assertTrue(orchestrator_enabled._calc_applies_to_import(svg_calc, "schema")) + @patch('aggregation.orchestrator.AggregationDeleter') + @patch('aggregation.orchestrator.LinkedEdgeGenerator') + def test_global_topic_list_edges_execution(self, mock_linked_gen, mock_deleter, mock_executor): + """Verifies global topic list edge consolidation runs when generate_topic_list_edges is True.""" + mock_job = MagicMock() + mock_job.job_id = "job-topic-1" + mock_linked_gen.return_value.run_topic_list_edges.return_value = mock_job + + orchestrator = AggregationOrchestrator(OrchestratorConfig( + connection_id="conn", + project_id="proj", + instance_id="inst", + database_id="db", + config_file_path=self.config_path, + generate_topic_list_edges=True + )) + orchestrator.executor = MagicMock() + orchestrator.executor.get_jobs_status.return_value = {"status": "DONE"} + + result = orchestrator.run(active_imports=[], dry_run=False) + self.assertTrue(result.success) + self.assertIn("GLOBAL", result.import_results) + self.assertTrue(result.import_results["GLOBAL"].success) + mock_deleter.return_value.delete_topic_list_edges.assert_called_once() + mock_linked_gen.return_value.run_topic_list_edges.assert_called_once() + class TestConfigSanity(unittest.TestCase): """Sanity checks for calculation configurations and metadata.""" diff --git a/pipeline/workflow/aggregation-helper/aggregation/schema.json b/pipeline/workflow/aggregation-helper/aggregation/schema.json index e2a07d55..02d3f01d 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/schema.json +++ b/pipeline/workflow/aggregation-helper/aggregation/schema.json @@ -43,7 +43,8 @@ "stat_var_calculation": { "type": "object" }, "embedding_generation": { "type": "object" }, "disabled": { "type": "boolean" }, - "should_prune_single_child_svgs": { "type": "boolean" } + "should_prune_single_child_svgs": { "type": "boolean" }, + "generate_topic_list_edges": { "type": "boolean" } }, "allOf": [ { diff --git a/pipeline/workflow/aggregation-helper/aggregation/stat_var_group_generator.py b/pipeline/workflow/aggregation-helper/aggregation/stat_var_group_generator.py index a3f89e80..e8086745 100644 --- a/pipeline/workflow/aggregation-helper/aggregation/stat_var_group_generator.py +++ b/pipeline/workflow/aggregation-helper/aggregation/stat_var_group_generator.py @@ -88,7 +88,7 @@ def run_stat_var_group(self) -> Optional[bigquery.job.QueryJob]: needed_predicates = ['populationType', 'measuredProperty', 'constraintProperties'] + constraint_props # Format into a SQL-safe string for Spanner injection - sv_predicates = [f"'{p.replace('\'', '')}'" for p in needed_predicates] + sv_predicates = ["'" + p.replace("'", "") + "'" for p in needed_predicates] sv_predicates_sql = ", ".join(sv_predicates) logging.info(f"Optimizing Spanner fetch. Pulling {len(sv_predicates)} specific predicates.") # ===================================================================== diff --git a/pipeline/workflow/aggregation-helper/main.py b/pipeline/workflow/aggregation-helper/main.py index 96801da4..3a5c26db 100644 --- a/pipeline/workflow/aggregation-helper/main.py +++ b/pipeline/workflow/aggregation-helper/main.py @@ -19,7 +19,6 @@ import logging import os import sys -from dataclasses import dataclass from typing import List, Optional from aggregation import AggregationOrchestrator, OrchestratorConfig @@ -90,6 +89,7 @@ def create_orchestrator_config( enable_embeddings=enable_embeddings, bq_dataset_id=bq_dataset_id, generate_stat_var_groups=args.generate_stat_var_groups, + generate_topic_list_edges=getattr(args, "generate_topic_list_edges", False), max_parallel_imports=args.max_parallel_imports, ) @@ -119,6 +119,12 @@ def main(): default=True, help="Whether to auto-generate StatVarGroup hierarchy tree (default: True, use --no-generate_stat_var_groups to disable)." ) + parser.add_argument( + "--generate_topic_list_edges", + action=argparse.BooleanOptionalAction, + default=False, + help="Whether to materialize topic relevantVariableList and memberList edges (default: False, use --generate_topic_list_edges to enable)." + ) parser.add_argument( "--dry_run", action=argparse.BooleanOptionalAction, diff --git a/pipeline/workflow/aggregation-helper/pyproject.toml b/pipeline/workflow/aggregation-helper/pyproject.toml index e6fc911b..80b8c901 100644 --- a/pipeline/workflow/aggregation-helper/pyproject.toml +++ b/pipeline/workflow/aggregation-helper/pyproject.toml @@ -46,6 +46,7 @@ dev = [ [tool.pytest.ini_options] pythonpath = ["."] +addopts = "--ignore=aggregation/e2e_tests" [[tool.uv.index]] name = "pypi"