1. add toYYYYMM[DD[hhmmss]] function 2. fix str(queryset.query) when default database is not clickhouse

Author: linzhiwen

clickhouse_backend/models/__init__.py

@@ -1,14 +1,17 @@
 from .base import ClickhouseModel
 from .engines import *
-from .engines import __all__ as engine_all  # NOQA
+from .engines import __all__ as engines_all  # NOQA
 from .fields import *
 from .fields import __all__ as fields_all  # NOQA
+from .functions import *
+from .functions import __all__ as fucntions_all  # NOQA
 from .indexes import *
-from .indexes import __all__ as index_all  # NOQA
+from .indexes import __all__ as indexes_all  # NOQA
 
 __all__ = [
     'ClickhouseModel',
+    *engines_all,
     *fields_all,
-    *engine_all,
-    *index_all,
+    *fucntions_all,
+    *indexes_all,
 ]

clickhouse_backend/models/base.py

@@ -58,14 +58,14 @@ def _do_update(self, base_qs, using, pk_val, values, update_fields, forced_updat
             # 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())
+            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())
         )

clickhouse_backend/models/functions.py

@@ -0,0 +1,36 @@
+from django.conf import settings
+from django.db.models import Func, Value
+
+__all__ = [
+    'toYYYYMM',
+    'toYYYYMMDD',
+    'toYYYYMMDDhhmmss',
+]
+
+
+class toYYYYMM(Func):
+    function = 'toYYYYMM'
+
+    def __init__(self, *expressions, output_field=None, **extra):
+        arity = len(expressions)
+        if arity < 1 or arity > 2:
+            raise TypeError(
+                "'%s' takes 1 or 2 arguments (%s given)"
+                % (
+                    self.__class__.__name__,
+                    len(expressions),
+                )
+            )
+        if arity == 2 and isinstance(expressions[1], str):
+            expressions = (expressions[0], Value(expressions[1]))
+        elif settings.USE_TZ:
+            expressions = (expressions[0], Value(settings.TIME_ZONE))
+        super().__init__(*expressions, output_field=output_field, **extra)
+
+
+class toYYYYMMDD(toYYYYMM):
+    function = 'toYYYYMMDD'
+
+
+class toYYYYMMDDhhmmss(toYYYYMM):
+    function = 'toYYYYMMDDhhmmss'

clickhouse_backend/models/sql/query.py

@@ -1,5 +1,6 @@
 from collections import namedtuple
 
+from django.db import router
 from django.db.models.sql import query
 from django.db.models.sql import subqueries
 
@@ -16,6 +17,10 @@ def __init__(self, model, where=query.WhereNode, alias_cols=True):
             super().__init__(model, where, alias_cols)
         self.setting_info = {}
 
+    def sql_with_params(self):
+        """Choose the right db when database router is used."""
+        return self.get_compiler(router.db_for_read(self.model)).as_sql()
+
     def clone(self):
         obj = super().clone()
         obj.setting_info = self.setting_info.copy()

example/testapp/models.py

@@ -25,7 +25,7 @@ class Meta:
         db_table = 'event'
         engine = models.ReplacingMergeTree(
             order_by=['id'],
-            partition_by=Func('timestamp', function='toYYYYMMDD')
+            partition_by=models.toYYYYMMDD('timestamp')
         )
         indexes = [
             models.Index(