get_or_404 passes kwargs to session.get

Author: jdimmerman

CHANGES.rst

@@ -4,6 +4,7 @@ Version 3.1.0
 Unreleased
 
 -   Remove previously deprecated code.
+-   Pass extra keyword arguments from ``get_or_404`` to ``session.get``. :issue:`1149`
 
 
 Version 3.0.3

src/flask_sqlalchemy/extension.py

@@ -597,18 +597,27 @@ def engine(self) -> sa.engine.Engine:
         return self.engines[None]
 
     def get_or_404(
-        self, entity: t.Type[t.Any], ident: t.Any, *, description: str | None = None
+        self,
+        entity: t.Type[t.Any],
+        ident: t.Any,
+        *,
+        description: str | None = None,
+        **kwargs: t.Any,
     ) -> t.Any:
         """Like :meth:`session.get() <sqlalchemy.orm.Session.get>` but aborts with a
         ``404 Not Found`` error instead of returning ``None``.
 
         :param entity: The model class to query.
         :param ident: The primary key to query.
         :param description: A custom message to show on the error page.
+        :param kwargs: Extra arguments passed to ``session.get()``.
+
+        .. versionchanged:: 3.1
+            Pass extra keyword arguments to ``session.get()``.
 
         .. versionadded:: 3.0
         """
-        value = self.session.get(entity, ident)
+        value = self.session.get(entity, ident, **kwargs)
 
         if value is None:
             abort(404, description=description)

tests/test_extension_object.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import typing as t
+
+import pytest
+import sqlalchemy as sa
+from flask import Flask
+from sqlalchemy.orm import joinedload
+from werkzeug.exceptions import NotFound
+
+from flask_sqlalchemy import SQLAlchemy
+from flask_sqlalchemy.record_queries import get_recorded_queries
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_get_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    item = Todo()
+    db.session.add(item)
+    db.session.commit()
+    assert db.get_or_404(Todo, 1) is item
+
+    with pytest.raises(NotFound):
+        db.get_or_404(Todo, 2)
+
+
+def test_get_or_404_kwargs(app: Flask) -> None:
+    app.config["SQLALCHEMY_RECORD_QUERIES"] = True
+    db = SQLAlchemy(app)
+
+    class User(db.Model):
+        id = sa.Column(db.Integer, primary_key=True)
+
+    class Todo(db.Model):
+        id = sa.Column(sa.Integer, primary_key=True)
+        user_id = sa.Column(sa.ForeignKey(User.id))
+        user = db.relationship(User)
+
+    with app.app_context():
+        db.create_all()
+        db.session.add(Todo(user=User()))
+        db.session.commit()
+
+    with app.app_context():
+        item = db.get_or_404(Todo, 1, options=[joinedload(Todo.user)])
+        assert item.user.id == 1
+        # one query with join, no second query when accessing relationship
+        assert len(get_recorded_queries()) == 1