feat: adopt Alembic for database schema migrations #4408
Conversation
Summary of ChangesHello @achilleasatha, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request modernizes the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Ignored Files
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
Adopt **Alembic** as the migration framework for `DatabaseSessionService`, replacing manual `ALTER TABLE` operations. This enables automatic schema tracking, upgrade/downgrade support, and Kubernetes Helm hook integration for production deployments. Addresses google#3343 --- - **AlembicMigrationRunner** - Wraps Alembic’s command API with programmatic config (instead of `alembic.ini`) - Ships migrations inside the package - **V1 baseline migration** - Creates all current schema tables - **V0 → V1 in-place bootstrap** - Migrates pickle-based databases to JSON - No separate destination database required - `env.py` and `script.py.mako` - Support both online and offline migration modes --- - `upgrade` — Apply pending migrations with auto-bootstrap for existing DBs - `downgrade` — Roll back by revision spec - `check` — Exit `0` if up-to-date, `1` if pending (CI/CD gate) - `stamp` — Bootstrap Alembic tracking for existing databases - `generate` — Create new migration scripts (autogenerate or template) --- - `_ensure_alembic_tracking()` stamps baseline after table creation - V0 databases auto-migrate to V1 on first access - `ADK_AUTO_MIGRATE_DB=true` - Runs pending migrations on startup - `ADK_AUTO_MIGRATE_DB=false` (default) - Raises `RuntimeError` if schema is behind --- - `get_alembic_revision()` - `is_alembic_managed()` - `is_alembic_managed_from_url()` --- - unit tests for `AlembicMigrationRunner` (SQLite) - CLI command tests covering all subcommands and edge cases - integration tests across SQLite / PostgreSQL / MySQL - Updated `test_database_schema.py` for V0 auto-migration behavior - GHA workflow: - PostgreSQL 15–17 - MySQL 8.0 / 8.4 / 9.2 - Python 3.10–3.14 - `docker-compose.yml` for local integration testing / development --- - `migration_guide.md` — User-facing upgrade and CLI reference - `helm_migration_guide.md` — Kubernetes Job and Cloud SQL Proxy examples - `migration/README.md` — Contributor workflow for adding schema versions
adk-bot
added
the
services
achilleasatha
force-pushed
the
feat/issue-3343-adopt-alembic-for-database-migrations
branch
from
530c7ee to
6066b3f
Compare
![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 is an excellent and comprehensive pull request that introduces Alembic for database schema migrations. The implementation is robust, with a new AlembicMigrationRunner to handle migrations programmatically, new CLI commands for managing migrations, and automatic migration capabilities within DatabaseSessionService. The addition of extensive unit and integration tests across multiple database backends is particularly impressive and ensures the reliability of this new system. The documentation updates are also very clear and helpful for both users and contributors. I have a few minor suggestions for the documentation and one security note regarding the use of pickle during the V0 to V1 data migration.
| if actions_raw is not None: | ||
| try: | ||
| if isinstance(actions_raw, bytes): | ||
| obj = pickle.loads(actions_raw) # noqa: S301 |
There was a problem hiding this comment.
Using pickle.loads() on data read from the database can introduce a security vulnerability (Insecure Deserialization). If an attacker can write a malicious pickle object to the database, they could achieve remote code execution when this migration script is run.
While I understand this is necessary for migrating away from the insecure pickle format, it's important to acknowledge and mitigate the risk. Please ensure that access to the database is strictly controlled, especially during the migration process.
There was a problem hiding this comment.
Hm not sure what gemini is trying to get out of this. It already acknowledges this is a requirement so we can move away from this. So we will inevitably need to perform the unsafe operation to do the V0->V1 migration. Best I can think of is to explicitly block dangerous types 🤷
Addressed in commit: 45d77e0
|
Note in this draft version I've included in Including it here for now so that this is self-contained. I am not addressing as part of this PR the locking mechanism discussed in #4387 |
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
| return | ||
|
|
||
| db_url = str(self.db_engine.url) | ||
| runner = AlembicMigrationRunner(db_url) |
There was a problem hiding this comment.
It would be better to have the runner use self.db_engine directly instead of using the url to create new engines. With that change, kwargs used to create the engine in the db session service init (which can be required for connection) would still be used.
3 participants
Adopt Alembic as the migration framework for
DatabaseSessionService, replacing manualALTER TABLEoperations.This enables automatic schema tracking, upgrade/downgrade support, and Kubernetes Helm hook integration for production deployments.
Addresses #3343
alembic.ini)env.pyandscript.py.makoupgrade— Apply pending migrations with auto-bootstrap for existing DBsdowngrade— Roll back by revision speccheck— Exit0if up-to-date,1if pending (CI/CD gate)stamp— Bootstrap Alembic tracking for existing databasesgenerate— Create new migration scripts (autogenerate or template)_ensure_alembic_tracking()stamps baseline after table creationADK_AUTO_MIGRATE_DB=trueADK_AUTO_MIGRATE_DB=false(default)RuntimeErrorif schema is behindget_alembic_revision()is_alembic_managed()is_alembic_managed_from_url()AlembicMigrationRunner(SQLite)test_database_schema.pyfor V0 auto-migration behaviordocker-compose.ymlfor local integration testing / developmentmigration_guide.md— User-facing upgrade and CLI referencehelm_migration_guide.md— Kubernetes Job and Cloud SQL Proxy examplesmigration/README.md— Contributor workflow for adding schema versionsPlease ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
If applicable, please follow the issue templates to provide as much detail as
possible.
Problem:
A clear and concise description of what the problem is.
Solution:
A clear and concise description of what you want to happen and why you choose
this solution.
Testing Plan
Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes.
Unit Tests:
Please include a summary of passed
pytestresults.Manual End-to-End (E2E) Tests:
Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.
Checklist
Additional context
Add any other context or screenshots about the feature request here.