Merge pull request #4 from M1hacka/qs-concat

Qs concat
M1ha-Shvn · Oct 28, 2018 · 7f1830d · 7f1830d
2 parents 6dfb05b + d085aed
commit 7f1830d
Show file tree

Hide file tree

Showing 4 changed files with 111 additions and 7 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -4,6 +4,13 @@ language: python
 cache: pip
 services:
   - postgresql
+addons:
+  postgresql: "10"
+  apt:
+    packages:
+    - postgresql-10
+    - postgresql-client-10
+
 
 python:
   - 2.7
@@ -18,31 +25,43 @@ env:
   - DJANGO=1.7 PG=9.4
   - DJANGO=1.7 PG=9.5
   - DJANGO=1.7 PG=9.6
+  - DJANGO=1.7 PG=10
   - DJANGO=1.8 PG=9.2
   - DJANGO=1.8 PG=9.3
   - DJANGO=1.8 PG=9.4
   - DJANGO=1.8 PG=9.5
   - DJANGO=1.8 PG=9.6
+  - DJANGO=1.8 PG=10
   - DJANGO=1.9 PG=9.2
   - DJANGO=1.9 PG=9.3
   - DJANGO=1.9 PG=9.4
   - DJANGO=1.9 PG=9.5
   - DJANGO=1.9 PG=9.6
+  - DJANGO=1.9 PG=10
   - DJANGO=1.10 PG=9.2
   - DJANGO=1.10 PG=9.3
   - DJANGO=1.10 PG=9.4
   - DJANGO=1.10 PG=9.5
   - DJANGO=1.10 PG=9.6
+  - DJANGO=1.10 PG=10
   - DJANGO=1.11 PG=9.2
   - DJANGO=1.11 PG=9.3
   - DJANGO=1.11 PG=9.4
   - DJANGO=1.11 PG=9.5
   - DJANGO=1.11 PG=9.6
+  - DJANGO=1.11 PG=10
   - DJANGO=2.0 PG=9.2
   - DJANGO=2.0 PG=9.3
   - DJANGO=2.0 PG=9.4
   - DJANGO=2.0 PG=9.5
   - DJANGO=2.0 PG=9.6
+  - DJANGO=2.0 PG=10
+  - DJANGO=2.1 PG=9.2
+  - DJANGO=2.1 PG=9.3
+  - DJANGO=2.1 PG=9.4
+  - DJANGO=2.1 PG=9.5
+  - DJANGO=2.1 PG=9.6
+  - DJANGO=2.1 PG=10
 
 matrix:
   exclude:
@@ -57,6 +76,20 @@ matrix:
     env: DJANGO=2.0 PG=9.5
   - python: 2.7
     env: DJANGO=2.0 PG=9.6
+  - python: 2.7
+    env: DJANGO=2.0 PG=10
+  - python: 2.7
+    env: DJANGO=2.1 PG=9.2
+  - python: 2.7
+    env: DJANGO=2.1 PG=9.3
+  - python: 2.7
+    env: DJANGO=2.1 PG=9.4
+  - python: 2.7
+    env: DJANGO=2.1 PG=9.5
+  - python: 2.7
+    env: DJANGO=2.1 PG=9.6
+  - python: 2.7
+    env: DJANGO=2.1 PG=10
 
   # Django 1.9+ doesn't support python 3.3
   - python: 3.3
@@ -69,6 +102,8 @@ matrix:
     env: DJANGO=1.9 PG=9.5
   - python: 3.3
     env: DJANGO=1.9 PG=9.6
+  - python: 3.3
+    env: DJANGO=1.9 PG=10
   - python: 3.3
     env: DJANGO=1.10 PG=9.2
   - python: 3.3
@@ -79,6 +114,8 @@ matrix:
     env: DJANGO=1.10 PG=9.5
   - python: 3.3
     env: DJANGO=1.10 PG=9.6
+  - python: 3.3
+    env: DJANGO=1.10 PG=10
   - python: 3.3
     env: DJANGO=1.11 PG=9.2
   - python: 3.3
@@ -89,6 +126,8 @@ matrix:
     env: DJANGO=1.11 PG=9.5
   - python: 3.3
     env: DJANGO=1.11 PG=9.6
+  - python: 3.3
+    env: DJANGO=1.11 PG=10
   - python: 3.3
     env: DJANGO=2.0 PG=9.2
   - python: 3.3
@@ -99,6 +138,35 @@ matrix:
     env: DJANGO=2.0 PG=9.5
   - python: 3.3
     env: DJANGO=2.0 PG=9.6
+  - python: 3.3
+    env: DJANGO=2.0 PG=10
+  - python: 3.3
+    env: DJANGO=2.1 PG=9.2
+  - python: 3.3
+    env: DJANGO=2.1 PG=9.3
+  - python: 3.3
+    env: DJANGO=2.1 PG=9.4
+  - python: 3.3
+    env: DJANGO=2.1 PG=9.5
+  - python: 3.3
+    env: DJANGO=2.1 PG=9.6
+  - python: 3.3
+    env: DJANGO=2.1 PG=10
+
+  # Django 2.1 doesn't support python 3.4
+  - python: 3.4
+    env: DJANGO=2.1 PG=9.2
+  - python: 3.4
+    env: DJANGO=2.1 PG=9.3
+  - python: 3.4
+    env: DJANGO=2.1 PG=9.4
+  - python: 3.4
+    env: DJANGO=2.1 PG=9.5
+  - python: 3.4
+    env: DJANGO=2.1 PG=9.6
+  - python: 3.4
+    env: DJANGO=2.1 PG=10
+
 
   # Django 1.7 doesn't support python 3.5+
   - python: 3.5
@@ -111,6 +179,8 @@ matrix:
     env: DJANGO=1.7 PG=9.5
   - python: 3.5
     env: DJANGO=1.7 PG=9.6
+  - python: 3.5
+    env: DJANGO=1.7 PG=10
   - python: 3.6
     env: DJANGO=1.7 PG=9.2
   - python: 3.6
@@ -121,8 +191,15 @@ matrix:
     env: DJANGO=1.7 PG=9.5
   - python: 3.6
     env: DJANGO=1.7 PG=9.6
+  - python: 3.6
+    env: DJANGO=1.7 PG=10
 
 before_install:
+  # Use default PostgreSQL 10 port
+  - sudo sed -i 's/port = 5433/port = 5432/' /etc/postgresql/10/main/postgresql.conf
+  - sudo cp /etc/postgresql/{9.6,10}/main/pg_hba.conf
+
+  # Start PostgreSQL version we need
   - sudo service postgresql stop && sudo service postgresql start $PG
 
 install:

diff --git a/setup.py b/setup.py
@@ -13,7 +13,7 @@
 
 setup(
     name='django-pg-returning',
-    version='1.0.1',
+    version='1.0.2',
     packages=['django_pg_returning'],
     package_dir={'': 'src'},
     url='https://github.com/M1hacka/django-pg-returning',

diff --git a/src/django_pg_returning/queryset.py b/src/django_pg_returning/queryset.py
@@ -1,5 +1,7 @@
 from collections import namedtuple
+from itertools import chain
 
+from copy import deepcopy
 from django.db.models.query import RawQuerySet
 from django.db import router
 from typing import Any, Union, List, Dict, Tuple
@@ -15,14 +17,15 @@ class ReturningQuerySet(RawQuerySet):
     4) Returns Database to write, not to read as .db
     """
     def __init__(self, *args, **kwargs):
-        # A list of fileds, fetched by returning statement, in order to form values_list
-        self._fields = kwargs.pop('fields')
+        # A list of fields, fetched by returning statement, in order to form values_list
+        self._fields = kwargs.pop('fields', [])
 
         super(ReturningQuerySet, self).__init__(*args, **kwargs)
 
         # HACK Using methods create a new RawQuerySet, based on current data without calling it.
         # I use it here in order to iterate with super().__iter__ method.
-        self._result_cache = list(super(ReturningQuerySet, self).using(self.db))
+        # If raw_query is empty, I think it's an empty QuerySet creation
+        self._result_cache = list(super(ReturningQuerySet, self).using(self.db)) if self.raw_query else []
 
     def __len__(self):
         return len(self._result_cache)
@@ -33,10 +36,22 @@ def __iter__(self):
     def __getitem__(self, k):
         return self._result_cache[k]
 
+    def __add__(self, other):
+        if self.fields != other.fields:
+            raise ValueError("Querysets with different fields can't be concatenated")
+
+        res = deepcopy(self)
+        res._result_cache = list(chain(res, other))
+        return res
+
     @property
     def db(self):  # type: () -> str
         return self._db or router.db_for_write(self.model, **self._hints)
 
+    @property
+    def fields(self):
+        return self._fields
+
     def using(self, alias):
         """
         This method is used with "lazy" RawQuerySet.

diff --git a/tests/test_returning_query_set.py b/tests/test_returning_query_set.py
@@ -1,6 +1,7 @@
-from django.db.models import F, ManyToOneRel
+from django.db.models import F
 from django.test import TestCase
 
+from django_pg_returning import ReturningQuerySet
 from tests.models import TestModel
 
 
@@ -48,5 +49,16 @@ def test_values_list(self):
         with self.assertRaises(ValueError):
             result.values_list('int_field', 'name', invalid=True)
 
-    def test_no_duplicate_query(self):
-        pass
+    def test_concat(self):
+        result = TestModel.objects.filter(id__gt=2, id__lte=5).update_returning(int_field=21)
+        result2 = TestModel.objects.filter(id__gt=5, id__lte=6).update_returning(int_field=21)
+        r = result + result2
+        self.assertSetEqual({3, 4, 5, 6}, set(r.values_list('id', flat=True)))
+
+        result3 = TestModel.objects.filter(id__gt=5, id__lte=6).only('id').update_returning(int_field=21)
+        with self.assertRaises(ValueError):
+            r = result + result3
+
+    def test_empty(self):
+        qs = ReturningQuerySet(None)
+        self.assertListEqual([], list(qs))