Update main.py

main.py

@@ -9,6 +9,7 @@ from google import genai
 from google.genai import types
 
 from pymysql.err import OperationalError
+import threading
 warnings.filterwarnings("ignore")
 
 # ───────────────────────────────────────────────────────────────────────────────
@@ -52,20 +53,13 @@ app = Flask(__name__)
 CORS(app)
 
 # ───────────────────────────────────────────────────────────────────────────────
-# DB CONNECTION (autocommit + TLS + auto-reconnect)
+# DB CONNECTION (refactored for better connection management)
 # ───────────────────────────────────────────────────────────────────────────────
-_CONN = None
+_connection_lock = threading.Lock()
 
-def _connect():
-    """Enhanced connection with better error handling"""
-    global _CONN
-    try:
-        if _CONN:
-            _CONN.close()
-    except Exception:
-        pass
-    
-    _CONN = pymysql.connect(
+def _create_connection():
+    """Create a new database connection with optimized settings"""
+    return pymysql.connect(
         host=TIDB_HOST,
         port=TIDB_PORT,
         user=TIDB_USER,
@@ -77,36 +71,90 @@ def _connect():
         autocommit=True,
         charset="utf8mb4",
         cursorclass=pymysql.cursors.DictCursor,
-        # Add these timeouts:
         connect_timeout=10,
-        read_timeout=30,
+        read_timeout=60,  # Increased for vector operations
         write_timeout=30,
+        # TiDB-specific optimizations:
+        init_command="SET SESSION sql_mode='STRICT_TRANS_TABLES,NO_ZERO_DATE,NO_ZERO_IN_DATE,ERROR_FOR_DIVISION_BY_ZERO'",
+        client_flag=pymysql.constants.CLIENT.MULTI_STATEMENTS,
     )
 
-def _ensure_conn():
-    global _CONN
+@contextmanager
+def cursor():
+    """Create a fresh connection for each request context with retry logic"""
+    conn = None
     max_retries = 3
+    
     for attempt in range(max_retries):
         try:
-            if _CONN is None:
-                _connect()
+            conn = _create_connection()
+            with conn.cursor() as cur:
+                yield cur
+            break
+        except (OperationalError, pymysql.err.InternalError) as e:
+            if conn:
+                try:
+                    conn.close()
+                except Exception:
+                    pass
+                conn = None
+            
+            if attempt == max_retries - 1:
+                log.error(f"Database connection failed after {max_retries} attempts: {e}")
+                raise
             else:
-                _CONN.ping(reconnect=False)  # Test without auto-reconnect first
-            return _CONN
+                log.warning(f"Database connection failed (attempt {attempt + 1}): {e}")
+                time.sleep(0.5 * (attempt + 1))  # Exponential backoff
         except Exception as e:
-            log.warning(f"Connection attempt {attempt + 1} failed: {e}")
-            _CONN = None
-            if attempt < max_retries - 1:
-                time.sleep(0.5 * (attempt + 1))  # Backoff
-            else:
-                raise
+            if conn:
+                try:
+                    conn.close()
+                except Exception:
+                    pass
+            log.error(f"Database connection failed: {e}")
+            raise
+        finally:
+            if conn:
+                try:
+                    conn.close()
+                except Exception:
+                    pass
+
+def with_db_retry(func):
+    """Decorator to retry database operations on connection failures"""
+    def wrapper(*args, **kwargs):
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                return func(*args, **kwargs)
+            except (OperationalError, pymysql.err.InternalError) as e:
+                if attempt == max_retries - 1:
+                    log.error(f"Database operation failed after {max_retries} attempts: {e}")
+                    raise
+                log.warning(f"Database operation failed (attempt {attempt + 1}): {e}")
+                time.sleep(0.5 * (attempt + 1))
+    return wrapper
 
-@contextmanager
-def cursor():
-    """DictCursor with auto-ping; use in each route."""
-    conn = _ensure_conn()
-    with conn.cursor() as cur:
-        yield cur
+# ───────────────────────────────────────────────────────────────────────────────
+# ERROR HANDLERS
+# ─────────────────────────────────────────────────��─────────────────────────────
+@app.errorhandler(OperationalError)
+def handle_db_error(e):
+    log.error(f"Database error: {e}")
+    return jsonify({
+        "ok": False, 
+        "error": "database_unavailable",
+        "message": "Database connection issue. Please try again."
+    }), 503
+
+@app.errorhandler(pymysql.err.InternalError)
+def handle_internal_error(e):
+    log.error(f"Database internal error: {e}")
+    return jsonify({
+        "ok": False,
+        "error": "database_error", 
+        "message": "Database operation failed. Please try again."
+    }), 500
 
 # ───────────────────────────────────────────────────────────────────────────────
 # EMBEDDINGS (lazy-load; same model as ingest; pad to 1536)
@@ -298,23 +346,41 @@ def root():
     return jsonify({"ok": True, "service": "provenance-radar-api", "device": _DEVICE_INFO})
 
 @app.get("/api/health")
+@with_db_retry
 def health():
     try:
+        start_time = time.time()
         with cursor() as cur:
             cur.execute("SELECT COUNT(*) AS c FROM objects");      objects = cur.fetchone()["c"]
             cur.execute("SELECT COUNT(*) AS c FROM provenance_sentences"); sentences = cur.fetchone()["c"]
             cur.execute("SELECT COUNT(*) AS c FROM risk_signals"); risks = cur.fetchone()["c"]
-        return jsonify({"ok": True, "device": _DEVICE_INFO, "counts": {
-            "objects": objects, "sentences": sentences, "risk_signals": risks}})
+        
+        db_latency = round((time.time() - start_time) * 1000, 2)
+        
+        return jsonify({
+            "ok": True, 
+            "device": _DEVICE_INFO, 
+            "db_latency_ms": db_latency,
+            "counts": {
+                "objects": objects, 
+                "sentences": sentences, 
+                "risk_signals": risks
+            }
+        })
     except Exception as e:
         log.exception("health failed")
-        return jsonify({"ok": False, "error": str(e)}), 500
+        return jsonify({
+            "ok": False, 
+            "error": str(e),
+            "db_status": "unavailable"
+        }), 503
 
 @app.get("/api/policy/windows")
 def policy_windows():
     return jsonify({"ok": True, "windows": POLICY_WINDOWS})
 
 @app.get("/api/leads")
+@with_db_retry
 def leads():
     limit = max(1, min(int(request.args.get("limit", 50)), 200))
     min_score = float(request.args.get("min_score", 0))
@@ -335,6 +401,7 @@ def leads():
     return jsonify({"ok": True, "data": rows})
 
 @app.get("/api/object/<int:object_id>")
+@with_db_retry
 def object_detail(object_id: int):
     with cursor() as cur:
         cur.execute("SELECT * FROM objects WHERE object_id=%s", (object_id,))
@@ -352,6 +419,7 @@ def object_detail(object_id: int):
     return jsonify({"ok": True, "object": obj, "sentences": sents, "events": events, "risks": risks})
 
 @app.get("/api/graph/<int:object_id>")
+@with_db_retry
 def graph(object_id: int):
     with cursor() as cur:
         cur.execute("SELECT object_id, source, title FROM objects WHERE object_id=%s", (object_id,))
@@ -365,6 +433,7 @@ def graph(object_id: int):
     return jsonify({"ok": True, **build_graph_from_events(obj, events)})
 
 @app.get("/api/timeline/<int:object_id>")
+@with_db_retry
 def timeline(object_id: int):
     with cursor() as cur:
         cur.execute("SELECT seq, sentence FROM provenance_sentences WHERE object_id=%s ORDER BY seq", (object_id,))
@@ -377,6 +446,7 @@ def timeline(object_id: int):
     return jsonify({"ok": True, "items": items})
 
 @app.get("/api/keyword")
+@with_db_retry
 def keyword_search():
     q = (request.args.get("q") or "").strip()
     limit = max(1, min(int(request.args.get("limit", 50)), 200))
@@ -396,6 +466,7 @@ def keyword_search():
 
 
 @app.post("/api/similar")
+@with_db_retry
 def similar_search():
     payload = request.get_json(force=True) or {}
     text = (payload.get("text") or "").strip()
@@ -471,6 +542,7 @@ def similar_search():
 
 
 @app.get("/api/vocab")
+@with_db_retry
 def vocab():
     field = (request.args.get("field") or "").strip().lower()
     limit = max(1, min(int(request.args.get("limit", 100)), 500))
@@ -490,6 +562,7 @@ def vocab():
 # ── Gemini-powered explanations ────────────────────────────────────────────────
 
 @app.get("/api/explain/object/<int:object_id>")
+@with_db_retry
 def explain_object(object_id: int):
     """Generate a concise, policy-aware research note for an object."""
     with cursor() as cur:
@@ -533,7 +606,7 @@ def explain_text():
         return jsonify({"ok": False, "error": "text required"}), 400
     sys = ("Explain this text as a provenance note for curators. "
            "Be precise and cautious; highlight possible red flags tied to 1933–1945 and post-1970 export rules.")
-    prompt = f"Explain and contextualize this provenance fragment:\n\n“{sentence}”."
+    prompt = f"Explain and contextualize this provenance fragment:\n\n"{sentence}"."
     text = gemini_explain(prompt, sys=sys)
     return jsonify({"ok": True, "model": EXPLAIN_MODEL, "explanation": text})
 
@@ -542,4 +615,4 @@ def explain_text():
 # ───────────────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", "7860"))
-    app.run(host="0.0.0.0", port=port, debug=False)
+    app.run(host="0.0.0.0", port=port, debug=False)