File size: 2,624 Bytes
f02c5b9
 
 
 
 
3b1f683
f02c5b9
3b1f683
f02c5b9
 
3b1f683
 
 
 
 
 
 
 
 
 
 
 
 
 
f02c5b9
 
 
 
3b1f683
 
 
 
 
 
 
ed1ea51
3b1f683
 
 
f02c5b9
 
ed1ea51
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8b37702
 
 
 
5ebe979
8b37702
 
 
 
 
 
 
 
 
ed1ea51
f02c5b9
 
3b1f683
 
 
f02c5b9
 
 
 
 
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
72
73
74
75
76
77
78
79
80
81
82
83
"""Database connection and session management."""
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from app.config import get_settings
from app.models import Base
import logging

logger = logging.getLogger(__name__)
settings = get_settings()

# Create engine with connection pooling
try:
    engine = create_engine(
        settings.DATABASE_URL,
        pool_pre_ping=True,
        pool_size=10,
        max_overflow=20
    )
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    logger.info("Database engine created successfully")
except Exception as e:
    logger.warning(f"Failed to create database engine: {e}")
    engine = None
    SessionLocal = None


def init_db():
    """Initialize database tables."""
    if engine is None:
        logger.warning("Database not configured - skipping initialization")
        return
    
    try:
        Base.metadata.create_all(bind=engine)
        logger.info("Database tables created successfully")
        _run_migrations()
    except Exception as e:
        logger.error(f"Failed to initialize database: {e}")
        raise


def _run_migrations():
    """Run any pending column migrations."""
    with engine.connect() as conn:
        # Add file_content column if it doesn't exist
        try:
            conn.execute(
                __import__('sqlalchemy').text(
                    "ALTER TABLE documents ADD COLUMN IF NOT EXISTS file_content BYTEA"
                )
            )
            conn.commit()
            logger.info("Migration: file_content column ensured")
        except Exception as e:
            logger.warning(f"Migration note: {e}")
            conn.rollback()

        # Add status and error_message columns
        for col_sql in [
            "ALTER TABLE documents ADD COLUMN IF NOT EXISTS status VARCHAR DEFAULT 'done'",
            "ALTER TABLE documents ADD COLUMN IF NOT EXISTS error_message TEXT",
            "ALTER TABLE folders ADD COLUMN IF NOT EXISTS parent_id VARCHAR REFERENCES folders(id)",
        ]:
            try:
                conn.execute(__import__('sqlalchemy').text(col_sql))
                conn.commit()
            except Exception as e:
                logger.warning(f"Migration note: {e}")
                conn.rollback()
        logger.info("Migration: status/error_message columns ensured")


def get_db() -> Session:
    """Dependency for getting database session."""
    if SessionLocal is None:
        raise RuntimeError("Database not configured")
    
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()