Merge "Add support for key-word based get()"

3 files changed, 103 insertions, 10 deletions

diff --git a/doc/build/changelog/unreleased_13/4316.rst b/doc/build/changelog/unreleased_13/4316.rst
new file mode 100644
index 000000000..dfab94043
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/4316.rst
@@ -0,0 +1,8 @@
+.. change::
+   :tags: feature, orm
+   :tickets: 4316
+
+   The :meth:`.Query.get` method can now accept a dictionary of attribute keys
+   and values as a means of indicating the primary key value to load; is
+   particularly useful for composite primary keys.  Pull request courtesy
+   Sanjana S.
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index f86cb9085..db8d8bcbe 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -883,6 +883,9 @@ class Query(object):
 
             some_object = session.query(VersionedFoo).get((5, 10))
 
+            some_object = session.query(VersionedFoo).get(
+                {"id": 5, "version_id": 10})
+
         :meth:`~.Query.get` is special in that it provides direct
         access to the identity map of the owning :class:`.Session`.
         If the given primary key identifier is present
@@ -918,14 +921,37 @@ class Query(object):
         before querying the database.  See :doc:`/orm/loading_relationships`
         for further details on relationship loading.
 
-        :param ident: A scalar or tuple value representing
-         the primary key.   For a composite primary key,
-         the order of identifiers corresponds in most cases
-         to that of the mapped :class:`.Table` object's
-         primary key columns.  For a :func:`.mapper` that
-         was given the ``primary key`` argument during
-         construction, the order of identifiers corresponds
-         to the elements present in this collection.
+        :param ident: A scalar, tuple, or dictionary representing the
+         primary key.  For a composite (e.g. multiple column) primary key,
+         a tuple or dictionary should be passed.
+
+         For a single-column primary key, the scalar calling form is typically
+         the most expedient.  If the primary key of a row is the value "5",
+         the call looks like::
+
+            my_object = query.get(5)
+
+         The tuple form contains primary key values typically in
+         the order in which they correspond to the mapped :class:`.Table`
+         object's primary key columns, or if the
+         :paramref:`.Mapper.primary_key` configuration parameter were used, in
+         the order used for that parameter. For example, if the primary key
+         of a row is represented by the integer
+         digits "5, 10" the call would look like::
+
+             my_object = query.get((5, 10))
+
+         The dictionary form should include as keys the mapped attribute names
+         corresponding to each element of the primary key.  If the mapped class
+         has the attributes ``id``, ``version_id`` as the attributes which
+         store the object's primary key value, the call would look like::
+
+            my_object = query.get({"id": 5, "version_id": 10})
+
+         .. versionadded:: 1.3 the :meth:`.Query.get` method now optionally
+            accepts a dictionary of attribute names to values in order to
+            indicate a primary key identifier.
+
 
         :return: The object instance, or ``None``.
 
@@ -991,10 +1017,12 @@ class Query(object):
         if hasattr(primary_key_identity, "__composite_values__"):
             primary_key_identity = primary_key_identity.__composite_values__()
 
-        primary_key_identity = util.to_list(primary_key_identity)
-
         mapper = self._only_full_mapper_zero("get")
 
+        is_dict = isinstance(primary_key_identity, dict)
+        if not is_dict:
+            primary_key_identity = util.to_list(primary_key_identity)
+
         if len(primary_key_identity) != len(mapper.primary_key):
             raise sa_exc.InvalidRequestError(
                 "Incorrect number of values in identifier to formulate "
@@ -1002,6 +1030,23 @@ class Query(object):
                 % ",".join("'%s'" % c for c in mapper.primary_key)
             )
 
+        if is_dict:
+            try:
+                primary_key_identity = list(
+                    primary_key_identity[prop.key]
+                    for prop in mapper._identity_key_props
+                )
+
+            except KeyError:
+                raise sa_exc.InvalidRequestError(
+                    "Incorrect names of values in identifier to formulate "
+                    "primary key for query.get(); primary key attribute names"
+                    " are %s" % ",".join(
+                        "'%s'" % prop.key
+                        for prop in mapper._identity_key_props
+                    )
+                )
+
         if (
             not self._populate_existing
             and not mapper.always_refresh
diff --git a/test/orm/test_query.py b/test/orm/test_query.py
index 01dfe204e..935e3e31c 100644
--- a/test/orm/test_query.py
+++ b/test/orm/test_query.py
@@ -633,6 +633,46 @@ class RawSelectTest(QueryTest, AssertsCompiledSQL):
 
 
 class GetTest(QueryTest):
+    def test_get_composite_pk_keyword_based_no_result(self):
+        CompositePk = self.classes.CompositePk
+
+        s = Session()
+        is_(s.query(CompositePk).get({"i": 100, "j": 100}), None)
+
+    def test_get_composite_pk_keyword_based_result(self):
+        CompositePk = self.classes.CompositePk
+
+        s = Session()
+        one_two = s.query(CompositePk).get({"i": 1, "j": 2})
+        eq_(one_two.i, 1)
+        eq_(one_two.j, 2)
+        eq_(one_two.k, 3)
+
+    def test_get_composite_pk_keyword_based_wrong_keys(self):
+        CompositePk = self.classes.CompositePk
+
+        s = Session()
+        q = s.query(CompositePk)
+        assert_raises(sa_exc.InvalidRequestError, q.get, {"i": 1, "k": 2})
+
+    def test_get_composite_pk_keyword_based_too_few_keys(self):
+        CompositePk = self.classes.CompositePk
+
+        s = Session()
+        q = s.query(CompositePk)
+        assert_raises(sa_exc.InvalidRequestError, q.get, {"i": 1})
+
+    def test_get_composite_pk_keyword_based_too_many_keys(self):
+        CompositePk = self.classes.CompositePk
+
+        s = Session()
+        q = s.query(CompositePk)
+        assert_raises(
+            sa_exc.InvalidRequestError,
+            q.get,
+            {"i": 1, "j": '2', "k": 3}
+        )
+
     def test_get(self):
         User = self.classes.User