SQL Alchemy Wrapper Decorator Code

# Decorator for managing the database session
def db_session_decorator(func):
    @wraps(func)
    async def session_wrapper(*args, **kwargs):
        async with AsyncSessionLocal() as session:
            try:
                # Pass the session to the wrapped function
                response = await func(session, *args, **kwargs)
                await session.commit()
                return response
            except IntegrityError as e:
                await session.rollback()
                logging.error("Integrity error detected.", exc_info=True)
                raise
            except (DataError, StatementError) as e:
                await session.rollback()
                logging.error("Data integrity or statement error detected.", exc_info=True)
                raise
            except (OperationalError, DisconnectionError, TimeoutError) as e:
                await session.rollback()
                logging.error("Connection or service availability issue.", exc_info=True)
                raise
            except ConnectionDoesNotExistError as e:
                await session.rollback()
                logging.error("Operation on closed connection attempted.", exc_info=True)
                raise
            except CannotConnectNowError as e:
                await session.rollback()
                logging.error("Cannot connect to the database.", exc_info=True)
                raise
            except MemoryError as e:
                await session.rollback()
                logging.critical("Memory or resource exhaustion detected.", exc_info=True)
                raise
            except PostgresError as e:  # Catch-all for Postgres related errors not caught by SQLAlchemy
                await session.rollback()
                logging.error("Postgres error detected.", exc_info=True)
                raise
            except Exception as e:
                # Rollback for any other exception not explicitly handled above
                await session.rollback()
                # Log the exception details
                logging.error("An unexpected error occurred.", exc_info=True)
                return {'error': str(e)}
    return session_wrapper

Example Usage of the Decorator

While this isn't the cleanest code I've written, it's an example I had on hand of how the decorator might be used.

async def get_object(modeltype, uuid, session=None):
    has_outside_session = session is not None
    # If no session was passed, create a new session for this call
    if not has_outside_session:
        session = AsyncSessionLocal()
    try:
        result = await session.execute(
            select(modeltype).filter_by(uuid=int(uuid))
        )
        instance = result.scalars().first()
        # If a new session was created, commit any changes (if any were made, although get_object is typically read-only) and close the session
        if not has_outside_session:
            await session.commit()
        return instance
    except Exception as e:
        # If a new session was created, rollback any changes due to an error
        if not has_outside_session:
            await session.rollback()
        raise
    finally:
        # If a new session was created, close it when done
        if not has_outside_session:
            await session.close()

@db_session_decorator
async def init_user(session, userid, username):
    user = await get_object(User, guildid, session)
    if not guild:
        parent = User(uuid=userid, registered_at=str(datetime.now()), guildname=username)
        session.add(parent)