File size: 2,734 Bytes
3b5d2e9 53ae71b 3b5d2e9 f986d90 3b5d2e9 5d1296c 3b5d2e9 5d1296c 3b5d2e9 786ac47 3b5d2e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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) |