sqlite is supported and now default

server/py/fastapi_server/main.py

@@ -8,6 +8,7 @@ import aiomysql
 import os
 import bcrypt
 import asyncio
+import aiosqlite
 import subprocess
 
 from yt_dlp import YoutubeDL
@@ -23,6 +24,14 @@ from yt_dlp.utils import DownloadError
 # So we can trust it
 
 class db:
+
+    mode = os.environ.get("DB_MODE", "sqlite")
+
+    # For SQLite usage
+
+    file = os.environ.get("SQLITE_FILE","./database.db")
+
+    # For MySQL usage
     pool = None
 
     host = os.environ.get("MYSQL_HOST","127.0.0.1")
@@ -33,21 +42,31 @@ class db:
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    db.pool = await aiomysql.create_pool(host=db.host,user=db.user,password=db.pwd,db=db.db_name, minsize=1, maxsize=10)
-    try:
-        yield
-    finally: 
-        db.pool.close()
-        await db.pool.wait_closed()
+    if db.mode == "mysql":
+        db.pool = await aiomysql.create_pool(host=db.host,user=db.user,password=db.pwd,db=db.db_name, minsize=1, maxsize=10)
+        try:
+            yield
+        finally: 
+            db.pool.close()
+            await db.pool.wait_closed()
+    else:
+        try:
+            yield
+        finally: pass # no action required now
 
 app = FastAPI(lifespan=lifespan)
 
 
-async def fetch_user_info(email: str):
-    async with db.pool.acquire() as conn:
-        async with conn.cursor(aiomysql.DictCursor) as cur:
-            await cur.execute("SELECT * FROM users WHERE email = %s", (email,))
-            return await cur.fetchone()
+async def fetch_user_info(arg: str, type: str):
+    if db.mode == "mysql":
+        async with db.pool.acquire() as conn:
+            async with conn.cursor(aiomysql.DictCursor) as cur:
+                await cur.execute("SELECT * FROM users WHERE %s = %s", (arg,type))
+                return await cur.fetchone()
+    elif db.mode == "sqlite":
+        async with aiosqlite.connect(db.file) as conn:
+            async with conn.execute("SELECT * FROM users WHERE ? = ?", (arg,type)) as cursor:
+                return await cursor.fetchone()
 
 def download_audio(video_url: str, output_folder: str) -> str:
     ydl_opts = {
@@ -132,11 +151,15 @@ async def create_user(req: Request):
     # Auto exceptions give more verbose to Nuxt.
 
     body["password"] = bcrypt.hashpw(body["password"].encode(), bcrypt.gensalt()).decode()
-    async with db.pool.acquire() as conn:
-        async with conn.cursor() as cur:
-            await cur.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (%s,%s,%s,%s)", (body.first_name, body.last_name, body.email, body.password))
-            await conn.commit()   
-
+    if db.mode == "mysql":
+        async with db.pool.acquire() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (%s,%s,%s,%s)", (body.first_name, body.last_name, body.email, body.password))
+                await conn.commit()   
+    elif db.mode == "sqlite":
+        async with aiosqlite.connect(db.file) as conn:
+            await conn.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (?, ?, ?, ?)", (body.first_name, body.last_name, body.email, body.password))
+            await conn.commit()
 
 
 @app.post("/login")
@@ -153,7 +176,7 @@ async def login_endpoint(req: Request):
             if (bcrypt.checkpw(body["password"].encode(), user_data.get("password").encode())):
                 return user_data
             else:
-                raise HTTPException(status_code=401)
+                raise HTTPException(status_code=401, detail="Wrong password")
     except KeyError:
         raise HTTPException(status_code=400)
     #except ValidationError:
@@ -168,9 +191,13 @@ async def me_get_endpoint(req: Request):
     Trusts Nuxt's backend to allow it or not.
     """
     args = dict(req.query_params)
-    #print(args)
+    me = None
     try:
-        me = await fetch_user_info(args["email"])
+    # Prefer id over e-mail
+        if "id" in args.keys(): me = fetch_user_info(args["id"], "id")
+        elif "email" in args.keys(): me = fetch_user_info(args["email"], "email")
+        else: raise HTTPException(400, "Provide either email or user id")
+    #print(args)
         if (me is None):
             raise HTTPException(404, "User doesn't exist")
         else: return me