Creating a "shortcut" base

[quentin-roche]

Hello everyone,

I have pretty much the same question as @indiVar0508 (can't find the discussion anymore). Basically he did what I am trying to achieve but he did it manually. I need to automate it and package it nicely because upgrading and downgrading happens a lot on our project.

Situation:

I have a CI jobs that are checking migration on every commit on different databases (AzureSQL/SQLite). To check the migration,
the CI creates an empty database and apply the migration from the base. Then it compares the database to the SQLAlchemy metadata. Finally, it checks if it is possible to downgrade to an empty database.

Now I want to add the possibility to verify the database for AzureSQL. I created a migration file to resolve this common issue with AzureSQL. The issue is that the CI job for AzureSQL is failing because for all the previous migrations, the database is not compatible with AzureSQL.

So I have created a new migration (migration50) that is fixing the compatibility issues with AzureSQL.

graph TD;
    database_V0:::nok   --migration0 -->  ...:::nok;
    ...:::nok           --migration49-->  database_V50:::nok;
    database_V50:::nok  --migration50-->  database_V51:::ok;

    classDef ok fill:#00AA00,color:white,stroke-width:2px;
    classDef nok fill:#DD5500,color:white,stroke-width:2px;

(orange: Incompatible with AzureSQL, green: Compatible with AzureSQL)

Issue:

The problem is that the CI job for AzureSQL is failing as expected because AzureSQL is not compatible from version V0 to V50. I could create fixes for all the migrations files but it is a lot of work and I would loose the main purpose of Alembic (to be able to downgrade/upgrade already existing databases).

Idea:

My idea would be to create some kind of switch between down revision based on SQLAlchemy dialect.

graph TD;
    database_V0:::nok   --migration0 -->  ...:::nok;
    ...:::nok           --migration49-->  database_V50:::nok;
    database_V50:::nok  --migration50-->  database_V51:::ok;
    database_V51:::ok:  --if MySQL or SQLite-->  switch?;
    
    database_V51_AzureSQL:::ok  --if AzureSQL -->  switch?;
    switch? --migration51--> database_V52:::ok;

    classDef ok fill:#00AA00,color:white,stroke-width:2px;
    classDef nok fill:#DD5500,color:white,stroke-width:2px;

What I created so far:

from typing import Sequence, Union
from common.database.connection import get_engine

if get_engine().dialect.name == "mssql":
    down_revision = '<database_V51>'  # From AzureSQL to MySQL/SQLite/AzureSQL
elif get_engine().dialect.name in ["sqlite", "mysql"]:
    down_revision = "<database_V51_AzureSQL>"  # From MySQL/SQLite to MySQL/SQLite/AzureSQL
else:
    raise Exception("Unsupported dialect at azuresql_switch migration")

# revision identifiers, used by Alembic.
revision: str = '<database_V52>'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

def upgrade() -> None:
    pass

def downgrade() -> None:
    pass

This is working but I need to specify the correct head branch since the branch not associated with the switch is still visible to Alembic. It would be ok if I was alone working on the project but many Data-scientist use the project, so i want migration management to be as simple as possible.

Do you have any idea how I could "hide" the branch ?
If you have a better idea on how to achieve automatic migration with two equivalent bases, I would be happy to hear it !

[CaselIT]

Hi,

If it's just on the CI that you have the issue you could use stamp to add a migration only to the migration table on the db without applying it.

So you could avoid the dynamic switch and instead do:

• stamp database_V51_AzureSQL when on sqlite
• stamp database_V50 when on AzureSQL

This would work for upgrade, but likely not for downgrade (unless you dynamically change where you want to downgrade to, instead of hard-coding base)

[quentin-roche]

That would work but it is not seamless for users.
What I did in the end is to create a decorator:

import logging
from functools import wraps

from alembic import context


def run_only_if(
    database_protocol_is: str | None = None,
    database_protocol_is_not: str | None = None,
):
    def decorator(func):
        @wraps(func)
        def wrapper():
            sqlalchemy_url = context.config.get_main_option("sqlalchemy.url")

            if ((
                database_protocol_is is not None
                and sqlalchemy_url.startswith(database_protocol_is)
            ) or (
                database_protocol_is_not is not None
                and not sqlalchemy_url.startswith(database_protocol_is_not)
            )):
                return func()
            else:
                logging.info(
                    "Skipping migration since the database "
                    f"protocol is {sqlalchemy_url.split('://')}."
                )

        return wrapper

    return decorator

And add it to upgrade and downgrade functions by doing either:

@run_only_if(database_protocol_is="mssql")
def upgrade():
    [...]
@run_only_if(database_protocol_is="mssql")
def downgrade():
    [...]

or

@run_only_if(database_protocol_is_not="mssql")
def upgrade():
    [...]
@run_only_if(database_protocol_is_not="mssql")
def downgrade():
    [...]

And I have this structure of migrations:

graph TD;
    empty_database:::ok  --create old tables-->  database_V0:::nok ;
    database_V0:::nok   --migration0 -->  ...:::nok;
    ...:::nok           --migration49-->  database_V50:::nok;
    database_V50:::nok  --migration50-->  database_V51:::ok;
    database_V51:::ok:  --merge      -->  database_V52:::ok;
    
    empty_database:::ok         --create tables-->  database_V51_AzureSQL:::ok ;
    database_V51_AzureSQL:::ok  --merge    -->  database_V52:::ok;
    
    classDef ok fill:#00AA00,color:white,stroke-width:2px;
    classDef nok fill:#DD5500,color:white,stroke-width:2px;