File size: 1,998 Bytes
603a641 5840851 ffd99f4 5840851 ffd99f4 603a641 5840851 603a641 5840851 603a641 5840851 603a641 5840851 603a641 5840851 381775d 5840851 | 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 | import os
from sqlalchemy import create_engine, event, text
from sqlalchemy.orm import sessionmaker
from dotenv import load_dotenv
from models import Base
load_dotenv()
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./stocking.db")
_is_postgres = DATABASE_URL.startswith("postgresql")
connect_args = {"sslmode": "require"} if _is_postgres else {"check_same_thread": False}
engine = create_engine(DATABASE_URL, connect_args=connect_args)
# SQLite 전용: WAL 모드 (PostgreSQL에는 적용 안 함)
if not _is_postgres:
@event.listens_for(engine, "connect")
def set_wal_mode(dbapi_connection, connection_record):
"""WAL 모드: 스케줄러 write + API read 동시 처리 가능."""
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA journal_mode=WAL")
cursor.close()
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def init_db():
# watchlist 테이블 스키마 마이그레이션: SQLite에만 적용 (PostgreSQL은 PRAGMA 없음)
if not _is_postgres:
with engine.connect() as conn:
cols = [row[1] for row in conn.execute(text("PRAGMA table_info(watchlist)")).fetchall()]
if cols and "user_id" not in cols:
conn.execute(text("DROP TABLE watchlist"))
conn.commit()
Base.metadata.create_all(bind=engine)
# annual_data 인덱스 마이그레이션: create_all은 기존 테이블에 인덱스를 추가하지 않으므로 직접 생성
with engine.connect() as conn:
conn.execute(text(
"CREATE INDEX IF NOT EXISTS ix_annual_data_ticker_year_is_estimate "
"ON annual_data (ticker, year, is_estimate)"
))
conn.execute(text(
"CREATE INDEX IF NOT EXISTS ix_annual_data_ticker_is_estimate "
"ON annual_data (ticker, is_estimate)"
))
conn.commit()
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
|