Extract shared primary key utility function

Consolidates duplicated PK inspection logic from BaseRepository._model_pk()
and result_presenters._pk_from_result_object() into a single
get_model_pk_name() function in common.py.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: febus982

sqlalchemy_bind_manager/_repository/base_repository.py

@@ -31,7 +31,7 @@
     Union,
 )
 
-from sqlalchemy import asc, desc, func, inspect, select
+from sqlalchemy import asc, desc, func, select
 from sqlalchemy.orm import Mapper, aliased, class_mapper, lazyload
 from sqlalchemy.orm.exc import UnmappedClassError
 from sqlalchemy.sql import Select
@@ -41,6 +41,7 @@
 from .common import (
     MODEL,
     CursorReference,
+    get_model_pk_name,
 )
 
 
@@ -342,11 +343,7 @@ def _model_pk(self) -> str:
 
         :return:
         """
-        primary_keys = inspect(self._model).primary_key  # type: ignore
-        if len(primary_keys) > 1:
-            raise NotImplementedError("Composite primary keys are not supported.")
-
-        return primary_keys[0].name
+        return get_model_pk_name(self._model)
 
     def _fail_if_invalid_models(self, objects: Iterable[MODEL]) -> None:
         if any(not isinstance(x, self._model) for x in objects):

sqlalchemy_bind_manager/_repository/common.py

@@ -18,15 +18,29 @@
 #  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 #  DEALINGS IN THE SOFTWARE.
 
-from typing import Generic, List, TypeVar, Union
+from typing import Generic, List, Type, TypeVar, Union
 from uuid import UUID
 
 from pydantic import BaseModel, StrictInt, StrictStr
+from sqlalchemy import inspect
 
 MODEL = TypeVar("MODEL")
 PRIMARY_KEY = Union[str, int, tuple, dict, UUID]
 
 
+def get_model_pk_name(model_class: Type) -> str:
+    """Retrieves the primary key column name from a SQLAlchemy model class.
+
+    :param model_class: A SQLAlchemy model class
+    :return: The name of the primary key column
+    :raises NotImplementedError: If the model has composite primary keys
+    """
+    primary_keys = inspect(model_class).primary_key  # type: ignore
+    if len(primary_keys) > 1:
+        raise NotImplementedError("Composite primary keys are not supported.")
+    return primary_keys[0].name
+
+
 class PageInfo(BaseModel):
     """
     Paginated query metadata.

sqlalchemy_bind_manager/_repository/result_presenters.py

@@ -21,15 +21,14 @@
 from math import ceil
 from typing import List, Union
 
-from sqlalchemy import inspect
-
 from .common import (
     MODEL,
     CursorPageInfo,
     CursorPaginatedResult,
     CursorReference,
     PageInfo,
     PaginatedResult,
+    get_model_pk_name,
 )
 
 
@@ -93,7 +92,7 @@ def _build_no_cursor_result(
         has_next_page = len(result_items) > items_per_page
         if has_next_page:
             result_items = result_items[0:items_per_page]
-        reference_column = _pk_from_result_object(result_items[0])
+        reference_column = get_model_pk_name(type(result_items[0]))
 
         return CursorPaginatedResult(
             items=result_items,
@@ -237,11 +236,3 @@ def build_result(
                 has_previous_page=has_previous_page,
             ),
         )
-
-
-def _pk_from_result_object(model) -> str:
-    primary_keys = inspect(type(model)).primary_key  # type: ignore
-    if len(primary_keys) > 1:
-        raise NotImplementedError("Composite primary keys are not supported.")
-
-    return primary_keys[0].name

tests/repository/result_presenters/test_composite_pk.py

@@ -2,15 +2,15 @@
 
 import pytest
 
-from sqlalchemy_bind_manager._repository.result_presenters import _pk_from_result_object
+from sqlalchemy_bind_manager._repository.common import get_model_pk_name
 
 
 def test_exception_raised_if_multiple_primary_keys():
     with (
         patch(
-            "sqlalchemy_bind_manager._repository.result_presenters.inspect",
+            "sqlalchemy_bind_manager._repository.common.inspect",
             return_value=Mock(primary_key=["1", "2"]),
         ),
         pytest.raises(NotImplementedError),
     ):
-        _pk_from_result_object("irrelevant")
+        get_model_pk_name(str)