Upload backend/database.py

backend/database.py

@@ -1,16 +1,24 @@
 """
 Bharat Tech Atlas — Database layer using SQLite with R-Tree spatial indexing.
+v3.2: Added connection pooling, schema_version table for migrations, and
+      proper cleanup on shutdown.
 """
 import sqlite3
 import os
-import json
 import math
+import threading
 
 DB_PATH = os.path.join(os.path.dirname(__file__), "..", "data", "bharattechatlas.db")
+SCHEMA_VERSION = 2  # Bump when schema changes
 
+# ─── Connection pool ──────────────────────────────────────────────────────────
+_pool_lock = threading.Lock()
+_pool = {}
 
-def get_db():
-    conn = sqlite3.connect(DB_PATH)
+
+def _create_connection():
+    """Create a new SQLite connection with optimized PRAGMAs."""
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
     conn.row_factory = sqlite3.Row
     conn.execute("PRAGMA journal_mode=WAL")
     conn.execute("PRAGMA foreign_keys=ON")
@@ -18,11 +26,34 @@ def get_db():
     conn.execute("PRAGMA mmap_size=67108864")
     conn.execute("PRAGMA synchronous=NORMAL")
     conn.execute("PRAGMA temp_store=MEMORY")
+    conn.execute("PRAGMA busy_timeout=5000")
     return conn
 
 
+def get_db():
+    """Get a connection from the pool (one per thread)."""
+    tid = threading.current_thread().ident
+    with _pool_lock:
+        if tid not in _pool or _pool[tid] is None:
+            _pool[tid] = _create_connection()
+        return _pool[tid]
+
+
+def close_all_connections():
+    """Close all pooled connections. Call on shutdown."""
+    with _pool_lock:
+        for tid, conn in list(_pool.items()):
+            try:
+                conn.close()
+            except Exception:
+                pass
+            _pool[tid] = None
+        _pool.clear()
+
+
 def init_db():
-    conn = get_db()
+    """Initialize database schema with migration tracking."""
+    conn = _create_connection()
 
     conn.executescript("""
         CREATE TABLE IF NOT EXISTS entities (
@@ -148,10 +179,42 @@ def init_db():
             ('wealthtech', 'WealthTech', '📈', '#7C3AED', 'dpiit_category');
     """)
 
+    # ─── Schema versioning / migrations ───────────────────────────────────────
+    conn.execute("""
+        CREATE TABLE IF NOT EXISTS schema_version (
+            version INTEGER PRIMARY KEY,
+            applied_at TEXT DEFAULT (datetime('now'))
+        )
+    """)
+    conn.execute("INSERT OR IGNORE INTO schema_version (version) VALUES (0)")
     conn.commit()
+
+    # Run migrations
+    current = conn.execute("SELECT version FROM schema_version ORDER BY version DESC LIMIT 1").fetchone()[0]
+    if current < 1:
+        conn.execute("UPDATE schema_version SET version = 1 WHERE version = 0")
+        conn.commit()
+    if current < 2:
+        try:
+            conn.execute("ALTER TABLE entities ADD COLUMN github_stars INTEGER DEFAULT 0")
+        except sqlite3.OperationalError:
+            pass  # Column already exists
+        conn.execute("UPDATE schema_version SET version = 2 WHERE version = 1")
+        conn.commit()
+
     conn.close()
 
 
+def is_seeded() -> bool:
+    """Check if database already has seed data."""
+    try:
+        conn = get_db()
+        count = conn.execute("SELECT COUNT(*) FROM entities WHERE is_active = 1").fetchone()[0]
+        return count > 100
+    except Exception:
+        return False
+
+
 def haversine_distance(lat1, lon1, lat2, lon2):
     R = 6371
     dlat = math.radians(lat2 - lat1)