You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
111 lines
3.0 KiB
Python
111 lines
3.0 KiB
Python
# src/database.py
|
|
import re
|
|
import operator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Annotated, Any
|
|
|
|
from fastapi import Depends
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
|
|
from sqlalchemy.ext.declarative import declarative_base, declared_attr
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
from starlette.requests import Request
|
|
|
|
from src.config import SQLALCHEMY_DATABASE_URI, SQLALCHEMY_COLLECTOR_URI
|
|
|
|
engine = create_async_engine(SQLALCHEMY_DATABASE_URI, echo=False, future=True)
|
|
collector_engine = create_async_engine(SQLALCHEMY_COLLECTOR_URI, echo=False, future=True)
|
|
|
|
async_session = async_sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autocommit=False,
|
|
autoflush=False,
|
|
)
|
|
|
|
async_collector_session = async_sessionmaker(
|
|
collector_engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autocommit=False,
|
|
autoflush=False,
|
|
)
|
|
|
|
async def get_collector_db(request: Request):
|
|
return request.state.collector_db
|
|
|
|
|
|
def get_db(request: Request):
|
|
return request.state.db
|
|
|
|
|
|
DbSession = Annotated[AsyncSession, Depends(get_db)]
|
|
CollectorDbSession = Annotated[AsyncSession, Depends(get_collector_db)]
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
@declared_attr.directive
|
|
def __tablename__(cls) -> str:
|
|
return resolve_table_name(cls.__name__)
|
|
|
|
def dict(self):
|
|
"""Returns a dict representation of a model."""
|
|
if hasattr(self, '__table__'):
|
|
return {c.name: operator.attrgetter(c.name)(self) for c in self.__table__.columns}
|
|
return {}
|
|
|
|
class CollectorBase(DeclarativeBase):
|
|
@declared_attr.directive
|
|
def __tablename__(cls) -> str:
|
|
return resolve_table_name(cls.__name__)
|
|
|
|
def dict(self):
|
|
"""Returns a dict representation of a model."""
|
|
if hasattr(self, '__table__'):
|
|
return {c.name: operator.attrgetter(c.name)(self) for c in self.__table__.columns}
|
|
return {}
|
|
|
|
@asynccontextmanager
|
|
async def get_main_session():
|
|
session = async_session()
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
@asynccontextmanager
|
|
async def get_collector_session():
|
|
session = async_collector_session()
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
def resolve_table_name(name):
|
|
"""Resolves table names to their mapped names."""
|
|
names = re.split("(?=[A-Z])", name) # noqa
|
|
return "_".join([x.lower() for x in names if x])
|
|
|
|
|
|
def get_class_by_tablename(table_fullname: str) -> Any:
|
|
"""Return class reference mapped to table."""
|
|
|
|
def _find_class(name):
|
|
for c in Base.registry._class_registry.values():
|
|
if hasattr(c, "__table__"):
|
|
if c.__table__.fullname.lower() == name.lower():
|
|
return c
|
|
|
|
mapped_name = resolve_table_name(table_fullname)
|
|
mapped_class = _find_class(mapped_name)
|
|
|
|
return mapped_class
|