DarkHelm.org/plex-playlist-old
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reformatted, and some adjustments to how sessionize operates.

Browse Source 
Signed-off-by: Cliff Hill <xlorep@darkhelm.org>
This commit is contained in:
Xlorep DarkHelm
2024-05-14 10:04:59 -04:00
parent 048707a2bd
commit 495d9d53a7
1 changed files with 62 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

151
backend/src/playlist/sql.py
View File

@@ -2,6 +2,7 @@
import collections
import collections.abc
import contextlib
import datetime
import functools
import inspect
@@ -15,15 +16,17 @@ import sqlalchemy.orm
from playlist import enums
from playlist import env
from playlist import models
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.ext.asyncio import create_async_engine
# Initialize asynchronous logger
logger = aiologger.Logger.with_default_handlers(name="sql_logger")
_engine = sqlalchemy.ext.asyncio.create_async_engine(env.DATABASE_URL, echo=True)
_engine = create_async_engine(env.DATABASE_URL, echo=True)
_async_session_maker = sqlalchemy.orm.sessionmaker(
_engine, class_=sqlalchemy.ext.asyncio.AsyncSession, expire_on_commit=False
_engine, class_=AsyncSession, expire_on_commit=False
)
@@ -33,7 +36,7 @@ class Sessionizable(typing.Protocol):
def __call__(
self,
*args: tuple[typing.Any, ...],
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
**kwargs: dict[str, typing.Any],
) -> (collections.abc.AsyncGenerator, collections.abc.Coroutine):
"""The signature for a sessionizable function.
@@ -52,72 +55,41 @@ class Sessionizable(typing.Protocol):
...
async def _validate_session_arg(name: str, kwargs) -> None:
@contextlib.asynccontextmanager
async def _db_logger(
func_name: str, session: AsyncSession
) -> collections.abc.AsyncGenerator[None, None]:
try:
if not (
kwargs["session"] is None
or isinstance(kwargs["session"], sqlalchemy.ext.asyncio.AsyncSession)
):
msg = f"{name} has an invalid session of type {type(kwargs['session'])}."
await logger.error(msg)
raise TypeError(msg)
yield
except sqlalchemy.exc.SQLAlchemyError as e:
await logger.exception(f"Database error in {func_name}: {e}")
await session.rollback()
raise
except Exception as e:
await logger.exception(f"Unexpected error in {func_name}: {e}")
await session.rollback()
raise
def _validate_signature(func: Sessionizable) -> None:
sig = inspect.signature(func)
try:
sig.parameters["session"]
except KeyError as e:
msg = f"{name} does not contain a session keyword argument."
await logger.error(msg)
msg = f"{func.__name__} does not have a session parameter."
raise UnboundLocalError(msg) from e
def _coro_db_logger(func: Sessionizable) -> Sessionizable:
@functools.wraps(func)
async def wrapper[
**P
](*args: P.args, **kwargs: P.kwargs) -> collections.abc.Coroutine:
try:
return await func(*args, **kwargs)
except sqlalchemy.exc.SQLAlchemyError as e:
await logger.exception(f"Database error in {func.__name__}: {e}")
await kwargs["session"].rollback()
raise
except Exception as e:
await logger.exception(f"Unexpected error in {func.__name__}: {e}")
await kwargs["session"].rollback()
raise
return wrapper
def _async_gen_db_logger(func: Sessionizable) -> Sessionizable:
@functools.wraps(func)
async def wrapper[
**P
](*args: P.args, **kwargs: P.kwargs) -> collections.abc.Coroutine:
try:
async for item in func(*args, **kwargs):
yield item
except sqlalchemy.exc.SQLAlchemyError as e:
await logger.exception(f"Database error in {func.__name__}: {e}")
await kwargs["session"].rollback()
raise
except Exception as e:
await logger.exception(f"Unexpected error in {func.__name__}: {e}")
await kwargs["session"].rollback()
raise
return wrapper
def sessionize(func: Sessionizable) -> Sessionizable:
"""Decorator that ensures a database session is available to an async function or generator.
This decorator automatically injects a `session` of type sqlalchemy.ext.asyncio.AsyncSession
This decorator automatically injects a `session` of type AsyncSession
into the decorated function if one is not provided. If the `session` keyword argument is
missing or set to None, a new session is created using the _async_session_maker and passed to
the function. If a session is already provided when the function is called, it uses the
missing or set to None, a new session is created using the `_async_session_maker` and passed
to the function. If a session is already provided when the function is called, it uses the
existing session.
The decorator is intended for use with coroutine functions and async generator functions that
@@ -144,7 +116,7 @@ def sessionize(func: Sessionizable) -> Sessionizable:
Example:
@sessionize
async def fetch_data(session: Optional[AsyncSession] = None):
async def fetch_data(session: AsyncSession | None = None):
# Function body using session
...
@@ -158,6 +130,8 @@ def sessionize(func: Sessionizable) -> Sessionizable:
- Sessionizable: The Protocol for sessionizable functions.
- AsyncSession: SQLAlchemy class used for asynchronous session management.
"""
_validate_signature(func)
ret: Sessionizable
match func:
case func if inspect.iscoroutinefunction(func):
@@ -166,14 +140,16 @@ def sessionize(func: Sessionizable) -> Sessionizable:
async def _coro_wrapper[
**P
](*args: P.args, **kwargs: P.kwargs) -> collections.abc.Coroutine:
await _validate_session_arg(func.__name__, kwargs)
logged_func = _coro_db_logger(func)
"""Wrap a sessionized coroutine function to inject the session if needed."""
if kwargs["session"] is None:
async with _async_session_maker() as kwargs["session"]:
return await logged_func(*args, **kwargs)
async with (
_async_session_maker() as kwargs["session"],
_db_logger(func.__name__, kwargs["session"]),
):
return await func(*args, **kwargs)
else:
return await logged_func(*args, **kwargs)
async with _db_logger(func.__name__, kwargs["session"]):
return await func(*args, **kwargs)
ret = _coro_wrapper
@@ -184,16 +160,17 @@ def sessionize(func: Sessionizable) -> Sessionizable:
**P
](*args: P.args, **kwargs: P.kwargs) -> collections.abc.AsyncGenerator:
"""Wrap a sessionized async generator function to inject the session if needed."""
await _validate_session_arg(func.__name__, kwargs)
logged_func = _async_gen_db_logger(func)
if kwargs["session"] is None:
async with _async_session_maker() as kwargs["session"]:
async for element in logged_func(*args, **kwargs):
async with (
_async_session_maker() as kwargs["session"],
_db_logger(func.__name__, kwargs["session"]),
):
async for element in func(*args, **kwargs):
yield element
else:
async for element in logged_func(*args, **kwargs):
yield element
async with _db_logger(func.__name__, kwargs["session"]):
async for element in func(*args, **kwargs):
yield element
ret = _asyncgen_wrapper
@@ -232,7 +209,7 @@ async def drop_db() -> None:
async def insert_or_update_track(
track_data: dict[str, typing.Any],
*,
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
) -> None:
"""Insert a new track or update an existing one in the database asynchronously.
@@ -263,9 +240,7 @@ async def insert_or_update_track(
@sessionize
async def cleanup_unused_genres(
*, session: sqlalchemy.ext.asyncio.AsyncSession = None
) -> None:
async def cleanup_unused_genres(*, session: AsyncSession | None = None) -> None:
"""Remove genres that are no longer used by any tracks.
Keyword Args:
@@ -292,7 +267,7 @@ async def set_timestamp(
event: enums.Event,
operation: enums.TimedEvent,
*,
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
) -> datetime.datetime:
"""Mark the start or end timestamps in the statistics model.
@@ -344,7 +319,7 @@ async def set_timestamp(
@sessionize
async def remove_holiday_by_id(
holiday_id: int, *, session: sqlalchemy.ext.asyncio.AsyncSession = None
holiday_id: int, *, session: AsyncSession | None = None
) -> None:
"""Remove a holiday from the database by ID.
@@ -383,7 +358,7 @@ async def update_statistics(
new_max_holiday_favorite_base: int,
new_max_holiday_general_base: int,
*,
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
) -> None:
"""Update the statistics model.
@@ -462,7 +437,7 @@ async def fetch_tracks_for_sublist(
limit: int,
playlist_entries: set[int],
*,
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
) -> collections.abc.AsyncGenerator[models.Track, None]:
"""Fetch tracks for a given sublist type within a category.
@@ -518,7 +493,7 @@ async def insert_into_playlist( # noqa: C901
sublist: enums.Sublist,
is_favorite: bool,
*,
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
) -> None:
"""Insert a track into the playlist.
@@ -552,7 +527,7 @@ async def insert_into_playlist( # noqa: C901
@sessionize
async def gen_all_category_ids(
*, session: sqlalchemy.ext.asyncio.AsyncSession = None
*, session: AsyncSession | None = None
) -> collections.abc.AsyncGenerator[int, None, None]:
"""Generate the current category ids from the database.
@@ -570,7 +545,7 @@ async def gen_all_category_ids(
@sessionize
async def gen_active_holiday_ids(
*, session: sqlalchemy.ext.asyncio.AsyncSession = None
*, session: AsyncSession | None = None
) -> collections.abc.AsyncGenerator[int, None, None]:
"""Generate the current active holiday ids from the database.
@@ -589,7 +564,7 @@ async def gen_active_holiday_ids(
@sessionize
async def get_existing_playlist_track_info(
*,
session: sqlalchemy.ext.asyncio.AsyncSession = None,
session: AsyncSession | None = None,
) -> tuple[set[int], dict[tuple[bool, int, bool, enums.Sublist], int]]:
"""Get the existing playlist track information.
@@ -623,9 +598,7 @@ async def get_existing_playlist_track_info(
@sessionize
async def get_statistics(
*, session: sqlalchemy.ext.asyncio.AsyncSession = None
) -> models.Statistics:
async def get_statistics(*, session: AsyncSession | None = None) -> models.Statistics:
"""Get the statistics object for the database.
Keyword Args:
@@ -641,7 +614,7 @@ async def get_statistics(
@sessionize
async def remove_track_from_playlist(
server_id: str, *, session: sqlalchemy.ext.asyncio.AsyncSession = None
server_id: str, *, session: AsyncSession | None = None
):
"""Directly removes all entries from PlaylistTrack that are associated with a given server_id.