Update database.py

database.py

@@ -5,13 +5,20 @@ from datetime import datetime
 import uuid
 import os
 
-# Use SQLite for reliable deployment on HF Spaces
-# The /data directory is persistent in HF Spaces
-DATABASE_URL = "sqlite:////data/chat_history.db"
+# Neon DB connection
+DATABASE_URL = os.getenv("DATABASE_URL")
+
+if not DATABASE_URL:
+    print("Warning: DATABASE_URL not set. Using sqlite fallback for local dev only.")
+    DATABASE_URL = "sqlite:///./local_chat.db"
+
+# Handle "postgres://" to "postgresql://" fix
+if DATABASE_URL and DATABASE_URL.startswith("postgres://"):
+    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)
 
 # Create engine
-# check_same_thread=False is needed for SQLite with FastAPI
-engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
+# sslmode is usually handled in the connection string itself
+engine = create_engine(DATABASE_URL, pool_pre_ping=True)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 Base = declarative_base()
 
@@ -40,12 +47,6 @@ class Message(Base):
 # Database functions
 def init_db():
     """Create all tables"""
-    # Ensure /data directory exists (for HF Spaces)
-    if DATABASE_URL.startswith("sqlite:////data"):
-        os.makedirs("/data", exist_ok=True)
-    elif DATABASE_URL.startswith("sqlite:///./data"):
-        os.makedirs("./data", exist_ok=True)
-        
     Base.metadata.create_all(bind=engine)
 
 def get_db():
@@ -54,7 +55,7 @@ def get_db():
     try:
         return db
     finally:
-        pass
+        db.close()
 
 def create_session(user_id: str, name: str = "New Chat"):
     """Create a new chat session"""