55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
|
import asyncio
|
||
|
|
import contextvars
|
||
|
|
import logging
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
_logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class PluginResourceProvider:
|
||
|
|
http_client: Any | None = None
|
||
|
|
db_session_factory: Any | None = None
|
||
|
|
redis_client: Any | None = None
|
||
|
|
config_provider: Any | None = None
|
||
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
async def cleanup(self) -> None:
|
||
|
|
for attr_name in ("http_client", "redis_client"):
|
||
|
|
resource = getattr(self, attr_name, None)
|
||
|
|
if resource is None:
|
||
|
|
continue
|
||
|
|
close_fn = getattr(resource, "close", None) or getattr(resource, "aclose", None)
|
||
|
|
if close_fn is None:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
result = close_fn()
|
||
|
|
if asyncio.iscoroutine(result):
|
||
|
|
await result
|
||
|
|
_logger.debug("PluginResourceProvider: closed %s", attr_name)
|
||
|
|
except Exception:
|
||
|
|
_logger.warning("PluginResourceProvider: failed to close %s", attr_name)
|
||
|
|
self.http_client = None
|
||
|
|
self.db_session_factory = None
|
||
|
|
self.redis_client = None
|
||
|
|
self.config_provider = None
|
||
|
|
self.metadata.clear()
|
||
|
|
|
||
|
|
|
||
|
|
_resources_ctx: contextvars.ContextVar[PluginResourceProvider] = contextvars.ContextVar(
|
||
|
|
"plugin_resources", default=PluginResourceProvider()
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def inject_resources(provider: PluginResourceProvider) -> None:
|
||
|
|
_resources_ctx.set(provider)
|
||
|
|
|
||
|
|
|
||
|
|
def get_resources() -> PluginResourceProvider:
|
||
|
|
return _resources_ctx.get()
|
||
|
|
|
||
|
|
|
||
|
|
def reset_resources() -> None:
|
||
|
|
_resources_ctx.set(PluginResourceProvider())
|