From 702ec28fe42377d18af699ccb8b49a0948abf186 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 12:23:51 -0700
Subject: [PATCH 01/11] add initial implementation

---
 .../datacommons_admin/admin_cli.py            | 113 ++++++-
 .../ingestion_helper_client.py                |  31 ++
 packages/datacommons-admin/pyproject.toml     |   2 +
 .../datacommons-admin/tests/test_admin_cli.py | 279 ++++++++++++++++-
 .../tests/test_ingestion_helper_client.py     | 289 ++++++++++++++++++
 uv.lock                                       |   7 +
 6 files changed, 700 insertions(+), 21 deletions(-)
 create mode 100644 packages/datacommons-admin/tests/test_ingestion_helper_client.py

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index 47eeaf57..f466c44e 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -28,6 +28,17 @@
     README_TEMPLATE,
     REMOTE_STATE_TEMPLATE,
 )
+from datacommons_admin.ingestion_helper_client import IngestionHelperClient
+from datacommons_admin.tf_utils import (
+    get_ingestion_service_url,
+    get_ingestion_workflow_service_account_email,
+    get_project_id,
+    get_spanner_database_id,
+    get_spanner_instance_id,
+)
+from datacommons_db.clients import SpannerClient
+from datacommons_db.migrations import MigrationRunner
+
 
 
 DEFAULT_BUCKET_LOCATION = "US"
@@ -508,34 +519,27 @@ def init(
     )
 
 
-def _setup_ingestion_client() -> Tuple[Any, str, str]:
+def _setup_ingestion_client() -> Tuple[Any, str, str, str]:
     click.secho(
         "Fetching ingestion service URL, workflow service account, and Spanner details from Terraform outputs...",
         fg="bright_black",
     )
 
-    from datacommons_admin.tf_utils import (
-        get_ingestion_service_url,
-        get_ingestion_workflow_service_account_email,
-        get_spanner_instance_id,
-        get_spanner_database_id,
-    )
-    from datacommons_admin.ingestion_helper_client import IngestionHelperClient
-
     url = get_ingestion_service_url()
     sa_email = get_ingestion_workflow_service_account_email()
+    project_id = get_project_id()
     instance_id = get_spanner_instance_id()
     database_id = get_spanner_database_id()
 
     click.secho(f"Found ingestion service URL: {url}", fg="green")
     click.secho(f"Found ingestion workflow service account: {sa_email}", fg="green")
     click.secho(
-        f"Found Spanner instance ID: {instance_id} / database ID: {database_id}",
+        f"Found Spanner details: project={project_id}, instance={instance_id}, database={database_id}",
         fg="green",
     )
 
     client = IngestionHelperClient(url, service_account_email=sa_email)
-    return client, instance_id, database_id
+    return client, project_id, instance_id, database_id
 
 
 def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
@@ -550,6 +554,85 @@ def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
         click.secho(f"Details: {message}", fg="bright_black")
 
 
+def _run_migrations(
+    client: Any, project_id: str, instance_id: str, database_id: str
+) -> None:
+    click.secho(
+        f"Checking schema migrations for Spanner database '{project_id}/{instance_id}/{database_id}'...",
+        fg="bright_black",
+    )
+
+    try:
+        spanner_client = SpannerClient(
+            project_id=project_id,
+            instance_id=instance_id,
+            database_id=database_id,
+        )
+        runner = MigrationRunner(spanner_client=spanner_client)
+        pending = runner.get_pending_migrations()
+    except Exception as e:
+        raise click.ClickException(f"Failed to check pending migrations: {e}") from e
+
+    if not pending:
+        click.secho(
+            "Database schema is already up-to-date. No migrations to apply.",
+            fg="green",
+        )
+        return
+
+    click.secho(f"Found {len(pending)} pending schema migration(s):", fg="cyan")
+    for m in pending:
+        click.echo(f"  - {m.creation_timestamp}: {m.description}")
+
+    click.secho(
+        "Acquiring database lock via the Ingestion Helper service...",
+        fg="bright_black",
+    )
+    try:
+        client.acquire_lock(workflow_id="schema-migration")
+    except Exception as e:
+        raise click.ClickException(
+            f"Could not acquire database lock: {e}\n"
+            "An ingestion workflow may currently be running. "
+            "Please wait for active ingestions to finish before running migrations."
+        ) from e
+
+    try:
+        click.secho("Applying pending schema migrations...", fg="bright_black")
+        results = runner.run_migrations()
+        for res in results:
+            click.secho(
+                f"  ✔ Applied migration {res.creation_timestamp}: {res.description}",
+                fg="green",
+            )
+        click.secho(
+            "Successfully applied all schema migrations!", fg="green", bold=True
+        )
+    except Exception as e:
+        raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
+    finally:
+        click.secho(
+            "Releasing database lock via the Ingestion Helper service...",
+            fg="bright_black",
+        )
+        try:
+            client.release_lock(workflow_id="schema-migration")
+        except Exception as e:
+            click.secho(
+                f"Warning: Failed to release database lock: {e}",
+                fg="yellow",
+            )
+
+
+
+@admin.command(name="migrate-db")
+def migrate_db() -> None:
+    """Apply pending schema migrations to the Spanner database."""
+    click.secho("Datacommons Admin Migrate-DB", fg="cyan", bold=True)
+    client, project_id, instance_id, database_id = _setup_ingestion_client()
+    _run_migrations(client, project_id, instance_id, database_id)
+
+
 @admin.command(name="init-db")
 @click.option(
     "--init-only", is_flag=True, help="Only initialize the database without seeding."
@@ -557,7 +640,7 @@ def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
 def init_db(init_only: bool) -> None:
     """Initialize (and by default seed) the Spanner database via the DCP Ingestion Helper service."""
     click.secho("Datacommons Admin Init-DB", fg="cyan", bold=True)
-    client, instance_id, database_id = _setup_ingestion_client()
+    client, project_id, instance_id, database_id = _setup_ingestion_client()
 
     click.secho(
         f"Initializing Spanner database '{instance_id}/{database_id}' via the Ingestion Helper service (this may take a few moments)...",
@@ -570,6 +653,8 @@ def init_db(init_only: bool) -> None:
     if message:
         click.secho(f"Details: {message}", fg="bright_black")
 
+    _run_migrations(client, project_id, instance_id, database_id)
+
     if not init_only:
         _run_seed_db(client, instance_id, database_id)
 
@@ -578,10 +663,12 @@ def init_db(init_only: bool) -> None:
 def seed_db() -> None:
     """Seed the Spanner database via the DCP Ingestion Helper service."""
     click.secho("Datacommons Admin Seed-DB", fg="cyan", bold=True)
-    client, instance_id, database_id = _setup_ingestion_client()
+    client, _project_id, instance_id, database_id = _setup_ingestion_client()
     _run_seed_db(client, instance_id, database_id)
 
 
+
 from datacommons_admin.ingest_cli import ingest
 
 admin.add_command(ingest)
+
diff --git a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
index 9be2a852..d663f102 100644
--- a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
+++ b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
@@ -120,3 +120,34 @@ def initialize_database(self) -> dict:
     def seed_database(self) -> dict:
         """Calls the seed_database endpoint on the ingestion helper service."""
         return self._call_endpoint("database/seed")
+
+    def acquire_lock(self, workflow_id: str, timeout: int = 300) -> dict:
+        """Acquires a distributed database lock via the ingestion helper service.
+
+        Args:
+            workflow_id: Identifier for the lock holder.
+            timeout: Maximum lock duration in seconds.
+
+        Returns:
+            API response dictionary from the Ingestion Helper service.
+        """
+        payload = {
+            "workflowId": workflow_id,
+            "timeout": timeout,
+        }
+        return self._call_endpoint("database/lock/acquire", payload=payload)
+
+    def release_lock(self, workflow_id: str) -> dict:
+        """Releases the distributed database lock via the ingestion helper service.
+
+        Args:
+            workflow_id: Identifier for the lock holder.
+
+        Returns:
+            API response dictionary from the Ingestion Helper service.
+        """
+        payload = {
+            "workflowId": workflow_id,
+        }
+        return self._call_endpoint("database/lock/release", payload=payload)
+
diff --git a/packages/datacommons-admin/pyproject.toml b/packages/datacommons-admin/pyproject.toml
index 874c93ae..547e27e0 100644
--- a/packages/datacommons-admin/pyproject.toml
+++ b/packages/datacommons-admin/pyproject.toml
@@ -12,8 +12,10 @@ dependencies = [
   "click>=8.1.7",
   "google-cloud-storage>=2.13.0",
   "pyopenssl>=24.0.0",
+  "datacommons-db",
 ]
 
+
 [build-system]
 requires = ["uv", "setuptools"]
 build-backend = "setuptools.build_meta"
diff --git a/packages/datacommons-admin/tests/test_admin_cli.py b/packages/datacommons-admin/tests/test_admin_cli.py
index 287fd33e..117bd2ee 100644
--- a/packages/datacommons-admin/tests/test_admin_cli.py
+++ b/packages/datacommons-admin/tests/test_admin_cli.py
@@ -13,12 +13,14 @@
 # limitations under the License.
 
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
-from click.testing import CliRunner
+import click
 import pytest
-
-from datacommons_admin.admin_cli import admin, init
+from click.testing import CliRunner
+from datacommons_admin.admin_cli import admin
+from datacommons_db.clients.spanner_client import ExecutionStatus
+from datacommons_db.migrations.migration_runner import MigrationResult
 
 
 @pytest.fixture
@@ -219,6 +221,7 @@ def test_init_db_terraform_error(
     assert "Failed to run 'terraform output'" in result.output
 
 
+@patch("datacommons_admin.admin_cli._run_migrations")
 @patch("datacommons_admin.tf_utils.shutil.which")
 @patch("datacommons_admin.tf_utils.subprocess.run")
 @patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
@@ -228,13 +231,14 @@ def test_init_db_success(
     mock_session: patch,
     mock_run: patch,
     mock_which: patch,
+    mock_run_migrations: patch,
     runner: CliRunner,
 ) -> None:
     mock_which.return_value = "terraform"
     from unittest.mock import MagicMock
 
     mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}}'
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
     mock_run.return_value = mock_proc
 
     mock_creds = MagicMock()
@@ -252,8 +256,10 @@ def test_init_db_success(
     assert "Successfully initialized Spanner database" in result.output
     assert "Details: DB Initialized" in result.output
     assert "Successfully seeded Spanner database" in result.output
+    mock_run_migrations.assert_called_once()
 
 
+@patch("datacommons_admin.admin_cli._run_migrations")
 @patch("datacommons_admin.tf_utils.shutil.which")
 @patch("datacommons_admin.tf_utils.subprocess.run")
 @patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
@@ -263,13 +269,14 @@ def test_init_db_success_no_details(
     mock_session: patch,
     mock_run: patch,
     mock_which: patch,
+    mock_run_migrations: patch,
     runner: CliRunner,
 ) -> None:
     mock_which.return_value = "terraform"
     from unittest.mock import MagicMock
 
     mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}}'
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
     mock_run.return_value = mock_proc
 
     mock_creds = MagicMock()
@@ -287,8 +294,10 @@ def test_init_db_success_no_details(
     assert "Successfully initialized Spanner database" in result.output
     assert "Details:" not in result.output
     assert "Successfully seeded Spanner database" in result.output
+    mock_run_migrations.assert_called_once()
 
 
+@patch("datacommons_admin.admin_cli._run_migrations")
 @patch("datacommons_admin.tf_utils.shutil.which")
 @patch("datacommons_admin.tf_utils.subprocess.run")
 @patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
@@ -298,13 +307,14 @@ def test_init_db_init_only(
     mock_session: patch,
     mock_run: patch,
     mock_which: patch,
+    mock_run_migrations: patch,
     runner: CliRunner,
 ) -> None:
     mock_which.return_value = "terraform"
     from unittest.mock import MagicMock
 
     mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}}'
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
     mock_run.return_value = mock_proc
 
     mock_creds = MagicMock()
@@ -321,6 +331,50 @@ def test_init_db_init_only(
     assert result.exit_code == 0
     assert "Successfully initialized Spanner database" in result.output
     assert "Seeding Spanner database" not in result.output
+    mock_run_migrations.assert_called_once()
+
+
+@patch("datacommons_admin.admin_cli._run_migrations")
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_init_db_migration_failure_halts_before_seed(
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    mock_run_migrations: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    from unittest.mock import MagicMock
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success", "message": "DB Initialized"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    mock_run_migrations.side_effect = click.ClickException("Migration failed")
+
+    result = runner.invoke(admin, ["init-db"])
+    assert result.exit_code != 0
+    assert "Successfully initialized Spanner database" in result.output
+    assert "Migration failed" in result.output
+    # Seeding should NOT be called if migrations fail
+    assert "Seeding Spanner database" not in result.output
+    assert "Successfully seeded Spanner database" not in result.output
+
+
 
 
 @patch("datacommons_admin.tf_utils.shutil.which")
@@ -338,7 +392,7 @@ def test_seed_db_success(
     from unittest.mock import MagicMock
 
     mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}}'
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
     mock_run.return_value = mock_proc
 
     mock_creds = MagicMock()
@@ -579,3 +633,212 @@ def test_ingest_start_with_imports_success(
     called_payload = called_args["json"]
     assert "argument" in called_payload
     assert json.loads(called_payload["argument"]) == expected_arg
+
+
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_no_pending(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_session.return_value = mock_session_inst
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = []
+    mock_runner_cls.return_value = mock_runner_inst
+
+    result = runner.invoke(admin, ["migrate-db"])
+    assert result.exit_code == 0
+    assert "Database schema is already up-to-date" in result.output
+    # Lock should not be acquired when there are no pending migrations
+    mock_session_inst.post.assert_not_called()
+
+
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_with_pending_success(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    mock_migration = MagicMock()
+    mock_migration.creation_timestamp = "20260817000000"
+    mock_migration.description = "Bootstrap migration"
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
+    mock_runner_inst.run_migrations.return_value = [
+        MigrationResult(
+            status=ExecutionStatus.SUCCESS,
+            creation_timestamp="20260817000000",
+            description="Bootstrap migration",
+        )
+    ]
+    mock_runner_cls.return_value = mock_runner_inst
+
+    result = runner.invoke(admin, ["migrate-db"])
+    assert result.exit_code == 0
+    assert "Found 1 pending schema migration" in result.output
+    assert "Applied migration 20260817000000: Bootstrap migration" in result.output
+    assert "Successfully applied all schema migrations!" in result.output
+
+    # Check lock acquired then released
+    assert mock_session_inst.post.call_count == 2
+    mock_session_inst.post.assert_any_call(
+        "https://mock-helper/database/lock/acquire",
+        json={"workflowId": "schema-migration", "timeout": 300},
+        timeout=300,
+    )
+    mock_session_inst.post.assert_any_call(
+        "https://mock-helper/database/lock/release",
+        json={"workflowId": "schema-migration"},
+        timeout=300,
+    )
+
+
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_failure_releases_lock(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    mock_migration = MagicMock()
+    mock_migration.creation_timestamp = "20260817000000"
+    mock_migration.description = "Bootstrap migration"
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
+    mock_runner_inst.run_migrations.side_effect = RuntimeError("DDL operation failed")
+    mock_runner_cls.return_value = mock_runner_inst
+
+    result = runner.invoke(admin, ["migrate-db"])
+    assert result.exit_code != 0
+    assert "Failed to apply schema migrations: DDL operation failed" in result.output
+
+    # Ensure release_lock was still called in finally block
+    mock_session_inst.post.assert_any_call(
+        "https://mock-helper/database/lock/release",
+        json={"workflowId": "schema-migration"},
+        timeout=300,
+    )
+
+
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_lock_busy_error(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = False
+    mock_resp.status_code = 503
+    mock_resp.json.return_value = {"detail": "Lock busy"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    mock_migration = MagicMock()
+    mock_migration.creation_timestamp = "20260817000000"
+    mock_migration.description = "Bootstrap migration"
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
+    mock_runner_cls.return_value = mock_runner_inst
+
+    result = runner.invoke(admin, ["migrate-db"])
+    assert result.exit_code != 0
+    assert "Ingestion Helper returned HTTP 503" in result.output
+    assert "Please wait for active ingestions to finish before running migrations" in result.output
+
+
+
diff --git a/packages/datacommons-admin/tests/test_ingestion_helper_client.py b/packages/datacommons-admin/tests/test_ingestion_helper_client.py
new file mode 100644
index 00000000..37a27033
--- /dev/null
+++ b/packages/datacommons-admin/tests/test_ingestion_helper_client.py
@@ -0,0 +1,289 @@
+# 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.
+
+from unittest.mock import MagicMock, patch
+
+import click
+import pytest
+import requests
+from datacommons_admin.ingestion_helper_client import IngestionHelperClient
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_acquire_lock_default_timeout(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success", "message": "Lock acquired"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    result = client.acquire_lock("schema-migration")
+
+    assert result == {"status": "success", "message": "Lock acquired"}
+    mock_session_inst.post.assert_called_once_with(
+        "https://mock-helper.a.run.app/database/lock/acquire",
+        json={"workflowId": "schema-migration", "timeout": 300},
+        timeout=300,
+    )
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_acquire_lock_custom_args(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success", "message": "Lock acquired"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    result = client.acquire_lock("custom-migration-123", timeout=600)
+
+    assert result == {"status": "success", "message": "Lock acquired"}
+    mock_session_inst.post.assert_called_once_with(
+        "https://mock-helper.a.run.app/database/lock/acquire",
+        json={"workflowId": "custom-migration-123", "timeout": 600},
+        timeout=300,
+    )
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_release_lock_success(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success", "message": "Lock released"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    result = client.release_lock("schema-migration")
+
+    assert result == {"status": "success", "message": "Lock released"}
+    mock_session_inst.post.assert_called_once_with(
+        "https://mock-helper.a.run.app/database/lock/release",
+        json={"workflowId": "schema-migration"},
+        timeout=300,
+    )
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_release_lock_custom_workflow_id(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success", "message": "Lock released"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    result = client.release_lock("custom-migration-123")
+
+    assert result == {"status": "success", "message": "Lock released"}
+    mock_session_inst.post.assert_called_once_with(
+        "https://mock-helper.a.run.app/database/lock/release",
+        json={"workflowId": "custom-migration-123"},
+        timeout=300,
+    )
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_lock_acquire_http_error(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = False
+    mock_resp.status_code = 503
+    mock_resp.json.return_value = {
+        "detail": "Database lock already held by workflow-456"
+    }
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    with pytest.raises(click.ClickException) as exc_info:
+        client.acquire_lock("schema-migration")
+
+    assert "Ingestion Helper returned HTTP 503" in str(exc_info.value)
+    assert "Database lock already held by workflow-456" in str(exc_info.value)
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_unauthorized_error(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = False
+    mock_resp.status_code = 401
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    with pytest.raises(click.ClickException) as exc_info:
+        client.acquire_lock("schema-migration")
+
+    assert "HTTP 401 Unauthorized when calling Ingestion Helper" in str(
+        exc_info.value
+    )
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("google.auth.impersonated_credentials.IDTokenCredentials")
+@patch("google.auth.impersonated_credentials.Credentials")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_network_request_exception(
+    mock_auth_default: patch,
+    mock_imp_creds: patch,
+    mock_id_token_creds: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_imp_creds.return_value = MagicMock()
+    mock_id_token_creds.return_value = MagicMock()
+
+    mock_session_inst = MagicMock()
+    mock_session_inst.post.side_effect = requests.exceptions.ConnectionError(
+        "Connection refused"
+    )
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient(
+        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
+    )
+    with pytest.raises(click.ClickException) as exc_info:
+        client.release_lock("schema-migration")
+
+    assert "Network or authentication error" in str(exc_info.value)
+
+
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.id_token.fetch_id_token")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+def test_id_token_fetch_when_no_service_account(
+    mock_auth_default: patch,
+    mock_fetch_id_token: patch,
+    mock_session: patch,
+) -> None:
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+    mock_fetch_id_token.return_value = "mock-id-token"
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    client = IngestionHelperClient("https://mock-helper.a.run.app")
+    result = client.initialize_database()
+
+    assert result == {"status": "success"}
+    mock_fetch_id_token.assert_called_once()
+
+
+def test_localhost_bypass_auth() -> None:
+    client = IngestionHelperClient("http://localhost:8080")
+    assert isinstance(client.session, requests.Session)
diff --git a/uv.lock b/uv.lock
index fdb58b1b..19347de6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -310,6 +310,7 @@ name = "datacommons-admin"
 source = { editable = "packages/datacommons-admin" }
 dependencies = [
     { name = "click" },
+    { name = "datacommons-db" },
     { name = "google-cloud-storage" },
     { name = "pyopenssl" },
 ]
@@ -317,6 +318,7 @@ dependencies = [
 [package.metadata]
 requires-dist = [
     { name = "click", specifier = ">=8.1.7" },
+    { name = "datacommons-db", editable = "packages/datacommons-db" },
     { name = "google-cloud-storage", specifier = ">=2.13.0" },
     { name = "pyopenssl", specifier = ">=24.0.0" },
 ]
@@ -706,6 +708,7 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/e8/2e1462c8fdbe0f210feb5ac7ad2d9029af8be3bf45bd9fa39765f821642f/greenlet-3.3.1-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:5fd23b9bc6d37b563211c6abbb1b3cab27db385a4449af5c32e932f93017080c", size = 274974, upload-time = "2026-01-23T15:31:02.891Z" },
     { url = "https://files.pythonhosted.org/packages/7e/a8/530a401419a6b302af59f67aaf0b9ba1015855ea7e56c036b5928793c5bd/greenlet-3.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:09f51496a0bfbaa9d74d36a52d2580d1ef5ed4fdfcff0a73730abfbbbe1403dd", size = 577175, upload-time = "2026-01-23T16:00:56.213Z" },
     { url = "https://files.pythonhosted.org/packages/8e/89/7e812bb9c05e1aaef9b597ac1d0962b9021d2c6269354966451e885c4e6b/greenlet-3.3.1-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb0feb07fe6e6a74615ee62a880007d976cf739b6669cce95daa7373d4fc69c5", size = 590401, upload-time = "2026-01-23T16:05:26.365Z" },
+    { url = "https://files.pythonhosted.org/packages/70/ae/e2d5f0e59b94a2269b68a629173263fa40b63da32f5c231307c349315871/greenlet-3.3.1-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:67ea3fc73c8cd92f42467a72b75e8f05ed51a0e9b1d15398c913416f2dafd49f", size = 601161, upload-time = "2026-01-23T16:15:53.456Z" },
     { url = "https://files.pythonhosted.org/packages/5c/ae/8d472e1f5ac5efe55c563f3eabb38c98a44b832602e12910750a7c025802/greenlet-3.3.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:39eda9ba259cc9801da05351eaa8576e9aa83eb9411e8f0c299e05d712a210f2", size = 590272, upload-time = "2026-01-23T15:32:49.411Z" },
     { url = "https://files.pythonhosted.org/packages/a8/51/0fde34bebfcadc833550717eade64e35ec8738e6b097d5d248274a01258b/greenlet-3.3.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e2e7e882f83149f0a71ac822ebf156d902e7a5d22c9045e3e0d1daf59cee2cc9", size = 1550729, upload-time = "2026-01-23T16:04:20.867Z" },
     { url = "https://files.pythonhosted.org/packages/16/c9/2fb47bee83b25b119d5a35d580807bb8b92480a54b68fef009a02945629f/greenlet-3.3.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:80aa4d79eb5564f2e0a6144fcc744b5a37c56c4a92d60920720e99210d88db0f", size = 1615552, upload-time = "2026-01-23T15:33:45.743Z" },
@@ -714,6 +717,7 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f9/c8/9d76a66421d1ae24340dfae7e79c313957f6e3195c144d2c73333b5bfe34/greenlet-3.3.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:7e806ca53acf6d15a888405880766ec84721aa4181261cd11a457dfe9a7a4975", size = 276443, upload-time = "2026-01-23T15:30:10.066Z" },
     { url = "https://files.pythonhosted.org/packages/81/99/401ff34bb3c032d1f10477d199724f5e5f6fbfb59816ad1455c79c1eb8e7/greenlet-3.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d842c94b9155f1c9b3058036c24ffb8ff78b428414a19792b2380be9cecf4f36", size = 597359, upload-time = "2026-01-23T16:00:57.394Z" },
     { url = "https://files.pythonhosted.org/packages/2b/bc/4dcc0871ed557792d304f50be0f7487a14e017952ec689effe2180a6ff35/greenlet-3.3.1-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:20fedaadd422fa02695f82093f9a98bad3dab5fcda793c658b945fcde2ab27ba", size = 607805, upload-time = "2026-01-23T16:05:28.068Z" },
+    { url = "https://files.pythonhosted.org/packages/3b/cd/7a7ca57588dac3389e97f7c9521cb6641fd8b6602faf1eaa4188384757df/greenlet-3.3.1-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c620051669fd04ac6b60ebc70478210119c56e2d5d5df848baec4312e260e4ca", size = 622363, upload-time = "2026-01-23T16:15:54.754Z" },
     { url = "https://files.pythonhosted.org/packages/cf/05/821587cf19e2ce1f2b24945d890b164401e5085f9d09cbd969b0c193cd20/greenlet-3.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14194f5f4305800ff329cbf02c5fcc88f01886cadd29941b807668a45f0d2336", size = 609947, upload-time = "2026-01-23T15:32:51.004Z" },
     { url = "https://files.pythonhosted.org/packages/a4/52/ee8c46ed9f8babaa93a19e577f26e3d28a519feac6350ed6f25f1afee7e9/greenlet-3.3.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7b2fe4150a0cf59f847a67db8c155ac36aed89080a6a639e9f16df5d6c6096f1", size = 1567487, upload-time = "2026-01-23T16:04:22.125Z" },
     { url = "https://files.pythonhosted.org/packages/8f/7c/456a74f07029597626f3a6db71b273a3632aecb9afafeeca452cfa633197/greenlet-3.3.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:49f4ad195d45f4a66a0eb9c1ba4832bb380570d361912fa3554746830d332149", size = 1636087, upload-time = "2026-01-23T15:33:47.486Z" },
@@ -722,6 +726,7 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/ab/d26750f2b7242c2b90ea2ad71de70cfcd73a948a49513188a0fc0d6fc15a/greenlet-3.3.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:7ab327905cabb0622adca5971e488064e35115430cec2c35a50fd36e72a315b3", size = 275205, upload-time = "2026-01-23T15:30:24.556Z" },
     { url = "https://files.pythonhosted.org/packages/10/d3/be7d19e8fad7c5a78eeefb2d896a08cd4643e1e90c605c4be3b46264998f/greenlet-3.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:65be2f026ca6a176f88fb935ee23c18333ccea97048076aef4db1ef5bc0713ac", size = 599284, upload-time = "2026-01-23T16:00:58.584Z" },
     { url = "https://files.pythonhosted.org/packages/ae/21/fe703aaa056fdb0f17e5afd4b5c80195bbdab701208918938bd15b00d39b/greenlet-3.3.1-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7a3ae05b3d225b4155bda56b072ceb09d05e974bc74be6c3fc15463cf69f33fd", size = 610274, upload-time = "2026-01-23T16:05:29.312Z" },
+    { url = "https://files.pythonhosted.org/packages/06/00/95df0b6a935103c0452dad2203f5be8377e551b8466a29650c4c5a5af6cc/greenlet-3.3.1-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:12184c61e5d64268a160226fb4818af4df02cfead8379d7f8b99a56c3a54ff3e", size = 624375, upload-time = "2026-01-23T16:15:55.915Z" },
     { url = "https://files.pythonhosted.org/packages/cb/86/5c6ab23bb3c28c21ed6bebad006515cfe08b04613eb105ca0041fecca852/greenlet-3.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6423481193bbbe871313de5fd06a082f2649e7ce6e08015d2a76c1e9186ca5b3", size = 612904, upload-time = "2026-01-23T15:32:52.317Z" },
     { url = "https://files.pythonhosted.org/packages/c2/f3/7949994264e22639e40718c2daf6f6df5169bf48fb038c008a489ec53a50/greenlet-3.3.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:33a956fe78bbbda82bfc95e128d61129b32d66bcf0a20a1f0c08aa4839ffa951", size = 1567316, upload-time = "2026-01-23T16:04:23.316Z" },
     { url = "https://files.pythonhosted.org/packages/8d/6e/d73c94d13b6465e9f7cd6231c68abde838bb22408596c05d9059830b7872/greenlet-3.3.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4b065d3284be43728dd280f6f9a13990b56470b81be20375a207cdc814a983f2", size = 1636549, upload-time = "2026-01-23T15:33:48.643Z" },
@@ -730,6 +735,7 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ae/fb/011c7c717213182caf78084a9bea51c8590b0afda98001f69d9f853a495b/greenlet-3.3.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:bd59acd8529b372775cd0fcbc5f420ae20681c5b045ce25bd453ed8455ab99b5", size = 275737, upload-time = "2026-01-23T15:32:16.889Z" },
     { url = "https://files.pythonhosted.org/packages/41/2e/a3a417d620363fdbb08a48b1dd582956a46a61bf8fd27ee8164f9dfe87c2/greenlet-3.3.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b31c05dd84ef6871dd47120386aed35323c944d86c3d91a17c4b8d23df62f15b", size = 646422, upload-time = "2026-01-23T16:01:00.354Z" },
     { url = "https://files.pythonhosted.org/packages/b4/09/c6c4a0db47defafd2d6bab8ddfe47ad19963b4e30f5bed84d75328059f8c/greenlet-3.3.1-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:02925a0bfffc41e542c70aa14c7eda3593e4d7e274bfcccca1827e6c0875902e", size = 658219, upload-time = "2026-01-23T16:05:30.956Z" },
+    { url = "https://files.pythonhosted.org/packages/e2/89/b95f2ddcc5f3c2bc09c8ee8d77be312df7f9e7175703ab780f2014a0e781/greenlet-3.3.1-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3e0f3878ca3a3ff63ab4ea478585942b53df66ddde327b59ecb191b19dbbd62d", size = 671455, upload-time = "2026-01-23T16:15:57.232Z" },
     { url = "https://files.pythonhosted.org/packages/80/38/9d42d60dffb04b45f03dbab9430898352dba277758640751dc5cc316c521/greenlet-3.3.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:34a729e2e4e4ffe9ae2408d5ecaf12f944853f40ad724929b7585bca808a9d6f", size = 660237, upload-time = "2026-01-23T15:32:53.967Z" },
     { url = "https://files.pythonhosted.org/packages/96/61/373c30b7197f9e756e4c81ae90a8d55dc3598c17673f91f4d31c3c689c3f/greenlet-3.3.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:aec9ab04e82918e623415947921dea15851b152b822661cce3f8e4393c3df683", size = 1615261, upload-time = "2026-01-23T16:04:25.066Z" },
     { url = "https://files.pythonhosted.org/packages/fd/d3/ca534310343f5945316f9451e953dcd89b36fe7a19de652a1dc5a0eeef3f/greenlet-3.3.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:71c767cf281a80d02b6c1bdc41c9468e1f5a494fb11bc8688c360524e273d7b1", size = 1683719, upload-time = "2026-01-23T15:33:50.61Z" },
@@ -738,6 +744,7 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/28/24/cbbec49bacdcc9ec652a81d3efef7b59f326697e7edf6ed775a5e08e54c2/greenlet-3.3.1-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:3e63252943c921b90abb035ebe9de832c436401d9c45f262d80e2d06cc659242", size = 282706, upload-time = "2026-01-23T15:33:05.525Z" },
     { url = "https://files.pythonhosted.org/packages/86/2e/4f2b9323c144c4fe8842a4e0d92121465485c3c2c5b9e9b30a52e80f523f/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:76e39058e68eb125de10c92524573924e827927df5d3891fbc97bd55764a8774", size = 651209, upload-time = "2026-01-23T16:01:01.517Z" },
     { url = "https://files.pythonhosted.org/packages/d9/87/50ca60e515f5bb55a2fbc5f0c9b5b156de7d2fc51a0a69abc9d23914a237/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c9f9d5e7a9310b7a2f416dd13d2e3fd8b42d803968ea580b7c0f322ccb389b97", size = 654300, upload-time = "2026-01-23T16:05:32.199Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/25/c51a63f3f463171e09cb586eb64db0861eb06667ab01a7968371a24c4f3b/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4b9721549a95db96689458a1e0ae32412ca18776ed004463df3a9299c1b257ab", size = 662574, upload-time = "2026-01-23T16:15:58.364Z" },
     { url = "https://files.pythonhosted.org/packages/1d/94/74310866dfa2b73dd08659a3d18762f83985ad3281901ba0ee9a815194fb/greenlet-3.3.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:92497c78adf3ac703b57f1e3813c2d874f27f71a178f9ea5887855da413cd6d2", size = 653842, upload-time = "2026-01-23T15:32:55.671Z" },
     { url = "https://files.pythonhosted.org/packages/97/43/8bf0ffa3d498eeee4c58c212a3905dd6146c01c8dc0b0a046481ca29b18c/greenlet-3.3.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ed6b402bc74d6557a705e197d47f9063733091ed6357b3de33619d8a8d93ac53", size = 1614917, upload-time = "2026-01-23T16:04:26.276Z" },
     { url = "https://files.pythonhosted.org/packages/89/90/a3be7a5f378fc6e84abe4dcfb2ba32b07786861172e502388b4c90000d1b/greenlet-3.3.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:59913f1e5ada20fde795ba906916aea25d442abcc0593fba7e26c92b7ad76249", size = 1676092, upload-time = "2026-01-23T15:33:52.176Z" },

From 2536015f9eac8a9e090c6f373e029d5b970d5755 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 15:31:20 -0700
Subject: [PATCH 02/11] refactor migrate functions to have extra checks

---
 .../datacommons_admin/admin_cli.py            | 153 ++++++++++++-----
 .../datacommons-admin/tests/test_admin_cli.py | 155 +++++++++++++++++-
 2 files changed, 268 insertions(+), 40 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index f466c44e..37178af2 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -519,7 +519,7 @@ def init(
     )
 
 
-def _setup_ingestion_client() -> Tuple[Any, str, str, str]:
+def _setup_ingestion_client() -> Tuple[IngestionHelperClient, str, str, str]:
     click.secho(
         "Fetching ingestion service URL, workflow service account, and Spanner details from Terraform outputs...",
         fg="bright_black",
@@ -538,58 +538,54 @@ def _setup_ingestion_client() -> Tuple[Any, str, str, str]:
         fg="green",
     )
 
-    client = IngestionHelperClient(url, service_account_email=sa_email)
-    return client, project_id, instance_id, database_id
+    ingestion_helper_client = IngestionHelperClient(
+        url, service_account_email=sa_email
+    )
+    return ingestion_helper_client, project_id, instance_id, database_id
 
 
-def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
+def _run_seed_db(
+    ingestion_helper_client: IngestionHelperClient,
+    instance_id: str,
+    database_id: str,
+) -> None:
     click.secho(
         f"Seeding Spanner database '{instance_id}/{database_id}' via the Ingestion Helper service (this may take a few moments)...",
         fg="bright_black",
     )
-    result = client.seed_database()
+    result = ingestion_helper_client.seed_database()
     click.secho("Successfully seeded Spanner database!", fg="green", bold=True)
     message = result.get("message")
     if message:
         click.secho(f"Details: {message}", fg="bright_black")
 
 
-def _run_migrations(
-    client: Any, project_id: str, instance_id: str, database_id: str
-) -> None:
-    click.secho(
-        f"Checking schema migrations for Spanner database '{project_id}/{instance_id}/{database_id}'...",
-        fg="bright_black",
-    )
-
+def _create_migration_runner(
+    project_id: str, instance_id: str, database_id: str
+) -> MigrationRunner:
+    """Initializes a SpannerClient and returns a MigrationRunner instance."""
     try:
         spanner_client = SpannerClient(
             project_id=project_id,
             instance_id=instance_id,
             database_id=database_id,
         )
-        runner = MigrationRunner(spanner_client=spanner_client)
-        pending = runner.get_pending_migrations()
+        return MigrationRunner(spanner_client=spanner_client)
     except Exception as e:
-        raise click.ClickException(f"Failed to check pending migrations: {e}") from e
-
-    if not pending:
-        click.secho(
-            "Database schema is already up-to-date. No migrations to apply.",
-            fg="green",
-        )
-        return
+        raise click.ClickException(f"Failed to initialize migration runner: {e}") from e
 
-    click.secho(f"Found {len(pending)} pending schema migration(s):", fg="cyan")
-    for m in pending:
-        click.echo(f"  - {m.creation_timestamp}: {m.description}")
 
+def _apply_migrations(
+    ingestion_helper_client: IngestionHelperClient, runner: MigrationRunner
+) -> None:
+    """Acquires a distributed database lock and applies all pending migrations."""
+    # Attempt to acquire Spanner database lock via the Ingestion Helper service.
     click.secho(
         "Acquiring database lock via the Ingestion Helper service...",
         fg="bright_black",
     )
     try:
-        client.acquire_lock(workflow_id="schema-migration")
+        ingestion_helper_client.acquire_lock(workflow_id="schema-migration")
     except Exception as e:
         raise click.ClickException(
             f"Could not acquire database lock: {e}\n"
@@ -597,6 +593,7 @@ def _run_migrations(
             "Please wait for active ingestions to finish before running migrations."
         ) from e
 
+    # Apply all pending migrations
     try:
         click.secho("Applying pending schema migrations...", fg="bright_black")
         results = runner.run_migrations()
@@ -610,13 +607,15 @@ def _run_migrations(
         )
     except Exception as e:
         raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
+
+    # Release database lock
     finally:
         click.secho(
             "Releasing database lock via the Ingestion Helper service...",
             fg="bright_black",
         )
         try:
-            client.release_lock(workflow_id="schema-migration")
+            ingestion_helper_client.release_lock(workflow_id="schema-migration")
         except Exception as e:
             click.secho(
                 f"Warning: Failed to release database lock: {e}",
@@ -624,13 +623,87 @@ def _run_migrations(
             )
 
 
+def _confirm_migration(num_pending: int, instance_id: str, database_id: str) -> bool:
+    """Displays a safety warning and prompts the user to confirm applying migrations."""
+    click.secho(
+        "\nWarning: Schema migrations will modify your Spanner database schema. "
+        "It is strongly recommended to create a database backup before proceeding in production environments.",
+        fg="yellow",
+    )
+    return _confirm(
+        f"Apply {num_pending} pending schema migration(s) to Spanner database '{instance_id}/{database_id}'?",
+        default=False,
+    )
+
+
+def _run_migrations(
+    ingestion_helper_client: IngestionHelperClient,
+    project_id: str,
+    instance_id: str,
+    database_id: str,
+    auto_approve: bool = False,
+) -> None:
+    """Checks, optionally confirms, and applies pending schema migrations to Spanner.
+
+    Args:
+        ingestion_helper_client: IngestionHelperClient instance.
+        project_id: GCP project ID hosting the Spanner database.
+        instance_id: Cloud Spanner instance ID.
+        database_id: Cloud Spanner database ID.
+        auto_approve: If False, prompts user for interactive confirmation before applying.
+    """
+    click.secho(
+        f"Checking schema migrations for Spanner database '{project_id}/{instance_id}/{database_id}'...",
+        fg="bright_black",
+    )
+    runner = _create_migration_runner(project_id, instance_id, database_id)
+
+    # Fetch pending migrations.
+    try:
+        pending = runner.get_pending_migrations()
+    except Exception as e:
+        raise click.ClickException(f"Failed to check pending migrations: {e}") from e
+
+    # Return early if there are no pending migrations.
+    if not pending:
+        click.secho(
+            "Database schema is already up-to-date. No migrations to apply.",
+            fg="green",
+        )
+        return
+
+    click.secho(f"Found {len(pending)} pending schema migration(s):", fg="cyan")
+    for m in pending:
+        click.echo(f"  - {m.creation_timestamp}: {m.description}")
+
+    # Ask user for confirmation if not auto-approved
+    if not auto_approve and not _confirm_migration(len(pending), instance_id, database_id):
+        click.secho("Migration cancelled.", fg="yellow")
+        return
+
+    # Apply migrations
+    _apply_migrations(ingestion_helper_client, runner)
+
 
 @admin.command(name="migrate-db")
-def migrate_db() -> None:
+@click.option(
+    "-y",
+    "--yes",
+    "auto_approve",
+    is_flag=True,
+    help="Automatically confirm and apply pending migrations without prompting.",
+)
+def migrate_db(auto_approve: bool) -> None:
     """Apply pending schema migrations to the Spanner database."""
     click.secho("Datacommons Admin Migrate-DB", fg="cyan", bold=True)
-    client, project_id, instance_id, database_id = _setup_ingestion_client()
-    _run_migrations(client, project_id, instance_id, database_id)
+    ingestion_helper_client, project_id, instance_id, database_id = _setup_ingestion_client()
+    _run_migrations(
+        ingestion_helper_client,
+        project_id,
+        instance_id,
+        database_id,
+        auto_approve=auto_approve,
+    )
 
 
 @admin.command(name="init-db")
@@ -640,31 +713,37 @@ def migrate_db() -> None:
 def init_db(init_only: bool) -> None:
     """Initialize (and by default seed) the Spanner database via the DCP Ingestion Helper service."""
     click.secho("Datacommons Admin Init-DB", fg="cyan", bold=True)
-    client, project_id, instance_id, database_id = _setup_ingestion_client()
+    ingestion_helper_client, project_id, instance_id, database_id = _setup_ingestion_client()
 
     click.secho(
         f"Initializing Spanner database '{instance_id}/{database_id}' via the Ingestion Helper service (this may take a few moments)...",
         fg="bright_black",
     )
-    result = client.initialize_database()
+    result = ingestion_helper_client.initialize_database()
 
     click.secho("Successfully initialized Spanner database!", fg="green", bold=True)
     message = result.get("message")
     if message:
         click.secho(f"Details: {message}", fg="bright_black")
 
-    _run_migrations(client, project_id, instance_id, database_id)
+    _run_migrations(
+        ingestion_helper_client,
+        project_id,
+        instance_id,
+        database_id,
+        auto_approve=True,
+    )
 
     if not init_only:
-        _run_seed_db(client, instance_id, database_id)
+        _run_seed_db(ingestion_helper_client, instance_id, database_id)
 
 
 @admin.command(name="seed-db")
 def seed_db() -> None:
     """Seed the Spanner database via the DCP Ingestion Helper service."""
     click.secho("Datacommons Admin Seed-DB", fg="cyan", bold=True)
-    client, _project_id, instance_id, database_id = _setup_ingestion_client()
-    _run_seed_db(client, instance_id, database_id)
+    ingestion_helper_client, _project_id, instance_id, database_id = _setup_ingestion_client()
+    _run_seed_db(ingestion_helper_client, instance_id, database_id)
 
 
 
diff --git a/packages/datacommons-admin/tests/test_admin_cli.py b/packages/datacommons-admin/tests/test_admin_cli.py
index 117bd2ee..4372a25f 100644
--- a/packages/datacommons-admin/tests/test_admin_cli.py
+++ b/packages/datacommons-admin/tests/test_admin_cli.py
@@ -721,9 +721,10 @@ def test_migrate_db_with_pending_success(
     ]
     mock_runner_cls.return_value = mock_runner_inst
 
-    result = runner.invoke(admin, ["migrate-db"])
+    result = runner.invoke(admin, ["migrate-db"], input="y\n")
     assert result.exit_code == 0
     assert "Found 1 pending schema migration" in result.output
+    assert "Warning: Schema migrations will modify your Spanner database schema" in result.output
     assert "Applied migration 20260817000000: Bootstrap migration" in result.output
     assert "Successfully applied all schema migrations!" in result.output
 
@@ -741,6 +742,153 @@ def test_migrate_db_with_pending_success(
     )
 
 
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_with_yes_flag(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_resp = MagicMock()
+    mock_resp.ok = True
+    mock_resp.json.return_value = {"status": "success"}
+    mock_session_inst.post.return_value = mock_resp
+    mock_session.return_value = mock_session_inst
+
+    mock_migration = MagicMock()
+    mock_migration.creation_timestamp = "20260817000000"
+    mock_migration.description = "Bootstrap migration"
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
+    mock_runner_inst.run_migrations.return_value = [
+        MigrationResult(
+            status=ExecutionStatus.SUCCESS,
+            creation_timestamp="20260817000000",
+            description="Bootstrap migration",
+        )
+    ]
+    mock_runner_cls.return_value = mock_runner_inst
+
+    result = runner.invoke(admin, ["migrate-db", "-y"])
+    assert result.exit_code == 0
+    assert "Found 1 pending schema migration" in result.output
+    assert "Applied migration 20260817000000: Bootstrap migration" in result.output
+    assert "Successfully applied all schema migrations!" in result.output
+    mock_runner_inst.run_migrations.assert_called_once()
+
+
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_user_cancels(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_session.return_value = mock_session_inst
+
+    mock_migration = MagicMock()
+    mock_migration.creation_timestamp = "20260817000000"
+    mock_migration.description = "Bootstrap migration"
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
+    mock_runner_cls.return_value = mock_runner_inst
+
+    result = runner.invoke(admin, ["migrate-db"], input="n\n")
+    assert result.exit_code == 0
+    assert "Found 1 pending schema migration" in result.output
+    assert "Warning: Schema migrations will modify your Spanner database schema" in result.output
+    assert "Migration cancelled." in result.output
+    # No lock or migrations should be run
+    mock_session_inst.post.assert_not_called()
+    mock_runner_inst.run_migrations.assert_not_called()
+
+
+@patch("datacommons_admin.tf_utils.shutil.which")
+@patch("datacommons_admin.tf_utils.subprocess.run")
+@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
+@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
+@patch("datacommons_admin.admin_cli.SpannerClient")
+@patch("datacommons_admin.admin_cli.MigrationRunner")
+def test_migrate_db_default_no_cancels(
+    mock_runner_cls: patch,
+    mock_spanner_cls: patch,
+    mock_auth_default: patch,
+    mock_session: patch,
+    mock_run: patch,
+    mock_which: patch,
+    runner: CliRunner,
+) -> None:
+    mock_which.return_value = "terraform"
+    mock_spanner_cls.return_value = MagicMock()
+
+    mock_proc = MagicMock()
+    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
+    mock_run.return_value = mock_proc
+
+    mock_creds = MagicMock()
+    mock_auth_default.return_value = (mock_creds, "test-project")
+
+    mock_session_inst = MagicMock()
+    mock_session.return_value = mock_session_inst
+
+    mock_migration = MagicMock()
+    mock_migration.creation_timestamp = "20260817000000"
+    mock_migration.description = "Bootstrap migration"
+
+    mock_runner_inst = MagicMock()
+    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
+    mock_runner_cls.return_value = mock_runner_inst
+
+    # Pressing Enter without typing 'y' should default to No and cancel
+    result = runner.invoke(admin, ["migrate-db"], input="\n")
+    assert result.exit_code == 0
+    assert "Found 1 pending schema migration" in result.output
+    assert "Warning: Schema migrations will modify your Spanner database schema" in result.output
+    assert "Migration cancelled." in result.output
+    mock_session_inst.post.assert_not_called()
+    mock_runner_inst.run_migrations.assert_not_called()
+
+
 @patch("datacommons_admin.tf_utils.shutil.which")
 @patch("datacommons_admin.tf_utils.subprocess.run")
 @patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
@@ -782,7 +930,7 @@ def test_migrate_db_failure_releases_lock(
     mock_runner_inst.run_migrations.side_effect = RuntimeError("DDL operation failed")
     mock_runner_cls.return_value = mock_runner_inst
 
-    result = runner.invoke(admin, ["migrate-db"])
+    result = runner.invoke(admin, ["migrate-db"], input="y\n")
     assert result.exit_code != 0
     assert "Failed to apply schema migrations: DDL operation failed" in result.output
 
@@ -835,10 +983,11 @@ def test_migrate_db_lock_busy_error(
     mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
     mock_runner_cls.return_value = mock_runner_inst
 
-    result = runner.invoke(admin, ["migrate-db"])
+    result = runner.invoke(admin, ["migrate-db"], input="y\n")
     assert result.exit_code != 0
     assert "Ingestion Helper returned HTTP 503" in result.output
     assert "Please wait for active ingestions to finish before running migrations" in result.output
 
 
 
+

From 4dbefc4016956030859348669b7c6c50e9612048 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 15:46:48 -0700
Subject: [PATCH 03/11] cleanup diffs and lint

---
 .../datacommons_admin/admin_cli.py            | 47 +++++++++----------
 .../ingestion_helper_client.py                |  1 -
 .../datacommons-admin/tests/test_admin_cli.py | 26 ++++++----
 .../tests/test_ingestion_helper_client.py     |  4 +-
 4 files changed, 38 insertions(+), 40 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index 37178af2..9075e219 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -40,7 +40,6 @@
 from datacommons_db.migrations import MigrationRunner
 
 
-
 DEFAULT_BUCKET_LOCATION = "US"
 GITHUB_RAW_BASE_URL = "https://raw.githubusercontent.com/datacommonsorg/datacommons"
 GITHUB_REPO_URL = "https://github.com/datacommonsorg/datacommons.git"
@@ -538,14 +537,12 @@ def _setup_ingestion_client() -> Tuple[IngestionHelperClient, str, str, str]:
         fg="green",
     )
 
-    ingestion_helper_client = IngestionHelperClient(
-        url, service_account_email=sa_email
-    )
-    return ingestion_helper_client, project_id, instance_id, database_id
+    client = IngestionHelperClient(url, service_account_email=sa_email)
+    return client, project_id, instance_id, database_id
 
 
 def _run_seed_db(
-    ingestion_helper_client: IngestionHelperClient,
+    client: Any,
     instance_id: str,
     database_id: str,
 ) -> None:
@@ -553,7 +550,7 @@ def _run_seed_db(
         f"Seeding Spanner database '{instance_id}/{database_id}' via the Ingestion Helper service (this may take a few moments)...",
         fg="bright_black",
     )
-    result = ingestion_helper_client.seed_database()
+    result = client.seed_database()
     click.secho("Successfully seeded Spanner database!", fg="green", bold=True)
     message = result.get("message")
     if message:
@@ -575,9 +572,7 @@ def _create_migration_runner(
         raise click.ClickException(f"Failed to initialize migration runner: {e}") from e
 
 
-def _apply_migrations(
-    ingestion_helper_client: IngestionHelperClient, runner: MigrationRunner
-) -> None:
+def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
     """Acquires a distributed database lock and applies all pending migrations."""
     # Attempt to acquire Spanner database lock via the Ingestion Helper service.
     click.secho(
@@ -585,7 +580,7 @@ def _apply_migrations(
         fg="bright_black",
     )
     try:
-        ingestion_helper_client.acquire_lock(workflow_id="schema-migration")
+        client.acquire_lock(workflow_id="schema-migration")
     except Exception as e:
         raise click.ClickException(
             f"Could not acquire database lock: {e}\n"
@@ -615,7 +610,7 @@ def _apply_migrations(
             fg="bright_black",
         )
         try:
-            ingestion_helper_client.release_lock(workflow_id="schema-migration")
+            client.release_lock(workflow_id="schema-migration")
         except Exception as e:
             click.secho(
                 f"Warning: Failed to release database lock: {e}",
@@ -637,7 +632,7 @@ def _confirm_migration(num_pending: int, instance_id: str, database_id: str) ->
 
 
 def _run_migrations(
-    ingestion_helper_client: IngestionHelperClient,
+    client: IngestionHelperClient,
     project_id: str,
     instance_id: str,
     database_id: str,
@@ -646,7 +641,7 @@ def _run_migrations(
     """Checks, optionally confirms, and applies pending schema migrations to Spanner.
 
     Args:
-        ingestion_helper_client: IngestionHelperClient instance.
+        client: IngestionHelperClient instance.
         project_id: GCP project ID hosting the Spanner database.
         instance_id: Cloud Spanner instance ID.
         database_id: Cloud Spanner database ID.
@@ -677,12 +672,14 @@ def _run_migrations(
         click.echo(f"  - {m.creation_timestamp}: {m.description}")
 
     # Ask user for confirmation if not auto-approved
-    if not auto_approve and not _confirm_migration(len(pending), instance_id, database_id):
+    if not auto_approve and not _confirm_migration(
+        len(pending), instance_id, database_id
+    ):
         click.secho("Migration cancelled.", fg="yellow")
         return
 
     # Apply migrations
-    _apply_migrations(ingestion_helper_client, runner)
+    _apply_migrations(client, runner)
 
 
 @admin.command(name="migrate-db")
@@ -696,9 +693,9 @@ def _run_migrations(
 def migrate_db(auto_approve: bool) -> None:
     """Apply pending schema migrations to the Spanner database."""
     click.secho("Datacommons Admin Migrate-DB", fg="cyan", bold=True)
-    ingestion_helper_client, project_id, instance_id, database_id = _setup_ingestion_client()
+    client, project_id, instance_id, database_id = _setup_ingestion_client()
     _run_migrations(
-        ingestion_helper_client,
+        client,
         project_id,
         instance_id,
         database_id,
@@ -713,13 +710,13 @@ def migrate_db(auto_approve: bool) -> None:
 def init_db(init_only: bool) -> None:
     """Initialize (and by default seed) the Spanner database via the DCP Ingestion Helper service."""
     click.secho("Datacommons Admin Init-DB", fg="cyan", bold=True)
-    ingestion_helper_client, project_id, instance_id, database_id = _setup_ingestion_client()
+    client, project_id, instance_id, database_id = _setup_ingestion_client()
 
     click.secho(
         f"Initializing Spanner database '{instance_id}/{database_id}' via the Ingestion Helper service (this may take a few moments)...",
         fg="bright_black",
     )
-    result = ingestion_helper_client.initialize_database()
+    result = client.initialize_database()
 
     click.secho("Successfully initialized Spanner database!", fg="green", bold=True)
     message = result.get("message")
@@ -727,7 +724,7 @@ def init_db(init_only: bool) -> None:
         click.secho(f"Details: {message}", fg="bright_black")
 
     _run_migrations(
-        ingestion_helper_client,
+        client,
         project_id,
         instance_id,
         database_id,
@@ -735,19 +732,17 @@ def init_db(init_only: bool) -> None:
     )
 
     if not init_only:
-        _run_seed_db(ingestion_helper_client, instance_id, database_id)
+        _run_seed_db(client, instance_id, database_id)
 
 
 @admin.command(name="seed-db")
 def seed_db() -> None:
     """Seed the Spanner database via the DCP Ingestion Helper service."""
     click.secho("Datacommons Admin Seed-DB", fg="cyan", bold=True)
-    ingestion_helper_client, _project_id, instance_id, database_id = _setup_ingestion_client()
-    _run_seed_db(ingestion_helper_client, instance_id, database_id)
-
+    client, _project_id, instance_id, database_id = _setup_ingestion_client()
+    _run_seed_db(client, instance_id, database_id)
 
 
 from datacommons_admin.ingest_cli import ingest
 
 admin.add_command(ingest)
-
diff --git a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
index d663f102..04cac3a7 100644
--- a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
+++ b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
@@ -150,4 +150,3 @@ def release_lock(self, workflow_id: str) -> dict:
             "workflowId": workflow_id,
         }
         return self._call_endpoint("database/lock/release", payload=payload)
-
diff --git a/packages/datacommons-admin/tests/test_admin_cli.py b/packages/datacommons-admin/tests/test_admin_cli.py
index 4372a25f..63ee81f9 100644
--- a/packages/datacommons-admin/tests/test_admin_cli.py
+++ b/packages/datacommons-admin/tests/test_admin_cli.py
@@ -375,8 +375,6 @@ def test_init_db_migration_failure_halts_before_seed(
     assert "Successfully seeded Spanner database" not in result.output
 
 
-
-
 @patch("datacommons_admin.tf_utils.shutil.which")
 @patch("datacommons_admin.tf_utils.subprocess.run")
 @patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
@@ -724,7 +722,10 @@ def test_migrate_db_with_pending_success(
     result = runner.invoke(admin, ["migrate-db"], input="y\n")
     assert result.exit_code == 0
     assert "Found 1 pending schema migration" in result.output
-    assert "Warning: Schema migrations will modify your Spanner database schema" in result.output
+    assert (
+        "Warning: Schema migrations will modify your Spanner database schema"
+        in result.output
+    )
     assert "Applied migration 20260817000000: Bootstrap migration" in result.output
     assert "Successfully applied all schema migrations!" in result.output
 
@@ -836,7 +837,10 @@ def test_migrate_db_user_cancels(
     result = runner.invoke(admin, ["migrate-db"], input="n\n")
     assert result.exit_code == 0
     assert "Found 1 pending schema migration" in result.output
-    assert "Warning: Schema migrations will modify your Spanner database schema" in result.output
+    assert (
+        "Warning: Schema migrations will modify your Spanner database schema"
+        in result.output
+    )
     assert "Migration cancelled." in result.output
     # No lock or migrations should be run
     mock_session_inst.post.assert_not_called()
@@ -883,7 +887,10 @@ def test_migrate_db_default_no_cancels(
     result = runner.invoke(admin, ["migrate-db"], input="\n")
     assert result.exit_code == 0
     assert "Found 1 pending schema migration" in result.output
-    assert "Warning: Schema migrations will modify your Spanner database schema" in result.output
+    assert (
+        "Warning: Schema migrations will modify your Spanner database schema"
+        in result.output
+    )
     assert "Migration cancelled." in result.output
     mock_session_inst.post.assert_not_called()
     mock_runner_inst.run_migrations.assert_not_called()
@@ -986,8 +993,7 @@ def test_migrate_db_lock_busy_error(
     result = runner.invoke(admin, ["migrate-db"], input="y\n")
     assert result.exit_code != 0
     assert "Ingestion Helper returned HTTP 503" in result.output
-    assert "Please wait for active ingestions to finish before running migrations" in result.output
-
-
-
-
+    assert (
+        "Please wait for active ingestions to finish before running migrations"
+        in result.output
+    )
diff --git a/packages/datacommons-admin/tests/test_ingestion_helper_client.py b/packages/datacommons-admin/tests/test_ingestion_helper_client.py
index 37a27033..bd9daada 100644
--- a/packages/datacommons-admin/tests/test_ingestion_helper_client.py
+++ b/packages/datacommons-admin/tests/test_ingestion_helper_client.py
@@ -223,9 +223,7 @@ def test_unauthorized_error(
     with pytest.raises(click.ClickException) as exc_info:
         client.acquire_lock("schema-migration")
 
-    assert "HTTP 401 Unauthorized when calling Ingestion Helper" in str(
-        exc_info.value
-    )
+    assert "HTTP 401 Unauthorized when calling Ingestion Helper" in str(exc_info.value)
 
 
 @patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")

From d7b9a38e3bf7c2f968c619b12ecb911d2299341f Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 15:53:00 -0700
Subject: [PATCH 04/11] more lint

---
 packages/datacommons-admin/datacommons_admin/admin_cli.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index 9075e219..fd98ff9e 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -541,11 +541,7 @@ def _setup_ingestion_client() -> Tuple[IngestionHelperClient, str, str, str]:
     return client, project_id, instance_id, database_id
 
 
-def _run_seed_db(
-    client: Any,
-    instance_id: str,
-    database_id: str,
-) -> None:
+def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
     click.secho(
         f"Seeding Spanner database '{instance_id}/{database_id}' via the Ingestion Helper service (this may take a few moments)...",
         fg="bright_black",

From 756408226cdbc5e7f727684045b79e27ac8a51f1 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 15:55:10 -0700
Subject: [PATCH 05/11] remove overkill jetski tests

---
 .../tests/test_ingestion_helper_client.py     | 287 ------------------
 1 file changed, 287 deletions(-)
 delete mode 100644 packages/datacommons-admin/tests/test_ingestion_helper_client.py

diff --git a/packages/datacommons-admin/tests/test_ingestion_helper_client.py b/packages/datacommons-admin/tests/test_ingestion_helper_client.py
deleted file mode 100644
index bd9daada..00000000
--- a/packages/datacommons-admin/tests/test_ingestion_helper_client.py
+++ /dev/null
@@ -1,287 +0,0 @@
-# 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.
-
-from unittest.mock import MagicMock, patch
-
-import click
-import pytest
-import requests
-from datacommons_admin.ingestion_helper_client import IngestionHelperClient
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_acquire_lock_default_timeout(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success", "message": "Lock acquired"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    result = client.acquire_lock("schema-migration")
-
-    assert result == {"status": "success", "message": "Lock acquired"}
-    mock_session_inst.post.assert_called_once_with(
-        "https://mock-helper.a.run.app/database/lock/acquire",
-        json={"workflowId": "schema-migration", "timeout": 300},
-        timeout=300,
-    )
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_acquire_lock_custom_args(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success", "message": "Lock acquired"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    result = client.acquire_lock("custom-migration-123", timeout=600)
-
-    assert result == {"status": "success", "message": "Lock acquired"}
-    mock_session_inst.post.assert_called_once_with(
-        "https://mock-helper.a.run.app/database/lock/acquire",
-        json={"workflowId": "custom-migration-123", "timeout": 600},
-        timeout=300,
-    )
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_release_lock_success(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success", "message": "Lock released"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    result = client.release_lock("schema-migration")
-
-    assert result == {"status": "success", "message": "Lock released"}
-    mock_session_inst.post.assert_called_once_with(
-        "https://mock-helper.a.run.app/database/lock/release",
-        json={"workflowId": "schema-migration"},
-        timeout=300,
-    )
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_release_lock_custom_workflow_id(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success", "message": "Lock released"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    result = client.release_lock("custom-migration-123")
-
-    assert result == {"status": "success", "message": "Lock released"}
-    mock_session_inst.post.assert_called_once_with(
-        "https://mock-helper.a.run.app/database/lock/release",
-        json={"workflowId": "custom-migration-123"},
-        timeout=300,
-    )
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_lock_acquire_http_error(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = False
-    mock_resp.status_code = 503
-    mock_resp.json.return_value = {
-        "detail": "Database lock already held by workflow-456"
-    }
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    with pytest.raises(click.ClickException) as exc_info:
-        client.acquire_lock("schema-migration")
-
-    assert "Ingestion Helper returned HTTP 503" in str(exc_info.value)
-    assert "Database lock already held by workflow-456" in str(exc_info.value)
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_unauthorized_error(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = False
-    mock_resp.status_code = 401
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    with pytest.raises(click.ClickException) as exc_info:
-        client.acquire_lock("schema-migration")
-
-    assert "HTTP 401 Unauthorized when calling Ingestion Helper" in str(exc_info.value)
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("google.auth.impersonated_credentials.IDTokenCredentials")
-@patch("google.auth.impersonated_credentials.Credentials")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_network_request_exception(
-    mock_auth_default: patch,
-    mock_imp_creds: patch,
-    mock_id_token_creds: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_imp_creds.return_value = MagicMock()
-    mock_id_token_creds.return_value = MagicMock()
-
-    mock_session_inst = MagicMock()
-    mock_session_inst.post.side_effect = requests.exceptions.ConnectionError(
-        "Connection refused"
-    )
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient(
-        "https://mock-helper.a.run.app", service_account_email="sa@mock.com"
-    )
-    with pytest.raises(click.ClickException) as exc_info:
-        client.release_lock("schema-migration")
-
-    assert "Network or authentication error" in str(exc_info.value)
-
-
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.id_token.fetch_id_token")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-def test_id_token_fetch_when_no_service_account(
-    mock_auth_default: patch,
-    mock_fetch_id_token: patch,
-    mock_session: patch,
-) -> None:
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-    mock_fetch_id_token.return_value = "mock-id-token"
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    client = IngestionHelperClient("https://mock-helper.a.run.app")
-    result = client.initialize_database()
-
-    assert result == {"status": "success"}
-    mock_fetch_id_token.assert_called_once()
-
-
-def test_localhost_bypass_auth() -> None:
-    client = IngestionHelperClient("http://localhost:8080")
-    assert isinstance(client.session, requests.Session)

From 41bbb2540e71f079d2c6ed9f56ac730bd6f80dbb Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 15:57:16 -0700
Subject: [PATCH 06/11] consolidate tests for admin cli

---
 .../datacommons-admin/tests/test_admin_cli.py | 376 ++++--------------
 1 file changed, 68 insertions(+), 308 deletions(-)

diff --git a/packages/datacommons-admin/tests/test_admin_cli.py b/packages/datacommons-admin/tests/test_admin_cli.py
index 63ee81f9..afc2336f 100644
--- a/packages/datacommons-admin/tests/test_admin_cli.py
+++ b/packages/datacommons-admin/tests/test_admin_cli.py
@@ -633,258 +633,83 @@ def test_ingest_start_with_imports_success(
     assert json.loads(called_payload["argument"]) == expected_arg
 
 
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
-def test_migrate_db_no_pending(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
-    runner: CliRunner,
-) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
+@pytest.fixture
+def mock_migration_setup():
+    with (
+        patch("datacommons_admin.admin_cli._setup_ingestion_client") as mock_setup,
+        patch(
+            "datacommons_admin.admin_cli._create_migration_runner"
+        ) as mock_runner_factory,
+    ):
+        mock_client = MagicMock()
+        mock_setup.return_value = (
+            mock_client,
+            "mock-proj",
+            "mock-instance",
+            "mock-db",
+        )
 
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
+        mock_runner = MagicMock()
+        mock_runner_factory.return_value = mock_runner
 
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
+        yield mock_client, mock_runner
 
-    mock_session_inst = MagicMock()
-    mock_session.return_value = mock_session_inst
 
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = []
-    mock_runner_cls.return_value = mock_runner_inst
+def test_migrate_db_no_pending(
+    mock_migration_setup: tuple[MagicMock, MagicMock], runner: CliRunner
+) -> None:
+    mock_client, mock_runner = mock_migration_setup
+    mock_runner.get_pending_migrations.return_value = []
 
     result = runner.invoke(admin, ["migrate-db"])
     assert result.exit_code == 0
     assert "Database schema is already up-to-date" in result.output
-    # Lock should not be acquired when there are no pending migrations
-    mock_session_inst.post.assert_not_called()
+    mock_client.acquire_lock.assert_not_called()
 
 
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
-def test_migrate_db_with_pending_success(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
+@pytest.mark.parametrize(("args", "input_str"), [(["-y"], None), ([], "y\n")])
+def test_migrate_db_apply_success(
+    mock_migration_setup: tuple[MagicMock, MagicMock],
     runner: CliRunner,
+    args: list[str],
+    input_str: str | None,
 ) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
-
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
-
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    mock_migration = MagicMock()
-    mock_migration.creation_timestamp = "20260817000000"
-    mock_migration.description = "Bootstrap migration"
-
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
-    mock_runner_inst.run_migrations.return_value = [
-        MigrationResult(
-            status=ExecutionStatus.SUCCESS,
-            creation_timestamp="20260817000000",
-            description="Bootstrap migration",
-        )
-    ]
-    mock_runner_cls.return_value = mock_runner_inst
-
-    result = runner.invoke(admin, ["migrate-db"], input="y\n")
-    assert result.exit_code == 0
-    assert "Found 1 pending schema migration" in result.output
-    assert (
-        "Warning: Schema migrations will modify your Spanner database schema"
-        in result.output
-    )
-    assert "Applied migration 20260817000000: Bootstrap migration" in result.output
-    assert "Successfully applied all schema migrations!" in result.output
-
-    # Check lock acquired then released
-    assert mock_session_inst.post.call_count == 2
-    mock_session_inst.post.assert_any_call(
-        "https://mock-helper/database/lock/acquire",
-        json={"workflowId": "schema-migration", "timeout": 300},
-        timeout=300,
-    )
-    mock_session_inst.post.assert_any_call(
-        "https://mock-helper/database/lock/release",
-        json={"workflowId": "schema-migration"},
-        timeout=300,
+    mock_client, mock_runner = mock_migration_setup
+    mock_migration = MagicMock(
+        creation_timestamp="20260817000000", description="Bootstrap migration"
     )
-
-
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
-def test_migrate_db_with_yes_flag(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
-    runner: CliRunner,
-) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
-
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
-
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    mock_migration = MagicMock()
-    mock_migration.creation_timestamp = "20260817000000"
-    mock_migration.description = "Bootstrap migration"
-
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
-    mock_runner_inst.run_migrations.return_value = [
+    mock_runner.get_pending_migrations.return_value = [mock_migration]
+    mock_runner.run_migrations.return_value = [
         MigrationResult(
             status=ExecutionStatus.SUCCESS,
             creation_timestamp="20260817000000",
             description="Bootstrap migration",
         )
     ]
-    mock_runner_cls.return_value = mock_runner_inst
 
-    result = runner.invoke(admin, ["migrate-db", "-y"])
+    result = runner.invoke(admin, ["migrate-db", *args], input=input_str)
     assert result.exit_code == 0
     assert "Found 1 pending schema migration" in result.output
     assert "Applied migration 20260817000000: Bootstrap migration" in result.output
     assert "Successfully applied all schema migrations!" in result.output
-    mock_runner_inst.run_migrations.assert_called_once()
+    mock_client.acquire_lock.assert_called_once_with(workflow_id="schema-migration")
+    mock_runner.run_migrations.assert_called_once()
+    mock_client.release_lock.assert_called_once_with(workflow_id="schema-migration")
 
 
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
+@pytest.mark.parametrize("input_str", ["n\n", "\n"])
 def test_migrate_db_user_cancels(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
+    mock_migration_setup: tuple[MagicMock, MagicMock],
     runner: CliRunner,
+    input_str: str,
 ) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
-
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
-
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-
-    mock_session_inst = MagicMock()
-    mock_session.return_value = mock_session_inst
-
-    mock_migration = MagicMock()
-    mock_migration.creation_timestamp = "20260817000000"
-    mock_migration.description = "Bootstrap migration"
-
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
-    mock_runner_cls.return_value = mock_runner_inst
-
-    result = runner.invoke(admin, ["migrate-db"], input="n\n")
-    assert result.exit_code == 0
-    assert "Found 1 pending schema migration" in result.output
-    assert (
-        "Warning: Schema migrations will modify your Spanner database schema"
-        in result.output
+    mock_client, mock_runner = mock_migration_setup
+    mock_migration = MagicMock(
+        creation_timestamp="20260817000000", description="Bootstrap migration"
     )
-    assert "Migration cancelled." in result.output
-    # No lock or migrations should be run
-    mock_session_inst.post.assert_not_called()
-    mock_runner_inst.run_migrations.assert_not_called()
+    mock_runner.get_pending_migrations.return_value = [mock_migration]
 
-
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
-def test_migrate_db_default_no_cancels(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
-    runner: CliRunner,
-) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
-
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
-
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-
-    mock_session_inst = MagicMock()
-    mock_session.return_value = mock_session_inst
-
-    mock_migration = MagicMock()
-    mock_migration.creation_timestamp = "20260817000000"
-    mock_migration.description = "Bootstrap migration"
-
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
-    mock_runner_cls.return_value = mock_runner_inst
-
-    # Pressing Enter without typing 'y' should default to No and cancel
-    result = runner.invoke(admin, ["migrate-db"], input="\n")
+    result = runner.invoke(admin, ["migrate-db"], input=input_str)
     assert result.exit_code == 0
     assert "Found 1 pending schema migration" in result.output
     assert (
@@ -892,105 +717,40 @@ def test_migrate_db_default_no_cancels(
         in result.output
     )
     assert "Migration cancelled." in result.output
-    mock_session_inst.post.assert_not_called()
-    mock_runner_inst.run_migrations.assert_not_called()
+    mock_client.acquire_lock.assert_not_called()
+    mock_runner.run_migrations.assert_not_called()
 
 
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
 def test_migrate_db_failure_releases_lock(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
-    runner: CliRunner,
+    mock_migration_setup: tuple[MagicMock, MagicMock], runner: CliRunner
 ) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
-
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
-
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = True
-    mock_resp.json.return_value = {"status": "success"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    mock_migration = MagicMock()
-    mock_migration.creation_timestamp = "20260817000000"
-    mock_migration.description = "Bootstrap migration"
-
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
-    mock_runner_inst.run_migrations.side_effect = RuntimeError("DDL operation failed")
-    mock_runner_cls.return_value = mock_runner_inst
+    mock_client, mock_runner = mock_migration_setup
+    mock_migration = MagicMock(
+        creation_timestamp="20260817000000", description="Bootstrap migration"
+    )
+    mock_runner.get_pending_migrations.return_value = [mock_migration]
+    mock_runner.run_migrations.side_effect = RuntimeError("DDL operation failed")
 
-    result = runner.invoke(admin, ["migrate-db"], input="y\n")
+    result = runner.invoke(admin, ["migrate-db", "-y"])
     assert result.exit_code != 0
     assert "Failed to apply schema migrations: DDL operation failed" in result.output
-
-    # Ensure release_lock was still called in finally block
-    mock_session_inst.post.assert_any_call(
-        "https://mock-helper/database/lock/release",
-        json={"workflowId": "schema-migration"},
-        timeout=300,
-    )
+    mock_client.acquire_lock.assert_called_once_with(workflow_id="schema-migration")
+    mock_client.release_lock.assert_called_once_with(workflow_id="schema-migration")
 
 
-@patch("datacommons_admin.tf_utils.shutil.which")
-@patch("datacommons_admin.tf_utils.subprocess.run")
-@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession")
-@patch("datacommons_admin.ingestion_helper_client.google.auth.default")
-@patch("datacommons_admin.admin_cli.SpannerClient")
-@patch("datacommons_admin.admin_cli.MigrationRunner")
 def test_migrate_db_lock_busy_error(
-    mock_runner_cls: patch,
-    mock_spanner_cls: patch,
-    mock_auth_default: patch,
-    mock_session: patch,
-    mock_run: patch,
-    mock_which: patch,
-    runner: CliRunner,
+    mock_migration_setup: tuple[MagicMock, MagicMock], runner: CliRunner
 ) -> None:
-    mock_which.return_value = "terraform"
-    mock_spanner_cls.return_value = MagicMock()
-
-    mock_proc = MagicMock()
-    mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "spanner_instance_id": {"value": "mock-instance"}, "spanner_database_id": {"value": "mock-db"}, "project_id": {"value": "mock-proj"}}'
-    mock_run.return_value = mock_proc
-
-    mock_creds = MagicMock()
-    mock_auth_default.return_value = (mock_creds, "test-project")
-
-    mock_session_inst = MagicMock()
-    mock_resp = MagicMock()
-    mock_resp.ok = False
-    mock_resp.status_code = 503
-    mock_resp.json.return_value = {"detail": "Lock busy"}
-    mock_session_inst.post.return_value = mock_resp
-    mock_session.return_value = mock_session_inst
-
-    mock_migration = MagicMock()
-    mock_migration.creation_timestamp = "20260817000000"
-    mock_migration.description = "Bootstrap migration"
-
-    mock_runner_inst = MagicMock()
-    mock_runner_inst.get_pending_migrations.return_value = [mock_migration]
-    mock_runner_cls.return_value = mock_runner_inst
+    mock_client, mock_runner = mock_migration_setup
+    mock_migration = MagicMock(
+        creation_timestamp="20260817000000", description="Bootstrap migration"
+    )
+    mock_runner.get_pending_migrations.return_value = [mock_migration]
+    mock_client.acquire_lock.side_effect = click.ClickException(
+        "Ingestion Helper returned HTTP 503"
+    )
 
-    result = runner.invoke(admin, ["migrate-db"], input="y\n")
+    result = runner.invoke(admin, ["migrate-db", "-y"])
     assert result.exit_code != 0
     assert "Ingestion Helper returned HTTP 503" in result.output
     assert (

From a2c9a53721adf85180559f9ef8ca9f8e067f3c32 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 16:16:45 -0700
Subject: [PATCH 07/11] disable built-in metrics in spanner client to avoid
 messy CLI output

---
 .../datacommons_db/clients/spanner_client.py             | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/packages/datacommons-db/datacommons_db/clients/spanner_client.py b/packages/datacommons-db/datacommons_db/clients/spanner_client.py
index 7ece9ffd..eeccdb0e 100644
--- a/packages/datacommons-db/datacommons_db/clients/spanner_client.py
+++ b/packages/datacommons-db/datacommons_db/clients/spanner_client.py
@@ -86,6 +86,8 @@ def __init__(
         instance_id: str,
         database_id: str,
         credentials: Credentials | None = None,
+        *,
+        disable_builtin_metrics: bool = True,
     ) -> None:
         """Initialize the SpannerClient.
 
@@ -94,6 +96,7 @@ def __init__(
             instance_id: Cloud Spanner instance ID.
             database_id: Cloud Spanner database ID.
             credentials: Optional Google Cloud credentials object.
+            disable_builtin_metrics: Whether to disable built-in Cloud Monitoring metrics export.
         """
         validate_resource_id("project_id", project_id)
         validate_resource_id("instance_id", instance_id)
@@ -103,7 +106,11 @@ def __init__(
         self.instance_id = instance_id
         self.database_id = database_id
 
-        self.client = spanner.Client(project=project_id, credentials=credentials)
+        self.client = spanner.Client(
+            project=project_id,
+            credentials=credentials,
+            disable_builtin_metrics=disable_builtin_metrics,
+        )
         self.instance = self.client.instance(self.instance_id)
         self.database = self.instance.database(self.database_id)
 

From 26a360e7b6af45aa9305400430eadb0fcd738773 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Tue, 18 Aug 2026 16:37:09 -0700
Subject: [PATCH 08/11] cleanup lock logic in apply_migrations

---
 .../datacommons_admin/admin_cli.py            | 53 ++++++++++---------
 .../tests/test_spanner_client.py              |  5 +-
 2 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index fd98ff9e..e8743812 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -570,22 +570,24 @@ def _create_migration_runner(
 
 def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
     """Acquires a distributed database lock and applies all pending migrations."""
-    # Attempt to acquire Spanner database lock via the Ingestion Helper service.
-    click.secho(
-        "Acquiring database lock via the Ingestion Helper service...",
-        fg="bright_black",
-    )
+    lock_acquired = False
     try:
-        client.acquire_lock(workflow_id="schema-migration")
-    except Exception as e:
-        raise click.ClickException(
-            f"Could not acquire database lock: {e}\n"
-            "An ingestion workflow may currently be running. "
-            "Please wait for active ingestions to finish before running migrations."
-        ) from e
+        # Attempt to acquire Spanner database lock via the Ingestion Helper service.
+        click.secho(
+            "Acquiring database lock via the Ingestion Helper service...",
+            fg="bright_black",
+        )
+        try:
+            client.acquire_lock(workflow_id="schema-migration")
+            lock_acquired = True
+        except Exception as e:
+            raise click.ClickException(
+                f"Could not acquire database lock: {e}\n"
+                "An ingestion workflow may currently be running. "
+                "Please wait for active ingestions to finish before running migrations."
+            ) from e
 
-    # Apply all pending migrations
-    try:
+        # Apply all pending migrations
         click.secho("Applying pending schema migrations...", fg="bright_black")
         results = runner.run_migrations()
         for res in results:
@@ -596,22 +598,23 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
         click.secho(
             "Successfully applied all schema migrations!", fg="green", bold=True
         )
+    except click.ClickException:
+        raise
     except Exception as e:
         raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
-
-    # Release database lock
     finally:
-        click.secho(
-            "Releasing database lock via the Ingestion Helper service...",
-            fg="bright_black",
-        )
-        try:
-            client.release_lock(workflow_id="schema-migration")
-        except Exception as e:
+        if lock_acquired:
             click.secho(
-                f"Warning: Failed to release database lock: {e}",
-                fg="yellow",
+                "Releasing database lock via the Ingestion Helper service...",
+                fg="bright_black",
             )
+            try:
+                client.release_lock(workflow_id="schema-migration")
+            except Exception as e:
+                click.secho(
+                    f"Warning: Failed to release database lock: {e}",
+                    fg="yellow",
+                )
 
 
 def _confirm_migration(num_pending: int, instance_id: str, database_id: str) -> bool:
diff --git a/packages/datacommons-db/tests/test_spanner_client.py b/packages/datacommons-db/tests/test_spanner_client.py
index 0cab04e5..f71cbdc3 100644
--- a/packages/datacommons-db/tests/test_spanner_client.py
+++ b/packages/datacommons-db/tests/test_spanner_client.py
@@ -194,8 +194,11 @@ def fake_spanner_db():
         "datacommons_db.clients.spanner_client.spanner.Client"
     ) as mock_client_cls:
 
-        def fake_client_factory(project: str, credentials: object = None) -> MagicMock:
+        def fake_client_factory(
+            project: str, credentials: object = None, **kwargs: object
+        ) -> MagicMock:
             _ = credentials
+            _ = kwargs
             mock_client = MagicMock()
             mock_client.project = project
             mock_instance = MagicMock()

From 02e63e3f12e6d76180028505d5c541b3f08d68e2 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Wed, 19 Aug 2026 10:16:00 -0700
Subject: [PATCH 09/11] clean up try except logic

---
 .../datacommons-admin/datacommons_admin/admin_cli.py   | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index e8743812..e53e571a 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -589,7 +589,11 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
 
         # Apply all pending migrations
         click.secho("Applying pending schema migrations...", fg="bright_black")
-        results = runner.run_migrations()
+        try:
+            results = runner.run_migrations()
+        except Exception as e:
+            raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
+
         for res in results:
             click.secho(
                 f"  ✔ Applied migration {res.creation_timestamp}: {res.description}",
@@ -598,10 +602,6 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
         click.secho(
             "Successfully applied all schema migrations!", fg="green", bold=True
         )
-    except click.ClickException:
-        raise
-    except Exception as e:
-        raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
     finally:
         if lock_acquired:
             click.secho(

From 4ac134d30be7dbf8d14d6ae3c958fe7d591fe9d9 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Wed, 19 Aug 2026 10:47:43 -0700
Subject: [PATCH 10/11] add bool output and make docstrings explicit

---
 .../datacommons_admin/admin_cli.py            | 71 ++++++++++++++++---
 1 file changed, 60 insertions(+), 11 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index e53e571a..7e82bec4 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -556,7 +556,19 @@ def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
 def _create_migration_runner(
     project_id: str, instance_id: str, database_id: str
 ) -> MigrationRunner:
-    """Initializes a SpannerClient and returns a MigrationRunner instance."""
+    """Initializes a SpannerClient and returns a MigrationRunner instance.
+
+    Args:
+        project_id: GCP project ID hosting the Spanner database.
+        instance_id: Cloud Spanner instance ID.
+        database_id: Cloud Spanner database ID.
+
+    Returns:
+        A MigrationRunner instance initialized with a SpannerClient.
+
+    Raises:
+        click.ClickException: If initialization of the SpannerClient or MigrationRunner fails.
+    """
     try:
         spanner_client = SpannerClient(
             project_id=project_id,
@@ -568,8 +580,19 @@ def _create_migration_runner(
         raise click.ClickException(f"Failed to initialize migration runner: {e}") from e
 
 
-def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
-    """Acquires a distributed database lock and applies all pending migrations."""
+def _apply_migrations(client: Any, runner: MigrationRunner) -> bool:
+    """Acquires a distributed database lock and applies all pending migrations.
+
+    Args:
+        client: IngestionHelperClient instance used for database lock management.
+        runner: MigrationRunner instance used to execute schema migrations.
+
+    Returns:
+        True if all migrations were successfully applied.
+
+    Raises:
+        click.ClickException: If acquiring the database lock or applying migrations fails.
+    """
     lock_acquired = False
     try:
         # Attempt to acquire Spanner database lock via the Ingestion Helper service.
@@ -602,6 +625,7 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
         click.secho(
             "Successfully applied all schema migrations!", fg="green", bold=True
         )
+        return True
     finally:
         if lock_acquired:
             click.secho(
@@ -618,7 +642,16 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> None:
 
 
 def _confirm_migration(num_pending: int, instance_id: str, database_id: str) -> bool:
-    """Displays a safety warning and prompts the user to confirm applying migrations."""
+    """Displays a safety warning and prompts the user to confirm applying migrations.
+
+    Args:
+        num_pending: Number of pending schema migrations.
+        instance_id: Cloud Spanner instance ID.
+        database_id: Cloud Spanner database ID.
+
+    Returns:
+        True if the user confirms the migration prompt, False otherwise.
+    """
     click.secho(
         "\nWarning: Schema migrations will modify your Spanner database schema. "
         "It is strongly recommended to create a database backup before proceeding in production environments.",
@@ -636,7 +669,7 @@ def _run_migrations(
     instance_id: str,
     database_id: str,
     auto_approve: bool = False,
-) -> None:
+) -> bool:
     """Checks, optionally confirms, and applies pending schema migrations to Spanner.
 
     Args:
@@ -645,6 +678,12 @@ def _run_migrations(
         instance_id: Cloud Spanner instance ID.
         database_id: Cloud Spanner database ID.
         auto_approve: If False, prompts user for interactive confirmation before applying.
+
+    Returns:
+        True if migrations were applied or database is already up-to-date, False if cancelled by the user.
+
+    Raises:
+        click.ClickException: If checking pending migrations, acquiring the database lock, or applying migrations fails.
     """
     click.secho(
         f"Checking schema migrations for Spanner database '{project_id}/{instance_id}/{database_id}'...",
@@ -664,7 +703,7 @@ def _run_migrations(
             "Database schema is already up-to-date. No migrations to apply.",
             fg="green",
         )
-        return
+        return True
 
     click.secho(f"Found {len(pending)} pending schema migration(s):", fg="cyan")
     for m in pending:
@@ -675,10 +714,10 @@ def _run_migrations(
         len(pending), instance_id, database_id
     ):
         click.secho("Migration cancelled.", fg="yellow")
-        return
+        return False
 
     # Apply migrations
-    _apply_migrations(client, runner)
+    return _apply_migrations(client, runner)
 
 
 @admin.command(name="migrate-db")
@@ -689,11 +728,21 @@ def _run_migrations(
     is_flag=True,
     help="Automatically confirm and apply pending migrations without prompting.",
 )
-def migrate_db(auto_approve: bool) -> None:
-    """Apply pending schema migrations to the Spanner database."""
+def migrate_db(auto_approve: bool) -> bool:
+    """Apply pending schema migrations to the Spanner database.
+
+    Args:
+        auto_approve: If True, automatically confirms and applies pending migrations without prompting.
+
+    Returns:
+        True if migrations were applied or database is already up-to-date, False if cancelled by the user.
+
+    Raises:
+        click.ClickException: If reading Terraform outputs, checking pending migrations, acquiring lock, or applying migrations fails.
+    """
     click.secho("Datacommons Admin Migrate-DB", fg="cyan", bold=True)
     client, project_id, instance_id, database_id = _setup_ingestion_client()
-    _run_migrations(
+    return _run_migrations(
         client,
         project_id,
         instance_id,

From 0d40b910a54baab37ae941370c9a81020f237799 Mon Sep 17 00:00:00 2001
From: Julia Wu <juliawu@google.com>
Date: Wed, 19 Aug 2026 11:23:01 -0700
Subject: [PATCH 11/11] move try except into acquire and release

---
 .../datacommons_admin/admin_cli.py            | 50 ++++++++-----------
 .../ingestion_helper_client.py                | 22 +++++++-
 .../datacommons-admin/tests/test_admin_cli.py |  6 ++-
 3 files changed, 45 insertions(+), 33 deletions(-)

diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py
index 7e82bec4..1cb2471d 100644
--- a/packages/datacommons-admin/datacommons_admin/admin_cli.py
+++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py
@@ -593,29 +593,17 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> bool:
     Raises:
         click.ClickException: If acquiring the database lock or applying migrations fails.
     """
-    lock_acquired = False
-    try:
-        # Attempt to acquire Spanner database lock via the Ingestion Helper service.
-        click.secho(
-            "Acquiring database lock via the Ingestion Helper service...",
-            fg="bright_black",
-        )
-        try:
-            client.acquire_lock(workflow_id="schema-migration")
-            lock_acquired = True
-        except Exception as e:
-            raise click.ClickException(
-                f"Could not acquire database lock: {e}\n"
-                "An ingestion workflow may currently be running. "
-                "Please wait for active ingestions to finish before running migrations."
-            ) from e
+    # Attempt to acquire Spanner database lock via the Ingestion Helper service.
+    click.secho(
+        "Acquiring database lock via the Ingestion Helper service...",
+        fg="bright_black",
+    )
+    client.acquire_lock(workflow_id="schema-migration")
 
+    try:
         # Apply all pending migrations
         click.secho("Applying pending schema migrations...", fg="bright_black")
-        try:
-            results = runner.run_migrations()
-        except Exception as e:
-            raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
+        results = runner.run_migrations()
 
         for res in results:
             click.secho(
@@ -626,19 +614,21 @@ def _apply_migrations(client: Any, runner: MigrationRunner) -> bool:
             "Successfully applied all schema migrations!", fg="green", bold=True
         )
         return True
+    except Exception as e:
+        raise click.ClickException(f"Failed to apply schema migrations: {e}") from e
     finally:
-        if lock_acquired:
+        # Always attempt to release the database lock after migration attempt
+        click.secho(
+            "Releasing database lock via the Ingestion Helper service...",
+            fg="bright_black",
+        )
+        try:
+            client.release_lock(workflow_id="schema-migration")
+        except Exception as e:
             click.secho(
-                "Releasing database lock via the Ingestion Helper service...",
-                fg="bright_black",
+                f"Warning: {e}",
+                fg="yellow",
             )
-            try:
-                client.release_lock(workflow_id="schema-migration")
-            except Exception as e:
-                click.secho(
-                    f"Warning: Failed to release database lock: {e}",
-                    fg="yellow",
-                )
 
 
 def _confirm_migration(num_pending: int, instance_id: str, database_id: str) -> bool:
diff --git a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
index 04cac3a7..d1514c53 100644
--- a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
+++ b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py
@@ -130,12 +130,22 @@ def acquire_lock(self, workflow_id: str, timeout: int = 300) -> dict:
 
         Returns:
             API response dictionary from the Ingestion Helper service.
+
+        Raises:
+            click.ClickException: If acquiring the database lock fails.
         """
         payload = {
             "workflowId": workflow_id,
             "timeout": timeout,
         }
-        return self._call_endpoint("database/lock/acquire", payload=payload)
+        try:
+            return self._call_endpoint("database/lock/acquire", payload=payload)
+        except click.ClickException as e:
+            raise click.ClickException(
+                f"Could not acquire database lock: {e.format_message()}\n"
+                "An ingestion workflow may currently be running. "
+                "Please wait for active ingestions to finish before running migrations."
+            ) from e
 
     def release_lock(self, workflow_id: str) -> dict:
         """Releases the distributed database lock via the ingestion helper service.
@@ -145,8 +155,16 @@ def release_lock(self, workflow_id: str) -> dict:
 
         Returns:
             API response dictionary from the Ingestion Helper service.
+
+        Raises:
+            click.ClickException: If releasing the database lock fails.
         """
         payload = {
             "workflowId": workflow_id,
         }
-        return self._call_endpoint("database/lock/release", payload=payload)
+        try:
+            return self._call_endpoint("database/lock/release", payload=payload)
+        except click.ClickException as e:
+            raise click.ClickException(
+                f"Could not release database lock: {e.format_message()}"
+            ) from e
diff --git a/packages/datacommons-admin/tests/test_admin_cli.py b/packages/datacommons-admin/tests/test_admin_cli.py
index afc2336f..b1ee023a 100644
--- a/packages/datacommons-admin/tests/test_admin_cli.py
+++ b/packages/datacommons-admin/tests/test_admin_cli.py
@@ -19,6 +19,7 @@
 import pytest
 from click.testing import CliRunner
 from datacommons_admin.admin_cli import admin
+from datacommons_admin.ingestion_helper_client import IngestionHelperClient
 from datacommons_db.clients.spanner_client import ExecutionStatus
 from datacommons_db.migrations.migration_runner import MigrationResult
 
@@ -747,7 +748,9 @@ def test_migrate_db_lock_busy_error(
     )
     mock_runner.get_pending_migrations.return_value = [mock_migration]
     mock_client.acquire_lock.side_effect = click.ClickException(
-        "Ingestion Helper returned HTTP 503"
+        "Could not acquire database lock: Ingestion Helper returned HTTP 503\n"
+        "An ingestion workflow may currently be running. "
+        "Please wait for active ingestions to finish before running migrations."
     )
 
     result = runner.invoke(admin, ["migrate-db", "-y"])
@@ -757,3 +760,4 @@ def test_migrate_db_lock_busy_error(
         "Please wait for active ingestions to finish before running migrations"
         in result.output
     )
+    mock_client.release_lock.assert_not_called()