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.
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
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 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):
|
|
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):
|
|
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):
|
|
names = re.split('(?=[A-Z])', name)
|
|
return '_'.join([x.lower() for x in names if x])
|
|
|
|
def get_class_by_tablename(table_fullname: str) -> Any:
|
|
|
|
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
|