Upload core/database.py with huggingface_hub

core/database.py

@@ -16,39 +16,34 @@ from sqlalchemy.engine import Engine
 from sqlalchemy.pool import QueuePool
 
 from core.logging import logger
-from core.models import (
+from core.models import (  # Database utilities; Other models needed by routers; Entity models; Compliance models; User models
     AccessReview,
     AuditLog,
-    # Database utilities
     Base,
     Case,
     CaseActivity,
-    # Other models needed by routers
     CaseNote,
     CaseStatus,
     CaseType,
     ComplianceAuditLog,
-    RegulatoryReport,
-    SecurityIncident,
-    # Entity models
     Entity,
     Evidence,
-    # Compliance models
     FraudAlert,
     GraphSnapshot,
     IdentityNode,
     Project,
     ReconciliationType,
+    RegulatoryReport,
     Relationship,
     RookieChecklist,
+    SecurityIncident,
     SessionLocal,
     Team,
+    TrainingRecord,
     Transaction,
-    # User models
     User,
-    TrainingRecord,
-    UserOnboardingState,
     UserDevice,
+    UserOnboardingState,
     UserRole,
     create_engine_and_session,
     create_tables,
@@ -123,17 +118,25 @@ class DatabaseOptimizer:
 
         @event.listens_for(self.engine, "after_execute")
         def after_execute(conn, clauseelement, multiparams, params, result):
-            execution_time = time.time() - conn.info.get("query_start_time", time.time())
+            execution_time = time.time() - conn.info.get(
+                "query_start_time", time.time()
+            )
             query_stats["query_count"] += 1
             query_stats["total_execution_time"] += execution_time
-            query_stats["avg_query_time"] = query_stats["total_execution_time"] / query_stats["query_count"]
+            query_stats["avg_query_time"] = (
+                query_stats["total_execution_time"] / query_stats["query_count"]
+            )
 
             # Log slow queries (>100ms)
             if execution_time > 0.1:
                 query_str = str(clauseelement)
                 query_stats["slow_queries"].append(
                     {
-                        "query": (query_str[:200] + "..." if len(query_str) > 200 else query_str),
+                        "query": (
+                            query_str[:200] + "..."
+                            if len(query_str) > 200
+                            else query_str
+                        ),
                         "execution_time": execution_time,
                         "timestamp": time.time(),
                     }
@@ -221,7 +224,9 @@ def optimize_database_performance():
         try:
             conn.execute(text("ANALYZE;"))  # Update table statistics
             conn.commit()
-            optimization_results["performance_improvements"]["statistics_updated"] = True
+            optimization_results["performance_improvements"]["statistics_updated"] = (
+                True
+            )
         except Exception as e:
             logger.warning(f"Failed to update statistics: {e}")
 
@@ -345,7 +350,8 @@ def after_execute(conn, clauseelement, multiparams, params, result):
         # Log slow queries
         if duration > 2.0:
             logger.warning(
-                f"Slow database query: {duration:.3f}s", extra={"query_duration": duration, "query_type": "database"}
+                f"Slow database query: {duration:.3f}s",
+                extra={"query_duration": duration, "query_type": "database"},
             )
         elif duration > 1.0:
             logger.info(f"Database query completed in {duration:.3f}s")