Support queryset slicing

CHANGELOG.md

@@ -5,6 +5,7 @@ Unreleased
 ----------
 - Changed license from PSF to BSD
 - Nullable fields support (yamiou)
+- Support for queryset slicing
 
 v0.9.2
 ------

docs/class_reference.md

@@ -539,8 +539,8 @@ infi.clickhouse_orm.query
 ### QuerySet
 
 
-A queryset is an object that represents a database query using a specific `Model`. 
-It is lazy, meaning that it does not hit the database until you iterate over its 
+A queryset is an object that represents a database query using a specific `Model`.
+It is lazy, meaning that it does not hit the database until you iterate over its
 matching rows (model instances).
 
 #### QuerySet(model_cls, database)

docs/querysets.md

@@ -4,7 +4,7 @@ Querysets
 A queryset is an object that represents a database query using a specific Model. It is lazy, meaning that it does not hit the database until you iterate over its matching rows (model instances). To create a base queryset for a model class, use:
 
     qs = Person.objects_in(database)
-    
+
 This queryset matches all Person instances in the database. You can get these instances using iteration:
 
     for person in qs:
@@ -19,7 +19,7 @@ The `filter` and `exclude` methods are used for filtering the matching instances
     >>> qs = qs.filter(first_name__startswith='V').exclude(birthday__lt='2000-01-01')
     >>> qs.conditions_as_sql()
     u"first_name LIKE 'V%' AND NOT (birthday < '2000-01-01')"
-    
+
 It is possible to specify several fields to filter or exclude by:
 
     >>> qs = Person.objects_in(database).filter(last_name='Smith', height__gt=1.75)
@@ -57,7 +57,7 @@ For example if we want to select only people with Irish last names:
 
     # A list of simple values
     qs = Person.objects_in(database).filter(last_name__in=["Murphy", "O'Sullivan"])
-    
+
     # A string
     subquery = "SELECT name from $db.irishlastname"
     qs = Person.objects_in(database).filter(last_name__in=subquery)
@@ -72,7 +72,7 @@ Counting and Checking Existence
 Use the `count` method to get the number of matches:
 
     Person.objects_in(database).count()
-    
+
 To check if there are any matches at all, you can use any of the following equivalent options:
 
     if qs.count(): ...
@@ -85,7 +85,7 @@ Ordering
 The sorting order of the results can be controlled using the `order_by` method:
 
     qs = Person.objects_in(database).order_by('last_name', 'first_name')
-    
+
 The default order is ascending. To use descending order, add a minus sign before the field name:
 
     qs = Person.objects_in(database).order_by('-height')
@@ -100,6 +100,25 @@ When some of the model fields aren't needed, it is more efficient to omit them f
     qs = Person.objects_in(database).only('first_name', 'birthday')
 
 
+Slicing
+-------
+
+It is possible to get a specific item from the queryset by index.
+
+      qs = Person.objects_in(database).order_by('last_name', 'first_name')
+      first = qs[0]
+
+It is also possible to get a range a instances using a slice. This returns a queryset,
+that you can either iterate over or convert to a list.
+
+      qs = Person.objects_in(database).order_by('last_name', 'first_name')
+      first_ten_people = list(qs[:10])
+      next_ten_people  = list(qs[10:20])
+
+You should use `order_by` to ensure a consistent ordering of the results.
+
+Trying to use negative indexes or a slice with a step (e.g. [0:100:2]) is not supported and will raise an `AssertionError`.
+
 ---
 
 [<< Models and Databases](models_and_databases.md) | [Table of Contents](toc.md) | [Field Types >>](field_types.md)

docs/toc.md

@@ -20,6 +20,7 @@
       * [Counting and Checking Existence](querysets.md#counting-and-checking-existence)
       * [Ordering](querysets.md#ordering)
       * [Omitting Fields](querysets.md#omitting-fields)
+      * [Slicing](querysets.md#slicing)
 
    * [Field Types](field_types.md#field-types)
       * [DateTimeField and Time Zones](field_types.md#datetimefield-and-time-zones)

src/infi/clickhouse_orm/query.py

@@ -59,7 +59,7 @@ class InOperator(Operator):
 class LikeOperator(Operator):
     """
     A LIKE operator that matches the field to a given pattern. Can be
-    case sensitive or insensitive. 
+    case sensitive or insensitive.
     """
 
     def __init__(self, pattern, case_sensitive=True):
@@ -168,8 +168,8 @@ class Q(object):
 
 class QuerySet(object):
     """
-    A queryset is an object that represents a database query using a specific `Model`. 
-    It is lazy, meaning that it does not hit the database until you iterate over its 
+    A queryset is an object that represents a database query using a specific `Model`.
+    It is lazy, meaning that it does not hit the database until you iterate over its
     matching rows (model instances).
     """
 
@@ -183,10 +183,11 @@ class QuerySet(object):
         self._order_by = []
         self._q = []
         self._fields = []
+        self._limits = None
 
     def __iter__(self):
         """
-        Iterates over the model instances matching this queryset 
+        Iterates over the model instances matching this queryset
         """
         return self._database.select(self.as_sql(), self._model_cls)
 
@@ -201,7 +202,25 @@ class QuerySet(object):
 
     def __unicode__(self):
         return self.as_sql()
-        
+
+    def __getitem__(self, s):
+        if isinstance(s, six.integer_types):
+            # Single index
+            assert s >= 0, 'negative indexes are not supported'
+            qs = copy(self)
+            qs._limits = (s, 1)
+            return iter(qs).next()
+        else:
+            # Slice
+            assert s.step in (None, 1), 'step is not supported in slices'
+            start = s.start or 0
+            stop = s.stop or 2**63 - 1
+            assert start >= 0 and stop >= 0, 'negative indexes are not supported'
+            assert start <= stop, 'start of slice cannot be smaller than its end'
+            qs = copy(self)
+            qs._limits = (start, stop - start)
+            return qs
+
     def as_sql(self):
         """
         Returns the whole query as a SQL string.
@@ -210,8 +229,10 @@ class QuerySet(object):
         if self._fields:
             fields = ', '.join('`%s`' % field for field in self._fields)
         ordering = '\nORDER BY ' + self.order_by_as_sql() if self._order_by else ''
-        params = (fields, self._database.db_name, self._model_cls.table_name(), self.conditions_as_sql(), ordering)
-        return u'SELECT %s\nFROM `%s`.`%s`\nWHERE %s%s' % params
+        limit = '\nLIMIT %d, %d' % self._limits if self._limits else ''
+        params = (fields, self._model_cls.table_name(),
+                  self.conditions_as_sql(), ordering, limit)
+        return u'SELECT %s\nFROM `%s`\nWHERE %s%s%s' % params
 
     def order_by_as_sql(self):
         """
@@ -236,7 +257,7 @@ class QuerySet(object):
         Returns the number of matching model instances.
         """
         return self._database.count(self._model_cls, self.conditions_as_sql())
-        
+
     def order_by(self, *field_names):
         """
         Returns a new `QuerySet` instance with the ordering changed.

tests/test_querysets.py

@@ -18,11 +18,11 @@ class QuerySetTestCase(TestCaseWithData):
     def setUp(self):
         super(QuerySetTestCase, self).setUp()
         self.database.insert(self._sample_data())
-        
+
     def _test_qs(self, qs, expected_count):
         logging.info(qs.as_sql())
         for instance in qs:
-            logging.info('\t%s' % instance.to_dict()) 
+            logging.info('\t%s' % instance.to_dict())
         self.assertEquals(qs.count(), expected_count)
 
     def test_no_filtering(self):
@@ -138,6 +138,30 @@ class QuerySetTestCase(TestCaseWithData):
         self._test_qs(qs.filter(num__in=(1, 2, 3)), 3)
         self._test_qs(qs.filter(num__in=range(1, 4)), 3)
 
+    def test_slicing(self):
+        db = Database('system')
+        numbers = range(100)
+        qs = Numbers.objects_in(db)
+        self.assertEquals(qs[0].number, numbers[0])
+        self.assertEquals(qs[5].number, numbers[5])
+        self.assertEquals([row.number for row in qs[:1]], numbers[:1])
+        self.assertEquals([row.number for row in qs[:10]], numbers[:10])
+        self.assertEquals([row.number for row in qs[3:10]], numbers[3:10])
+        self.assertEquals([row.number for row in qs[9:10]], numbers[9:10])
+        self.assertEquals([row.number for row in qs[10:10]], numbers[10:10])
+
+    def test_invalid_slicing(self):
+        db = Database('system')
+        qs = Numbers.objects_in(db)
+        with self.assertRaises(AssertionError):
+            qs[3:10:2]
+        with self.assertRaises(AssertionError):
+            qs[-5]
+        with self.assertRaises(AssertionError):
+            qs[:-5]
+        with self.assertRaises(AssertionError):
+            qs[50:1]
+
 
 Color = Enum('Color', u'red blue green yellow brown white black')
 
@@ -149,4 +173,9 @@ class SampleModel(Model):
     num = Int32Field()
     color = Enum8Field(Color)
 
-    engine = MergeTree('materialized_date', ('materialized_date',))
+    engine = MergeTree('materialized_date', ('materialized_date',))
+
+
+class Numbers(Model):
+
+    number = UInt64Field()