From 4b51e9a7eeeb219e031e7df235ae3c62f38d331b Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 4 Feb 2023 16:35:21 -0500
Subject: coerce elements in mapper.primary_key, process in __mapper_args__

Repaired ORM Declarative mappings to allow for the
:paramref:`_orm.Mapper.primary_key` parameter to be specified within
``__mapper_args__`` when using :func:`_orm.mapped_column`. Despite this
usage being directly in the 2.0 documentation, the :class:`_orm.Mapper` was
not accepting the :func:`_orm.mapped_column` construct in this context. Ths
feature was already working for the :paramref:`_orm.Mapper.version_id_col`
and :paramref:`_orm.Mapper.polymorphic_on` parameters.

As part of this change, the ``__mapper_args__`` attribute may be specified
without using :func:`_orm.declared_attr` on a non-mapped mixin class,
including a ``"primary_key"`` entry that refers to :class:`_schema.Column`
or :func:`_orm.mapped_column` objects locally present on the mixin;
Declarative will also translate these columns into the correct ones for a
particular mapped class. This again was working already for the
:paramref:`_orm.Mapper.version_id_col` and
:paramref:`_orm.Mapper.polymorphic_on` parameters.  Additionally,
elements within ``"primary_key"`` may be indicated as string names of
existing mapped properties.

Fixes: #9240
Change-Id: Ie2000273289fa23e0af21ef9c6feb3962a8b848c
---
 lib/sqlalchemy/orm/mapper.py | 57 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 51 insertions(+), 6 deletions(-)

(limited to 'lib/sqlalchemy/orm/mapper.py')

diff --git a/lib/sqlalchemy/orm/mapper.py b/lib/sqlalchemy/orm/mapper.py
index a3b209e4a..660c61691 100644
--- a/lib/sqlalchemy/orm/mapper.py
+++ b/lib/sqlalchemy/orm/mapper.py
@@ -83,6 +83,7 @@ from ..sql import util as sql_util
 from ..sql import visitors
 from ..sql.cache_key import MemoizedHasCacheKey
 from ..sql.elements import KeyedColumnElement
+from ..sql.schema import Column
 from ..sql.schema import Table
 from ..sql.selectable import LABEL_STYLE_TABLENAME_PLUS_COL
 from ..util import HasMemoized
@@ -112,7 +113,6 @@ if TYPE_CHECKING:
     from ..sql.base import ReadOnlyColumnCollection
     from ..sql.elements import ColumnClause
     from ..sql.elements import ColumnElement
-    from ..sql.schema import Column
     from ..sql.selectable import FromClause
     from ..util import OrderedSet
 
@@ -650,11 +650,15 @@ class Mapper(
                :ref:`orm_mapping_classes_toplevel`
 
         :param primary_key: A list of :class:`_schema.Column`
-           objects which define
+           objects, or alternatively string names of attribute names which
+           refer to :class:`_schema.Column`, which define
            the primary key to be used against this mapper's selectable unit.
            This is normally simply the primary key of the ``local_table``, but
            can be overridden here.
 
+           .. versionchanged:: 2.0.2 :paramref:`_orm.Mapper.primary_key`
+              arguments may be indicated as string attribute names as well.
+
            .. seealso::
 
                 :ref:`mapper_primary_key` - background and example use
@@ -1557,6 +1561,29 @@ class Mapper(
 
         self.__dict__.pop("_configure_failed", None)
 
+    def _str_arg_to_mapped_col(self, argname: str, key: str) -> Column[Any]:
+        try:
+            prop = self._props[key]
+        except KeyError as err:
+            raise sa_exc.ArgumentError(
+                f"Can't determine {argname} column '{key}' - "
+                "no attribute is mapped to this name."
+            ) from err
+        try:
+            expr = prop.expression
+        except AttributeError as ae:
+            raise sa_exc.ArgumentError(
+                f"Can't determine {argname} column '{key}'; "
+                "property does not refer to a single mapped Column"
+            ) from ae
+        if not isinstance(expr, Column):
+            raise sa_exc.ArgumentError(
+                f"Can't determine {argname} column '{key}'; "
+                "property does not refer to a single "
+                "mapped Column"
+            )
+        return expr
+
     def _configure_pks(self) -> None:
         self.tables = sql_util.find_tables(self.persist_selectable)
 
@@ -1585,10 +1612,28 @@ class Mapper(
                 all_cols
             )
 
+        if self._primary_key_argument:
+
+            coerced_pk_arg = [
+                self._str_arg_to_mapped_col("primary_key", c)
+                if isinstance(c, str)
+                else c
+                for c in (
+                    coercions.expect(  # type: ignore
+                        roles.DDLConstraintColumnRole,
+                        coerce_pk,
+                        argname="primary_key",
+                    )
+                    for coerce_pk in self._primary_key_argument
+                )
+            ]
+        else:
+            coerced_pk_arg = None
+
         # if explicit PK argument sent, add those columns to the
         # primary key mappings
-        if self._primary_key_argument:
-            for k in self._primary_key_argument:
+        if coerced_pk_arg:
+            for k in coerced_pk_arg:
                 if k.table not in self._pks_by_table:
                     self._pks_by_table[k.table] = util.OrderedSet()
                 self._pks_by_table[k.table].add(k)
@@ -1625,12 +1670,12 @@ class Mapper(
             # determine primary key from argument or persist_selectable pks
             primary_key: Collection[ColumnElement[Any]]
 
-            if self._primary_key_argument:
+            if coerced_pk_arg:
                 primary_key = [
                     cc if cc is not None else c
                     for cc, c in (
                         (self.persist_selectable.corresponding_column(c), c)
-                        for c in self._primary_key_argument
+                        for c in coerced_pk_arg
                     )
                 ]
             else:
-- 
cgit v1.2.1