From 5c3eb8e6f23742239bace4dd721b10dfe0890f76 Mon Sep 17 00:00:00 2001
From: Adam Holmberg <Adam.Holmberg@datastax.com>
Date: Tue, 22 Mar 2016 14:18:16 -0500
Subject: [PATCH] cqle: make NameTable produce part. keys for Token function

PYTHON-260
---
 cassandra/cqlengine/columns.py   |  9 ---------
 cassandra/cqlengine/functions.py |  2 --
 cassandra/cqlengine/named.py     | 20 ++++++++++++++++++++
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/cassandra/cqlengine/columns.py b/cassandra/cqlengine/columns.py
index 3c353d90..1981287d 100644
--- a/cassandra/cqlengine/columns.py
+++ b/cassandra/cqlengine/columns.py
@@ -881,12 +881,3 @@ class _PartitionKeysToken(Column):
     @property
     def db_field_name(self):
         return 'token({0})'.format(', '.join(['"{0}"'.format(c.db_field_name) for c in self.partition_columns]))
-
-    def to_database(self, value):
-        from cqlengine.functions import Token
-        assert isinstance(value, Token)
-        value.set_columns(self.partition_columns)
-        return value
-
-    def get_cql(self):
-        return "token({0})".format(", ".join(c.cql for c in self.partition_columns))
diff --git a/cassandra/cqlengine/functions.py b/cassandra/cqlengine/functions.py
index 0200f1fe..ccfe9de9 100644
--- a/cassandra/cqlengine/functions.py
+++ b/cassandra/cqlengine/functions.py
@@ -100,14 +100,12 @@ class MaxTimeUUID(TimeUUIDQueryFunction):
     format_string = 'MaxTimeUUID(%({0})s)'
 
 
-
 class Token(BaseQueryFunction):
     """
     compute the token for a given partition key
 
     http://cassandra.apache.org/doc/cql3/CQL.html#tokenFun
     """
-
     def __init__(self, *values):
         if len(values) == 1 and isinstance(values[0], (list, tuple)):
             values = values[0]
diff --git a/cassandra/cqlengine/named.py b/cassandra/cqlengine/named.py
index f2f26aed..90a0d8fd 100644
--- a/cassandra/cqlengine/named.py
+++ b/cassandra/cqlengine/named.py
@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from cassandra.util import OrderedDict
+
 from cassandra.cqlengine import CQLEngineException
+from cassandra.cqlengine.columns import Column
+from cassandra.cqlengine.connection import get_cluster
 from cassandra.cqlengine.query import AbstractQueryableColumn, SimpleQuerySet
 from cassandra.cqlengine.query import DoesNotExist as _DoesNotExist
 from cassandra.cqlengine.query import MultipleObjectsReturned as _MultipleObjectsReturned
@@ -78,6 +82,8 @@ class NamedTable(object):
 
     objects = QuerySetDescriptor()
 
+    __partition_keys = None
+
     class DoesNotExist(_DoesNotExist):
         pass
 
@@ -88,6 +94,20 @@ class NamedTable(object):
         self.keyspace = keyspace
         self.name = name
 
+    @property
+    def _partition_keys(self):
+        if not self.__partition_keys:
+            self._get_partition_keys()
+        return self.__partition_keys
+
+    def _get_partition_keys(self):
+        try:
+            table_meta = get_cluster().metadata.keyspaces[self.keyspace].tables[self.name]
+            self.__partition_keys = OrderedDict((pk.name, Column(primary_key=True, partition_key=True, db_field=pk.name)) for pk in table_meta.partition_key)
+        except Exception as e:
+            raise CQLEngineException("Failed inspecting partition keys for {0}."
+                                     "Ensure cqlengine is connected before attempting this with NamedTable.".format(self.column_family_name()))
+
     def column(self, name):
         return NamedColumn(name)