Merge pull request #620 from datastax/584

PYTHON-584: Fix pk__token equality filter

cassandra/cqlengine/query.py

@@ -545,7 +545,7 @@ class AbstractQuerySet(object):
         if len(statement) == 1:
             return arg, None
         elif len(statement) == 2:
-            return statement[0], statement[1]
+            return (statement[0], statement[1]) if arg != 'pk__token' else (arg, None)
         else:
             raise QueryException("Can't parse '{0}'".format(arg))
 
@@ -954,7 +954,8 @@ class ModelQuerySet(AbstractQuerySet):
     def _validate_select_where(self):
         """ Checks that a filterset will not create invalid select statement """
         # check that there's either a =, a IN or a CONTAINS (collection) relationship with a primary key or indexed field
-        equal_ops = [self.model._get_column_by_db_name(w.field) for w in self._where if isinstance(w.operator, EqualsOperator)]
+        equal_ops = [self.model._get_column_by_db_name(w.field) \
+                     for w in self._where if isinstance(w.operator, EqualsOperator) and not isinstance(w.value, Token)]
         token_comparison = any([w for w in self._where if isinstance(w.value, Token)])
         if not any(w.primary_key or w.index for w in equal_ops) and not token_comparison and not self._allow_filtering:
             raise QueryException(('Where clauses require either  =, a IN or a CONTAINS (collection) '

tests/integration/cqlengine/query/test_queryoperators.py

@@ -72,7 +72,7 @@ class TestTokenFunction(BaseCassEngTestCase):
         super(TestTokenFunction, self).tearDown()
         drop_table(TokenTestModel)
 
-    @execute_count(14)
+    @execute_count(15)
     def test_token_function(self):
         """ Tests that token functions work properly """
         assert TokenTestModel.objects().count() == 0
@@ -91,6 +91,10 @@ class TestTokenFunction(BaseCassEngTestCase):
         assert len(seen_keys) == 10
         assert all([i in seen_keys for i in range(10)])
 
+        # pk__token equality
+        r = TokenTestModel.objects(pk__token=functions.Token(last_token))
+        self.assertEqual(len(r), 1)
+
     def test_compound_pk_token_function(self):
 
         class TestModel(Model):