Fix get_one_or_create IntegrityError handler (PP-1838) (#2134)

Fix an issue with the Exception handler for an IntegrityError in the get_one_or_create function.

If one of the parameters for this function is a sqlalchemy class, then the logging function can cause a PendingRollbackError, since the rollback is happening after the log message.

Update the function to use db.begin_nested() as a context manager, and write a test for this case.

I added type hints to a couple related functions while I was debugging this before i narrowed in on exactly the cause. These hints are included in this PR. They aren't really related to this fix, but they seem harmless to come in with this one. I can break them out to separate PR if desired.

src/palace/manager/api/overdrive.py

@@ -904,7 +904,7 @@ def patron_request(
             return response
 
     def get_patron_credential(
-        self, patron: Patron, pin: str | None, is_fulfillment=False
+        self, patron: Patron, pin: str | None, is_fulfillment: bool = False
     ) -> Credential:
         """Create an OAuth token for the given patron.
 
@@ -913,7 +913,7 @@ def get_patron_credential(
         :param is_fulfillment: Boolean indicating whether we need a fulfillment credential.
         """
 
-        def refresh(credential):
+        def refresh(credential: Credential) -> Credential:
             return self.refresh_patron_access_token(
                 credential, patron, pin, is_fulfillment=is_fulfillment
             )

src/palace/manager/sqlalchemy/model/credential.py

@@ -2,7 +2,8 @@
 
 import datetime
 import uuid
-from typing import TYPE_CHECKING
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
 from sqlalchemy import Column, DateTime, ForeignKey, Index, Integer, String
 from sqlalchemy.orm import Mapped, relationship
@@ -120,15 +121,15 @@ def _filter_invalid_credential(
     @classmethod
     def lookup(
         cls,
-        _db,
-        data_source,
-        token_type,
-        patron,
-        refresher_method,
-        allow_persistent_token=False,
-        allow_empty_token=False,
-        collection=None,
-        force_refresh=False,
+        _db: Session,
+        data_source: DataSource | str,
+        token_type: str,
+        patron: Patron | None,
+        refresher_method: Callable[[Credential], Any] | None,
+        allow_persistent_token: bool = False,
+        allow_empty_token: bool = False,
+        collection: Collection | None = None,
+        force_refresh: bool = False,
     ) -> Credential:
         from palace.manager.sqlalchemy.model.datasource import DataSource
 

src/palace/manager/sqlalchemy/util.py

@@ -11,6 +11,8 @@
 from sqlalchemy.exc import IntegrityError, MultipleResultsFound, NoResultFound
 from sqlalchemy.orm import Session
 
+from palace.manager.util.log import logger_for_function
+
 # This is the lock ID used to ensure that only one circulation manager
 # initializes or migrates the database at a time.
 LOCK_ID_DB_INIT = 1000000001
@@ -121,25 +123,23 @@ def get_one_or_create(
     if one:
         return one, False
     else:
-        __transaction = db.begin_nested()
         try:
-            # These kwargs are supported by get_one() but not by create().
-            get_one_keys = ["on_multiple", "constraint"]
-            for key in get_one_keys:
-                if key in kwargs:
-                    del kwargs[key]
-            obj = create(db, model, create_method, create_method_kwargs, **kwargs)
-            __transaction.commit()
-            return obj
+            with db.begin_nested():
+                # These kwargs are supported by get_one() but not by create().
+                get_one_keys = ["on_multiple", "constraint"]
+                for key in get_one_keys:
+                    if key in kwargs:
+                        del kwargs[key]
+                obj = create(db, model, create_method, create_method_kwargs, **kwargs)
+                return obj
         except IntegrityError as e:
-            logging.info(
+            logger_for_function().debug(
                 "INTEGRITY ERROR on %r %r, %r: %r",
                 model,
                 create_method_kwargs,
                 kwargs,
                 e,
             )
-            __transaction.rollback()
             return db.query(model).filter_by(**kwargs).one(), False
 
 

tests/manager/sqlalchemy/test_util.py

@@ -1,8 +1,15 @@
+import functools
+from unittest.mock import patch
+
 import pytest
 from psycopg2.extras import NumericRange
 from sqlalchemy import not_
 from sqlalchemy.orm.exc import MultipleResultsFound
 
+from palace.manager.service.logging.configuration import LogLevel
+from palace.manager.sqlalchemy import util
+from palace.manager.sqlalchemy.model.credential import Credential
+from palace.manager.sqlalchemy.model.datasource import DataSource
 from palace.manager.sqlalchemy.model.edition import Edition
 from palace.manager.sqlalchemy.util import (
     get_one,
@@ -45,6 +52,46 @@ def test_get_one(self, db: DatabaseTransactionFixture):
         result = get_one(db.session, Edition, constraint=constraint)
         assert None == result
 
+    def test_get_one_or_create(
+        self, db: DatabaseTransactionFixture, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        data_source = DataSource.lookup(db.session, "Test", autocreate=True)
+        patron = db.patron()
+        collection = db.collection()
+
+        get_one_or_create = functools.partial(
+            util.get_one_or_create,
+            db.session,
+            Credential,
+            data_source=data_source,
+            type="Test token",
+            patron=patron,
+            collection=collection,
+        )
+
+        # If it doesn't exist... the function will create it
+        credential, is_new = get_one_or_create()
+        assert is_new
+        assert isinstance(credential, Credential)
+
+        # If it already exists, that gets returned instead
+        credential_existing, is_new = get_one_or_create()
+        assert not is_new
+        assert isinstance(credential_existing, Credential)
+        assert credential_existing == credential
+
+        # If there is a race condition, and an IntegrityError happens
+        # it will be handled, an error message logged, and the existing
+        # object returned
+        caplog.clear()
+        caplog.set_level(LogLevel.debug)
+        with patch.object(util, "get_one", return_value=None):
+            credential_existing, is_new = get_one_or_create()
+        assert not is_new
+        assert isinstance(credential_existing, Credential)
+        assert credential_existing == credential
+        assert "INTEGRITY ERROR" in caplog.text
+
 
 class TestNumericRangeConversion:
     """Test the helper functions that convert between tuples and NumericRange