fix #9

Author: jayvynl

clickhouse_backend/driver/client.py

@@ -5,7 +5,7 @@
 
 from .escape import escape_params
 
-insert_pattern = re.compile(r'\s*insert\s*into', flags=re.IGNORECASE)
+insert_pattern = re.compile(r'\s*insert\s+into', flags=re.IGNORECASE)
 
 
 class Client(client.Client):

clickhouse_backend/driver/connection.py

@@ -1,10 +1,14 @@
+import re
+
 from clickhouse_driver.dbapi import connection
 from clickhouse_driver.dbapi import cursor
 from clickhouse_driver.dbapi import errors
 from clickhouse_pool.pool import ChPoolError
 
 from .pool import ClickhousePool
 
+update_pattern = re.compile(r'\s*alter\s+table\s+(.+)\s+update.+where\s+(.+)', flags=re.IGNORECASE)
+
 
 class Cursor(cursor.Cursor):
     def close(self):
@@ -28,6 +32,19 @@ def __del__(self):
             except ChPoolError:
                 pass
 
+    def execute(self, operation, parameters=None):
+        """fix https://github.com/jayvynl/django-clickhouse-backend/issues/9"""
+        if update_pattern.match(operation):
+            query = self._client.substitute_params(
+                    operation, parameters, self._client.connection.context
+                )
+            table, where = update_pattern.match(query).groups()
+            super().execute(f'select count(*) from {table} where {where}')
+            rowcount, = self.fetchone()
+            self._reset_state()
+            self._rowcount = rowcount
+        super().execute(operation, parameters)
+
 
 class Connection(connection.Connection):
     """Connection class with support for connection pool."""

clickhouse_backend/models/base.py

@@ -25,7 +25,9 @@ def as_clickhouse(self, compiler, connection, **extra_context):
 functions.Random.as_clickhouse = as_clickhouse
 
 
-class ClickhouseManager(BaseManager.from_queryset(QuerySet)):
+class ClickhouseManager(models.Manager):
+    _queryset_class = QuerySet
+
     def get_queryset(self):
         """
         User defined Query and QuerySet class that support clickhouse particular query.
@@ -46,9 +48,9 @@ class Meta:
 
     def _do_update(self, base_qs, using, pk_val, values, update_fields, forced_update):
         """
-         Try to update the model. Return True if the model was updated (if an
-         update query was done and a matching row was found in the DB).
-         """
+        Try to update the model. Return True if the model was updated (if an
+        update query was done and a matching row was found in the DB).
+        """
         filtered = base_qs.filter(pk=pk_val)
         if not values:
             # We can end up here when saving a model in inheritance chain where
@@ -57,15 +59,4 @@ def _do_update(self, base_qs, using, pk_val, values, update_fields, forced_updat
             # is a model with just PK - in that case check that the PK still
             # exists.
             return update_fields is not None or filtered.exists()
-        return (
-            filtered.exists()
-            and
-            # It may happen that the object is deleted from the DB right after
-            # this check, causing the subsequent UPDATE to return zero matching
-            # rows. The same result can occur in some rare cases when the
-            # database returns zero despite the UPDATE being executed
-            # successfully (a row is matched and updated). In order to
-            # distinguish these two cases, the object's existence in the
-            # database is again checked for if the UPDATE query returns 0.
-            (filtered._update(values) > 0 or filtered.exists())
-        )
+        return filtered._update(values) > 0

tests/basic/models.py

@@ -53,3 +53,14 @@ class PrimaryKeyWithDefault(ClickhouseModel):
 
 class ChildPrimaryKeyWithDefault(PrimaryKeyWithDefault):
     pass
+
+
+class DjangoArticle(models.Model):
+    headline = models.CharField(max_length=100, default="Default headline")
+    pub_date = models.DateTimeField()
+
+    class Meta:
+        ordering = ("pub_date", "headline")
+
+    def __str__(self):
+        return self.headline

tests/basic/tests.py

@@ -4,24 +4,22 @@
 
 from django.core.exceptions import MultipleObjectsReturned, ObjectDoesNotExist
 from django.db import DEFAULT_DB_ALIAS, DatabaseError, connections, models
-from django.db.models.manager import BaseManager
 from django.db.models.query import MAX_GET_RESULTS, EmptyQuerySet
 from django.test import (
-    SimpleTestCase,
     TestCase,
     TransactionTestCase,
     skipUnlessDBFeature,
 )
 from django.utils.translation import gettext_lazy
 
-from clickhouse_backend import compat
 from .models import (
     Article,
     ArticleSelectOnSave,
     ChildPrimaryKeyWithDefault,
     FeaturedArticle,
     PrimaryKeyWithDefault,
     SelfRef,
+    DjangoArticle
 )
 
 
@@ -632,12 +630,12 @@ def deleter():
 class SelectOnSaveTests(TestCase):
     def test_select_on_save(self):
         a1 = Article.objects.create(pub_date=datetime.now())
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(1):
             a1.save()
         asos = ArticleSelectOnSave.objects.create(pub_date=datetime.now())
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(1):
             asos.save()
-        with self.assertNumQueries(3):
+        with self.assertNumQueries(1):
             asos.save(force_update=True)
         Article.objects.all().delete()
         with self.assertRaisesMessage(
@@ -671,7 +669,7 @@ def _update(self, *args, **kwargs):
         try:
             Article._base_manager._queryset_class = FakeQuerySet
             asos = ArticleSelectOnSave.objects.create(pub_date=datetime.now())
-            with self.assertNumQueries(3):
+            with self.assertNumQueries(2):
                 asos.save()
                 self.assertTrue(FakeQuerySet.called)
         finally:
@@ -808,3 +806,17 @@ def test_prefetched_cache_cleared(self):
         a2_prefetched.refresh_from_db(fields=["selfref_set"])
         # Cache was cleared and new results are available.
         self.assertCountEqual(a2_prefetched.selfref_set.all(), [s])
+
+
+class DjangoModelTests(TestCase):
+    def test_saving_an_object_again_does_not_create_a_new_object(self):
+        a = DjangoArticle(headline="original", pub_date=datetime(2014, 5, 16))
+        a.save()
+        current_id = a.id
+
+        a.save()
+        self.assertEqual(a.id, current_id)
+
+        a.headline = "Updated headline"
+        a.save()
+        self.assertEqual(a.id, current_id)