36 lines
1.1 KiB
Python
36 lines
1.1 KiB
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":
|
|
try:
|
|
print("Going with mysql")
|
|
global mysql_pool
|
|
mysql_pool = await aiomysql.create_pool(host=mysql_host,user=mysql_user,password=mysql_pwd,db=mysql_db_name, minsize=1, maxsize=10)
|
|
yield
|
|
finally:
|
|
mysql_pool.close()
|
|
await mysql_pool.wait_closed()
|
|
elif mode == "sqlite":
|
|
try:
|
|
print("Going with sqlite")
|
|
yield
|
|
finally: pass # no action required now
|
|
else:
|
|
raise EnvironmentError("DB_MODE must either be mysql or sqlite") |