|
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextvars
|
|
import importlib
|
|
from typing import Any
|
|
|
|
_DB_MODULES = (
|
|
"schema",
|
|
"content",
|
|
"pagination",
|
|
"stats",
|
|
"forks",
|
|
"awards",
|
|
"customization",
|
|
"soft_delete",
|
|
"usage",
|
|
"email",
|
|
"follows",
|
|
"settings",
|
|
"seo_meta",
|
|
"deepsearch",
|
|
"activity",
|
|
"attachments_data",
|
|
"users",
|
|
"notifications",
|
|
"comments",
|
|
"ranking",
|
|
"engagement",
|
|
"relations",
|
|
)
|
|
|
|
_EXTERNAL_DB_MODULES = (
|
|
"devplacepy.services.audit.store",
|
|
"devplacepy.services.backup.store",
|
|
"devplacepy.services.openai_gateway.routing",
|
|
"devplacepy.services.openai_gateway.usage",
|
|
"devplacepy.services.openai_gateway.analytics",
|
|
"devplacepy.services.base",
|
|
"devplacepy.services.jobs.queue",
|
|
"devplacepy.services.containers.store",
|
|
"devplacepy.services.devii.store",
|
|
"devplacepy.services.devii.hub",
|
|
"devplacepy.attachments",
|
|
"devplacepy.project_files",
|
|
)
|
|
|
|
# The read pool runs on asyncio.to_thread worker threads while the write
|
|
# lane runs on the main event loop task - a plain module-global "current
|
|
# connection" swapped in/out per call is a data race between them (a
|
|
# concurrent read could reassign the global to a read-only connection
|
|
# while a write is still mid-flight, raising "attempt to write a readonly
|
|
# database"). asyncio.to_thread propagates a COPY of the calling task's
|
|
# contextvars.Context into the new thread, so storing the active
|
|
# connection in a ContextVar gives each call its own isolated view with
|
|
# no cross-talk, while every db-layer module keeps a single shared proxy
|
|
# object as its permanent `db` binding.
|
|
_CURRENT_DB: contextvars.ContextVar[Any] = contextvars.ContextVar("current_db", default=None)
|
|
|
|
# Fallback for any code path that touches the db outside an explicit
|
|
# use_db()-scoped call (e.g. startup's own init_db()) - write-once at
|
|
# broker startup, read-only afterward, so it carries no race of its own.
|
|
_fallback_db: Any = None
|
|
|
|
|
|
def _resolve() -> Any:
|
|
db = _CURRENT_DB.get()
|
|
if db is not None:
|
|
return db
|
|
if _fallback_db is not None:
|
|
return _fallback_db
|
|
raise RuntimeError("no database connection configured for this context")
|
|
|
|
|
|
class _ContextDb:
|
|
def __getattr__(self, name: str) -> Any:
|
|
return getattr(_resolve(), name)
|
|
|
|
def __getitem__(self, name: str) -> Any:
|
|
return _resolve()[name]
|
|
|
|
def __enter__(self):
|
|
return _resolve().__enter__()
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return _resolve().__exit__(exc_type, exc, tb)
|
|
|
|
|
|
_CONTEXT_DB = _ContextDb()
|
|
|
|
|
|
def patch_all_db() -> None:
|
|
import devplacepy.database as database_pkg
|
|
import devplacepy.database.core as core
|
|
|
|
core.db = _CONTEXT_DB
|
|
database_pkg.db = _CONTEXT_DB
|
|
for name in _DB_MODULES:
|
|
mod = importlib.import_module(f"devplacepy.database.{name}")
|
|
if hasattr(mod, "db"):
|
|
mod.db = _CONTEXT_DB
|
|
for name in _EXTERNAL_DB_MODULES:
|
|
mod = importlib.import_module(name)
|
|
if hasattr(mod, "db"):
|
|
mod.db = _CONTEXT_DB
|
|
|
|
|
|
def set_fallback_db(db_conn) -> None:
|
|
global _fallback_db
|
|
_fallback_db = db_conn
|
|
|
|
|
|
def use_db(db_conn) -> contextvars.Token:
|
|
return _CURRENT_DB.set(db_conn)
|
|
|
|
|
|
def reset_db(token: contextvars.Token) -> None:
|
|
_CURRENT_DB.reset(token)
|