diff --git a/warehouse/db.py b/warehouse/db.py index 5f111b8626b5..79b1b68f5118 100644 --- a/warehouse/db.py +++ b/warehouse/db.py @@ -27,7 +27,14 @@ from sqlalchemy import event, func, inspect from sqlalchemy.dialects.postgresql import UUID as PG_UUID from sqlalchemy.exc import IntegrityError, OperationalError -from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker +from sqlalchemy.orm import ( + DeclarativeBase, + Mapped, + Session as BaseSession, + mapped_column, + raiseload, + sessionmaker, +) from warehouse.metrics import IMetricsService from warehouse.utils.attrs import make_repr @@ -102,10 +109,18 @@ class Model(ModelBase): ) +# Custom Session to prevent lazy-loading +class StrictSession(BaseSession): + def query(self, *entities, **kwargs): + query = super().query(*entities) + query = query.options(raiseload("*")) + return query + + # Create our session class here, this will stay stateless as we'll bind the # engine to each new state we create instead of binding it to the session # class. -Session = sessionmaker() +Session = sessionmaker(class_=StrictSession) def listens_for(target, identifier, *args, **kwargs):