48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
from contextlib import contextmanager
|
|
from typing import Optional
|
|
from app.database.session import sessionmanager
|
|
from sqlalchemy import MetaData
|
|
from app.database.models import Base
|
|
|
|
|
|
@contextmanager
|
|
def with_db(tenant_schema: Optional[str] = None):
|
|
"""
|
|
Gerencia uma sessão do banco de dados com suporte a schema_translate_map.
|
|
|
|
Args:
|
|
tenant_schema (Optional[str]): Nome do schema do tenant (opcional).
|
|
|
|
Yields:
|
|
AsyncSession: Sessão configurada para o schema especificado.
|
|
"""
|
|
if not sessionmanager.is_initialized():
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
|
|
# Configura o schema_translate_map para o inquilino
|
|
if tenant_schema:
|
|
schema_translate_map = {None: tenant_schema}
|
|
else:
|
|
schema_translate_map = None
|
|
|
|
# Configura a conexão com o schema correto
|
|
connectable = sessionmanager.get_engine().execution_options(schema_translate_map=schema_translate_map)
|
|
|
|
# Cria uma sessão vinculada ao connectable configurado
|
|
session = sessionmanager.get_sessionmaker(bind=connectable)
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
def get_tenant_specific_metadata():
|
|
meta = MetaData(schema=None)
|
|
for table in Base.metadata.tables.values():
|
|
if table.schema is None:
|
|
table.to_metadata(meta)
|
|
return meta
|