move compile_sql_or_scalar meth to decorators.py

This commit is contained in:
Jakub Miazek 2024-10-17 09:20:28 +02:00
parent adfb0496b9
commit b18f150b6b
2 changed files with 43 additions and 17 deletions

View File

@ -1,6 +1,5 @@
import uuid import uuid
from sqlalchemy.dialects import postgresql
from sqlalchemy import String, select, ForeignKey from sqlalchemy import String, select, ForeignKey
from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
@ -8,22 +7,7 @@ from sqlalchemy.orm import mapped_column, Mapped, relationship, joinedload
from app.models.base import Base from app.models.base import Base
from app.models.nonsense import Nonsense from app.models.nonsense import Nonsense
from app.utils.decorators import compile_sql_or_scalar
from functools import wraps
def compile_sql_or_scalar(func):
@wraps(func)
async def wrapper(cls, db_session, name, compile_sql=False, *args, **kwargs):
stmt = await func(cls, db_session, name, *args, **kwargs)
if compile_sql:
return stmt.compile(
dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True}
)
result = await db_session.execute(stmt)
return result.scalars().first()
return wrapper
class Stuff(Base): class Stuff(Base):

42
app/utils/decorators.py Normal file
View File

@ -0,0 +1,42 @@
from sqlalchemy.dialects import postgresql
from sqlalchemy.ext.asyncio import AsyncSession
from functools import wraps
def compile_sql_or_scalar(func):
"""
A decorator that compiles an SQL statement or executes it and returns the first scalar result.
Args:
func (Callable): The function to be decorated. It should return an SQLAlchemy statement.
Returns:
Callable: A wrapper function that either compiles the SQL statement or executes it and returns the first scalar result.
"""
@wraps(func)
async def wrapper(cls, db_session, name, compile_sql=False, *args, **kwargs):
"""
Wrapper function that either compiles the SQL statement or executes it.
Args:
cls (Type): The class on which the method is called.
db_session (AsyncSession): The SQLAlchemy async session.
name (str): The name to be used in the SQL statement.
compile_sql (bool, optional): If True, compiles the SQL statement. Defaults to False.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
Compiled SQL statement or the first scalar result of the executed statement.
"""
stmt = await func(cls, db_session, name, *args, **kwargs)
if compile_sql:
return stmt.compile(
dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True}
)
result = await db_session.execute(stmt)
return result.scalars().first()
return wrapper