55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Literal, Protocol
|
|
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
|
|
DatastoreMode = Literal["shared", "schema", "database"]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class TenantDatastoreRef:
|
|
tenant_id: str
|
|
mode: DatastoreMode = "shared"
|
|
datastore_id: str | None = None
|
|
schema_name: str | None = None
|
|
database_url_secret_ref: str | None = None
|
|
|
|
|
|
class TenantDatastore(Protocol):
|
|
ref: TenantDatastoreRef
|
|
|
|
@contextmanager
|
|
def session_for(self, module_id: str) -> Generator[Session, None, None]:
|
|
...
|
|
|
|
|
|
class DatastoreResolver(Protocol):
|
|
def for_tenant(self, tenant_id: str) -> TenantDatastore:
|
|
...
|
|
|
|
|
|
class SharedTenantDatastore:
|
|
def __init__(self, ref: TenantDatastoreRef, session_factory: sessionmaker[Session]) -> None:
|
|
self.ref = ref
|
|
self._session_factory = session_factory
|
|
|
|
@contextmanager
|
|
def session_for(self, module_id: str) -> Generator[Session, None, None]:
|
|
del module_id
|
|
with self._session_factory() as session:
|
|
yield session
|
|
|
|
|
|
class SharedDatastoreResolver:
|
|
def __init__(self, session_factory: sessionmaker[Session]) -> None:
|
|
self._session_factory = session_factory
|
|
|
|
def for_tenant(self, tenant_id: str) -> TenantDatastore:
|
|
return SharedTenantDatastore(TenantDatastoreRef(tenant_id=tenant_id), self._session_factory)
|
|
|