-
Notifications
You must be signed in to change notification settings - Fork 1
/
database.py
34 lines (26 loc) · 1.15 KB
/
database.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from types import SimpleNamespace
from typing import Dict
from starlette.applications import ASGIApp
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from ..config.models import Database as DatabaseConfig
from ..database import DatabaseProvider, Database
class DatabaseMiddleware(BaseHTTPMiddleware):
instances: Dict[str, DatabaseProvider]
@staticmethod
def get(r: Request) -> Database:
return r.state.databases
@staticmethod
async def default(r: Request) -> Database:
async with r.state.databases.default() as database:
yield database
def __init__(self, app: ASGIApp, databases: Dict[str, DatabaseConfig]):
super(DatabaseMiddleware, self).__init__(app)
self.instances = {
database: DatabaseProvider(config)
for database, config in databases.items()
}
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
request.state.databases = SimpleNamespace(**self.instances)
return await call_next(request)