db.py

"""
db.py — Database abstraction layer.

- Production (HF Spaces): set DATABASE_URL env var → uses Neon PostgreSQL via psycopg2
- Local dev: no DATABASE_URL → falls back to SQLite

Usage in other modules:
    from db import get_conn, PH, IntegrityError

    conn = get_conn()
    cursor = conn.cursor()
    cursor.execute(f"SELECT * FROM users WHERE username = {PH}", (username,))
    row = cursor.fetchone()
    conn.close()
"""

import os
import sqlite3
import logging

logger = logging.getLogger(__name__)

# Connection string from environment (set as HF Space Secret)
DATABASE_URL = os.environ.get('DATABASE_URL', '')

# Are we using PostgreSQL?
USE_POSTGRES = bool(DATABASE_URL)

if USE_POSTGRES:
    import psycopg2
    import psycopg2.extras
    import psycopg2.errors
    # Placeholder style for PostgreSQL
    PH = '%s'
    # Exception alias
    IntegrityError = psycopg2.errors.UniqueViolation
    logger.info("DB: Using Neon PostgreSQL")
else:
    # Placeholder style for SQLite
    PH = '?'
    # Exception alias
    IntegrityError = sqlite3.IntegrityError
    logger.info("DB: Using local SQLite (dev mode)")

# SQLite path (only used in local dev)
_SQLITE_PATH = os.path.join(os.path.dirname(__file__), 'data', 'users.db')


def get_conn():
    """
    Returns an open database connection.
    Caller is responsible for calling conn.close() (use finally blocks).
    """
    if USE_POSTGRES:
        if 'sslmode=' not in DATABASE_URL:
            conn = psycopg2.connect(DATABASE_URL, sslmode='require')
        else:
            conn = psycopg2.connect(DATABASE_URL)
        conn.autocommit = False
        return conn
    else:
        os.makedirs(os.path.dirname(_SQLITE_PATH), exist_ok=True)
        conn = sqlite3.connect(_SQLITE_PATH, timeout=10)
        conn.execute('PRAGMA busy_timeout=5000')
        return conn