Update services/auth.py

services/auth.py

@@ -1,11 +1,8 @@
-# services/auth.py
-
 from passlib.context import CryptContext
 from sqlmodel import Session, select
 from typing import Optional
 
-# Import User directly to avoid potential circularity with models.__init__ if not careful
-from models.user import User, UserCreate
+from models.user import User, UserCreate # Direct import
 from models.db import get_session_context
 from services.logger import app_logger
 
@@ -17,21 +14,29 @@ def verify_password(plain_password: str, hashed_password: str) -> bool:
 def get_password_hash(password: str) -> str:
     return pwd_context.hash(password)
 
-def get_user(db: Session, username: str) -> Optional[User]:
+def get_user_by_username(db: Session, username: str) -> Optional[User]:
+    """Fetches a user by username."""
     statement = select(User).where(User.username == username)
     user = db.exec(statement).first()
     return user
 
 def create_user_in_db(user_data: UserCreate) -> Optional[User]:
+    """
+    Creates a new user in the database.
+    Returns the User object with basic attributes loaded, suitable for immediate
+    use even after the session creating it has closed.
+    """
     hashed_password = get_password_hash(user_data.password)
     
-    # Create the Python object
-    db_user = User(
+    # Create the Python object instance.
+    # Initialize relationships like chat_sessions to an empty list if appropriate,
+    # as they won't be populated from the DB for a new user yet.
+    new_user_instance = User(
         username=user_data.username,
         hashed_password=hashed_password,
         email=user_data.email,
         full_name=user_data.full_name,
-        chat_sessions=[] # Initialize relationships if they are expected to be non-None
+        chat_sessions=[] # Good practice to initialize for new instances
     )
     
     created_user_id: Optional[int] = None
@@ -39,61 +44,89 @@ def create_user_in_db(user_data: UserCreate) -> Optional[User]:
     try:
         with get_session_context() as db:
             # Check if user already exists
-            existing_user_check = db.exec(select(User).where(User.username == user_data.username)).first()
-            if existing_user_check:
+            existing_user = get_user_by_username(db, user_data.username)
+            if existing_user:
                 app_logger.warning(f"User {user_data.username} already exists.")
-                return None
+                return None # User already exists
             
-            db.add(db_user)
-            db.commit()  # Commit to save the user and generate ID
-            db.refresh(db_user) # Refresh to get all attributes, especially ID
-            created_user_id = db_user.id # Store the ID
+            db.add(new_user_instance)
+            db.commit() # Commit to save the user and allow DB to generate ID
+            db.refresh(new_user_instance) # Refresh to get all attributes, especially the DB-generated ID
             
-            # Explicitly load the username to be safe, though it should be loaded.
-            # This is more of a "just in case" step.
-            _ = db_user.username 
-            _ = db_user.email # if you plan to use it immediately after
+            # Store the ID so we can re-fetch a "clean" instance if needed,
+            # or to ensure we are working with the persisted state.
+            created_user_id = new_user_instance.id
             
-        # After session closes, db_user is detached.
-        # If we need to return the full User object, and it might be used in a way
-        # that requires re-attaching or accessing lazy-loaded fields, this can be an issue.
-        # For simple display like username, it should be fine.
-        # The error suggests something is trying to "refresh" it.
+            # To ensure the returned object is safe for use after this session closes,
+            # especially for attributes that might be displayed immediately,
+            # we can load them explicitly or re-fetch.
+            # For simple attributes, refresh should be enough.
+            # We will return a freshly fetched instance to be absolutely sure.
+
+    except Exception as e:
+        app_logger.error(f"Error during database operation for user {user_data.username}: {e}")
+        return None # DB operation failed
 
-        # To be absolutely safe for the immediate use case (displaying username),
-        # we can fetch a "clean" version of the user by ID in a new, short-lived session,
-        # though this is usually overkill for just displaying a username.
-        # This ensures the returned object is "fresh" but also detached.
-        if created_user_id is not None:
+    # If user was created, fetch a fresh instance to return.
+    # This new instance will be detached but will have its basic attributes loaded.
+    if created_user_id is not None:
+        try:
             with get_session_context() as db:
-                # Re-fetch the user to ensure it's a clean, detached object
-                # with its basic attributes loaded.
-                user_to_return = db.get(User, created_user_id)
-                if user_to_return:
-                    # Access attributes needed immediately to ensure they are loaded
-                    # before the session closes again.
-                    _ = user_to_return.username
-                    _ = user_to_return.email
-                    return user_to_return # This user object is now detached.
-        return None # Should not happen if creation was successful
+                # .get() is efficient for fetching by primary key
+                final_user_to_return = db.get(User, created_user_id)
+                if final_user_to_return:
+                    # "Touch" attributes that will be used immediately after this function returns
+                    # to ensure they are loaded before this (new) session closes.
+                    _ = final_user_to_return.id
+                    _ = final_user_to_return.username
+                    _ = final_user_to_return.email
+                    _ = final_user_to_return.full_name
+                    # Do NOT try to access final_user_to_return.chat_sessions here unless you
+                    # intend to load them, which might be a heavy operation.
+                    # They will be lazy-loaded when accessed within an active session later.
+                    return final_user_to_return
+        except Exception as e:
+            app_logger.error(f"Error re-fetching created user {created_user_id}: {e}")
+            return None # Failed to re-fetch
+
+    return None # Fallback if user was not created or couldn't be re-fetched
 
-    except Exception as e:
-        app_logger.error(f"Error creating user {user_data.username}: {e}")
-        return None
 
 def authenticate_user(username: str, password: str) -> Optional[User]:
-    with get_session_context() as db:
-        user = get_user(db, username)
-        if not user:
-            return None
-        if not verify_password(password, user.hashed_password):
-            return None
-        # Before returning, ensure any attributes needed by the caller are loaded
-        # if they might be lazy-loaded and the user object will be used after this session.
-        # For st.session_state.authenticated_user, this is important.
-        _ = user.id
-        _ = user.username
-        _ = user.email # etc.
-        # For relationships, you might need to configure eager loading (e.g. selectinload)
-        # or accept that they will be lazy-loaded (requiring a session) or not accessible.
-        return user
+    """
+    Authenticates a user.
+    Returns the User object with basic attributes loaded.
+    """
+    try:
+        with get_session_context() as db:
+            user = get_user_by_username(db, username)
+            if not user:
+                return None # User not found
+            if not verify_password(password, user.hashed_password):
+                return None # Invalid password
+            
+            # User is authenticated. Before returning the user object,
+            # ensure any attributes that will be immediately accessed by the caller
+            # (e.g., when setting st.session_state.authenticated_user) are loaded.
+            # This helps avoid DetachedInstanceError if Streamlit or other parts
+            # of the app try to access these attributes after this session closes.
+            
+            user_id_to_return = user.id # Get ID to re-fetch
+
+        # Re-fetch the user in a new session context to return a "clean" detached object
+        # This is a robust way to prevent issues with detached instances being used across session boundaries.
+        if user_id_to_return is not None:
+            with get_session_context() as db:
+                authenticated_user = db.get(User, user_id_to_return)
+                if authenticated_user:
+                    # "Touch" attributes
+                    _ = authenticated_user.id
+                    _ = authenticated_user.username
+                    _ = authenticated_user.email
+                    # Do not touch relationships unless intended for eager load here.
+                    return authenticated_user
+        return None # Should not happen if authenticated
+
+    except Exception as e:
+        app_logger.error(f"Error during authentication for user {username}: {e}")
+        return None