db.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  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. # Try to autoload a local .env file if python-dotenv is available.
  12. # Python does not read .env files automatically, so without this the app
  13. # will only see environment variables explicitly set in the process.
  14. try:
  15. from dotenv import load_dotenv
  16. # Prefer project-level .env (FastApi/.env) then repository root .env
  17. base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
  18. env_paths = [os.path.join(base_dir, ".env"), os.path.join(os.path.dirname(base_dir), ".env")]
  19. for p in env_paths:
  20. if os.path.exists(p):
  21. load_dotenv(p, override=False)
  22. break
  23. except Exception:
  24. # dotenv not installed; skip silently
  25. pass
  26. from sqlalchemy import create_engine
  27. from sqlalchemy.engine import URL
  28. from sqlalchemy.orm import sessionmaker, declarative_base
  29. # Read connection info from environment variables
  30. DATABASE_URL = os.getenv("DATABASE_URL")
  31. DB_USER = os.getenv("DB_USER", "root")
  32. DB_PASS = os.getenv("DB_PASS")
  33. DB_HOST = os.getenv("DB_HOST", "192.168.1.177")
  34. DB_PORT = int(os.getenv("DB_PORT", "3306"))
  35. DB_NAME = os.getenv("DB_NAME", "wishing_tree")
  36. # Query params for pymysql / SQLAlchemy. Keep only supported options (charset).
  37. QUERY = {
  38. "charset": "utf8mb4",
  39. }
  40. if DATABASE_URL:
  41. # If the user set a full SQLAlchemy URL, prefer it
  42. engine_url = DATABASE_URL
  43. else:
  44. # Build a SQLAlchemy URL object (handles proper quoting)
  45. engine_url = URL.create(
  46. drivername="mysql+pymysql",
  47. username=DB_USER,
  48. password=DB_PASS,
  49. host=DB_HOST,
  50. port=DB_PORT,
  51. database=DB_NAME,
  52. query=QUERY,
  53. )
  54. # Create engine and session factory
  55. # Tune pool parameters as needed for your environment
  56. engine = create_engine(
  57. engine_url,
  58. pool_pre_ping=True,
  59. pool_size=10,
  60. max_overflow=20,
  61. )
  62. SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
  63. Base = declarative_base()
  64. # Dependency for FastAPI endpoints
  65. def get_db() -> Generator:
  66. db = SessionLocal()
  67. try:
  68. yield db
  69. finally:
  70. db.close()