🔧 Fix SQLAlchemy Session Issues for Sample Data Loading

✅ Root Issue: SQLAlchemy session management problems with in-memory DB
• Scoped sessions were causing connection pool issues
• Thread safety problems in HF Spaces environment
• Missing error handling and debugging information

⚡ Complete Fix:
• Use SessionLocal() instead of scoped Session() for better control
• Added database connection test before loading sample data
• Enhanced error handling with full traceback logging
• Better debugging output showing exact counts and status
• Proper session management with guaranteed cleanup

🎯 New debugging features:
• test_database_connection() - verifies DB connectivity
• Detailed environment variable logging (SPACE_ID, etc.)
• Verification counts after data insertion
• Full exception tracebacks for troubleshooting

🚀 Expected result:
• Sample data should now load correctly in HF Spaces
• Better error messages if something goes wrong
• Robust session handling prevents connection pool issues

backend/app.py

@@ -76,7 +76,7 @@ async def startup_event():
     # 🗄️ Initialize database on startup
     try:
         from backend.database.init_db import init_database
-        from backend.database import init_db, add_sample_data_for_hf
+        from backend.database import init_db, test_database_connection, add_sample_data_for_hf
         
         init_database(reset=False, force=False)
         init_db()  # Create tables using SQLAlchemy
@@ -86,7 +86,11 @@ async def startup_event():
             logger.info("🔒 HF Spaces: Using in-memory database for user privacy")
             logger.info("📝 Note: Data will be cleared when container restarts")
             
+            # Test database connection first
+            logger.info("🔍 Testing database connection...")
+            
             # Add sample data for HF Spaces
+            logger.info("📊 Loading sample data for demonstration...")
             add_sample_data_for_hf()
         else:
             logger.info("💾 Local development: Using persistent database")
@@ -94,6 +98,9 @@ async def startup_event():
         logger.info("🗄️ Database initialized successfully")
     except Exception as e:
         logger.error(f"❌ Database initialization failed: {e}")
+        import traceback
+        logger.error("Full traceback:")
+        logger.error(traceback.format_exc())
         # Don't fail startup - continue with empty database
     
     logger.info("🚀 Backend API available at: http://0.0.0.0:7860")

backend/database/__init__.py

@@ -60,23 +60,58 @@ def init_db():
     """Initialize the database by creating all tables."""
     Base.metadata.create_all(bind=engine)
 
+def test_database_connection():
+    """Test if database connection is working."""
+    try:
+        session = SessionLocal()
+        # Try to create tables
+        Base.metadata.create_all(bind=engine)
+        
+        # Test a simple query
+        result = session.execute("SELECT 1").fetchone()
+        print(f"✅ Database connection test successful: {result}")
+        
+        session.close()
+        return True
+    except Exception as e:
+        print(f"❌ Database connection test failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
 def add_sample_data_for_hf():
     """Add sample data for HF Spaces using SQLAlchemy."""
     if not os.getenv("SPACE_ID"):
         return  # Only run on HF Spaces
+        
+    # First test database connection
+    if not test_database_connection():
+        print("❌ Database connection failed, skipping sample data")
+        return
     
     import json
     import uuid
     import hashlib
     from backend.database.models import Trace, KnowledgeGraph, Entity, Relation
     
-    session = Session()
+    # Use SessionLocal instead of scoped Session for better control
+    session = SessionLocal()
     try:
+        # Ensure tables exist first
+        Base.metadata.create_all(bind=engine)
+        
         # Check if data already exists
         existing_traces = session.query(Trace).count()
         existing_kgs = session.query(KnowledgeGraph).count()
         
+        print(f"🔍 HF Spaces environment check:")
+        print(f"   • SPACE_ID: {os.getenv('SPACE_ID')}")
+        print(f"   • Database URL: {DATABASE_URL}")
+        print(f"   • Existing traces: {existing_traces}")
+        print(f"   • Existing KGs: {existing_kgs}")
+        
         if existing_traces > 0 or existing_kgs > 0:
+            print("📊 Sample data already exists, skipping...")
             return  # Data already exists
         
         print("🎯 HF Spaces: Adding sample data to in-memory database...")
@@ -296,13 +331,24 @@ def add_sample_data_for_hf():
         # Commit all data
         session.commit()
         
+        # Verify data was added
+        final_traces = session.query(Trace).count()
+        final_kgs = session.query(KnowledgeGraph).count()
+        final_entities = session.query(Entity).count()
+        final_relations = session.query(Relation).count()
+        
         print("✅ Sample data added successfully to HF Spaces in-memory database!")
-        print(f"   • 1 sample trace: Multi-Agent Customer Service Resolution")
-        print(f"   • 1 knowledge graph with {len(entities_data)} entities and {len(relations_data)} relations")
-        print(f"   • Demonstrates: Multi-agent coordination, error handling, optimization opportunities")
+        print(f"   • Traces added: {final_traces}")
+        print(f"   • Knowledge graphs added: {final_kgs}")
+        print(f"   • Entities added: {final_entities}")
+        print(f"   • Relations added: {final_relations}")
+        print(f"   • Scenario: Multi-Agent Customer Service with error handling and optimization")
         
     except Exception as e:
         print(f"❌ Failed to add sample data: {e}")
+        import traceback
+        print("Full error traceback:")
+        traceback.print_exc()
         session.rollback()
     finally:
         session.close()
@@ -311,6 +357,7 @@ __all__ = [
     'get_db',
     'models',
     'init_db',
+    'test_database_connection',
     'add_sample_data_for_hf',
     'save_knowledge_graph',
     'update_knowledge_graph_status',