import concurrent.futures import io import json import tempfile import threading import unittest from datetime import datetime, timezone from pathlib import Path from unittest import mock from botocore.exceptions import ClientError import packed_backup_watermark as worker def s3_client_error(code): return ClientError( {"Error": {"Code": code, "Message": code}}, "HeadObject", ) class FakeCursor: def __init__(self, connection): self.connection = connection self.rowcount = 0 def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): return False def execute(self, query, parameters): self.connection.queries.append((query, parameters)) if query.startswith("UPDATE"): self.rowcount = self.connection.update_rowcount def fetchall(self): return self.connection.rows class FakeConnection: def __init__(self, watermark, update_rowcount=1, rows=None): self.watermark = watermark self.rows = [(watermark,)] if rows is None else rows self.update_rowcount = update_rowcount self.queries = [] self.commits = 0 self.rollbacks = 0 self.closed = False def cursor(self): return FakeCursor(self) def commit(self): self.commits += 1 def rollback(self): self.rollbacks += 1 def close(self): self.closed = True class WatermarkWorkerTests(unittest.TestCase): def test_main_succeeds_with_unverified_marker_count(self): def reconcile(_args, result): result["packed_clusters"] = 1 result["unverified_complete_marker_clusters"] = 1 stdout = io.StringIO() with tempfile.TemporaryDirectory() as directory: termination_path = Path(directory) / "termination-log" with ( mock.patch.object(worker, "reconcile", side_effect=reconcile), mock.patch.object(worker.sys, "stdout", stdout), mock.patch.object( worker, "TERMINATION_MESSAGE_PATH", termination_path, ), ): return_code = worker.main( [ "--watermark-instance", "watermark-run", "--export-instance", "export-run", "--tasks-json", "/unused/tasks.json", "--out-prefix", "s3://bucket/export", "--backup-time", "2026-07-31T00:00:00Z", "--mysql-defaults-file", "/unused/mysql.cnf", "--aws-region", "us-east-1", ] ) payload = json.loads(termination_path.read_text(encoding="utf-8")) self.assertEqual(return_code, 0) self.assertEqual(stdout.getvalue(), "") self.assertEqual(payload["status"], "succeeded") self.assertEqual(payload["unverified_complete_marker_clusters"], 1) self.assertNotIn("error", payload) def test_complete_marker_body_may_contain_output_pointer(self): client = mock.Mock() client.head_object.return_value = { "ContentLength": 64, "Metadata": { "job-instance-name": "export-run", "backup-time": "2026-07-31T00:00:00Z", }, } verified, has_newer_backup = worker.verify_complete_markers( client, "s3://bucket/export", {"cluster-1": ["cluster-1_default"]}, "export-run", datetime(2026, 7, 31, tzinfo=timezone.utc), ) self.assertEqual(verified, 1) self.assertFalse(has_newer_backup) def test_missing_marker_is_counted_without_failing_reconcile(self): client = mock.Mock() def head_object(*, Bucket, Key): self.assertEqual(Bucket, "bucket") if "cluster-missing_default" in Key: raise s3_client_error("404") return { "Metadata": { "job-instance-name": "export-run", "backup-time": "2026-07-31T00:00:00Z", }, } client.head_object.side_effect = head_object args = mock.Mock( tasks_json="/unused/tasks.json", aws_region="us-east-1", parallel=1, out_prefix="s3://bucket/export", export_instance="export-run", mysql_defaults_file="/unused/mysql.cnf", backup_time=datetime(2026, 7, 31, tzinfo=timezone.utc), ) result = {} connections = [] def open_connection(_defaults_file): connection = FakeConnection(datetime(2026, 7, 30)) connections.append(connection) return connection with ( mock.patch.object(worker, "load_tasks", return_value=[{}]), mock.patch.object(worker, "create_s3_client", return_value=client), mock.patch.object( worker, "packed_meta_identities", return_value=[ "cluster-complete_default", "cluster-missing_default", ], ), mock.patch.object( worker, "open_mysql_connection", side_effect=open_connection, ), mock.patch.object(worker, "log") as log_mock, ): worker.reconcile(args, result) self.assertEqual(result["mysql_updated_clusters"], 1) self.assertEqual(result["unverified_complete_marker_clusters"], 1) self.assertEqual(len(connections), 1) self.assertEqual(len(connections[0].queries), 2) self.assertEqual( connections[0].queries[0][1][1], "cluster-complete", ) self.assertTrue(connections[0].queries[1][0].startswith("INSERT IGNORE")) self.assertEqual( connections[0].queries[1][1], ( "cluster-missing", "br_export", "export_instance=export-run, " "backup_time=2026-07-31T00:00:00Z", ), ) self.assertTrue( any( "cluster-missing" in call.args[0] and "complete marker does not exist" in call.args[0] for call in log_mock.call_args_list ) ) def test_one_missing_meta_marker_blocks_the_whole_cluster_watermark(self): client = mock.Mock() def head_object(*, Bucket, Key): self.assertEqual(Bucket, "bucket") if "cluster-1_second" in Key: raise s3_client_error("NoSuchKey") return { "Metadata": { "job-instance-name": "export-run", "backup-time": "2026-07-31T00:00:00Z", }, } client.head_object.side_effect = head_object connection = FakeConnection(None) mysql_connection = mock.Mock(return_value=connection) advance = worker.advance_verified_cluster_watermark( client, "s3://bucket/export", "cluster-1", ["cluster-1_first", "cluster-1_second"], "export-run", mysql_connection, datetime(2026, 7, 31, tzinfo=timezone.utc), ) self.assertEqual(advance.updated_clusters, 0) self.assertEqual(advance.unverified_clusters, 1) mysql_connection.assert_called_once_with() self.assertEqual( connection.queries[0][1], ( "cluster-1", "br_export", "export_instance=export-run, " "backup_time=2026-07-31T00:00:00Z", ), ) def test_mismatched_marker_is_blocklisted_without_updating_watermark(self): client = mock.Mock() client.head_object.return_value = { "Metadata": { "job-instance-name": "another-export-run", "backup-time": "2026-07-31T00:00:00Z", }, } connection = FakeConnection(None) mysql_connection = mock.Mock(return_value=connection) with mock.patch.object(worker, "log") as log_mock: advance = worker.advance_verified_cluster_watermark( client, "s3://bucket/export", "cluster-1", ["cluster-1_default"], "export-run", mysql_connection, datetime(2026, 7, 31, tzinfo=timezone.utc), ) self.assertEqual(advance.updated_clusters, 0) self.assertEqual(advance.unverified_clusters, 1) mysql_connection.assert_called_once_with() self.assertTrue(connection.queries[0][0].startswith("INSERT IGNORE")) self.assertEqual( connection.queries[0][1], ( "cluster-1", "br_export", "export_instance=export-run, " "backup_time=2026-07-31T00:00:00Z", ), ) self.assertIn( "complete marker job instance does not match", log_mock.call_args.args[0], ) def test_older_marker_is_blocklisted_even_for_current_export_instance(self): client = mock.Mock() client.head_object.return_value = { "Metadata": { "job-instance-name": "export-run", "backup-time": "2026-07-30T00:00:00Z", }, } connection = FakeConnection(None) mysql_connection = mock.Mock(return_value=connection) advance = worker.advance_verified_cluster_watermark( client, "s3://bucket/export", "cluster-1", ["cluster-1_default"], "export-run", mysql_connection, datetime(2026, 7, 31, tzinfo=timezone.utc), ) self.assertEqual(advance.unverified_clusters, 1) self.assertEqual(connection.queries[0][1][1], "br_export") self.assertIn("export_instance=export-run", connection.queries[0][1][2]) self.assertIn( "backup_time=2026-07-31T00:00:00Z", connection.queries[0][1][2], ) def test_marker_without_metadata_is_blocklisted(self): client = mock.Mock() client.head_object.return_value = {"Metadata": {}} connection = FakeConnection(None) mysql_connection = mock.Mock(return_value=connection) advance = worker.advance_verified_cluster_watermark( client, "s3://bucket/export", "cluster-1", ["cluster-1_default"], "export-run", mysql_connection, datetime(2026, 7, 31, tzinfo=timezone.utc), ) self.assertEqual(advance.unverified_clusters, 1) self.assertEqual(connection.queries[0][1][1], "br_export") self.assertIn("export_instance=export-run", connection.queries[0][1][2]) self.assertIn( "backup_time=2026-07-31T00:00:00Z", connection.queries[0][1][2], ) def test_newer_marker_does_not_add_cluster_to_blocklist(self): client = mock.Mock() client.head_object.return_value = { "Metadata": { "job-instance-name": "another-export-run", "backup-time": "2026-08-01T00:00:00Z", }, } mysql_connection = mock.Mock() advance = worker.advance_verified_cluster_watermark( client, "s3://bucket/export", "cluster-1", ["cluster-1_default"], "export-run", mysql_connection, datetime(2026, 7, 31, tzinfo=timezone.utc), ) self.assertEqual(advance.newer_backup_clusters, 1) self.assertEqual(advance.unverified_clusters, 0) mysql_connection.assert_not_called() def test_non_missing_marker_s3_error_still_fails(self): client = mock.Mock() client.head_object.side_effect = s3_client_error("AccessDenied") mysql_connection = mock.Mock() with self.assertRaisesRegex( worker.WatermarkError, "code=AccessDenied", ): worker.advance_verified_cluster_watermark( client, "s3://bucket/export", "cluster-1", ["cluster-1_default"], "export-run", mysql_connection, datetime(2026, 7, 31, tzinfo=timezone.utc), ) mysql_connection.assert_not_called() def test_meta_identity_is_split_from_right_without_cluster_list(self): mapping = worker.map_packed_clusters( [ "cluster_with_underscores_keyspaceB", "cluster-2_default", "cluster_with_underscores_keyspaceA", ] ) self.assertEqual( mapping, { "cluster-2": ["cluster-2_default"], "cluster_with_underscores": [ "cluster_with_underscores_keyspaceA", "cluster_with_underscores_keyspaceB", ], }, ) def test_invalid_meta_identity_is_rejected(self): for identity in ("no-separator", "_keyspace", "cluster_"): with self.subTest(identity=identity), self.assertRaisesRegex( worker.WatermarkError, "_", ): worker.map_packed_clusters([identity]) def test_newer_mysql_watermark_is_never_regressed(self): connection = FakeConnection( datetime(2026, 8, 1), update_rowcount=0, ) target = datetime(2026, 7, 31, tzinfo=timezone.utc) has_newer = worker.update_cluster_export_watermark( connection, "cluster-1", target, ) self.assertTrue(has_newer) self.assertEqual(connection.commits, 1) self.assertEqual(connection.rollbacks, 1) self.assertEqual(len(connection.queries), 2) self.assertTrue(connection.queries[0][0].startswith("UPDATE")) self.assertTrue(connection.queries[1][0].startswith("SELECT")) def test_update_is_conditional_on_the_current_watermark(self): connection = FakeConnection(datetime(2026, 7, 30)) target = datetime(2026, 7, 31, tzinfo=timezone.utc) has_newer = worker.update_cluster_export_watermark( connection, "cluster-1", target, ) self.assertFalse(has_newer) self.assertEqual(connection.commits, 1) self.assertEqual(connection.rollbacks, 0) self.assertEqual(len(connection.queries), 1) update_query = connection.queries[0][0] self.assertIn("s3_data_updated_at < %s", update_query) def test_missing_cluster_is_queried_after_update_misses(self): connection = FakeConnection( None, update_rowcount=0, rows=[], ) target = datetime(2026, 7, 31, tzinfo=timezone.utc) with self.assertRaisesRegex( worker.WatermarkError, "expected one row", ): worker.update_cluster_export_watermark( connection, "missing-cluster", target, ) self.assertEqual( [query.split(None, 1)[0] for query, _ in connection.queries], ["UPDATE", "SELECT"], ) def test_mysql_update_failure_still_fails_reconcile(self): client = mock.Mock() client.head_object.return_value = { "Metadata": { "job-instance-name": "export-run", "backup-time": "2026-07-31T00:00:00Z", }, } args = mock.Mock( tasks_json="/unused/tasks.json", aws_region="us-east-1", parallel=1, out_prefix="s3://bucket/export", export_instance="export-run", mysql_defaults_file="/unused/mysql.cnf", backup_time=datetime(2026, 7, 31, tzinfo=timezone.utc), ) result = {} connection = FakeConnection(None, update_rowcount=0, rows=[]) with ( mock.patch.object(worker, "load_tasks", return_value=[{}]), mock.patch.object(worker, "create_s3_client", return_value=client), mock.patch.object( worker, "packed_meta_identities", return_value=["cluster-missing_default"], ), mock.patch.object( worker, "open_mysql_connection", return_value=connection, ), self.assertRaisesRegex( worker.WatermarkError, "1 cluster.*failed to update", ), ): worker.reconcile(args, result) self.assertEqual(result["verified_complete_markers"], 1) self.assertEqual(result["unverified_complete_marker_clusters"], 0) self.assertEqual(result["mysql_update_failed_clusters"], 1) def test_cluster_reconciliation_has_bounded_inflight_work(self): real_executor = concurrent.futures.ThreadPoolExecutor parallel = 3 cluster_count = 10 release = threading.Event() workers_started = threading.Event() state_lock = threading.Lock() active = 0 max_active = 0 opened_connections = [] class TrackingExecutor: submitted = 0 submit_lock = threading.Lock() def __init__(self, max_workers): self.executor = real_executor(max_workers=max_workers) def __enter__(self): self.executor.__enter__() return self def __exit__(self, exc_type, exc_value, traceback): return self.executor.__exit__(exc_type, exc_value, traceback) def submit(self, function, *args): with self.submit_lock: type(self).submitted += 1 return self.executor.submit(function, *args) def verify_markers(*_args): nonlocal active, max_active with state_lock: active += 1 max_active = max(max_active, active) if active == parallel: workers_started.set() try: if not release.wait(timeout=5): raise AssertionError("test did not release workers") finally: with state_lock: active -= 1 return 1, False def open_connection(_defaults_file): connection = FakeConnection(datetime(2026, 7, 30)) with state_lock: opened_connections.append(connection) return connection clusters = { f"cluster-{index:02d}": [f"cluster-{index:02d}_default"] for index in range(cluster_count) } result = [] errors = [] def run_reconciliation(): try: result.append( worker.advance_verified_cluster_watermarks( object(), "s3://bucket/export", clusters, "export-run", "/unused/mysql.cnf", datetime(2026, 7, 31, tzinfo=timezone.utc), parallel, ) ) except BaseException as exc: errors.append(exc) with ( mock.patch.object( worker.concurrent.futures, "ThreadPoolExecutor", TrackingExecutor, ), mock.patch.object( worker, "verify_complete_markers", side_effect=verify_markers, ), mock.patch.object( worker, "open_mysql_connection", side_effect=open_connection, ), ): reconciliation = threading.Thread(target=run_reconciliation) reconciliation.start() try: self.assertTrue(workers_started.wait(timeout=5)) self.assertEqual(TrackingExecutor.submitted, parallel) finally: release.set() reconciliation.join(timeout=5) self.assertFalse(reconciliation.is_alive()) if errors: raise errors[0] self.assertEqual(TrackingExecutor.submitted, cluster_count) self.assertLessEqual(max_active, parallel) self.assertEqual(result[0].verified_markers, cluster_count) self.assertEqual(result[0].updated_clusters, cluster_count) self.assertLessEqual(len(opened_connections), parallel) self.assertTrue(opened_connections) self.assertTrue( all(connection.closed for connection in opened_connections) ) if __name__ == "__main__": unittest.main()