Update app.py

app.py

@@ -42,6 +42,7 @@ from fastapi.templating import Jinja2Templates
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.security import APIKeyHeader, OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from starlette.middleware.cors import CORSMiddleware
+import asyncpg
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(filename)s - %(lineno)d - %(message)s')
 logger = logging.getLogger(__name__)
@@ -54,17 +55,6 @@ ACCESS_TOKEN_EXPIRE_MINUTES = 30
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
-conn = sqlite3.connect('users.db', check_same_thread=False)
-cursor = conn.cursor()
-cursor.execute('''
-    CREATE TABLE IF NOT EXISTS users (
-        username TEXT PRIMARY KEY,
-        email TEXT UNIQUE,
-        hashed_password TEXT
-    )
-''')
-conn.commit()
-
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 API_KEY = os.getenv("API_KEY")
 api_key_header = APIKeyHeader(name="X-API-Key")
@@ -77,6 +67,7 @@ HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 TEMP_DIR = "/tmp"
 STATIC_DIR = "static"
 TEMPLATES = Jinja2Templates(directory="templates")
+DATABASE_URL = os.getenv("DATABASE_URL")
 
 app = FastAPI()
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
@@ -568,73 +559,56 @@ class Token(BaseModel):
     access_token: str
     token_type: str
 
+async def get_db():
+    async with asyncpg.create_pool(DATABASE_URL) as pool:
+        async with pool.acquire() as conn:
+            yield conn
+
+async def authenticate_user(username, password, conn):
+    row = await conn.fetchrow("SELECT * FROM users WHERE username = $1", username)
+    if row is not None and pwd_context.verify(password, row["hashed_password"]):
+        return {"username": username}
+    return None
+
 @app.post("/token", response_model=Token)
-async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
-    user = authenticate_user(form_data.username, form_data.password)
+async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), conn = Depends(get_db)):
+    user = await authenticate_user(form_data.username, form_data.password, conn)
     if not user:
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect username or password", headers={"WWW-Authenticate": "Bearer"})
     access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
     access_token = create_access_token(data={"sub": user["username"]}, expires_delta=access_token_expires)
     return {"access_token": access_token, "token_type": "bearer"}
 
-def authenticate_user(username: str, password: str):
-    cursor.execute("SELECT * FROM users WHERE username = ?", (username,))
-    user = cursor.fetchone()
-    if user and pwd_context.verify(password, user[2]):
-        return {"username": username}
-    return None
-
-def create_access_token(data: Dict[str, Any], expires_delta: timedelta = None):
-    to_encode = data.copy()
-    if expires_delta:
-        expire = datetime.utcnow() + expires_delta
-    else:
-        expire = datetime.utcnow() + timedelta(minutes=15)
-    to_encode.update({"exp": expire})
-    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
-    return encoded_jwt
-
 
-@app.get("/users/me")
-async def read_users_me(current_user: str = Depends(get_current_user)):
-    return {"username": current_user}
-
-async def get_current_user(token: str = Depends(oauth2_scheme)):
+async def get_current_user(token: str = Depends(oauth2_scheme), conn = Depends(get_db)):
     credentials_exception = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"})
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
         username: str = payload.get("sub")
         if username is None:
             raise credentials_exception
+        user = await conn.fetchrow("SELECT * FROM users WHERE username = $1", username)
+        if user is None:
+            raise credentials_exception
+        return username
     except JWTError:
         raise credentials_exception
-    cursor.execute("SELECT * FROM users WHERE username = ?", (username,))
-    user = cursor.fetchone()
-    if user is None:
-        raise credentials_exception
-    return username
-
 
 @app.post("/register", response_model=User, status_code=status.HTTP_201_CREATED)
-async def create_user(user: User):
+async def create_user(user: User, conn = Depends(get_db)):
+    hashed_password = pwd_context.hash(user.password)
     try:
-        hashed_password = pwd_context.hash(user.password)
-        cursor.execute("INSERT INTO users (username, email, hashed_password) VALUES (?, ?, ?)", (user.username, user.email, hashed_password))
-        conn.commit()
+        await conn.execute("INSERT INTO users (username, email, hashed_password) VALUES ($1, $2, $3)", user.username, user.email, hashed_password)
         return user
-    except sqlite3.IntegrityError:
+    except asyncpg.exceptions.UniqueViolationError:
         raise HTTPException(status_code=400, detail="Username or email already exists")
-    except Exception as e:
-        logger.error(f"Error creating user: {e}")
-        raise HTTPException(status_code=500, detail=f"Error creating user: {e}")
 
 
 @app.put("/users/{username}", response_model=User, dependencies=[Depends(get_current_user)])
-async def update_user_data(username: str, user: User):
+async def update_user_data(username: str, user: User, conn = Depends(get_db)):
+    hashed_password = pwd_context.hash(user.password)
     try:
-        hashed_password = pwd_context.hash(user.password)
-        cursor.execute("UPDATE users SET email = ?, hashed_password = ? WHERE username = ?", (user.email, hashed_password, username))
-        conn.commit()
+        await conn.execute("UPDATE users SET email = $1, hashed_password = $2 WHERE username = $3", user.email, hashed_password, username)
         return user
     except Exception as e:
         logger.error(f"Error updating user: {e}")
@@ -642,10 +616,9 @@ async def update_user_data(username: str, user: User):
 
 
 @app.delete("/users/{username}", dependencies=[Depends(get_current_user)])
-async def delete_user_account(username: str):
+async def delete_user_account(username: str, conn = Depends(get_db)):
     try:
-        cursor.execute("DELETE FROM users WHERE username = ?", (username,))
-        conn.commit()
+        await conn.execute("DELETE FROM users WHERE username = $1", username)
         return JSONResponse({"message": "User deleted successfully."}, status_code=200)
     except Exception as e:
         logger.error(f"Error deleting user: {e}")
@@ -653,10 +626,17 @@ async def delete_user_account(username: str):
 
 
 @app.get("/users", dependencies=[Depends(get_current_user)])
-async def get_all_users_route():
-    cursor.execute("SELECT username, email FROM users")
-    users = cursor.fetchall()
-    return [{"username": user[0], "email": user[1]} for user in users]
+async def get_all_users_route(conn = Depends(get_db)):
+    rows = await conn.fetch("SELECT username, email FROM users")
+    return [{"username": row["username"], "email": row["email"]} for row in rows]
+
+
+@app.get("/users/me", dependencies=[Depends(get_current_user)])  # Requires authentication
+async def read_users_me(current_user: str = Depends(get_current_user), conn=Depends(get_db)):
+    user = await conn.fetchrow("SELECT username, email FROM users WHERE username = $1", current_user)
+    if user:
+        return {"username": user["username"], "email": user["email"]}
+    raise HTTPException(status_code=404, detail="User not found")
 
 
 @app.exception_handler(RequestValidationError)
@@ -667,5 +647,16 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
     )
 
 
+def create_access_token(data: Dict[str, Any], expires_delta: timedelta = None):
+    to_encode = data.copy()
+    if expires_delta:
+        expire = datetime.utcnow() + expires_delta
+    else:
+        expire = datetime.utcnow() + timedelta(minutes=15)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+
 if __name__ == "__main__":
     uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)