104 lines
3.1 KiB
Python
104 lines
3.1 KiB
Python
|
"""Persist explicit unknown point-in-time sector membership snapshots."""
|
||
|
|
|
||
|
|
from collections.abc import Sequence
|
||
|
|
|
||
|
|
import sqlalchemy as sa
|
||
|
|
from alembic import op
|
||
|
|
|
||
|
|
revision: str = "0006_membership_unknown"
|
||
|
|
down_revision: str | None = "0005_radar_daily_aggregate"
|
||
|
|
branch_labels: str | Sequence[str] | None = None
|
||
|
|
depends_on: str | Sequence[str] | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def upgrade() -> None:
|
||
|
|
"""Allow one null-stock marker for an explicitly empty sector partition."""
|
||
|
|
|
||
|
|
op.add_column(
|
||
|
|
"sector_radar_membership",
|
||
|
|
sa.Column("membership_key", sa.String(32), nullable=True),
|
||
|
|
)
|
||
|
|
op.execute("UPDATE sector_radar_membership SET membership_key = stock_code")
|
||
|
|
op.drop_constraint(
|
||
|
|
"sector_radar_membership_pkey",
|
||
|
|
"sector_radar_membership",
|
||
|
|
type_="primary",
|
||
|
|
)
|
||
|
|
op.drop_constraint(
|
||
|
|
"ck_sector_radar_membership_status",
|
||
|
|
"sector_radar_membership",
|
||
|
|
type_="check",
|
||
|
|
)
|
||
|
|
op.alter_column(
|
||
|
|
"sector_radar_membership",
|
||
|
|
"membership_key",
|
||
|
|
existing_type=sa.String(32),
|
||
|
|
nullable=False,
|
||
|
|
)
|
||
|
|
op.alter_column(
|
||
|
|
"sector_radar_membership",
|
||
|
|
"stock_code",
|
||
|
|
existing_type=sa.String(12),
|
||
|
|
nullable=True,
|
||
|
|
)
|
||
|
|
op.alter_column(
|
||
|
|
"sector_radar_membership",
|
||
|
|
"stock_name",
|
||
|
|
existing_type=sa.String(128),
|
||
|
|
nullable=True,
|
||
|
|
)
|
||
|
|
op.create_primary_key(
|
||
|
|
"sector_radar_membership_pkey",
|
||
|
|
"sector_radar_membership",
|
||
|
|
["source_snapshot_id", "sector_code", "membership_key"],
|
||
|
|
)
|
||
|
|
op.create_check_constraint(
|
||
|
|
"ck_sector_radar_membership_status",
|
||
|
|
"sector_radar_membership",
|
||
|
|
"(membership_status = 'available' "
|
||
|
|
"AND stock_code IS NOT NULL AND stock_name IS NOT NULL "
|
||
|
|
"AND membership_key = stock_code) OR "
|
||
|
|
"(membership_status = 'membership_unknown' "
|
||
|
|
"AND stock_code IS NULL AND stock_name IS NULL "
|
||
|
|
"AND membership_key = '__membership_unknown__')",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def downgrade() -> None:
|
||
|
|
"""Discard unknown markers and restore the available-member-only schema."""
|
||
|
|
|
||
|
|
op.execute("DELETE FROM sector_radar_membership WHERE membership_status = 'membership_unknown'")
|
||
|
|
op.drop_constraint(
|
||
|
|
"ck_sector_radar_membership_status",
|
||
|
|
"sector_radar_membership",
|
||
|
|
type_="check",
|
||
|
|
)
|
||
|
|
op.drop_constraint(
|
||
|
|
"sector_radar_membership_pkey",
|
||
|
|
"sector_radar_membership",
|
||
|
|
type_="primary",
|
||
|
|
)
|
||
|
|
op.alter_column(
|
||
|
|
"sector_radar_membership",
|
||
|
|
"stock_code",
|
||
|
|
existing_type=sa.String(12),
|
||
|
|
nullable=False,
|
||
|
|
)
|
||
|
|
op.alter_column(
|
||
|
|
"sector_radar_membership",
|
||
|
|
"stock_name",
|
||
|
|
existing_type=sa.String(128),
|
||
|
|
nullable=False,
|
||
|
|
)
|
||
|
|
op.drop_column("sector_radar_membership", "membership_key")
|
||
|
|
op.create_primary_key(
|
||
|
|
"sector_radar_membership_pkey",
|
||
|
|
"sector_radar_membership",
|
||
|
|
["source_snapshot_id", "sector_code", "stock_code"],
|
||
|
|
)
|
||
|
|
op.create_check_constraint(
|
||
|
|
"ck_sector_radar_membership_status",
|
||
|
|
"sector_radar_membership",
|
||
|
|
"membership_status = 'available'",
|
||
|
|
)
|