# 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`
1. 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:

```python
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':`)

```python
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 ###

```
