From 46bc71f16a70912f535c9d4924be08f7c476d681 Mon Sep 17 00:00:00 2001
From: Rainy21 <simonzheng950@gmail.com>
Date: Tue, 10 Mar 2026 23:45:53 -0400
Subject: [PATCH] add database but need to test it

---
 db/crud.py       | 18 ++++++++++++++++++
 db/database.py   |  9 +++++++++
 db/models.py     | 14 ++++++++++++++
 db/schemas.py    |  9 +++++++++
 main.py          | 39 +++++++++++++++++++++++++--------------
 requirements.txt |  1 +
 6 files changed, 76 insertions(+), 14 deletions(-)
 create mode 100644 db/crud.py
 create mode 100644 db/database.py
 create mode 100644 db/models.py
 create mode 100644 db/schemas.py

diff --git a/db/crud.py b/db/crud.py
new file mode 100644
index 0000000..214d65d
--- /dev/null
+++ b/db/crud.py
@@ -0,0 +1,18 @@
+from pydantic import Secret
+import sqlalchemy.orm import Session
+from . import models
+
+def create_url(db:Session, short_code:str , long_url:str):
+    db_url = models.URL(
+        short_code=short_code,
+        long_url=long_url
+    )
+    db.add(db_url)
+    db.commit()
+    db.refresh(db_url)
+    return db_url
+
+def get_url(db:Session , short_code:str):
+    return db.query(models.URL).filter(
+        models.URL.short_code == short_code
+    ).first()
diff --git a/db/database.py b/db/database.py
new file mode 100644
index 0000000..12f4d82
--- /dev/null
+++ b/db/database.py
@@ -0,0 +1,9 @@
+from sqlalchemy import create_engine, engine
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+DATABASE_URL = "sqlite:///./urls.db"
+
+engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+Base = declarative_base()
diff --git a/db/models.py b/db/models.py
new file mode 100644
index 0000000..87b4ebf
--- /dev/null
+++ b/db/models.py
@@ -0,0 +1,14 @@
+from sqlalchemy import Column, Integer, String, DateTime
+from datetime import datetime
+
+from sqlalchemy.sql.traversals import ColIdentityComparatorStrategy
+from .database import Base
+
+
+class URL(Base):
+    __tablename__ = "urls"
+
+    id = Column(Integer, primary_key=True, index=True)
+    short_code = Column(String, unique=True, index=True)
+    long_url = Column(String)
+    create_at = Column(DateTime, default=datetime.utcnow)
diff --git a/db/schemas.py b/db/schemas.py
new file mode 100644
index 0000000..4a0641b
--- /dev/null
+++ b/db/schemas.py
@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+
+class URLCreate(BaseModel):
+    long_url: str
+
+
+class URLResponse(BaseModel):
+    short_url: str
diff --git a/main.py b/main.py
index 004e1c7..4ac2412 100644
--- a/main.py
+++ b/main.py
@@ -1,17 +1,24 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
 from fastapi.responses import RedirectResponse
-from pydantic import BaseModel
+
+# the sqlalchemy
+from sqlalchemy.orm import Session
 import string
 import random
 
+from db.database import SessionLocal, engine
+from db import models, crud, schemas
+
+models.Base.metadata.create_call(bind=engine)
 app = FastAPI()
 
-# temp storage
-url_db = {}
 
-
-class URLRequest(BaseModel):
-    long_url: str
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
 
 
 def generate_short_code(length=6):
@@ -19,15 +26,19 @@ def generate_short_code(length=6):
     return "".join(random.choice(chars) for _ in range(length))
 
 
-@app.post("/shorten")
-def shorten_url(request: URLRequest):
+@app.post("/shorten", response_model=schemas.URLResponse)
+def shorten_url(request: schemas.URLCreate, db: Session = Depends(get_db)):
     short_code = generate_short_code()
-    url_db[short_code] = request.long_url
+
+    crud.create_url(db, short_code, request.long_url)
+
     return {"short_url": f"http://localhost:8000/{short_code}"}
 
 
 @app.get("/{short_code}")
-def redirect_url(short_code: str):
-    if short_code in url_db:
-        return RedirectResponse(url=url_db[short_code])
-    raise HTTPException(status_code=404, detail="URL not found")
+def redirect_url(short_code: str, db: Session = Depends(get_db)):
+    db_url = crud.get_url(db, short_code)
+
+    if not db_url:
+        raise HTTPException(status_code=404, detail="URL not found")
+    return {"long_url": db_url.long_url}
diff --git a/requirements.txt b/requirements.txt
index 97dc7cd..89bae1f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 fastapi
 uvicorn
+sqlalchemy