apps/omyfish_api/db/engine.py

import os
import uuid
from contextlib import contextmanager
from pathlib import Path

DATABASE_URL = os.getenv("DATABASE_URL", "")
IS_POSTGIS = DATABASE_URL.startswith("postgresql")

if not IS_POSTGIS:
    _db_path = Path(__file__).resolve().parent.parent.parent.parent / "data" / "observations.db"
    DATABASE_URL = f"sqlite:///{_db_path}"

from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker

_connect_args = {} if IS_POSTGIS else {"check_same_thread": False}
engine = create_engine(DATABASE_URL, pool_pre_ping=True, connect_args=_connect_args)
Session = sessionmaker(bind=engine)

_db_initialized = False


def init_db():
    """Create tables via raw DDL. Used for SQLite; PostGIS uses Alembic migrations."""
    with engine.connect() as conn:
        if IS_POSTGIS:
            conn.execute(text("CREATE EXTENSION IF NOT EXISTS postgis"))
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS observations (
                    id               UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    species_name     TEXT,
                    scientific_name  TEXT,
                    confidence       FLOAT,
                    timestamp        TIMESTAMPTZ DEFAULT now(),
                    latitude         FLOAT,
                    longitude        FLOAT,
                    geom             GEOGRAPHY(POINT, 4326),
                    image_url        TEXT,
                    user_id          TEXT,
                    source           TEXT DEFAULT 'upload'
                )
            """))
            conn.execute(text(
                "CREATE INDEX IF NOT EXISTS observations_geom_idx "
                "ON observations USING GIST(geom)"
            ))
        else:
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS observations (
                    id               TEXT PRIMARY KEY,
                    species_name     TEXT,
                    scientific_name  TEXT,
                    confidence       REAL,
                    timestamp        TEXT DEFAULT (datetime('now')),
                    latitude         REAL,
                    longitude        REAL,
                    image_url        TEXT,
                    user_id          TEXT,
                    source           TEXT DEFAULT 'upload'
                )
            """))
            conn.execute(text("""
                CREATE TABLE IF NOT EXISTS users (
                    id               TEXT PRIMARY KEY,
                    email            TEXT UNIQUE NOT NULL,
                    hashed_password  TEXT NOT NULL,
                    role             TEXT NOT NULL DEFAULT 'user',
                    is_active        INTEGER NOT NULL DEFAULT 1,
                    created_at       TEXT DEFAULT (datetime('now'))
                )
            """))
        conn.commit()


def run_migrations():
    """Run Alembic migrations. PostGIS only."""
    from alembic.config import Config
    from alembic import command

    alembic_cfg = Config(str(Path(__file__).resolve().parent / "alembic.ini"))
    command.upgrade(alembic_cfg, "head")


def ensure_db() -> bool:
    """Idempotent DB setup: Alembic for PostGIS, DDL shim for SQLite."""
    global _db_initialized
    if not _db_initialized:
        try:
            if IS_POSTGIS:
                run_migrations()
            else:
                init_db()
            _db_initialized = True
        except Exception:
            pass
    return _db_initialized


def db_ready() -> bool:
    return _db_initialized


def new_id() -> str:
    """Generate a UUID string (works for both backends)."""
    return str(uuid.uuid4())


@contextmanager
def get_db():
    session = Session()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()