From 84d02f7c8a83e65e7723dd36683d8d845e14613c Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Sat, 6 Dec 2014 10:21:14 +0200
Subject: [PATCH] Add loose typed column support for get_column_key

---
 CHANGES.rst                            |  8 ++++++++
 sqlalchemy_utils/functions/orm.py      | 21 ++++++++++++++-------
 tests/functions/test_get_column_key.py | 14 ++++++++++++--
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index e749aa5..83dcdc3 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -3,6 +3,14 @@ Changelog
 
 Here you can see the full list of changes between each SQLAlchemy-Utils release.
 
+
+0.27.11 (2014-12-06)
+^^^^^^^^^^^^^^^^^^^^
+
+- Added loose typed column checking support for get_column_key
+- Made get_column_key throw UnmappedColumnError to be consistent with SQLAlchemy
+
+
 0.27.10 (2014-12-03)
 ^^^^^^^^^^^^^^^^^^^^
 
diff --git a/sqlalchemy_utils/functions/orm.py b/sqlalchemy_utils/functions/orm.py
index 954ddb5..05a70e3 100644
--- a/sqlalchemy_utils/functions/orm.py
+++ b/sqlalchemy_utils/functions/orm.py
@@ -35,14 +35,21 @@ def get_column_key(model, column):
         get_column_key(User, User.__table__.c.name)  # 'name'
 
     .. versionadded: 0.26.5
+
+    .. versionchanged: 0.27.11
+        Throws UnmappedColumnError instead of ValueError when no property was
+        found for given column. This is consistent with how SQLAlchemy works.
     """
-    for key, c in sa.inspect(model).columns.items():
-        if c is column:
-            return key
-    raise ValueError(
-        "Class %s doesn't have a column '%s'",
-        model.__name__,
-        column
+    mapper = sa.inspect(model)
+    try:
+        return mapper.get_property_by_column(column).key
+    except sa.orm.exc.UnmappedColumnError:
+        for key, c in mapper.columns.items():
+            if c.name == column.name and c.table is column.table:
+                return key
+    raise sa.orm.exc.UnmappedColumnError(
+        'No column %s is configured on mapper %s...' %
+        (column, mapper)
     )
 
 
diff --git a/tests/functions/test_get_column_key.py b/tests/functions/test_get_column_key.py
index 8211d28..c7ed02f 100644
--- a/tests/functions/test_get_column_key.py
+++ b/tests/functions/test_get_column_key.py
@@ -1,3 +1,4 @@
+from copy import copy
 from pytest import raises
 
 import sqlalchemy as sa
@@ -15,7 +16,12 @@ class TestGetColumnKey(object):
             id = sa.Column(sa.Integer, primary_key=True)
             name = sa.Column('_name', sa.Unicode(255))
 
+        class Movie(Base):
+            __tablename__ = 'movie'
+            id = sa.Column(sa.Integer, primary_key=True)
+
         self.Building = Building
+        self.Movie = Movie
 
     def test_supports_aliases(self):
         assert (
@@ -29,6 +35,10 @@ class TestGetColumnKey(object):
             'name'
         )
 
+    def test_supports_vague_matching_of_column_objects(self):
+        column = copy(self.Building.__table__.c._name)
+        assert get_column_key(self.Building, column) == 'name'
+
     def test_throws_value_error_for_unknown_column(self):
-        with raises(ValueError):
-            get_column_key(self.Building, 'unknown')
+        with raises(sa.orm.exc.UnmappedColumnError):
+            get_column_key(self.Building, self.Movie.__table__.c.id)