69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
|
"""add reference_sites and migrate references.site to FK
|
||
|
|
|
||
|
|
Revision ID: d7e3b19f0a2c
|
||
|
|
Revises: c3f8a1e924d5
|
||
|
|
Create Date: 2026-06-09
|
||
|
|
"""
|
||
|
|
from alembic import op
|
||
|
|
import sqlalchemy as sa
|
||
|
|
from sqlalchemy.sql import table, column, select
|
||
|
|
|
||
|
|
revision = "d7e3b19f0a2c"
|
||
|
|
down_revision = "c3f8a1e924d5"
|
||
|
|
branch_labels = None
|
||
|
|
depends_on = None
|
||
|
|
|
||
|
|
|
||
|
|
def upgrade() -> None:
|
||
|
|
# 1. Create reference_sites table
|
||
|
|
op.create_table(
|
||
|
|
"reference_sites",
|
||
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
||
|
|
sa.Column("name", sa.String(), nullable=False, unique=True),
|
||
|
|
)
|
||
|
|
|
||
|
|
# 2. Add site_id FK column (nullable)
|
||
|
|
op.add_column(
|
||
|
|
"references",
|
||
|
|
sa.Column("site_id", sa.Integer(),
|
||
|
|
sa.ForeignKey("reference_sites.id", ondelete="SET NULL"),
|
||
|
|
nullable=True),
|
||
|
|
)
|
||
|
|
|
||
|
|
# 3. Migrate existing site string data
|
||
|
|
conn = op.get_bind()
|
||
|
|
refs_t = table("references", column("id", sa.Integer), column("site", sa.String), column("site_id", sa.Integer))
|
||
|
|
sites_t = table("reference_sites", column("id", sa.Integer), column("name", sa.String))
|
||
|
|
|
||
|
|
# Collect distinct non-null site strings
|
||
|
|
existing = conn.execute(
|
||
|
|
select(refs_t.c.site).where(refs_t.c.site.isnot(None)).distinct()
|
||
|
|
).fetchall()
|
||
|
|
|
||
|
|
for (name,) in existing:
|
||
|
|
# Insert into reference_sites (ignore duplicates)
|
||
|
|
result = conn.execute(
|
||
|
|
sites_t.insert().returning(sites_t.c.id).values(name=name)
|
||
|
|
)
|
||
|
|
site_id = result.fetchone()[0]
|
||
|
|
# Update references rows with this site string
|
||
|
|
conn.execute(
|
||
|
|
refs_t.update().where(refs_t.c.site == name).values(site_id=site_id)
|
||
|
|
)
|
||
|
|
|
||
|
|
# 4. Drop old site string column
|
||
|
|
op.drop_column("references", "site")
|
||
|
|
|
||
|
|
|
||
|
|
def downgrade() -> None:
|
||
|
|
op.add_column("references", sa.Column("site", sa.String(), nullable=True))
|
||
|
|
# best-effort: restore site name from FK
|
||
|
|
conn = op.get_bind()
|
||
|
|
refs_t = table("references", column("id", sa.Integer), column("site", sa.String), column("site_id", sa.Integer))
|
||
|
|
sites_t = table("reference_sites", column("id", sa.Integer), column("name", sa.String))
|
||
|
|
rows = conn.execute(select(refs_t.c.id, sites_t.c.name).join(sites_t, refs_t.c.site_id == sites_t.c.id)).fetchall()
|
||
|
|
for ref_id, name in rows:
|
||
|
|
conn.execute(refs_t.update().where(refs_t.c.id == ref_id).values(site=name))
|
||
|
|
op.drop_column("references", "site_id")
|
||
|
|
op.drop_table("reference_sites")
|