🗃️ Use of dedicated db sessions

This commit is contained in:
2023-05-07 12:18:12 +02:00
parent 5505209760
commit b713042359
5 changed files with 201 additions and 185 deletions

View File

@@ -1,5 +1,6 @@
from __future__ import annotations
from logging import getLogger
from typing import Iterable, Self, TYPE_CHECKING
from sqlalchemy import select
@@ -9,31 +10,36 @@ from sqlalchemy.orm import DeclarativeBase
if TYPE_CHECKING:
from .db import Database
logger = getLogger(__name__)
class Base(DeclarativeBase):
db: Database | None = None
@classmethod
async def add(cls, stops: Self | Iterable[Self]) -> bool:
async def add(cls, objs: Self | Iterable[Self]) -> bool:
if cls.db is not None and (session := await cls.db.get_session()) is not None:
async with session.begin():
try:
if isinstance(stops, Iterable):
cls.db.session.add_all(stops) # type: ignore
if isinstance(objs, Iterable):
session.add_all(objs)
else:
cls.db.session.add(stops) # type: ignore
await cls.db.session.commit() # type: ignore
session.add(objs)
except (AttributeError, IntegrityError) as err:
print(err)
logger.error(err)
return False
return True
@classmethod
async def get_by_id(cls, id_: int | str) -> Self | None:
try:
stmt = select(cls).where(cls.id == id_) # type: ignore
res = await cls.db.session.execute(stmt) # type: ignore
element = res.scalar_one_or_none()
except AttributeError as err:
print(err)
element = None
return element
if cls.db is not None and (session := await cls.db.get_session()) is not None:
async with session.begin():
stmt = select(cls).where(cls.id == id_)
res = await session.execute(stmt)
return res.scalar_one_or_none()
return None

View File

@@ -1,5 +1,10 @@
from logging import getLogger
from typing import Annotated, AsyncIterator
from fastapi import Depends
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import (
async_sessionmaker,
AsyncEngine,
@@ -10,42 +15,47 @@ from sqlalchemy.ext.asyncio import (
from .base_class import Base
logger = getLogger(__name__)
class Database:
def __init__(self) -> None:
self._engine: AsyncEngine | None = None
self._session_maker: async_sessionmaker[AsyncSession] | None = None
self._session: AsyncSession | None = None
self._async_engine: AsyncEngine | None = None
self._async_session_local: async_sessionmaker[AsyncSession] | None = None
@property
def session(self) -> AsyncSession | None:
if self._session is None and (session_maker := self._session_maker) is not None:
self._session = session_maker()
return self._session
async def get_session(self) -> AsyncSession | None:
try:
return self._async_session_local() # type: ignore
async def connect(self, db_path: str, clear_static_data: bool = False) -> bool:
except (SQLAlchemyError, AttributeError) as e:
logger.exception(e)
return None
# TODO: Preserve UserLastStopSearchResults table from drop.
self._engine = create_async_engine(db_path)
if self._engine is not None:
SQLAlchemyInstrumentor().instrument(engine=self._engine.sync_engine)
self._session_maker = async_sessionmaker(
self._engine, expire_on_commit=False, class_=AsyncSession
async def connect(self, db_path: str, clear_static_data: bool = False) -> bool:
self._async_engine = create_async_engine(
db_path, pool_pre_ping=True, pool_size=10, max_overflow=20
)
if (session := self.session) is not None:
await session.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm;"))
async with self._engine.begin() as conn:
if self._async_engine is not None:
SQLAlchemyInstrumentor().instrument(engine=self._async_engine.sync_engine)
self._async_session_local = async_sessionmaker(
bind=self._async_engine,
# autoflush=False,
expire_on_commit=False,
class_=AsyncSession,
)
async with self._async_engine.begin() as session:
await session.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm;"))
if clear_static_data:
await conn.run_sync(Base.metadata.drop_all)
await conn.run_sync(Base.metadata.create_all)
await session.run_sync(Base.metadata.drop_all)
await session.run_sync(Base.metadata.create_all)
return True
async def disconnect(self) -> None:
if self._session is not None:
await self._session.close()
self._session = None
if self._engine is not None:
await self._engine.dispose()
if self._async_engine is not None:
await self._async_engine.dispose()

View File

@@ -15,7 +15,7 @@ from aiohttp import ClientSession
from msgspec import ValidationError
from msgspec.json import Decoder
from pyproj import Transformer
from shapefile import Reader as ShapeFileReader, ShapeRecord
from shapefile import Reader as ShapeFileReader, ShapeRecord # type: ignore
from ..db import Database
from ..models import ConnectionArea, Line, LinePicto, Stop, StopArea, StopShape
@@ -357,7 +357,6 @@ class IdfmInterface:
fields = line.fields
picto_id = fields.picto.id_ if fields.picto is not None else None
picto = await LinePicto.get_by_id(picto_id) if picto_id else None
ret.append(
Line(
@@ -384,7 +383,6 @@ class IdfmInterface:
fields.audiblesigns_available.value
),
picto_id=fields.picto.id_ if fields.picto is not None else None,
picto=picto,
record_id=line.recordid,
record_ts=int(line.record_timestamp.timestamp()),
)

View File

@@ -94,10 +94,9 @@ class Line(Base):
async def get_by_name(
cls, name: str, operator_name: None | str = None
) -> Sequence[Self] | None:
session = cls.db.session
if session is None:
return None
if (session := await cls.db.get_session()) is not None:
async with session.begin():
filters = {"name": name}
if operator_name is not None:
filters["operator_name"] = operator_name
@@ -112,6 +111,8 @@ class Line(Base):
return lines
return None
@classmethod
async def _add_picto_to_line(cls, line: str | Self, picto: LinePicto) -> None:
formatted_line: Self | None = None
@@ -133,23 +134,25 @@ class Line(Base):
@classmethod
async def add_pictos(cls, line_to_pictos: Iterable[tuple[str, LinePicto]]) -> bool:
session = cls.db.session
if session is None:
return False
if (session := await cls.db.get_session()) is not None:
async with session.begin():
await asyncio_gather(
*[cls._add_picto_to_line(line, picto) for line, picto in line_to_pictos]
*[
cls._add_picto_to_line(line, picto)
for line, picto in line_to_pictos
]
)
await session.commit()
return True
return False
@classmethod
async def add_stops(cls, line_to_stop_ids: Iterable[tuple[str, str, int]]) -> int:
session = cls.db.session
if session is None:
return 0
if (session := await cls.db.get_session()) is not None:
async with session.begin():
line_names_ops, stop_ids = set(), set()
for line_name, operator_name, stop_id in line_to_stop_ids:
@@ -166,13 +169,17 @@ class Line(Base):
for line in lines_res.scalars():
lines[(line.name, line.operator_name)].append(line)
stops_res = await session.execute(select(_Stop).where(_Stop.id.in_(stop_ids)))
stops_res = await session.execute(
select(_Stop).where(_Stop.id.in_(stop_ids))
)
stops = {stop.id: stop for stop in stops_res.scalars()}
found = 0
for line_name, operator_name, stop_id in line_to_stop_ids:
if (stop := stops.get(stop_id)) is not None:
if (stop_lines := lines.get((line_name, operator_name))) is not None:
if (
stop_lines := lines.get((line_name, operator_name))
) is not None:
for stop_line in stop_lines:
stop_line.stops.append(stop)
found += 1
@@ -184,6 +191,6 @@ class Line(Base):
f"(used by {line_name}/{operator_name})"
)
await session.commit()
return found
return 0

View File

@@ -1,6 +1,7 @@
from __future__ import annotations
from typing import Iterable, Sequence, TYPE_CHECKING
from logging import getLogger
from typing import Annotated, Iterable, Sequence, TYPE_CHECKING
from sqlalchemy import (
BigInteger,
@@ -22,7 +23,6 @@ from sqlalchemy.orm import (
Mapped,
relationship,
selectinload,
with_polymorphic,
)
from sqlalchemy.schema import Index
from sqlalchemy_utils.types.ts_vector import TSVectorType
@@ -34,6 +34,8 @@ if TYPE_CHECKING:
from .line import Line
logger = getLogger(__name__)
stop_area_stop_association_table = Table(
"stop_area_stop_association_table",
Base.metadata,
@@ -91,35 +93,24 @@ class _Stop(Base):
),
)
# TODO: Test https://www.cybertec-postgresql.com/en/postgresql-more-performance-for-like-and-ilike-statements/
# TODO: Should be able to remove with_polymorphic ?
@classmethod
async def get_by_name(cls, name: str) -> Sequence[type[_Stop]] | None:
session = cls.db.session
if session is None:
return None
async def get_by_name(cls, name: str) -> Sequence[_Stop] | None:
if (session := await cls.db.get_session()) is not None:
stop_stop_area = with_polymorphic(_Stop, [Stop, StopArea])
match_stmt = stop_stop_area.names_tsv.match(name, postgresql_regconfig="french")
async with session.begin():
match_stmt = cls.names_tsv.match(name, postgresql_regconfig="french")
ranking_stmt = func.ts_rank_cd(
stop_stop_area.names_tsv, func.plainto_tsquery("french", name)
)
stmt = (
select(stop_stop_area)
.filter(match_stmt)
.order_by(desc(ranking_stmt))
.options(
selectinload(stop_stop_area.areas),
selectinload(stop_stop_area.lines),
)
cls.names_tsv, func.plainto_tsquery("french", name)
)
stmt = select(cls).filter(match_stmt).order_by(desc(ranking_stmt))
res = await session.execute(stmt)
stops = res.scalars().all()
return stops
return None
class Stop(_Stop):
@@ -160,9 +151,9 @@ class StopArea(_Stop):
async def add_stops(
cls, stop_area_to_stop_ids: Iterable[tuple[int, int]]
) -> int | None:
session = cls.db.session
if session is None:
return None
if (session := await cls.db.get_session()) is not None:
async with session.begin():
stop_area_ids, stop_ids = set(), set()
for stop_area_id, stop_id in stop_area_to_stop_ids:
@@ -178,7 +169,9 @@ class StopArea(_Stop):
stop_area.id: stop_area for stop_area in stop_areas_res.all()
}
stop_res = await session.execute(select(Stop).where(Stop.id.in_(stop_ids)))
stop_res = await session.execute(
select(Stop).where(Stop.id.in_(stop_ids))
)
stops: dict[int, Stop] = {stop.id: stop for stop in stop_res.scalars()}
found = 0
@@ -192,10 +185,10 @@ class StopArea(_Stop):
else:
print(f"No stop area found for {stop_area_id}")
await session.commit()
return found
return None
class StopShape(Base):
@@ -235,9 +228,9 @@ class ConnectionArea(Base):
async def add_stops(
cls, conn_area_to_stop_ids: Iterable[tuple[int, int]]
) -> int | None:
session = cls.db.session
if session is None:
return None
if (session := await cls.db.get_session()) is not None:
async with session.begin():
conn_area_ids, stop_ids = set(), set()
for conn_area_id, stop_id in conn_area_to_stop_ids:
@@ -253,8 +246,10 @@ class ConnectionArea(Base):
conn.id: conn for conn in conn_area_res.scalars()
}
stop_res = await session.execute(select(_Stop).where(_Stop.id.in_(stop_ids)))
stops: dict[int, _Stop] = {stop.id: stop for stop in stop_res.scalars()}
stop_res = await session.execute(
select(Stop).where(Stop.id.in_(stop_ids))
)
stops: dict[int, Stop] = {stop.id: stop for stop in stop_res.scalars()}
found = 0
for conn_area_id, stop_id in conn_area_to_stop_ids:
@@ -267,6 +262,6 @@ class ConnectionArea(Base):
else:
print(f"No connection area found for {conn_area_id}")
await session.commit()
return found
return None