zhixing-system/zhixing-server/migrations/versions/0005_radar_daily_aggregate.py

"""Persist publication-owned sector daily aggregates for exact metric replay."""

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

revision: str = "0005_radar_daily_aggregate"
down_revision: str | None = "0004_sector_radar"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    """Create source recovery links and exact multi-day metric inputs."""

    op.create_table(
        "sector_radar_publication_source",
        sa.Column(
            "publication_id",
            sa.String(64),
            sa.ForeignKey("sector_radar_publication.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("source_group", sa.String(32), nullable=False),
        sa.Column("source_order", sa.Integer(), nullable=False),
        sa.Column(
            "refresh_on_retry",
            sa.Boolean(),
            nullable=False,
            server_default=sa.false(),
        ),
        sa.Column(
            "source_snapshot_id",
            sa.String(64),
            sa.ForeignKey("sector_radar_source_snapshot.id", ondelete="RESTRICT"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("publication_id", "source_group", "source_order"),
        sa.UniqueConstraint(
            "publication_id",
            "source_group",
            "source_snapshot_id",
            name="uq_sector_radar_publication_source_snapshot",
        ),
        sa.CheckConstraint(
            "source_group IN ('calendar', 'concept_indices', 'industry_indices', "
            "'members', 'stock_basics', 'suspensions', 'daily', 'moneyflow_dc')",
            name="ck_sector_radar_publication_source_group",
        ),
        sa.CheckConstraint(
            "source_order >= 0",
            name="ck_sector_radar_publication_source_order",
        ),
    )
    op.create_index(
        "ix_sector_radar_publication_source_snapshot",
        "sector_radar_publication_source",
        ["source_snapshot_id"],
    )

    op.create_unique_constraint(
        "uq_sector_radar_publication_id_date",
        "sector_radar_publication",
        ["id", "target_trade_date"],
    )
    op.create_foreign_key(
        "fk_sector_radar_ranking_publication_date",
        "sector_radar_ranking",
        "sector_radar_publication",
        ["publication_id", "trade_date"],
        ["id", "target_trade_date"],
        ondelete="CASCADE",
    )

    op.create_table(
        "sector_radar_daily_aggregate",
        sa.Column("publication_id", sa.String(64), nullable=False),
        sa.Column("trade_date", sa.Date(), nullable=False),
        sa.Column("sector_type", sa.String(16), nullable=False),
        sa.Column("sector_code", sa.String(16), nullable=False),
        sa.Column("sector_name", sa.String(128), nullable=False),
        sa.Column("member_count", sa.Integer(), nullable=False),
        sa.Column("valid_sample_count", sa.Integer(), nullable=False),
        sa.Column("net_amount_yuan", sa.Numeric(28, 6)),
        sa.Column("turnover_yuan", sa.Numeric(28, 6)),
        sa.Column("membership_coverage", sa.Numeric(8, 6), nullable=False),
        sa.Column("moneyflow_coverage", sa.Numeric(8, 6), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.PrimaryKeyConstraint("publication_id", "sector_type", "sector_code"),
        sa.ForeignKeyConstraint(
            ["publication_id", "trade_date"],
            ["sector_radar_publication.id", "sector_radar_publication.target_trade_date"],
            name="fk_sector_radar_daily_aggregate_publication_date",
            ondelete="CASCADE",
        ),
        sa.CheckConstraint(
            "sector_type IN ('concept', 'industry')",
            name="ck_sector_radar_daily_aggregate_type",
        ),
        sa.CheckConstraint(
            "member_count >= 0 AND valid_sample_count >= 0 AND valid_sample_count <= member_count",
            name="ck_sector_radar_daily_aggregate_counts",
        ),
        sa.CheckConstraint(
            "membership_coverage >= 0 AND membership_coverage <= 1 "
            "AND moneyflow_coverage >= 0 AND moneyflow_coverage <= 1",
            name="ck_sector_radar_daily_aggregate_coverage",
        ),
        sa.CheckConstraint(
            "turnover_yuan IS NULL OR (turnover_yuan >= 0 AND "
            "turnover_yuan NOT IN ('NaN'::numeric, 'Infinity'::numeric))",
            name="ck_sector_radar_daily_aggregate_turnover",
        ),
        sa.CheckConstraint(
            "net_amount_yuan IS NULL OR net_amount_yuan NOT IN "
            "('NaN'::numeric, 'Infinity'::numeric, '-Infinity'::numeric)",
            name="ck_sector_radar_daily_aggregate_net_amount",
        ),
    )
    op.create_index(
        "ix_sector_radar_daily_aggregate_history",
        "sector_radar_daily_aggregate",
        ["trade_date", "sector_type", "sector_code"],
    )


def downgrade() -> None:
    """Drop only the replay aggregate extension."""

    op.drop_index(
        "ix_sector_radar_daily_aggregate_history",
        table_name="sector_radar_daily_aggregate",
    )
    op.drop_table("sector_radar_daily_aggregate")
    op.drop_constraint(
        "fk_sector_radar_ranking_publication_date",
        "sector_radar_ranking",
        type_="foreignkey",
    )
    op.drop_constraint(
        "uq_sector_radar_publication_id_date",
        "sector_radar_publication",
        type_="unique",
    )
    op.drop_index(
        "ix_sector_radar_publication_source_snapshot",
        table_name="sector_radar_publication_source",
    )
    op.drop_table("sector_radar_publication_source")