src/core/database.py

import uuid
from datetime import datetime
from typing import AsyncGenerator
import os

from fastapi import Depends
from fastapi_users.db import SQLAlchemyBaseUserTableUUID, SQLAlchemyUserDatabase
from sqlalchemy import Column, String, Integer, DateTime, ForeignKey, Boolean
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.ext.declarative import DeclarativeMeta, declarative_base
from sqlalchemy.orm import relationship

# Use the DATABASE_URL from environment variables, with a fallback to SQLite for local development
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./knowledge_assistant.db")

Base: DeclarativeMeta = declarative_base()


class User(SQLAlchemyBaseUserTableUUID, Base):
    """User model extending FastAPI-Users base table"""
    __tablename__ = "users"
    
    # Use the standard UUID type for PostgreSQL
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    
    # Additional fields beyond the base user table
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class DocumentMetadata(Base):
    """Document metadata model for tracking user uploads"""
    __tablename__ = "documents"
    
    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    user_id = Column(UUID(as_uuid=True), ForeignKey("users.id", ondelete="CASCADE"), nullable=False)
    filename = Column(String(255), nullable=False)
    original_size = Column(Integer)
    chunks_count = Column(Integer)
    upload_date = Column(DateTime, default=datetime.utcnow)
    file_hash = Column(String(64), unique=True)  # Prevent duplicate uploads
    
    # Relationship to user
    user = relationship("User", back_populates="documents")


# Add relationship to User model
User.documents = relationship("DocumentMetadata", back_populates="user", cascade="all, delete-orphan")


# Database engine and session configuration
engine = create_async_engine(DATABASE_URL, connect_args={"statement_cache_size": 0})
async_session_maker = async_sessionmaker(engine, expire_on_commit=False)


async def create_db_and_tables():
    """Create database tables"""
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    """Get async database session"""
    async with async_session_maker() as session:
        yield session


async def get_user_db(session: AsyncSession = Depends(get_async_session)):
    """Get user database instance for FastAPI-Users"""
    yield SQLAlchemyUserDatabase(session, User)