31 lines
910 B
Python
31 lines
910 B
Python
import aiomysql, aiosqlite, os
|
|
from fastapi import FastAPI
|
|
from contextlib import asynccontextmanager
|
|
|
|
# Either "mysql" or "sqlite"
|
|
mode = os.environ.get("DB_MODE", "sqlite")
|
|
# For SQLite usage
|
|
sqlite_file = os.environ.get("SQLITE_FILE","./database.db")
|
|
|
|
# For MySQL usage
|
|
mysql_pool = None
|
|
|
|
mysql_host = os.environ.get("MYSQL_HOST","127.0.0.1")
|
|
mysql_user = os.environ.get("MYSQL_USER","yt2mp3")
|
|
mysql_db_name = os.environ.get("MYSQL_DB","yt2mp3")
|
|
mysql_pwd = os.environ.get("MYSQL_PASSWORD","")
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
if mode == "mysql":
|
|
pool = await aiomysql.create_pool(host=mysql_host,user=mysql_user,password=mysql_pwd,db=mysql_db_name, minsize=1, maxsize=10)
|
|
try:
|
|
yield
|
|
finally:
|
|
pool.close()
|
|
await pool.wait_closed()
|
|
else:
|
|
try:
|
|
yield
|
|
finally: pass # no action required now |