ORM missed attributes fix

clickhouse_sqlalchemy/drivers/compilers/sqlcompiler.py

@@ -8,6 +8,31 @@
 
 
 class ClickHouseSQLCompiler(compiler.SQLCompiler):
+    def visit_select(
+        self,
+        select_stmt,
+        **kwargs,
+    ):
+        orig_compile_state_factory = select_stmt._compile_state_factory
+
+        def f(self, *args, **kwargs):
+            tmp = orig_compile_state_factory(self, *args, **kwargs)
+
+            if hasattr(tmp, 'select_statement'):
+                # Fix missed attributes
+                for attr in ["_with_cube", "_with_rollup", "_with_totals", "_final_clause", "_sample_clause", "_limit_by_clause", "_array_join"]:
+                    val = getattr(tmp.select_statement, attr, None)
+
+                    if val:
+                        setattr(tmp.statement, attr, val)
+
+            return tmp
+
+        select_stmt._compile_state_factory = f
+        tmp_select = super().visit_select(select_stmt=select_stmt, **kwargs)
+
+        return tmp_select
+
     def visit_mod_binary(self, binary, operator, **kw):
         return self.process(binary.left, **kw) + ' %% ' + \
             self.process(binary.right, **kw)

clickhouse_sqlalchemy/orm/query.py

@@ -19,6 +19,21 @@ class Query(BaseQuery):
     _limit_by = None
     _array_join = None
 
+    def _statement_20(self, for_statement=False, use_legacy_query_style=True):
+        orig_smt = super(Query, self)._statement_20(for_statement=for_statement,
+                                                    use_legacy_query_style=use_legacy_query_style)
+
+        orig_smt._with_cube = self._with_cube
+        orig_smt._with_rollup = self._with_rollup
+        orig_smt._with_totals = self._with_totals
+        orig_smt._final_clause = self._final
+        orig_smt._sample_clause = sample_clause(self._sample)
+        orig_smt._limit_by_clause = self._limit_by
+        orig_smt._array_join = self._array_join
+
+        return orig_smt
+
+
     def _compile_context(self, *args, **kwargs):
         context = super(Query, self)._compile_context(*args, **kwargs)
         query = context.query