From cfb7ddedd3fc6b4e5dc497511a8c8f6040e40c94 Mon Sep 17 00:00:00 2001
From: mxr612 <work@mxr612.io>
Date: Tue, 17 Jun 2025 15:44:10 +0800
Subject: [PATCH] feat: implement user identity middleware for session
 management

- Added UserIdentityMiddleware to manage user sessions by generating a unique user ID stored in a cookie.
- Implemented logic to update the last seen timestamp for existing users in the database.
- Enhanced ScaleResult model to associate responses with users, improving data tracking and user experience.
---
 app.py      | 44 +++++++++++++++++++++++++++++++++++++++++++-
 database.py | 27 +++++++++++++++++++++++++--
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/app.py b/app.py
index 4245460..3d046b3 100644
--- a/app.py
+++ b/app.py
@@ -9,7 +9,7 @@ import uvicorn
 from datetime import datetime
 from xml.etree import ElementTree as ET
 from sqlalchemy.orm import Session
-from database import get_db, ScaleResult
+from database import get_db, ScaleResult, new_user, User
 import geoip2.database
 from datetime import datetime, UTC
 import csv
@@ -100,8 +100,49 @@ class LanguageMiddleware(BaseHTTPMiddleware):
         
         return response
 
+class UserIdentityMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp):
+        super().__init__(app)
+
+    async def dispatch(self, request: Request, call_next):
+        # Get user_id from cookie
+        user_id = request.cookies.get("user_id")
+        
+        # If no user_id in cookie, generate a new one
+        if not user_id:
+            user_id = new_user()
+        else:
+            # Update last_seen for existing user
+            db = next(get_db())
+            try:
+                user = db.query(User).filter(User.id == int(user_id)).first()
+                if user:
+                    user.last_seen = datetime.now(UTC)
+                    db.commit()
+            finally:
+                db.close()
+        
+        # Add user_id to request state
+        request.state.user_id = user_id
+        
+        # Continue processing the request
+        response = await call_next(request)
+        
+        # Set cookie if it's not already set
+        if not request.cookies.get("user_id"):
+            response.set_cookie(
+                key="user_id",
+                value=user_id,
+                max_age=None,  # Cookie will never expire
+                httponly=True,
+                samesite="lax"
+            )
+        
+        return response
+
 app = FastAPI()
 app.add_middleware(LanguageMiddleware)
+app.add_middleware(UserIdentityMiddleware)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = {}
 for lang in os.listdir("templates"):
@@ -200,6 +241,7 @@ async def result(request: Request, scale_id: str, db: Session = Depends(get_db))
             location = get_location_from_ip(ip)# Get location information
             db_response = ScaleResult(
                 scale_id=scale_id,
+                user_id=request.state.user_id,
                 user_agent=request.headers.get("user-agent", "Unknown"),
                 ip_address=ip,
                 location=location,
diff --git a/database.py b/database.py
index 5a20eff..a9f8821 100644
--- a/database.py
+++ b/database.py
@@ -1,7 +1,8 @@
-from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, ForeignKey, JSON
+from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, ForeignKey, JSON, func
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, relationship
 import json
+from datetime import datetime, UTC
 
 SQLALCHEMY_DATABASE_URL = "sqlite:///./psychoscales.db"
 
@@ -14,9 +15,17 @@ SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 Base = declarative_base()
 
+class User(Base):
+    __tablename__ = "users"
+    id = Column(Integer, primary_key=True, index=True)
+    created_at = Column(DateTime)
+    last_seen = Column(DateTime)
+    responses = relationship("ScaleResult", back_populates="user")
+
 class ScaleResult(Base):
     __tablename__ = "responses"
     id = Column(Integer, primary_key=True, index=True)
+    user_id = Column(Integer, ForeignKey("users.id"))
     scale_id = Column(String, index=True)
     user_agent = Column(String)
     ip_address = Column(String)
@@ -25,6 +34,7 @@ class ScaleResult(Base):
     sum_response = Column(JSON)
     avg_response = Column(JSON)
     created_at = Column(DateTime)
+    user = relationship("User", back_populates="responses")
 
 # Create tables
 Base.metadata.create_all(bind=engine)
@@ -35,4 +45,17 @@ def get_db():
     try:
         yield db
     finally:
-        db.close() 
\ No newline at end of file
+        db.close()
+
+def new_user() -> int:
+    db = SessionLocal()
+    try:
+        with db.begin():
+            user = User()
+            user.last_seen = user.created_at = datetime.now(UTC)
+            db.add(user)
+            db.flush()
+            return user.id
+    finally:
+        db.close()
+