feat: support prewhere, finish test for django 4.2

Author: jayvynl

README.md

@@ -28,6 +28,7 @@ Read [Documentation](https://github.com/jayvynl/django-clickhouse-backend/blob/m
 - Support creating test database and table, working with django TestCase and pytest-django.
 - Support most clickhouse data types.
 - Support [SETTINGS in SELECT Query](https://clickhouse.com/docs/en/sql-reference/statements/select/#settings-in-select-query).
+- Support [PREWHERE clause](https://clickhouse.com/docs/en/sql-reference/statements/select/prewhere).
 
 **Notes:**
 

clickhouse_backend/models/base.py

@@ -26,6 +26,9 @@ def get_queryset(self):
             model=self.model, query=Query(self.model), using=self._db, hints=self._hints
         )
 
+    def prewhere(self, *args, **kwargs):
+        return self.get_queryset().prewhere(*args, **kwargs)
+
 
 class ClickhouseModel(models.Model):
     objects = ClickhouseManager()

clickhouse_backend/models/query.py

@@ -1,6 +1,6 @@
-from django.db.models import query
+from django.db.models import Q, query
 
-from .sql import Query
+from clickhouse_backend.models import sql
 
 
 class QuerySet(query.QuerySet):
@@ -14,7 +14,7 @@ def explain(self, *, format=None, type=None, **settings):
 
     def settings(self, **kwargs):
         clone = self._chain()
-        if isinstance(clone.query, Query):
+        if isinstance(clone.query, sql.Query):
             clone.query.setting_info.update(kwargs)
         return clone
 
@@ -27,5 +27,5 @@ def prewhere(self, *args, **kwargs):
         if (args or kwargs) and self.query.is_sliced:
             raise TypeError("Cannot prewhere a query once a slice has been taken.")
         clone = self._chain()
-        clone._query.add_prewhere(query.Q(*args, **kwargs))
+        clone._query.add_prewhere(Q(*args, **kwargs))
         return clone

clickhouse_backend/models/sql/compiler.py

@@ -8,6 +8,7 @@
 from clickhouse_backend import compat
 from clickhouse_backend.idworker import id_worker
 from clickhouse_backend.models import engines
+from clickhouse_backend.models.sql import Query
 
 if compat.dj_ge42:
     from django.core.exceptions import FullResultSet
@@ -75,13 +76,24 @@ def pre_sql_setup(self, with_col_aliases=False):
             ) = self.query.where.split_having_qualify(
                 must_group_by=self.query.group_by is not None
             )
-            (
-                self.prewhere,
-                prehaving,
-                prequalify,
-            ) = self.query.prewhere.split_having_qualify(
-                must_group_by=self.query.group_by is not None
-            )
+            if isinstance(self.query, Query):
+                (
+                    self.prewhere,
+                    prehaving,
+                    prequalify,
+                ) = self.query.prewhere.split_having_qualify(
+                    must_group_by=self.query.group_by is not None
+                )
+            else:
+                (
+                    self.prewhere,
+                    prehaving,
+                    prequalify,
+                ) = (
+                    None,
+                    None,
+                    None,
+                )
             # Check before ClickHouse complain.
             # DB::Exception: Window function is found in PREWHERE in query. (ILLEGAL_AGGREGATION)
             if prequalify:

clickhouse_backend/models/sql/query.py

@@ -3,6 +3,7 @@
 from django.db import router
 from django.db.models.sql import query, subqueries
 from django.db.models.sql.constants import INNER
+from django.db.models.sql.datastructures import BaseTable, Join
 from django.db.models.sql.where import AND
 
 from clickhouse_backend import compat
@@ -26,6 +27,7 @@ def sql_with_params(self):
     def clone(self):
         obj = super().clone()
         obj.setting_info = self.setting_info.copy()
+        obj.prewhere = self.prewhere.clone()
         return obj
 
     def explain(self, using, format=None, type=None, **settings):
@@ -36,15 +38,8 @@ def explain(self, using, format=None, type=None, **settings):
 
     def add_prewhere(self, q_object):
         """
-        A preprocessor for the internal _add_q(). Responsible for doing final
-        join promotion.
+        refer add_q
         """
-        # For join promotion this case is doing an AND for the added q_object
-        # and existing conditions. So, any existing inner join forces the join
-        # type to remain inner. Existing outer joins can however be demoted.
-        # (Consider case where rel_a is LOUTER and rel_a__col=1 is added - if
-        # rel_a doesn't produce any rows, then the whole condition must fail.
-        # So, demotion is OK.
         existing_inner = {
             a for a in self.alias_map if self.alias_map[a].join_type == INNER
         }
@@ -59,6 +54,82 @@ def add_prewhere(self, q_object):
         def is_sliced(self):
             return self.low_mark != 0 or self.high_mark is not None
 
+    def resolve_expression(self, query, *args, **kwargs):
+        clone = self.clone()
+        # Subqueries need to use a different set of aliases than the outer query.
+        clone.bump_prefix(query)
+        clone.subquery = True
+        clone.where.resolve_expression(query, *args, **kwargs)
+        clone.prewhere.resolve_expression(query, *args, **kwargs)
+        # Resolve combined queries.
+        if clone.combinator:
+            clone.combined_queries = tuple(
+                [
+                    combined_query.resolve_expression(query, *args, **kwargs)
+                    for combined_query in clone.combined_queries
+                ]
+            )
+        for key, value in clone.annotations.items():
+            resolved = value.resolve_expression(query, *args, **kwargs)
+            if hasattr(resolved, "external_aliases"):
+                resolved.external_aliases.update(clone.external_aliases)
+            clone.annotations[key] = resolved
+        # Outer query's aliases are considered external.
+        for alias, table in query.alias_map.items():
+            clone.external_aliases[alias] = (
+                isinstance(table, Join)
+                and table.join_field.related_model._meta.db_table != alias
+            ) or (
+                isinstance(table, BaseTable) and table.table_name != table.table_alias
+            )
+        return clone
+
+    def change_aliases(self, change_map):
+        """
+        Change the aliases in change_map (which maps old-alias -> new-alias),
+        relabelling any references to them in select columns and the where
+        clause.
+        """
+        # If keys and values of change_map were to intersect, an alias might be
+        # updated twice (e.g. T4 -> T5, T5 -> T6, so also T4 -> T6) depending
+        # on their order in change_map.
+        assert set(change_map).isdisjoint(change_map.values())
+
+        # 1. Update references in "select" (normal columns plus aliases),
+        # "group by" and "where".
+        self.where.relabel_aliases(change_map)
+        self.prewhere.relabel_aliases(change_map)
+        if isinstance(self.group_by, tuple):
+            self.group_by = tuple(
+                [col.relabeled_clone(change_map) for col in self.group_by]
+            )
+        self.select = tuple([col.relabeled_clone(change_map) for col in self.select])
+        self.annotations = self.annotations and {
+            key: col.relabeled_clone(change_map)
+            for key, col in self.annotations.items()
+        }
+
+        # 2. Rename the alias in the internal table/alias datastructures.
+        for old_alias, new_alias in change_map.items():
+            if old_alias not in self.alias_map:
+                continue
+            alias_data = self.alias_map[old_alias].relabeled_clone(change_map)
+            self.alias_map[new_alias] = alias_data
+            self.alias_refcount[new_alias] = self.alias_refcount[old_alias]
+            del self.alias_refcount[old_alias]
+            del self.alias_map[old_alias]
+
+            table_aliases = self.table_map[alias_data.table_name]
+            for pos, alias in enumerate(table_aliases):
+                if alias == old_alias:
+                    table_aliases[pos] = new_alias
+                    break
+        self.external_aliases = {
+            # Table is aliased or it's being changed and thus is aliased.
+            change_map.get(alias, alias): (aliased or alias in change_map)
+            for alias, aliased in self.external_aliases.items()
+        }
+
 
 def clone_decorator(cls):
     old_clone = cls.clone

tests/clickhouse_queries/__init__.py

@@ -0,0 +1,18 @@
+from django.db.models import CASCADE, ForeignKey
+
+from clickhouse_backend import models
+
+
+class Author(models.ClickhouseModel):
+    name = models.StringField(max_length=10)
+    num = models.UInt32Field()
+
+
+class Book(models.ClickhouseModel):
+    author = ForeignKey(Author, on_delete=CASCADE, related_name="books")
+    name = models.StringField(max_length=10)
+
+
+class Article(models.ClickhouseModel):
+    title = models.StringField(max_length=10)
+    book = models.Int64Field()

tests/clickhouse_queries/models.py

@@ -0,0 +1,75 @@
+from django.db import NotSupportedError
+from django.db.models import Count, Window
+from django.db.models.functions import Rank
+from django.test import TestCase
+
+from . import models
+
+
+class QueriesTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.a1, cls.a2 = models.Author.objects.bulk_create(
+            [models.Author(name="a1", num=1001), models.Author(name="a2", num=2002)]
+        )
+        cls.b1, cls.b2, cls.b3, cls.b4 = models.Book.objects.bulk_create(
+            [
+                models.Book(name="b1", author=cls.a1),
+                models.Book(name="b2", author=cls.a1),
+                models.Book(name="b3", author=cls.a2),
+                models.Book(name="b4", author=cls.a2),
+            ]
+        )
+        models.Article.objects.bulk_create(
+            [
+                models.Article(title="t1", book=cls.b1.id),
+                models.Article(title="t2", book=cls.b2.id),
+            ]
+        )
+
+    def test_prewhere(self):
+        qs = models.Author.objects.prewhere(name="a1")
+        self.assertIn("PREWHERE", str(qs.query))
+        self.assertEqual(qs[0].name, "a1")
+
+    def test_prewhere_fk(self):
+        self.assertQuerySetEqual(
+            models.Book.objects.filter(author__name=self.a1.name)
+            .prewhere(author_id=self.a1.id)
+            .order_by("name"),
+            [self.b1, self.b2],
+        )
+
+    # clickhouse backend will generate suitable query, but clickhouse will raise exception.
+    # clickhouse 23.11
+    # DB::Exception: Missing columns: 'clickhouse_queries_article.book' while processing query: 'SELECT name FROM clickhouse_queries_book AS U0 PREWHERE id = clickhouse_queries_article.book', required columns: 'name' 'id' 'clickhouse_queries_article.book', maybe you meant: 'name' or 'id': While processing (SELECT U0.name FROM clickhouse_queries_book AS U0 PREWHERE U0.id = clickhouse_queries_article.book) AS book_name.
+    # clickhouse 24.6
+    # DB::Exception: Resolve identifier 'clickhouse_queries_article.book' from parent scope only supported for constants and CTE. Actual test_default.clickhouse_queries_article.book node type COLUMN. In scope (SELECT U0.name FROM clickhouse_queries_book AS U0 PREWHERE U0.id = clickhouse_queries_article.book) AS book_name.
+    # def test_prewhere_subquery(self):
+    #     a = models.Article.objects.annotate(
+    #         book_name=Subquery(
+    #             models.Book.objects.prewhere(id=OuterRef("book")).values("name")
+    #         )
+    #     ).get(title="t1")
+    #     self.assertEqual(a.book_name, self.b1.name)
+
+    def test_prewhere_agg(self):
+        with self.assertRaisesMessage(
+            NotSupportedError,
+            "Aggregate function is disallowed in the prewhere clause.",
+        ):
+            list(
+                models.Author.objects.annotate(count=Count("books")).prewhere(
+                    count__gt=0
+                )
+            )
+
+    def test_prewhere_window(self):
+        with self.assertRaisesMessage(
+            NotSupportedError, "Window function is disallowed in the prewhere clause."
+        ):
+            list(
+                models.Book.objects.annotate(
+                    rank=Window(Rank(), partition_by="author", order_by="name")
+                ).prewhere(rank__gt=1)
+            )