db.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. """SQLAlchemy database setup for MySQL.
  2. Reads connection info from environment variables to avoid committing secrets.
  3. - DATABASE_URL: full SQLAlchemy URL (optional)
  4. - DB_USER, DB_PASS, DB_HOST, DB_PORT, DB_NAME: used when DATABASE_URL is not set
  5. Example to set DATABASE_URL locally (Windows PowerShell):
  6. $env:DATABASE_URL = "mysql+pymysql://root:%40%40qinys12346..@192.168.101.7:3306/wishing_tree?charset=utf8mb4&serverTimezone=Asia/Shanghai&use_unicode=true&characterEncoding=utf-8"
  7. Install dependencies: pip install sqlalchemy pymysql
  8. """
  9. import os
  10. from typing import Generator
  11. from sqlalchemy import create_engine
  12. from sqlalchemy.engine import URL
  13. from sqlalchemy.orm import sessionmaker, declarative_base
  14. # Read connection info from environment variables
  15. DATABASE_URL = os.getenv("DATABASE_URL")
  16. DB_USER = os.getenv("DB_USER", "root")
  17. DB_PASS = os.getenv("DB_PASS")
  18. DB_HOST = os.getenv("DB_HOST", "192.168.1.177")
  19. DB_PORT = int(os.getenv("DB_PORT", "3306"))
  20. DB_NAME = os.getenv("DB_NAME", "wishing_tree")
  21. # Query params for pymysql / SQLAlchemy. Keep only supported options (charset).
  22. QUERY = {
  23. "charset": "utf8mb4",
  24. }
  25. if DATABASE_URL:
  26. # If the user set a full SQLAlchemy URL, prefer it
  27. engine_url = DATABASE_URL
  28. else:
  29. # Build a SQLAlchemy URL object (handles proper quoting)
  30. engine_url = URL.create(
  31. drivername="mysql+pymysql",
  32. username=DB_USER,
  33. password=DB_PASS,
  34. host=DB_HOST,
  35. port=DB_PORT,
  36. database=DB_NAME,
  37. query=QUERY,
  38. )
  39. # Create engine and session factory
  40. # Tune pool parameters as needed for your environment
  41. engine = create_engine(
  42. engine_url,
  43. pool_pre_ping=True,
  44. pool_size=10,
  45. max_overflow=20,
  46. )
  47. SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
  48. Base = declarative_base()
  49. # Dependency for FastAPI endpoints
  50. def get_db() -> Generator:
  51. db = SessionLocal()
  52. try:
  53. yield db
  54. finally:
  55. db.close()