diff --git a/app/main.py b/app/main.py index 08de4b8..2bb4b69 100644 --- a/app/main.py +++ b/app/main.py @@ -5,29 +5,26 @@ from app.utils import generate_short_code from dotenv import load_dotenv import os import sqlite3 +from contextlib import asynccontextmanager -load_dotenv() - -app = FastAPI() - -DB_PATH = "urls.db" - -def get_db_connection(): - conn = sqlite3.connect(DB_PATH) - return conn - -# URL schema -class URLRequest(BaseModel): - url: str - -@app.on_event("startup") -def startup(): - conn = get_db_connection() +@asynccontextmanager +async def lifespan(app: FastAPI): + conn = sqlite3.connect('urls.db') conn.execute('''CREATE TABLE IF NOT EXISTS urls (id INTEGER PRIMARY KEY, short_code TEXT, original_url TEXT, clicks INTEGER)''') conn.commit() + yield conn.close() +app = FastAPI(lifespan=lifespan) + +def get_db_connection(): + conn = sqlite3.connect('urls.db') + return conn + +class URLRequest(BaseModel): + url: str + @app.get("/") def home(): return {"message": "URL Shortener API"} @@ -35,9 +32,8 @@ def home(): @app.post("/shorten") def shorten_url(request: URLRequest): short_code = generate_short_code() - conn = get_db_connection() - + while conn.execute("SELECT 1 FROM urls WHERE short_code = ?", (short_code,)).fetchone(): short_code = generate_short_code() @@ -52,12 +48,11 @@ def shorten_url(request: URLRequest): def redirect_to_url(short_code: str): conn = get_db_connection() url_data = conn.execute("SELECT original_url, clicks FROM urls WHERE short_code = ?", (short_code,)).fetchone() - + if url_data is None: raise HTTPException(status_code=404, detail="Shortened URL not found") - - original_url, clicks = url_data + original_url, clicks = url_data conn.execute("UPDATE urls SET clicks = ? WHERE short_code = ?", (clicks + 1, short_code)) conn.commit() conn.close() diff --git a/urls.db b/urls.db index 7ef956d..d2b6c00 100644 Binary files a/urls.db and b/urls.db differ