Skip to content

Latest commit

 

History

History
85 lines (60 loc) · 2.5 KB

README.md

File metadata and controls

85 lines (60 loc) · 2.5 KB

Database migrations

We use alembic to perform DB migrations. alembic tutorial: https://alembic.sqlalchemy.org/en/latest/tutorial.html

Before running commands

  1. Open the file mage-ai/mage_ai/orchestration/db/alembic.ini
  2. Change the value sqlalchemy.url to the URL of your database.
    • e.g. sqlalchemy.url = postgresql+psycopg2://postgres:postgres@host.docker.internal:5432/demo
    • e.g. postgresql://postgres:postgres@host.docker.internal/mage_development

Create migration scripts

Change directory into:

cd mage_ai/orchestration/db

alembic revision --autogenerate -m "Migration message"

If you don’t see your changes in the migration file, add your new models module in this file: mage-ai/mage_ai/orchestration/db/migrations/env.py

For example:

import mage_ai.orchestration.db.models.tags  # noqa: E402, F401

Run migration

alembic upgrade head

Rollback migration

alembic downgrade -1

Clean up

  1. Change the sqlalchemy.url back to the original value.

Running migrations on multiple databases

After you run the migration initially on SQLite or PostgreSQL, you must run the migration targeting the other database (e.g. if you initially ran it on SQLite, you must run it on PostgreSQL).

If it fails, create the migration file again for that database: alembic revision --autogenerate -m "Migration message". Find the differences and then add conditional statements for the commands that work for a database and doesn’t for the other.

Example: (Note if bind.engine.name == 'postgresql':)

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    bind = op.get_bind()
    if bind.engine.name == 'postgresql':
        with op.get_context().autocommit_block():
            op.execute("ALTER TYPE blockrunstatus ADD VALUE 'CONDITION_FAILED'")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    bind = op.get_bind()
    if bind.engine.name == 'postgresql':
        op.execute("ALTER TYPE blockrunstatus RENAME TO blockrunstatus_old")
        op.execute("CREATE TYPE blockrunstatus AS ENUM('INITIAL', 'QUEUED', 'RUNNING', 'COMPLETED', 'FAILED', 'CANCELLED', 'UPSTREAM_FAILED')")
        op.execute((
            "ALTER TABLE block_run ALTER COLUMN status TYPE blockrunstatus USING "
            "status::text::blockrunstatus"
        ))
        op.execute("DROP TYPE blockrunstatus_old")
    # ### end Alembic commands ###