bikeslop/db.py

import json
import os
import uuid
from datetime import datetime, timezone

from sqlalchemy import (
    Boolean,
    Column,
    DateTime,
    Float,
    ForeignKey,
    Integer,
    String,
    Text,
    create_engine,
)
from sqlalchemy.orm import DeclarativeBase, Session, relationship

DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///bikeslop.db")

engine = create_engine(
    DATABASE_URL,
    connect_args={"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {},
)


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "users"

    id = Column(String, primary_key=True, default=lambda: uuid.uuid4().hex)
    oidc_sub = Column(String, unique=True, nullable=False)
    email = Column(String, nullable=True)
    display_name = Column(String, nullable=True)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))

    rides = relationship("Ride", back_populates="user", cascade="all, delete-orphan")
    insights = relationship("Insight", back_populates="user", cascade="all, delete-orphan")


class Ride(Base):
    __tablename__ = "rides"

    id = Column(String, primary_key=True, default=lambda: uuid.uuid4().hex)
    user_id = Column(String, ForeignKey("users.id"), nullable=False)
    s3_key = Column(String, nullable=False)
    filename = Column(String, nullable=False)
    name = Column(String, nullable=True)
    activity_type = Column(String, nullable=True)
    start_time = Column(String, nullable=True)
    end_time = Column(String, nullable=True)
    duration_s = Column(Integer, default=0)
    distance_m = Column(Integer, default=0)
    distance_km = Column(Float, default=0.0)
    avg_speed_kmh = Column(Float, default=0.0)
    elevation_gain_m = Column(Integer, default=0)
    elevation_loss_m = Column(Integer, default=0)
    avg_hr = Column(Integer, nullable=True)
    max_hr = Column(Integer, nullable=True)
    hr_zones_json = Column(Text, default="{}")
    has_gps = Column(Boolean, default=True)

    user = relationship("User", back_populates="rides")
    insight = relationship("Insight", back_populates="ride", uselist=False, cascade="all, delete-orphan")

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "s3_key": self.s3_key,
            "filename": self.filename,
            "name": self.name,
            "activity_type": self.activity_type,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_s": self.duration_s,
            "distance_m": self.distance_m,
            "distance_km": self.distance_km,
            "avg_speed_kmh": self.avg_speed_kmh,
            "elevation_gain_m": self.elevation_gain_m,
            "elevation_loss_m": self.elevation_loss_m,
            "avg_hr": self.avg_hr,
            "max_hr": self.max_hr,
            "hr_zones": json.loads(self.hr_zones_json or "{}"),
            "has_gps": self.has_gps,
        }


class Insight(Base):
    __tablename__ = "insights"

    id = Column(String, primary_key=True, default=lambda: uuid.uuid4().hex)
    ride_id = Column(String, ForeignKey("rides.id"), nullable=False, unique=True)
    user_id = Column(String, ForeignKey("users.id"), nullable=False)
    text = Column(Text, nullable=False)
    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))

    ride = relationship("Ride", back_populates="insight")
    user = relationship("User", back_populates="insights")


def create_tables() -> None:
    Base.metadata.create_all(engine)


def get_session() -> Session:
    return Session(engine)


def upsert_user(oidc_sub: str, email: str | None, display_name: str | None) -> User:
    with get_session() as s:
        user = s.query(User).filter_by(oidc_sub=oidc_sub).first()
        if user is None:
            user = User(oidc_sub=oidc_sub, email=email, display_name=display_name)
            s.add(user)
        else:
            if email:
                user.email = email
            if display_name:
                user.display_name = display_name
        s.commit()
        s.refresh(user)
        return User(
            id=user.id,
            oidc_sub=user.oidc_sub,
            email=user.email,
            display_name=user.display_name,
            created_at=user.created_at,
        )


def get_user_by_id(user_id: str) -> User | None:
    with get_session() as s:
        user = s.query(User).filter_by(id=user_id).first()
        if user is None:
            return None
        return User(
            id=user.id,
            oidc_sub=user.oidc_sub,
            email=user.email,
            display_name=user.display_name,
            created_at=user.created_at,
        )