Security_AI_integrated/db/models.py

import os
import sys
from datetime import datetime
from typing import List, Optional

from sqlalchemy import (
    Boolean,
    DateTime,
    Float,
    ForeignKey,
    Integer,
    String,
    Text,
    create_engine,
    event,
)
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
    relationship,
    sessionmaker,
)

from config import get_config


class Base(DeclarativeBase):
    pass


class Camera(Base):
    __tablename__ = "cameras"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    cloud_id: Mapped[Optional[int]] = mapped_column(Integer, unique=True, nullable=True)
    name: Mapped[str] = mapped_column(String(64), nullable=False)
    rtsp_url: Mapped[str] = mapped_column(Text, nullable=False)
    enabled: Mapped[bool] = mapped_column(Boolean, default=True)
    fps_limit: Mapped[int] = mapped_column(Integer, default=30)
    process_every_n_frames: Mapped[int] = mapped_column(Integer, default=3)
    pending_sync: Mapped[bool] = mapped_column(Boolean, default=False)
    sync_failed_at: Mapped[Optional[datetime]] = mapped_column(DateTime)
    sync_retry_count: Mapped[int] = mapped_column(Integer, default=0)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, default=datetime.utcnow, onupdate=datetime.utcnow
    )

    rois: Mapped[List["ROI"]] = relationship(
        "ROI", back_populates="camera", cascade="all, delete-orphan"
    )
    status: Mapped[Optional["CameraStatus"]] = relationship(
        "CameraStatus", back_populates="camera", uselist=False
    )
    alarms: Mapped[List["Alarm"]] = relationship(
        "Alarm", back_populates="camera", cascade="all, delete-orphan"
    )


class CameraStatus(Base):
    __tablename__ = "camera_status"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    camera_id: Mapped[int] = mapped_column(
        Integer, ForeignKey("cameras.id"), unique=True, nullable=False
    )
    is_running: Mapped[bool] = mapped_column(Boolean, default=False)
    last_frame_time: Mapped[Optional[datetime]] = mapped_column(DateTime)
    fps: Mapped[float] = mapped_column(Float, default=0.0)
    error_message: Mapped[Optional[str]] = mapped_column(Text)
    last_check_time: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)

    camera: Mapped["Camera"] = relationship("Camera", back_populates="status")


class ROI(Base):
    __tablename__ = "rois"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    cloud_id: Mapped[Optional[int]] = mapped_column(Integer, unique=True, nullable=True)
    camera_id: Mapped[int] = mapped_column(
        Integer, ForeignKey("cameras.id"), nullable=False
    )
    roi_id: Mapped[str] = mapped_column(String(64), nullable=False)
    name: Mapped[str] = mapped_column(String(128), nullable=False)
    roi_type: Mapped[str] = mapped_column(String(32), nullable=False)
    points: Mapped[str] = mapped_column(Text, nullable=False)
    rule_type: Mapped[str] = mapped_column(String(32), nullable=False)
    direction: Mapped[Optional[str]] = mapped_column(String(32))
    stay_time: Mapped[Optional[int]] = mapped_column(Integer)
    enabled: Mapped[bool] = mapped_column(Boolean, default=True)
    threshold_sec: Mapped[int] = mapped_column(Integer, default=360)
    confirm_sec: Mapped[int] = mapped_column(Integer, default=30)
    return_sec: Mapped[int] = mapped_column(Integer, default=5)
    pending_sync: Mapped[bool] = mapped_column(Boolean, default=False)
    sync_version: Mapped[int] = mapped_column(Integer, default=0)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, default=datetime.utcnow, onupdate=datetime.utcnow
    )

    camera: Mapped["Camera"] = relationship("Camera", back_populates="rois")


class Alarm(Base):
    __tablename__ = "alarms"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    cloud_id: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
    camera_id: Mapped[int] = mapped_column(
        Integer, ForeignKey("cameras.id"), nullable=False
    )
    roi_id: Mapped[Optional[str]] = mapped_column(String(64))
    event_type: Mapped[str] = mapped_column(String(32), nullable=False)
    confidence: Mapped[float] = mapped_column(Float, default=0.0)
    snapshot_path: Mapped[Optional[str]] = mapped_column(Text)
    region_data: Mapped[Optional[str]] = mapped_column(Text)
    upload_status: Mapped[str] = mapped_column(String(32), default='pending_upload')
    upload_retry_count: Mapped[int] = mapped_column(Integer, default=0)
    error_message: Mapped[Optional[str]] = mapped_column(Text)
    llm_checked: Mapped[bool] = mapped_column(Boolean, default=False)
    llm_result: Mapped[Optional[str]] = mapped_column(Text)
    processed: Mapped[bool] = mapped_column(Boolean, default=False)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)

    camera: Mapped["Camera"] = relationship("Camera", back_populates="alarms")


_engine = None
_SessionLocal = None


def get_engine():
    global _engine
    if _engine is None:
        config = get_config()
        _engine = create_engine(
            config.database.url,
            echo=config.database.echo,
            pool_pre_ping=True,
            pool_recycle=3600,
        )
    return _engine


def get_session_factory():
    global _SessionLocal
    if _SessionLocal is None:
        _SessionLocal = sessionmaker(
            autocommit=False, autoflush=False, bind=get_engine()
        )
    return _SessionLocal


def get_db():
    SessionLocal = get_session_factory()
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db():
    config = get_config()
    engine = get_engine()

    Base.metadata.create_all(bind=engine)

    if config.database.dialect == "sqlite":
        for table in [Camera, ROI, Alarm]:
            table.__table__.create(engine, checkfirst=True)


def reset_engine():
    global _engine, _SessionLocal
    _engine = None
    _SessionLocal = None