From a60a65f10ee9bf631b07d3c60de7dbdccdaa54ff Mon Sep 17 00:00:00 2001 From: alanjds Date: Thu, 20 Aug 2015 14:31:29 +0000 Subject: [PATCH 1/8] Debug for automatic pagination machinery --- cassandra/cluster.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 13afecd245..5cc4aa6f6f 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -2829,7 +2829,9 @@ def start_fetching_next_page(self): .. versionadded:: 2.0.0 """ + log.debug('Pagination: Fetching next page') if not self._paging_state: + log.debug('Pagination: No more pages to fetch') raise QueryExhausted() self._make_query_plan() From ddff533310cfb18e0aa01b3a67b5f15f0a823b06 Mon Sep 17 00:00:00 2001 From: alanjds Date: Thu, 20 Aug 2015 14:34:09 +0000 Subject: [PATCH 2/8] AbstractQuerySet.__iter__ is no more greedy. Fetches on demand. The __iter__ now yields until the end of actual page, than start to fetch the next page. --- cassandra/cqlengine/query.py | 41 ++++++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py index 46918e0eb4..da4cba4fa8 100644 --- a/cassandra/cqlengine/query.py +++ b/cassandra/cqlengine/query.py @@ -289,8 +289,12 @@ def __init__(self, model): # results cache self._result_cache = None self._result_idx = None + self._result_generator = None self._distinct_fields = None + + self._count = None + self._batch = None self._ttl = getattr(model, '__default_ttl__', None) self._consistency = None @@ -341,7 +345,7 @@ def __deepcopy__(self, memo): def __len__(self): self._execute_query() - return len(self._result_cache) + return self.count() # ----query generation / execution---- @@ -374,7 +378,8 @@ def _execute_query(self): if self._batch: raise CQLEngineException("Only inserts, updates, and deletes are available in batch mode") if self._result_cache is None: - self._result_cache = list(self._execute(self._select_query())) + self._result_generator = (i for i in self._execute(self._select_query())) + self._result_cache = [] self._construct_result = self._get_result_constructor() def _fill_result_cache_to_idx(self, idx): @@ -388,13 +393,24 @@ def _fill_result_cache_to_idx(self, idx): else: for idx in range(qty): self._result_idx += 1 - self._result_cache[self._result_idx] = self._construct_result(self._result_cache[self._result_idx]) + while True: + try: + self._result_cache[self._result_idx] = self._construct_result(self._result_cache[self._result_idx]) + break + except IndexError: + self._result_cache.append(next(self._result_generator)) def __iter__(self): self._execute_query() - for idx in range(len(self._result_cache)): - instance = self._result_cache[idx] + for idx in range(self.count()): + while True: + try: + instance = self._result_cache[idx] + break + except IndexError: + self._result_cache.append(next(self._result_generator)) + if isinstance(instance, dict): self._fill_result_cache_to_idx(idx) yield self._result_cache[idx] @@ -402,7 +418,7 @@ def __iter__(self): def __getitem__(self, s): self._execute_query() - num_results = len(self._result_cache) + num_results = self.count() if isinstance(s, slice): # calculate the amount of results that need to be loaded @@ -615,10 +631,10 @@ def get(self, *args, **kwargs): return self.filter(*args, **kwargs).get() self._execute_query() - if len(self._result_cache) == 0: + if self.count() == 0: raise self.model.DoesNotExist - elif len(self._result_cache) > 1: - raise self.model.MultipleObjectsReturned('{0} objects found'.format(len(self._result_cache))) + elif self.count() > 1: + raise self.model.MultipleObjectsReturned('{0} objects found'.format(self.count())) else: return self[0] @@ -679,13 +695,12 @@ def count(self): if self._batch: raise CQLEngineException("Only inserts, updates, and deletes are available in batch mode") - if self._result_cache is None: + if self._count is None: query = self._select_query() query.count = True result = self._execute(query) - return result[0]['count'] - else: - return len(self._result_cache) + self._count = result[0]['count'] + return self._count def distinct(self, distinct_fields=None): """ From 2feb57fe96140b67cbcd7278012bb2480e84e8f0 Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Wed, 2 Mar 2016 16:18:30 -0500 Subject: [PATCH 3/8] perf improvements on cqlengine queryset __iter__ and __getitem__ --- cassandra/cqlengine/query.py | 53 ++++++++++++++++++------------- cassandra/cqlengine/statements.py | 9 ++++-- 2 files changed, 37 insertions(+), 25 deletions(-) diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py index da4cba4fa8..2d05f7de12 100644 --- a/cassandra/cqlengine/query.py +++ b/cassandra/cqlengine/query.py @@ -328,7 +328,7 @@ def __call__(self, *args, **kwargs): def __deepcopy__(self, memo): clone = self.__class__(self.model) for k, v in self.__dict__.items(): - if k in ['_con', '_cur', '_result_cache', '_result_idx']: # don't clone these + if k in ['_con', '_cur', '_result_cache', '_result_idx', '_result_generator']: # don't clone these clone.__dict__[k] = None elif k == '_batch': # we need to keep the same batch instance across @@ -403,45 +403,53 @@ def _fill_result_cache_to_idx(self, idx): def __iter__(self): self._execute_query() - for idx in range(self.count()): - while True: + idx = 0 + while True: + if len(self._result_cache) <= idx: try: - instance = self._result_cache[idx] - break - except IndexError: self._result_cache.append(next(self._result_generator)) + except StopIteration: + break + instance = self._result_cache[idx] if isinstance(instance, dict): self._fill_result_cache_to_idx(idx) yield self._result_cache[idx] + idx += 1 + def __getitem__(self, s): self._execute_query() - num_results = self.count() - if isinstance(s, slice): # calculate the amount of results that need to be loaded - end = num_results if s.step is None else s.step - if end < 0: - end += num_results - else: - end -= 1 - self._fill_result_cache_to_idx(end) + end = s.stop + if s.start < 0 or s.stop < 0: + end = self.count() + + try: + self._fill_result_cache_to_idx(end) + except StopIteration: + pass + return self._result_cache[s.start:s.stop:s.step] else: - # return the object at this index - s = int(s) + try: + s = int(s) + except (ValueError, TypeError): + raise TypeError('QuerySet indices must be integers') - # handle negative indexing + # Using negative indexing is costly since we have to execute a count() if s < 0: + num_results = self.count() s += num_results - if s >= num_results: - raise IndexError - else: + try: self._fill_result_cache_to_idx(s) - return self._result_cache[s] + except StopIteration: + raise IndexError + + return self._result_cache[s] def _get_result_constructor(self): """ @@ -699,7 +707,8 @@ def count(self): query = self._select_query() query.count = True result = self._execute(query) - self._count = result[0]['count'] + count_row = result[0].popitem() + self._count = count_row[1] return self._count def distinct(self, distinct_fields=None): diff --git a/cassandra/cqlengine/statements.py b/cassandra/cqlengine/statements.py index 39f09307d0..7bf6f27fcf 100644 --- a/cassandra/cqlengine/statements.py +++ b/cassandra/cqlengine/statements.py @@ -577,10 +577,13 @@ def __init__(self, def __unicode__(self): qs = ['SELECT'] - if self.count: + if self.distinct_fields: + if self.count: + qs += ['DISTINCT COUNT({0})'.format(', '.join(['"{0}"'.format(f) for f in self.distinct_fields]))] + else: + qs += ['DISTINCT {0}'.format(', '.join(['"{0}"'.format(f) for f in self.distinct_fields]))] + elif self.count: qs += ['COUNT(*)'] - elif self.distinct_fields: - qs += ['DISTINCT {0}'.format(', '.join(['"{0}"'.format(f) for f in self.distinct_fields]))] else: qs += [', '.join(['"{0}"'.format(f) for f in self.fields]) if self.fields else '*'] qs += ['FROM', self.table] From 58f5e2e033138c1389b003064be8ab0cecfdd08c Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Wed, 2 Mar 2016 16:33:01 -0500 Subject: [PATCH 4/8] Add the ability to disable a query limit with .limit(None) --- cassandra/cqlengine/query.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py index 2d05f7de12..38918ced2d 100644 --- a/cassandra/cqlengine/query.py +++ b/cassandra/cqlengine/query.py @@ -758,7 +758,11 @@ def limit(self, v): for user in User.objects().limit(100): print(user) """ - if not (v is None or isinstance(v, six.integer_types)): + + if v is None: + v = 0 + + if not isinstance(v, six.integer_types): raise TypeError if v == self._limit: return self From 3e379abf8fc2a26d2907eb043b6f08769060362e Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Thu, 3 Mar 2016 13:22:52 -0500 Subject: [PATCH 5/8] Add QuerySet.fetch_size() support --- cassandra/cqlengine/connection.py | 3 +-- cassandra/cqlengine/query.py | 28 +++++++++++++++++++++++++++- cassandra/cqlengine/statements.py | 10 +++++++--- 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/cassandra/cqlengine/connection.py b/cassandra/cqlengine/connection.py index 7d53af6917..1e043822b5 100644 --- a/cassandra/cqlengine/connection.py +++ b/cassandra/cqlengine/connection.py @@ -162,8 +162,7 @@ def execute(query, params=None, consistency_level=None, timeout=NOT_SET): elif isinstance(query, BaseCQLStatement): params = query.get_context() - query = str(query) - query = SimpleStatement(query, consistency_level=consistency_level) + query = SimpleStatement(str(query), consistency_level=consistency_level, fetch_size=query.fetch_size) elif isinstance(query, six.string_types): query = SimpleStatement(query, consistency_level=consistency_level) diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py index 38918ced2d..963c4dceb9 100644 --- a/cassandra/cqlengine/query.py +++ b/cassandra/cqlengine/query.py @@ -302,6 +302,7 @@ def __init__(self, model): self._if_not_exists = False self._timeout = connection.NOT_SET self._if_exists = False + self._fetch_size = None @property def column_family_name(self): @@ -369,7 +370,8 @@ def _select_query(self): order_by=self._order, limit=self._limit, allow_filtering=self._allow_filtering, - distinct_fields=self._distinct_fields + distinct_fields=self._distinct_fields, + fetch_size=self._fetch_size ) # ----Reads------ @@ -774,6 +776,30 @@ def limit(self, v): clone._limit = v return clone + def fetch_size(self, v): + """ + Sets the number of rows that are fetched at a time. + + *Note that driver's default fetch size is 5000. + + .. code-block:: python + + for user in User.objects().fetch_size(500): + print(user) + """ + + if not isinstance(v, six.integer_types): + raise TypeError + if v == self._fetch_size: + return self + + if v < 1: + raise QueryException("fetch size less than 1 is not allowed") + + clone = copy.deepcopy(self) + clone._fetch_size = v + return clone + def allow_filtering(self): """ Enables the (usually) unwise practive of querying on a clustering key without also defining a partition key diff --git a/cassandra/cqlengine/statements.py b/cassandra/cqlengine/statements.py index 7bf6f27fcf..73f9197ed3 100644 --- a/cassandra/cqlengine/statements.py +++ b/cassandra/cqlengine/statements.py @@ -16,6 +16,7 @@ import time import six +from cassandra.query import FETCH_SIZE_UNSET from cassandra.cqlengine import UnicodeMixin from cassandra.cqlengine.functions import QueryValue from cassandra.cqlengine.operators import BaseWhereOperator, InOperator @@ -470,13 +471,14 @@ def __unicode__(self): class BaseCQLStatement(UnicodeMixin): """ The base cql statement class """ - def __init__(self, table, consistency=None, timestamp=None, where=None): + def __init__(self, table, consistency=None, timestamp=None, where=None, fetch_size=None): super(BaseCQLStatement, self).__init__() self.table = table self.consistency = consistency self.context_id = 0 self.context_counter = self.context_id self.timestamp = timestamp + self.fetch_size = fetch_size if fetch_size else FETCH_SIZE_UNSET self.where_clauses = [] for clause in where or []: @@ -556,7 +558,8 @@ def __init__(self, order_by=None, limit=None, allow_filtering=False, - distinct_fields=None): + distinct_fields=None, + fetch_size=None): """ :param where @@ -565,7 +568,8 @@ def __init__(self, super(SelectStatement, self).__init__( table, consistency=consistency, - where=where + where=where, + fetch_size=fetch_size ) self.fields = [fields] if isinstance(fields, six.string_types) else (fields or []) From 84c9c8c04b043a64b65dbee75697c12d79207c63 Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Thu, 3 Mar 2016 14:41:49 -0500 Subject: [PATCH 6/8] Add some tests for cqlengine queryset slicing --- .../cqlengine/query/test_queryset.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/integration/cqlengine/query/test_queryset.py b/tests/integration/cqlengine/query/test_queryset.py index e0cc6ab4b0..a1609f3c1f 100644 --- a/tests/integration/cqlengine/query/test_queryset.py +++ b/tests/integration/cqlengine/query/test_queryset.py @@ -554,16 +554,31 @@ def test_negative_indexing_works_properly(self): def test_slicing_works_properly(self): q = TestModel.objects(test_id=0).order_by('attempt_id') expected_order = [0, 1, 2, 3] + for model, expect in zip(q[1:3], expected_order[1:3]): - assert model.attempt_id == expect + self.assertEqual(model.attempt_id, expect) + + for model, expect in zip(q[0:3:2], expected_order[0:3:2]): + self.assertEqual(model.attempt_id, expect) def test_negative_slicing(self): q = TestModel.objects(test_id=0).order_by('attempt_id') expected_order = [0, 1, 2, 3] + for model, expect in zip(q[-3:], expected_order[-3:]): - assert model.attempt_id == expect + self.assertEqual(model.attempt_id, expect) + for model, expect in zip(q[:-1], expected_order[:-1]): - assert model.attempt_id == expect + self.assertEqual(model.attempt_id, expect) + + for model, expect in zip(q[1:-1], expected_order[1:-1]): + self.assertEqual(model.attempt_id, expect) + + for model, expect in zip(q[-3:-1], expected_order[-3:-1]): + self.assertEqual(model.attempt_id, expect) + + for model, expect in zip(q[-3:-1:2], expected_order[-3:-1:2]): + self.assertEqual(model.attempt_id, expect) class TestQuerySetValidation(BaseQuerySetUsage): From dabb8d412b27aa91b772d625e41c6711eed1b4a1 Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Thu, 3 Mar 2016 15:17:04 -0500 Subject: [PATCH 7/8] Add tests for distcint with count --- tests/integration/cqlengine/query/test_queryset.py | 7 +++++++ .../cqlengine/statements/test_select_statement.py | 3 +++ 2 files changed, 10 insertions(+) diff --git a/tests/integration/cqlengine/query/test_queryset.py b/tests/integration/cqlengine/query/test_queryset.py index a1609f3c1f..7fffe8137b 100644 --- a/tests/integration/cqlengine/query/test_queryset.py +++ b/tests/integration/cqlengine/query/test_queryset.py @@ -485,6 +485,13 @@ def test_zero_result(self): q = TestModel.objects.distinct(['test_id']).filter(test_id__in=[52]) self.assertEqual(len(q), 0) + def test_distinct_with_explicit_count(self): + q = TestModel.objects.distinct(['test_id']) + self.assertEqual(q.count(), 3) + + q = TestModel.objects.distinct(['test_id']).filter(test_id__in=[1, 2]) + self.assertEqual(q.count(), 2) + class TestQuerySetOrdering(BaseQuerySetUsage): diff --git a/tests/integration/cqlengine/statements/test_select_statement.py b/tests/integration/cqlengine/statements/test_select_statement.py index 33b1ef07e9..3f1c371c9e 100644 --- a/tests/integration/cqlengine/statements/test_select_statement.py +++ b/tests/integration/cqlengine/statements/test_select_statement.py @@ -64,6 +64,9 @@ def test_distinct(self): ss = SelectStatement('table', distinct_fields=['field1', 'field2']) self.assertEqual(six.text_type(ss), 'SELECT DISTINCT "field1", "field2" FROM table') + ss = SelectStatement('table', distinct_fields=['field1'], count=True) + self.assertEqual(six.text_type(ss), 'SELECT DISTINCT COUNT("field1") FROM table') + def test_context(self): ss = SelectStatement('table') ss.add_where_clause(WhereClause('a', EqualsOperator(), 'b')) From a4d329eceb1bc3854d14023f6a0ba57dbcfa1d64 Mon Sep 17 00:00:00 2001 From: Alan Boudreault Date: Thu, 3 Mar 2016 15:35:38 -0500 Subject: [PATCH 8/8] Add tests for limit(None) and fetch_size(...) --- cassandra/cluster.py | 2 -- .../cqlengine/statements/test_base_statement.py | 12 ++++++++++++ .../cqlengine/statements/test_select_statement.py | 12 ++++++++++++ 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 5cc4aa6f6f..13afecd245 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -2829,9 +2829,7 @@ def start_fetching_next_page(self): .. versionadded:: 2.0.0 """ - log.debug('Pagination: Fetching next page') if not self._paging_state: - log.debug('Pagination: No more pages to fetch') raise QueryExhausted() self._make_query_plan() diff --git a/tests/integration/cqlengine/statements/test_base_statement.py b/tests/integration/cqlengine/statements/test_base_statement.py index 2a9bb80fae..1bda588eff 100644 --- a/tests/integration/cqlengine/statements/test_base_statement.py +++ b/tests/integration/cqlengine/statements/test_base_statement.py @@ -16,6 +16,7 @@ except ImportError: import unittest # noqa +from cassandra.query import FETCH_SIZE_UNSET from cassandra.cqlengine.statements import BaseCQLStatement, StatementException @@ -26,3 +27,14 @@ def test_where_clause_type_checking(self): stmt = BaseCQLStatement('table', []) with self.assertRaises(StatementException): stmt.add_where_clause('x=5') + + def test_fetch_size(self): + """ tests that fetch_size is correctly set """ + stmt = BaseCQLStatement('table', None, fetch_size=1000) + self.assertEqual(stmt.fetch_size, 1000) + + stmt = BaseCQLStatement('table', None, fetch_size=None) + self.assertEqual(stmt.fetch_size, FETCH_SIZE_UNSET) + + stmt = BaseCQLStatement('table', None) + self.assertEqual(stmt.fetch_size, FETCH_SIZE_UNSET) diff --git a/tests/integration/cqlengine/statements/test_select_statement.py b/tests/integration/cqlengine/statements/test_select_statement.py index 3f1c371c9e..6612333809 100644 --- a/tests/integration/cqlengine/statements/test_select_statement.py +++ b/tests/integration/cqlengine/statements/test_select_statement.py @@ -96,3 +96,15 @@ def test_additional_rendering(self): self.assertIn('ORDER BY x, y', qstr) self.assertIn('ALLOW FILTERING', qstr) + def test_limit_rendering(self): + ss = SelectStatement('table', None, limit=10) + qstr = six.text_type(ss) + self.assertIn('LIMIT 10', qstr) + + ss = SelectStatement('table', None, limit=0) + qstr = six.text_type(ss) + self.assertNotIn('LIMIT', qstr) + + ss = SelectStatement('table', None, limit=None) + qstr = six.text_type(ss) + self.assertNotIn('LIMIT', qstr)