diff --git a/clickhouse_sqlalchemy/drivers/compilers/sqlcompiler.py b/clickhouse_sqlalchemy/drivers/compilers/sqlcompiler.py index b38d63ca..d1011591 100644 --- a/clickhouse_sqlalchemy/drivers/compilers/sqlcompiler.py +++ b/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) diff --git a/clickhouse_sqlalchemy/orm/query.py b/clickhouse_sqlalchemy/orm/query.py index 3b486e2f..611b4f71 100644 --- a/clickhouse_sqlalchemy/orm/query.py +++ b/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