144 lines
4.2 KiB
Python
144 lines
4.2 KiB
Python
# Importações de bibliotecas padrão
|
|
import contextlib
|
|
from typing import AsyncIterator, Optional
|
|
|
|
from fastapi import Depends
|
|
# Importações de bibliotecas de terceiros
|
|
from sqlalchemy import MetaData
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncConnection,
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
AsyncAttrs,
|
|
)
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
from app.config import ECHO
|
|
import traceback
|
|
|
|
EXTENSIONS = ["uuid-ossp", "postgis", "postgis_topology"]
|
|
|
|
naming_convention = {
|
|
"ix": "ix_ct_%(table_name)s_%(column_0_N_name)s",
|
|
"uq": "uq_ct_%(table_name)s_%(column_0_N_name)s",
|
|
"ck": "ck_ct_%(table_name)s_%(constraint_name)s",
|
|
"fk": "fk_ct_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
|
|
"pk": "pk_ct_%(table_name)s",
|
|
}
|
|
|
|
|
|
class Base(DeclarativeBase, AsyncAttrs):
|
|
metadata = MetaData(naming_convention=naming_convention)
|
|
|
|
|
|
class DatabaseSessionManager:
|
|
def __init__(self):
|
|
self._engine: AsyncEngine | None = None
|
|
self._sessionmaker: async_sessionmaker | None = None
|
|
|
|
def init(self, host: str):
|
|
self._engine = create_async_engine(host, echo=ECHO)
|
|
self._sessionmaker = async_sessionmaker(autocommit=False, bind=self._engine)
|
|
|
|
async def close(self):
|
|
if self._engine is None:
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
await self._engine.dispose()
|
|
self._engine = None
|
|
self._sessionmaker = None
|
|
|
|
def is_initialized(self) -> bool:
|
|
"""Verifica se o engine foi inicializado"""
|
|
return self._engine is not None
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def connect(self) -> AsyncIterator[AsyncConnection]:
|
|
if self._engine is None:
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
|
|
async with self._engine.begin() as connection:
|
|
try:
|
|
yield connection
|
|
except Exception:
|
|
await connection.rollback()
|
|
raise
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def session(self) -> AsyncIterator[AsyncSession]:
|
|
if self._sessionmaker is None:
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
|
|
session = self._sessionmaker()
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def session_with_tenant(self, tenant_schema: Optional[str]) -> AsyncIterator[AsyncSession]:
|
|
|
|
if self._engine is None:
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
|
|
# Verifica e configura o schema_translate_map
|
|
if tenant_schema:
|
|
|
|
schema_engine = self._engine.execution_options(
|
|
schema_translate_map={None: str(tenant_schema)}
|
|
)
|
|
else:
|
|
|
|
schema_engine = self._engine.execution_options(
|
|
schema_translate_map=None
|
|
)
|
|
|
|
# Cria a sessão vinculada ao schema_engine
|
|
try:
|
|
session = self._sessionmaker(bind=schema_engine)
|
|
|
|
|
|
except Exception as e:
|
|
print(f"Erro ao configurar o bind da sessão: {e}")
|
|
raise
|
|
|
|
try:
|
|
yield session
|
|
except Exception as e:
|
|
await session.rollback()
|
|
print(f"Erro encontrado: {e}. Traceback: {traceback.format_exc()}")
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
# print("Sessão encerrada")
|
|
|
|
async def create_all(self, connection: AsyncConnection):
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
|
|
async def drop_all(self, connection: AsyncConnection):
|
|
await connection.run_sync(Base.metadata.drop_all)
|
|
|
|
# Multi Tenant
|
|
def get_engine(self):
|
|
if not self.is_initialized():
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
return self._engine
|
|
|
|
def get_sessionmaker(self):
|
|
if self._sessionmaker is None:
|
|
raise Exception("DatabaseSessionManager is not initialized")
|
|
return self._sessionmaker
|
|
|
|
# Multi Tenant
|
|
|
|
|
|
sessionmanager = DatabaseSessionManager()
|
|
|
|
|
|
async def get_db():
|
|
async with sessionmanager.session() as session:
|
|
yield session
|