haunt-fm/alembic/versions/001_initial_schema.py

"""Initial schema

Revision ID: 001
Revises:
Create Date: 2026-02-22

"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from pgvector.sqlalchemy import Vector

revision: str = "001"
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Enable pgvector extension
    op.execute("CREATE EXTENSION IF NOT EXISTS vector")

    # Tracks
    op.create_table(
        "tracks",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("title", sa.Text, nullable=False),
        sa.Column("artist", sa.Text, nullable=False),
        sa.Column("album", sa.Text),
        sa.Column("fingerprint", sa.Text, unique=True, nullable=False),
        sa.Column("lastfm_url", sa.Text),
        sa.Column("itunes_track_id", sa.BigInteger),
        sa.Column("itunes_preview_url", sa.Text),
        sa.Column("apple_music_id", sa.Text),
        sa.Column("duration_ms", sa.Integer),
        sa.Column("genre", sa.Text),
        sa.Column("embedding_status", sa.Text, nullable=False, server_default="pending"),
        sa.Column("embedding_error", sa.Text),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
    )

    # Listen events
    op.create_table(
        "listen_events",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("track_id", sa.BigInteger, sa.ForeignKey("tracks.id"), nullable=False),
        sa.Column("source", sa.Text, nullable=False, server_default="music_assistant"),
        sa.Column("speaker_name", sa.Text),
        sa.Column("listened_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.Column("duration_played", sa.Integer),
        sa.Column("raw_payload", sa.dialects.postgresql.JSONB),
    )

    # Track embeddings (512-dim CLAP)
    op.create_table(
        "track_embeddings",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("track_id", sa.BigInteger, sa.ForeignKey("tracks.id"), unique=True, nullable=False),
        sa.Column("embedding", Vector(512), nullable=False),
        sa.Column("model_version", sa.Text, nullable=False, server_default="laion/larger_clap_music"),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
    )
    op.execute(
        "CREATE INDEX ix_track_embeddings_hnsw ON track_embeddings "
        "USING hnsw (embedding vector_cosine_ops)"
    )

    # Similarity links
    op.create_table(
        "similarity_links",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("source_track_id", sa.BigInteger, sa.ForeignKey("tracks.id"), nullable=False),
        sa.Column("target_track_id", sa.BigInteger, sa.ForeignKey("tracks.id"), nullable=False),
        sa.Column("lastfm_match", sa.REAL),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
        sa.UniqueConstraint("source_track_id", "target_track_id", name="uq_similarity_link"),
    )

    # Taste profiles
    op.create_table(
        "taste_profiles",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("name", sa.Text, unique=True, nullable=False, server_default="default"),
        sa.Column("embedding", Vector(512), nullable=False),
        sa.Column("track_count", sa.Integer, nullable=False),
        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
    )

    # Playlists
    op.create_table(
        "playlists",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("name", sa.Text),
        sa.Column("known_pct", sa.Integer, nullable=False),
        sa.Column("total_tracks", sa.Integer, nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
    )

    op.create_table(
        "playlist_tracks",
        sa.Column("id", sa.BigInteger, primary_key=True),
        sa.Column("playlist_id", sa.BigInteger, sa.ForeignKey("playlists.id", ondelete="CASCADE"), nullable=False),
        sa.Column("track_id", sa.BigInteger, sa.ForeignKey("tracks.id"), nullable=False),
        sa.Column("position", sa.Integer, nullable=False),
        sa.Column("is_known", sa.Boolean, nullable=False),
        sa.Column("similarity_score", sa.REAL),
    )


def downgrade() -> None:
    op.drop_table("playlist_tracks")
    op.drop_table("playlists")
    op.drop_table("taste_profiles")
    op.drop_table("similarity_links")
    op.execute("DROP INDEX IF EXISTS ix_track_embeddings_hnsw")
    op.drop_table("track_embeddings")
    op.drop_table("listen_events")
    op.drop_table("tracks")
    op.execute("DROP EXTENSION IF EXISTS vector")