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/querysets.md

@@ -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

@@ -183,6 +183,7 @@ class QuerySet(object):
         self._order_by = []
         self._q = []
         self._fields = []
+        self._limits = None
 
     def __iter__(self):
         """
@@ -202,6 +203,24 @@ 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):
         """

tests/test_querysets.py

@@ -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')
 
@@ -150,3 +174,8 @@ class SampleModel(Model):
     color = Enum8Field(Color)
 
     engine = MergeTree('materialized_date', ('materialized_date',))
+
+
+class Numbers(Model):
+
+    number = UInt64Field()