hardbanrecords/Metadata-Engine

1from sqlalchemy import create_engine, Column, Integer, String, DateTime, text2from sqlalchemy.types import JSON3from sqlalchemy.orm import DeclarativeBase, sessionmaker4import os5import logging6 7logger = logging.getLogger(__name__)8 9# Determine database path (use /data on HF Spaces for persistence if available)10PERSISTENT_DATA_PATH = "/data"11SQLITE_DB_NAME = "music_metadata.db"12 13if os.path.exists(PERSISTENT_DATA_PATH):14    DEFAULT_DB_URL = f"sqlite:///{PERSISTENT_DATA_PATH}/{SQLITE_DB_NAME}"15else:16    DEFAULT_DB_URL = f"sqlite:///./{SQLITE_DB_NAME}"17 18DATABASE_URL = os.getenv("DATABASE_URL", DEFAULT_DB_URL)19engine = create_engine(20    DATABASE_URL,21    connect_args={"check_same_thread": False} if "sqlite" in DATABASE_URL else {},22)23SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)24 25def get_db():26    db = SessionLocal()27    try:28        yield db29    finally:30        db.close()31 32class Base(DeclarativeBase):33    pass34 35class Job(Base):36    __tablename__ = "jobs"37    id = Column(String, primary_key=True, index=True)38    user_id = Column(String, index=True, nullable=True)39    status = Column(String, default="pending")40    file_name = Column(String)41    result = Column(JSON, nullable=True)42    error = Column(String, nullable=True)43    message = Column(String, nullable=True)44    duration = Column(Integer, nullable=True)45    structure = Column(JSON, nullable=True)46    coverArt = Column(String, nullable=True)47    ipfs_hash = Column(String, nullable=True)48    ipfs_url = Column(String, nullable=True)49    timestamp = Column(DateTime)50 51class AnalysisHistory(Base):52    __tablename__ = "analysis_history"53    id = Column(Integer, primary_key=True, index=True)54    user_id = Column(String, index=True)55    file_name = Column(String)56    result = Column(JSON)57    created_at = Column(DateTime, default=text('CURRENT_TIMESTAMP'))58 59# Helper for migrations60def run_migrations():61    if "sqlite" not in DATABASE_URL:62        return63    64    try:65        with engine.connect() as conn:66            # Check for missing columns in 'jobs'67            result = conn.execute(text("PRAGMA table_info(jobs)"))68            cols = [row[1] for row in result.fetchall()]69            70            if not cols:71                return # Table doesn't exist yet, create_all will handle it72                73            required = {74                "message": "TEXT",75                "duration": "INTEGER",76                "structure": "JSON",77                "coverArt": "TEXT",78                "ipfs_hash": "TEXT",79                "ipfs_url": "TEXT",80            }81            82            for col, col_type in required.items():83                if col not in cols:84                    logger.info(f"Adding missing column '{col}' to 'jobs' table...")85                    conn.execute(text(f"ALTER TABLE jobs ADD COLUMN {col} {col_type}"))86                    conn.commit()87    except Exception as e:88        logger.error(f"Migration error: {e}")89 90# Apply migrations and create tables91Base.metadata.create_all(bind=engine)92run_migrations()93