From 5d54a3054381817ccd08351cc71c506cce20d759 Mon Sep 17 00:00:00 2001 From: Dustin Ingram Date: Mon, 27 Mar 2023 21:31:14 +0000 Subject: [PATCH] Use `raiseload` option for all queries --- warehouse/db.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/warehouse/db.py b/warehouse/db.py index 13966f4a54a3..18d0924ffb87 100644 --- a/warehouse/db.py +++ b/warehouse/db.py @@ -23,7 +23,7 @@ from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.exc import IntegrityError, OperationalError from sqlalchemy.ext.declarative import declarative_base # type: ignore -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import sessionmaker, Session as BaseSession, raiseload from warehouse.metrics import IMetricsService from warehouse.utils.attrs import make_repr @@ -92,10 +92,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):