From beff77d6e49450baf78699511adfdc555ccd5759 Mon Sep 17 00:00:00 2001 From: Rohit Kumar Date: Mon, 8 Jun 2026 18:03:28 +0000 Subject: [PATCH 1/4] feat: add aggregation framework and implement run_aggregation_spec API endpoint --- .../aggregation_spec_utils.py | 165 +++++++++++++ .../aggregation_spec_utils_test.py | 217 ++++++++++++++++++ pipeline/workflow/ingestion-helper/main.py | 28 +++ .../workflow/ingestion-helper/main_test.py | 26 ++- 4 files changed, 435 insertions(+), 1 deletion(-) create mode 100644 pipeline/workflow/ingestion-helper/aggregation_spec_utils.py create mode 100644 pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py new file mode 100644 index 000000000..787507946 --- /dev/null +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py @@ -0,0 +1,165 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from dataclasses import dataclass +from typing import Any, Callable, Dict, Iterable, List, Optional + + +Row = Dict[str, Any] +Mapper = Callable[[Row], Optional[Row]] + + +@dataclass +class AggregationSpec: + name: str + source: Any + sink: Any + mapper: Optional[Mapper] = None + + +class SpannerSqlSource: + + def __init__(self, + database: Any, + sql: str, + params: Optional[Dict[str, Any]] = None, + param_types: Optional[Dict[str, Any]] = None) -> None: + self.database = database + self.sql = sql + self.params = params or {} + self.param_types = param_types or {} + + def read(self) -> Iterable[Row]: + with self.database.snapshot() as snapshot: + results = snapshot.execute_sql(self.sql, + params=self.params, + param_types=self.param_types) + field_names = None + for row in results: + if isinstance(row, dict): + yield row + continue + + if field_names is None: + field_names = [ + field.name for field in getattr(results, 'fields', []) + ] + + if field_names: + yield dict(zip(field_names, row)) + else: + yield {"value": row} + + +class LocalJsonlSource: + + def __init__(self, path: str) -> None: + self.path = path + + def read(self) -> Iterable[Row]: + with open(self.path, 'r', encoding='utf-8') as jsonl_file: + for line in jsonl_file: + line = line.strip() + if line: + yield json.loads(line) + + +class SpannerTableSink: + + def __init__(self, + database: Any, + table: str, + columns: List[str], + batch_size: int = 1000) -> None: + self.database = database + self.table = table + self.columns = columns + self.batch_size = batch_size + + def start(self) -> None: + pass + + def write_many(self, rows: List[Row]) -> None: + if not rows: + return + + values = [[row[column] for column in self.columns] for row in rows] + # Use a mutation pool here if this sink needs higher write throughput. + with self.database.batch() as batch: + batch.insert_or_update(table=self.table, + columns=self.columns, + values=values) + + def close(self) -> None: + pass + + +class LocalJsonlSink: + + def __init__(self, path: str, batch_size: int = 1000) -> None: + self.path = path + self.batch_size = batch_size + self._file = None + + def start(self) -> None: + directory = os.path.dirname(self.path) + if directory: + os.makedirs(directory, exist_ok=True) + self._file = open(self.path, 'w', encoding='utf-8') + + def write_many(self, rows: List[Row]) -> None: + if not self._file: + raise RuntimeError("LocalJsonlSink must be started before writing") + for row in rows: + self._file.write(json.dumps(row, sort_keys=True) + '\n') + + def close(self) -> None: + if self._file: + self._file.close() + self._file = None + + +class AggregationRunner: + + def run(self, spec: AggregationSpec) -> int: + mapper = spec.mapper or _identity_mapper + batch = [] + row_count = 0 + + spec.sink.start() + try: + for row in spec.source.read(): + mapped_row = mapper(row) + if mapped_row is None: + continue + + batch.append(mapped_row) + if len(batch) >= spec.sink.batch_size: + spec.sink.write_many(batch) + row_count += len(batch) + batch = [] + + if batch: + spec.sink.write_many(batch) + row_count += len(batch) + + return row_count + finally: + spec.sink.close() + + +def _identity_mapper(row: Row) -> Row: + return row diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py new file mode 100644 index 000000000..b2aad3de7 --- /dev/null +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py @@ -0,0 +1,217 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +from aggregation_spec_utils import AggregationRunner +from aggregation_spec_utils import AggregationSpec +from aggregation_spec_utils import LocalJsonlSink +from aggregation_spec_utils import LocalJsonlSource +from aggregation_spec_utils import SpannerSqlSource +from aggregation_spec_utils import SpannerTableSink + + +class MockField: + + def __init__(self, name): + self.name = name + + +class MockResults: + + def __init__(self, rows, field_names): + self.rows = rows + self.fields = [MockField(name) for name in field_names] + + def __iter__(self): + return iter(self.rows) + + +class DelayedFieldResults: + + def __init__(self, rows, field_names): + self.rows = rows + self.field_names = field_names + self._iteration_started = False + + @property + def fields(self): + if not self._iteration_started: + return [] + return [MockField(name) for name in self.field_names] + + def __iter__(self): + self._iteration_started = True + return iter(self.rows) + + +class TestAggregationSpecUtils(unittest.TestCase): + + def test_local_jsonl_source_to_local_jsonl_sink(self): + with tempfile.TemporaryDirectory() as temp_dir: + input_path = Path(temp_dir) / 'input.jsonl' + output_path = Path(temp_dir) / 'output.jsonl' + self._write_jsonl(input_path, [{ + 'subject_id': 'dc/1' + }, { + 'subject_id': 'dc/2' + }]) + + count = AggregationRunner().run( + AggregationSpec( + name='copy', + source=LocalJsonlSource(str(input_path)), + sink=LocalJsonlSink(str(output_path), batch_size=1), + )) + + self.assertEqual(count, 2) + self.assertEqual(self._read_jsonl(output_path), [{ + 'subject_id': 'dc/1' + }, { + 'subject_id': 'dc/2' + }]) + + def test_mapper_transforms_row_shape(self): + with tempfile.TemporaryDirectory() as temp_dir: + input_path = Path(temp_dir) / 'input.jsonl' + output_path = Path(temp_dir) / 'output.jsonl' + self._write_jsonl(input_path, [{'source': 'dc/1'}]) + + count = AggregationRunner().run( + AggregationSpec( + name='mapped', + source=LocalJsonlSource(str(input_path)), + sink=LocalJsonlSink(str(output_path)), + mapper=lambda row: {'subject_id': row['source']}, + )) + + self.assertEqual(count, 1) + self.assertEqual(self._read_jsonl(output_path), + [{'subject_id': 'dc/1'}]) + + def test_spanner_source_streams_named_rows_from_snapshot(self): + database = MagicMock() + snapshot = database.snapshot.return_value.__enter__.return_value + snapshot.execute_sql.return_value = MockResults( + rows=[('dc/1', 'Node 1'), ('dc/2', 'Node 2')], + field_names=['subject_id', 'name'], + ) + source = SpannerSqlSource(database, + 'SELECT subject_id, name FROM Node', + params={'limit': 2}, + param_types={'limit': 'INT64'}) + + rows = list(source.read()) + + self.assertEqual(rows, [{ + 'subject_id': 'dc/1', + 'name': 'Node 1' + }, { + 'subject_id': 'dc/2', + 'name': 'Node 2' + }]) + snapshot.execute_sql.assert_called_once_with( + 'SELECT subject_id, name FROM Node', + params={'limit': 2}, + param_types={'limit': 'INT64'}, + ) + + def test_spanner_source_reads_fields_after_stream_starts(self): + database = MagicMock() + snapshot = database.snapshot.return_value.__enter__.return_value + snapshot.execute_sql.return_value = DelayedFieldResults( + rows=[('dc/1', 'Node 1')], + field_names=['subject_id', 'name'], + ) + + rows = list( + SpannerSqlSource(database, + 'SELECT subject_id, name FROM Node').read()) + + self.assertEqual(rows, [{'subject_id': 'dc/1', 'name': 'Node 1'}]) + + def test_spanner_sink_writes_multiple_batches(self): + with tempfile.TemporaryDirectory() as temp_dir: + input_path = Path(temp_dir) / 'input.jsonl' + self._write_jsonl(input_path, [{ + 'subject_id': f'dc/{idx}' + } for idx in range(5)]) + + database = MagicMock() + batches = [] + + def make_batch_context(): + context = MagicMock() + batch = MagicMock() + context.__enter__.return_value = batch + batches.append(batch) + return context + + database.batch.side_effect = make_batch_context + sink = SpannerTableSink(database, + table='Node', + columns=['subject_id'], + batch_size=2) + + count = AggregationRunner().run( + AggregationSpec( + name='write_nodes', + source=LocalJsonlSource(str(input_path)), + sink=sink, + )) + + self.assertEqual(count, 5) + self.assertEqual(database.batch.call_count, 3) + self.assertEqual( + [call.kwargs['values'] for batch in batches + for call in batch.insert_or_update.call_args_list], + [[['dc/0'], ['dc/1']], [['dc/2'], ['dc/3']], [['dc/4']]], + ) + + def test_empty_source_writes_nothing(self): + with tempfile.TemporaryDirectory() as temp_dir: + input_path = Path(temp_dir) / 'input.jsonl' + input_path.write_text('', encoding='utf-8') + database = MagicMock() + + count = AggregationRunner().run( + AggregationSpec( + name='empty', + source=LocalJsonlSource(str(input_path)), + sink=SpannerTableSink(database, + table='Node', + columns=['subject_id'], + batch_size=2), + )) + + self.assertEqual(count, 0) + database.batch.assert_not_called() + + def _write_jsonl(self, path, rows): + path.write_text(''.join(json.dumps(row) + '\n' for row in rows), + encoding='utf-8') + + def _read_jsonl(self, path): + return [ + json.loads(line) + for line in path.read_text(encoding='utf-8').splitlines() + ] + + +if __name__ == '__main__': + unittest.main() diff --git a/pipeline/workflow/ingestion-helper/main.py b/pipeline/workflow/ingestion-helper/main.py index 5b60514bc..4b2eae8e0 100644 --- a/pipeline/workflow/ingestion-helper/main.py +++ b/pipeline/workflow/ingestion-helper/main.py @@ -8,6 +8,10 @@ import import_utils from flask import jsonify from aggregation_utils import AggregationUtils +from aggregation_spec_utils import AggregationRunner +from aggregation_spec_utils import AggregationSpec +from aggregation_spec_utils import LocalJsonlSink +from aggregation_spec_utils import SpannerSqlSource logging.getLogger().setLevel(logging.INFO) @@ -60,6 +64,19 @@ def _validate_params(request_json, required_params): return None +def _get_aggregation_spec(spec_name, spanner): + if spec_name == 'sample_nodes_to_local_jsonl': + return AggregationSpec( + name=spec_name, + source=SpannerSqlSource( + spanner.graph_database, + 'SELECT subject_id, name FROM Node LIMIT 10', + ), + sink=LocalJsonlSink('/tmp/aggregation_spec_nodes.jsonl'), + ) + return None + + @functions_framework.http def ingestion_helper(request): """ @@ -292,6 +309,17 @@ def ingestion_helper(request): return jsonify({'status': 'SUBMITTED', 'jobIds': job_ids}), 200 except Exception as e: return (f"Aggregation failed: {str(e)}", 500) + elif action_type == 'run_aggregation_spec': + spec_name = request_json.get('specName', 'sample_nodes_to_local_jsonl') + spec = _get_aggregation_spec(spec_name, spanner) + if not spec: + return (f'Unknown aggregation spec: {spec_name}', 400) + + try: + row_count = AggregationRunner().run(spec) + return (f'OK [Spec: {spec_name} Rows: {row_count}]', 200) + except Exception as e: + return (f"Aggregation spec failed: {str(e)}", 500) elif action_type == 'clear_redis_cache': logging.info("Action: clear_redis_cache") redis_host = os.environ.get("REDIS_HOST") diff --git a/pipeline/workflow/ingestion-helper/main_test.py b/pipeline/workflow/ingestion-helper/main_test.py index 69c87ebfa..4bf9b050b 100644 --- a/pipeline/workflow/ingestion-helper/main_test.py +++ b/pipeline/workflow/ingestion-helper/main_test.py @@ -141,6 +141,30 @@ def test_seed_database_success(self, mock_spanner_client_class): self.assertIn("OK", response) mock_spanner_client.seed_database.assert_called_once() + @patch('main.StorageClient') + @patch('main.AggregationRunner') + @patch('main.SpannerClient') + def test_run_aggregation_spec_success(self, mock_spanner_client_class, + mock_runner_class, + mock_storage_client_class): + mock_spanner_client = MagicMock() + mock_spanner_client_class.return_value = mock_spanner_client + mock_runner_class.return_value.run.return_value = 3 + + mock_request = MagicMock() + mock_request.get_json.return_value = { + "actionType": "run_aggregation_spec", + "specName": "sample_nodes_to_local_jsonl", + } + + response, status_code = main.ingestion_helper(mock_request) + + self.assertEqual(status_code, 200) + self.assertIn("Rows: 3", response) + spec = mock_runner_class.return_value.run.call_args.args[0] + self.assertEqual(spec.name, "sample_nodes_to_local_jsonl") + self.assertEqual(spec.sink.path, "/tmp/aggregation_spec_nodes.jsonl") + mock_storage_client_class.assert_called_once() + if __name__ == '__main__': unittest.main() - From e340ce5cebcb7f5c9c6ad4b25c496e953fd76556 Mon Sep 17 00:00:00 2001 From: Rohit Kumar Date: Mon, 8 Jun 2026 18:23:18 +0000 Subject: [PATCH 2/4] fix: handle None fields in SpannerSqlSource to prevent attribute errors during ingestion --- .../aggregation_spec_utils.py | 6 +++--- .../aggregation_spec_utils_test.py | 20 +++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py index 787507946..838d1290e 100644 --- a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py @@ -54,9 +54,9 @@ def read(self) -> Iterable[Row]: continue if field_names is None: - field_names = [ - field.name for field in getattr(results, 'fields', []) - ] + fields = getattr(results, "fields", None) + if fields: + field_names = [field.name for field in fields] if field_names: yield dict(zip(field_names, row)) diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py index b2aad3de7..a91cd794d 100644 --- a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py @@ -60,6 +60,15 @@ def __iter__(self): return iter(self.rows) +class NoneFieldResults: + def __init__(self, rows): + self.rows = rows + self.fields = None + + def __iter__(self): + return iter(self.rows) + + class TestAggregationSpecUtils(unittest.TestCase): def test_local_jsonl_source_to_local_jsonl_sink(self): @@ -145,6 +154,17 @@ def test_spanner_source_reads_fields_after_stream_starts(self): self.assertEqual(rows, [{'subject_id': 'dc/1', 'name': 'Node 1'}]) + def test_spanner_source_handles_none_fields(self): + database = MagicMock() + snapshot = database.snapshot.return_value.__enter__.return_value + snapshot.execute_sql.return_value = NoneFieldResults(rows=[("dc/1", "Node 1")]) + + rows = list( + SpannerSqlSource(database, "SELECT subject_id, name FROM Node").read() + ) + + self.assertEqual(rows, [{"value": ("dc/1", "Node 1")}]) + def test_spanner_sink_writes_multiple_batches(self): with tempfile.TemporaryDirectory() as temp_dir: input_path = Path(temp_dir) / 'input.jsonl' From a9a27a6c4b764ace5043d87ca1699ff022e500f7 Mon Sep 17 00:00:00 2001 From: Rohit Kumar Date: Mon, 8 Jun 2026 18:30:21 +0000 Subject: [PATCH 3/4] fix: add default batch size of 1000 for sinks that do not define one --- .../aggregation_spec_utils.py | 3 +- .../aggregation_spec_utils_test.py | 40 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py index 838d1290e..7a850eb88 100644 --- a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py @@ -138,6 +138,7 @@ def run(self, spec: AggregationSpec) -> int: mapper = spec.mapper or _identity_mapper batch = [] row_count = 0 + batch_size = getattr(spec.sink, "batch_size", 1000) spec.sink.start() try: @@ -147,7 +148,7 @@ def run(self, spec: AggregationSpec) -> int: continue batch.append(mapped_row) - if len(batch) >= spec.sink.batch_size: + if len(batch) >= batch_size: spec.sink.write_many(batch) row_count += len(batch) batch = [] diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py index a91cd794d..031012e36 100644 --- a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py @@ -61,6 +61,7 @@ def __iter__(self): class NoneFieldResults: + def __init__(self, rows): self.rows = rows self.fields = None @@ -69,6 +70,22 @@ def __iter__(self): return iter(self.rows) +class NoBatchSizeSink: + def __init__(self): + self.batches = [] + self.started = False + self.closed = False + + def start(self): + self.started = True + + def write_many(self, rows): + self.batches.append(rows) + + def close(self): + self.closed = True + + class TestAggregationSpecUtils(unittest.TestCase): def test_local_jsonl_source_to_local_jsonl_sink(self): @@ -222,6 +239,29 @@ def test_empty_source_writes_nothing(self): self.assertEqual(count, 0) database.batch.assert_not_called() + def test_runner_uses_default_batch_size_when_sink_omits_it(self): + with tempfile.TemporaryDirectory() as temp_dir: + input_path = Path(temp_dir) / "input.jsonl" + self._write_jsonl( + input_path, [{"subject_id": "dc/1"}, {"subject_id": "dc/2"}] + ) + sink = NoBatchSizeSink() + + count = AggregationRunner().run( + AggregationSpec( + name="default_batch_size", + source=LocalJsonlSource(str(input_path)), + sink=sink, + ) + ) + + self.assertEqual(count, 2) + self.assertTrue(sink.started) + self.assertTrue(sink.closed) + self.assertEqual( + sink.batches, [[{"subject_id": "dc/1"}, {"subject_id": "dc/2"}]] + ) + def _write_jsonl(self, path, rows): path.write_text(''.join(json.dumps(row) + '\n' for row in rows), encoding='utf-8') From cda8453ebedd9ff5663f1ddddd4e8c05d1f39d46 Mon Sep 17 00:00:00 2001 From: Rohit Kumar Date: Mon, 8 Jun 2026 18:40:20 +0000 Subject: [PATCH 4/4] feat: add SpannerDmlSink support and allow AggregationRunner to execute sinks without a source --- .../aggregation_spec_utils.py | 51 ++++++++++++- .../aggregation_spec_utils_test.py | 71 +++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py index 7a850eb88..94533aefa 100644 --- a/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils.py @@ -25,11 +25,16 @@ @dataclass class AggregationSpec: name: str - source: Any + source: Optional[Any] sink: Any mapper: Optional[Mapper] = None +class SpannerDmlMode: + TRANSACTIONAL = 'transactional' + PARTITIONED = 'partitioned' + + class SpannerSqlSource: def __init__(self, @@ -132,9 +137,53 @@ def close(self) -> None: self._file = None +@dataclass +class SpannerDmlSink: + database: Any + sql: str + params: Optional[Dict[str, Any]] = None + param_types: Optional[Dict[str, Any]] = None + mode: str = SpannerDmlMode.TRANSACTIONAL + timeout: Optional[int] = None + + def run(self) -> int: + if self.mode == SpannerDmlMode.TRANSACTIONAL: + return self._run_transactional() + if self.mode == SpannerDmlMode.PARTITIONED: + return self._run_partitioned() + raise ValueError(f"Unknown Spanner DML mode: {self.mode}") + + def _run_transactional(self) -> int: + + def _execute_dml(transaction): + kwargs = { + 'params': self.params or {}, + 'param_types': self.param_types or {}, + } + if self.timeout is not None: + kwargs['timeout'] = self.timeout + return transaction.execute_update(self.sql, **kwargs) + + return self.database.run_in_transaction(_execute_dml) + + def _run_partitioned(self) -> int: + if self.timeout is not None: + raise ValueError( + "Partitioned DML does not support timeout in this client") + + return self.database.execute_partitioned_dml( + self.sql, + params=self.params or {}, + param_types=self.param_types or {}, + ) + + class AggregationRunner: def run(self, spec: AggregationSpec) -> int: + if spec.source is None: + return spec.sink.run() + mapper = spec.mapper or _identity_mapper batch = [] row_count = 0 diff --git a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py index 031012e36..f3db25d3c 100644 --- a/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py +++ b/pipeline/workflow/ingestion-helper/aggregation_spec_utils_test.py @@ -22,6 +22,8 @@ from aggregation_spec_utils import AggregationSpec from aggregation_spec_utils import LocalJsonlSink from aggregation_spec_utils import LocalJsonlSource +from aggregation_spec_utils import SpannerDmlMode +from aggregation_spec_utils import SpannerDmlSink from aggregation_spec_utils import SpannerSqlSource from aggregation_spec_utils import SpannerTableSink @@ -262,6 +264,75 @@ def test_runner_uses_default_batch_size_when_sink_omits_it(self): sink.batches, [[{"subject_id": "dc/1"}, {"subject_id": "dc/2"}]] ) + def test_runner_runs_transactional_dml_sink_without_source(self): + database = MagicMock() + transaction = MagicMock() + transaction.execute_update.return_value = 3 + sql = "INSERT INTO Target (id) SELECT id FROM Source" + + def run_in_transaction_side_effect(callback): + return callback(transaction) + + database.run_in_transaction.side_effect = run_in_transaction_side_effect + + count = AggregationRunner().run( + AggregationSpec( + name="insert_select", + source=None, + sink=SpannerDmlSink( + database=database, + sql=sql, + params={"import_name": "dc/test"}, + param_types={"import_name": "STRING"}, + timeout=120, + ), + )) + + self.assertEqual(count, 3) + database.run_in_transaction.assert_called_once() + transaction.execute_update.assert_called_once_with( + sql, + params={"import_name": "dc/test"}, + param_types={"import_name": "STRING"}, + timeout=120, + ) + + def test_runner_runs_partitioned_dml_sink_without_source(self): + database = MagicMock() + database.execute_partitioned_dml.return_value = 4 + sql = "DELETE FROM Target WHERE import_name = @import_name" + + count = AggregationRunner().run( + AggregationSpec( + name="delete_target", + source=None, + sink=SpannerDmlSink( + database=database, + sql=sql, + params={"import_name": "dc/test"}, + param_types={"import_name": "STRING"}, + mode=SpannerDmlMode.PARTITIONED, + ), + )) + + self.assertEqual(count, 4) + database.execute_partitioned_dml.assert_called_once_with( + sql, + params={"import_name": "dc/test"}, + param_types={"import_name": "STRING"}, + ) + + def test_partitioned_dml_rejects_timeout(self): + sink = SpannerDmlSink( + database=MagicMock(), + sql="DELETE FROM Target WHERE import_name = @import_name", + mode=SpannerDmlMode.PARTITIONED, + timeout=120, + ) + + with self.assertRaisesRegex(ValueError, "does not support timeout"): + sink.run() + def _write_jsonl(self, path, rows): path.write_text(''.join(json.dumps(row) + '\n' for row in rows), encoding='utf-8')