Search code examples
pythonsqlalchemymigrationalembic

How to update tables schemes in Python, using only Python (No CLI)


I am working with SQL Alchemy, and wish to update the tables schemes automatically, without using CLI (The structure of my code prevent it) I don't need to store migrations files either. I have been looking for Alembic, which seems to be the best tool to do this. However, I cannot find a proper process or explanation on how to do it without using CLIs. The closest I have comes from this: How to run a migration with Python Alembic by code? However, the same error still appears, even with the solution. I was thinking about using autogenerate.produce_migrations(alembic_context, metadata) to get the operations (That works fine), but how to execute those, then? Thanks


Solution

  • I finally found it. Below the code

    from alembic import autogenerate, util
    from alembic.config import Config
    
    from alembic.runtime.environment import EnvironmentContext
    from alembic.script import ScriptDirectory
    config = Config()
    
    config.set_main_option("script_location", <script location>)
    
    alembic_script = ScriptDirectory.from_config(config)
    environment = EnvironmentContext(config, alembic_script)
    
    with metadata.bind.connect() as connection:
        environment.configure(
            connection=connection,
            target_metadata=metadata,
            fn=fn,
            include_object=include_object
        )
    
        migrations = autogenerate.produce_migrations(context=environment.get_context(),
                                                     metadata=metadata)
        alembic_script.generate_revision(util.rev_id(), "update table",
                                         upgrades=autogenerate.render_python_code(migrations.upgrade_ops))
    
        with environment.begin_transaction():
            environment.run_migrations()
    
    
    def fn(revision, context):
        script = ScriptDirectory.from_config(context.config)
        return script._upgrade_revs(script.get_heads(), revision)
    
    
    def include_object(object, name, type_, reflected, compare_to):
        if type_ == "table" and reflected and compare_to is None:
            return False
        else:
            return True
    

    Notes:

    • include_object option ensures that no table is deleted. It is then optional

    • replace <'script location'> with the path of alembic. This directory must contain a direction "versions" and a file script.py.mako. Such data can be automatically generated by calling once the cli alembic init <root path of data>

    • Such code doesn't handle mutliple branches and downgrade: I assume it wouldn't be that complicate to handle those

    • All this code can be called in a function, with a single metadata input (That must be bounded to an engine).