feat(datacommons-admin) Add CLI for schema migrations

packages/datacommons-admin/datacommons_admin/admin_cli.py

@@ -28,6 +28,16 @@
     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 +518,27 @@ def init(
     )
 
 
-def _setup_ingestion_client() -> Tuple[Any, 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",
     )
 
-    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,14 +553,160 @@ def _run_seed_db(client: Any, instance_id: str, database_id: str) -> None:
         click.secho(f"Details: {message}", 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,
+        )
+        return MigrationRunner(spanner_client=spanner_client)
+    except Exception as e:
+        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."""
+    # 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")
+    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:
+        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
+
+    # 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:
+            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:
+    """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(
+    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:
+        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(client, runner)
+
+
+@admin.command(name="migrate-db")
+@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,
+        auto_approve=auto_approve,
+    )
+
+
 @admin.command(name="init-db")
 @click.option(
     "--init-only", is_flag=True, help="Only initialize the database without seeding."
 )
 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 +719,14 @@ 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,
+        auto_approve=True,
+    )
+
     if not init_only:
         _run_seed_db(client, instance_id, database_id)
 
@@ -578,7 +735,7 @@ 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)
 
 

packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py

@@ -120,3 +120,33 @@ 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)

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"