"""add source_files table for raw data tracking

Revision ID: 597473a8001c
Revises: 7bbb8648030f
Create Date: 2026-01-24 23:44:58.892640

"""

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "597473a8001c"
down_revision = "7bbb8648030f"
branch_labels = None
depends_on = None


def upgrade() -> None:
  # ### commands auto generated by Alembic - please adjust! ###
  op.create_table(
    "source_files",
    sa.Column("id", sa.String(), nullable=False),
    sa.Column("graph_id", sa.String(), nullable=False),
    sa.Column("storage_key", sa.String(), nullable=False),
    sa.Column("file_type", sa.String(), nullable=False),
    sa.Column("file_size_bytes", sa.Integer(), nullable=True),
    sa.Column("source_id", sa.String(), nullable=True),
    sa.Column("status", sa.String(), nullable=False),
    sa.Column("error_reason", sa.Text(), nullable=True),
    sa.Column("attempts", sa.Integer(), nullable=False),
    sa.Column("partition_key", sa.String(), nullable=True),
    sa.Column("discovered_at", sa.DateTime(timezone=True), nullable=False),
    sa.Column("processed_at", sa.DateTime(timezone=True), nullable=True),
    sa.Column("last_attempt_at", sa.DateTime(timezone=True), nullable=True),
    sa.ForeignKeyConstraint(["graph_id"], ["graphs.graph_id"], ondelete="CASCADE"),
    sa.PrimaryKeyConstraint("id"),
    sa.UniqueConstraint("storage_key"),
  )
  op.create_index(
    "idx_source_files_file_type", "source_files", ["file_type"], unique=False
  )
  op.create_index(
    "idx_source_files_graph_id", "source_files", ["graph_id"], unique=False
  )
  op.create_index(
    "idx_source_files_partition_key", "source_files", ["partition_key"], unique=False
  )
  op.create_index("idx_source_files_status", "source_files", ["status"], unique=False)
  op.create_index(
    "idx_source_files_status_attempts",
    "source_files",
    ["status", "attempts"],
    unique=False,
  )
  # ### end Alembic commands ###


def downgrade() -> None:
  # ### commands auto generated by Alembic - please adjust! ###
  op.drop_index("idx_source_files_status_attempts", table_name="source_files")
  op.drop_index("idx_source_files_status", table_name="source_files")
  op.drop_index("idx_source_files_partition_key", table_name="source_files")
  op.drop_index("idx_source_files_graph_id", table_name="source_files")
  op.drop_index("idx_source_files_file_type", table_name="source_files")
  op.drop_table("source_files")
  # ### end Alembic commands ###
