main.py

import os
from alembic import command
from alembic.config import Config
from alembic.runtime.migration import MigrationContext
from alembic.script import ScriptDirectory
from sqlalchemy import Column, DateTime, create_engine
from sqlalchemy.orm import sessionmaker

from db.models import Base, Model, Tag


def get_alembic_cfg():
    return Config("alembic.ini")


def get_db_url(alembic_cfg):
    return alembic_cfg.get_main_option("sqlalchemy.url")


def get_db_name(alembic_cfg):
    return get_db_url(alembic_cfg).split("/")[-1]


def reset_db(alembic_cfg):
    db_name = get_db_name(alembic_cfg)
    if os.path.exists(db_name):
        print(f"Removing existing database: {db_name}")
        os.remove(db_name)


def init_db(alembic_cfg, revision="head"):
    db_url = get_db_url(alembic_cfg)
    print(f"Database URL: {db_url}")

    # Check if we need to upgrade
    engine = create_engine(db_url)
    conn = engine.connect()

    context = MigrationContext.configure(conn)
    current_rev = context.get_current_revision()

    script = ScriptDirectory.from_config(alembic_cfg)
    if revision == "head":
        target_rev = script.get_current_head()
    else:
        target_rev = revision

    print(f"Current revision: {current_rev}")
    print(f"Target revision: {target_rev}")

    if current_rev != target_rev:
        # Database is not up to date, run upgrade
        command.upgrade(alembic_cfg, target_rev)
        print(f"Database upgraded from {current_rev} to {target_rev}")
    else:
        print("Database is up to date")

    return engine


def register_date_added_column():
    col = Column(
        "date_added", DateTime
    )
    setattr(Model, col.name, col)


if __name__ == "__main__":
    alembic_cfg = get_alembic_cfg()
    reset_db(alembic_cfg)

    # Initialise database to first revision
    engine = init_db(alembic_cfg, "15d874f688e8")
    Session = sessionmaker(bind=engine)

    # Seed database
    with Session() as session:
        readme_model = Model(
            path="readme.md",
            name="Readme",
            description="Project readme",
            type="readme",
            architecture="markdown"
        )
        session.add(readme_model)

        session.add(Model(
            path="requirements.txt",
            name="Requirements",
            description="Project requirements",
            type="requirements",
            architecture="text"
        ))

        my_first_tag = Tag(name="my first tag")
        session.add(my_first_tag)

        session.commit()

        # Add a tag to the readme model
        readme_model.tags.append(my_first_tag)
        session.commit()

        tagged_models = (
            session.query(Model)
            .join(Model.tags)
            .filter(Tag.name == my_first_tag.name)
            .all()
        )

        print(f"Tagged models: {len(tagged_models)}")
        print(f"All models: {len(session.query(Model).all())}")

    # Dynamically register our date_added column, normally this would just be part of the model
    # but for this example we need to add it at runtime so the initial seeding doesnt try to set it
    register_date_added_column()

    # Upgrade to latest revision
    engine = init_db(alembic_cfg)
    Session = sessionmaker(bind=engine)

    # Print all models with their date_added, this should now be populated by the migration
    with Session() as session:
        for model in session.query(Model).all():
            print(model.path, model.date_added)