feat(datacommons-admin) Add CLI for schema migrations - #219
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces database schema migration capabilities to the datacommons-admin CLI by adding a new migrate-db command and integrating automatic migrations into the init-db command. It utilizes a distributed lock mechanism via the ingestion helper service to prevent concurrent migrations. The review feedback highlights a potential lock leak vulnerability in _apply_migrations if an exception occurs during lock acquisition, and recommends supporting service account impersonation for SpannerClient by passing the service account email to the migration runner.
| 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 |
There was a problem hiding this comment.
In environments where the user or runner relies on service account impersonation to access Spanner, initializing SpannerClient with default credentials will result in permission denied errors. To ensure Spanner operations run with the correct permissions, we should construct and pass impersonated credentials to SpannerClient when service_account_email is provided.
def _create_migration_runner(
project_id: str,
instance_id: str,
database_id: str,
service_account_email: str | None = None,
) -> MigrationRunner:
"""Initializes a SpannerClient and returns a MigrationRunner instance."""
try:
credentials = None
if service_account_email:
import google.auth
from google.auth import impersonated_credentials
base_credentials, _ = google.auth.default()
credentials = impersonated_credentials.Credentials(
source_credentials=base_credentials,
target_principal=service_account_email,
target_scopes=["https://www.googleapis.com/auth/cloud-platform"],
)
spanner_client = SpannerClient(
project_id=project_id,
instance_id=instance_id,
database_id=database_id,
credentials=credentials,
)
return MigrationRunner(spanner_client=spanner_client)
except Exception as e:
raise click.ClickException(f"Failed to initialize migration runner: {e}") from eThere was a problem hiding this comment.
I don't think this is necessary for now.
There was a problem hiding this comment.
can you remind me how's the auth for spanner writes is granted? It's through the CLI service account?
There was a problem hiding this comment.
Right now it's using ADC credentials acquired via gcloud auth application-default login.
Would we ever expect users to run migrations in a CI/CD automation and not manually by CLI? Given that it's performing potentially irreversible operations on the database, I assumed we'd be requiring users to migrate manually via CLI, so I haven't configured any service accounts.
| 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 |
There was a problem hiding this comment.
can you remind me how's the auth for spanner writes is granted? It's through the CLI service account?
| """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( |
There was a problem hiding this comment.
So I think migrate_db belongs in admin_cli, but nearly all things above (essentially _run_migrations and all other helpers) need to move out.
Admin_cli should be an orchestrator file, it's already racing to 1,000 lines.
The ingestion_helper_client I think sets the right pattern here. Perhaps we can have a migration_runner_client, or something similar.
In datacommosn_db, I believe you created a new subdir for it. How about reusing the same pattern in datacommons_admin?
| assert json.loads(called_payload["argument"]) == expected_arg | ||
|
|
||
|
|
||
| @pytest.fixture |
There was a problem hiding this comment.
A bit of a similar comment as above, I think organizing the test file as well will be helpful. It's growing very quickly. Curious to hear your thoughts on how we can improve test file structure a bit, and lay the foundation for future proliferation of tests, while keeping this organized :)
2 participants
Overview
This PR adds the
datacommons admin migrate-dbcommand and integrates automated schema migrations intodatacommons admin init-db.With this change, DCP admins can update their databases to the latest schema by running:
Any databases intialized via
init-dbwill automatically have all available migrations applied.Key Changes
Added Locking in IngestionHelperClient (ingestion_helper_client.py):
acquire_lock(): sends POST /database/lock/acquire.release_lock(): sends POST /database/lock/release.Admin CLI Migration Integration (admin_cli.py):
datacommons admin migrate-db:datacommons admin init-db:Added datacommons-db dependency to datacommons-admin.
Testing Strategy
Added new unit tests to
tests/test_admin_cli.pyManually verified workflow by creating a new DCP instance to test the init-db part, including:
init-dbcall runs migrations automaticallyinit-dbresults in no-op from migration standpointManually verified workflow by testing the migrate-db on an existing DCP instance, including:
migrate-dbon an existing instance does prompt before proceedingmigrate-dbon an up-to-date instance results in no-op from migration standpoint