Upload studio_generate.py

studio_generate.py

@@ -54,70 +54,79 @@ POLAR_ORG_ID   = os.environ.get("POLAR_ORG_ID", "")
 FIREBASE_KEY_PATH = os.environ.get("FIREBASE_KEY_PATH", "serviceAccountKey.json")
 ADMIN_UIDS     = set(os.environ.get("STUDIO_ADMIN_UIDS", "").split(","))
 
-# ── MONGODB (reuse main.py's connection if available, else create our own) ────
-# IMPORTANT: We use pymongo (sync) — NOT motor (async). Motor creates a separate
-# connection that fails on Render. Instead we reuse main.py's existing working
-# MongoClient via sys.modules, or create our own with the same SSL settings.
+# ── MONGODB (lazy — waits until main.py is loaded, then reuses its connection) ─
 import sys as _sys
 
+_db_instance = None
+
 def _get_db():
-    """Get the neuraprompt database. Reuses main.py's connection if available."""
+    """Get the neuraprompt database. LAZY: reuses main.py's connection.
+    Called on first DB operation, not at import time — this ensures main.py
+    has finished initializing its MongoClient before we try to reuse it."""
+    global _db_instance
+    if _db_instance is not None:
+        return _db_instance
+    # Try to reuse main.py's connection (it's already connected to Atlas)
     main_mod = _sys.modules.get("main")
     if main_mod and hasattr(main_mod, "neuraprompt_db"):
-        return main_mod.neuraprompt_db
+        _db_instance = main_mod.neuraprompt_db
+        print("[Studio v2] MongoDB: reusing main.py connection")
+        return _db_instance
     # Fallback: create our own pymongo client (same SSL settings as main.py)
+    print("[Studio v2] MongoDB: main.py not available, creating own connection")
     _client = MongoClient(
         MONGO_URL, ssl=True,
         tlsAllowInvalidCertificates=False,
         tlsCAFile="/etc/ssl/certs/ca-certificates.crt",
         server_api=ServerApi("1"),
     )
-    return _client["neuraprompt"]
+    _db_instance = _client["neuraprompt"]
+    return _db_instance
 
-db = _get_db()
-print(f"[Studio v2] MongoDB: {'reusing main.py connection' if _sys.modules.get('main') and hasattr(_sys.modules.get('main'), 'neuraprompt_db') else 'own connection'}")
+# NOTE: db is NOT set at import time — it's resolved lazily on first use.
+# All _db_* helpers call _get_db() internally.
 
 # ── ASYNC DB HELPER (pymongo is sync, so we wrap calls in asyncio.to_thread) ──
 async def _db_find_one(collection, query, projection=None):
     """Async wrapper for db.collection.find_one()"""
     return await asyncio.to_thread(
-        lambda: db[collection].find_one(query, projection) if projection else db[collection].find_one(query)
+        lambda: _get_db()[collection].find_one(query, projection) if projection else db[collection].find_one(query)
     )
 
 async def _db_update_one(collection, query, update, upsert=False):
-    return await asyncio.to_thread(lambda: db[collection].update_one(query, update, upsert=upsert))
+    return await asyncio.to_thread(lambda: _get_db()[collection].update_one(query, update, upsert=upsert))
 
 async def _db_insert_one(collection, doc):
-    return await asyncio.to_thread(lambda: db[collection].insert_one(doc))
+    return await asyncio.to_thread(lambda: _get_db()[collection].insert_one(doc))
 
 async def _db_count_documents(collection, query):
-    return await asyncio.to_thread(lambda: db[collection].count_documents(query))
+    return await asyncio.to_thread(lambda: _get_db()[collection].count_documents(query))
 
 async def _db_find(collection, query, projection=None, sort=None, limit=None):
     def _do():
-        cursor = db[collection].find(query, projection) if projection else db[collection].find(query)
+        cursor = _get_db()[collection].find(query, projection) if projection else db[collection].find(query)
         if sort: cursor = cursor.sort(*sort) if isinstance(sort, tuple) else cursor.sort(sort)
         if limit: cursor = cursor.limit(limit)
         return list(cursor)
     return await asyncio.to_thread(_do)
 
 async def _db_delete_one(collection, query):
-    return await asyncio.to_thread(lambda: db[collection].delete_one(query))
+    return await asyncio.to_thread(lambda: _get_db()[collection].delete_one(query))
 
 async def _db_delete_many(collection, query):
-    return await asyncio.to_thread(lambda: db[collection].delete_many(query))
+    return await asyncio.to_thread(lambda: _get_db()[collection].delete_many(query))
 
 async def _db_insert_many(collection, docs):
-    return await asyncio.to_thread(lambda: db[collection].insert_many(docs))
+    return await asyncio.to_thread(lambda: _get_db()[collection].insert_many(docs))
 
 async def _db_distinct(collection, field, query=None):
-    return await asyncio.to_thread(lambda: db[collection].distinct(field, query or {}))
+    return await asyncio.to_thread(lambda: _get_db()[collection].distinct(field, query or {}))
 
 async def _db_aggregate(collection, pipeline):
-    return await asyncio.to_thread(lambda: list(db[collection].aggregate(pipeline)))
+    return await asyncio.to_thread(lambda: list(_get_db()[collection].aggregate(pipeline)))
 
 async def _db_create_index(collection, keys, **kwargs):
-    return await asyncio.to_thread(lambda: db[collection].create_index(keys, **kwargs))
+    return await asyncio.to_thread(lambda: _get_db()[collection].create_index(keys, **kwargs))
 
 
 # ── CREDIT SYSTEM ─────────────────────────────────────────────────────────────