devplacepy/devplacepy_services/database/db_patch.py

# 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)