|
1 | 1 | from typing import Any, cast |
2 | 2 |
|
3 | 3 | from sqlalchemy import Table, and_ |
4 | | -from sqlalchemy.sql.selectable import FromClause, Select |
| 4 | +from sqlalchemy.sql.selectable import CTE, FromClause, HasCTE, Select |
5 | 5 |
|
| 6 | +from clickhouse_connect import driver_name |
6 | 7 | from clickhouse_connect.cc_sqlalchemy.sql.clauses import ArrayJoin, LimitByClause, PreWhereClause |
7 | 8 | from clickhouse_connect.cc_sqlalchemy.sql.clauses import array_join as _array_join_fromclause |
8 | 9 | from clickhouse_connect.cc_sqlalchemy.sql.clauses import ch_join as _ch_join_fromclause |
|
12 | 13 | # compiled-statement cache keys when FINAL/SAMPLE/PREWHERE/LIMIT BY are applied. |
13 | 14 | _CH_MODIFIER_DIALECT = "_ch_modifier" |
14 | 15 |
|
| 16 | +# SQLAlchemy renders CTE prefixes between the name and the body, which is exactly where |
| 17 | +# ClickHouse expects the materialization keyword: WITH <name> AS MATERIALIZED (...). |
| 18 | +_MATERIALIZED_KEYWORD = "MATERIALIZED" |
| 19 | + |
15 | 20 |
|
16 | 21 | def full_table(table_name: str, schema: str | None = None) -> str: |
17 | 22 | if table_name.startswith("(") or not schema: |
@@ -158,6 +163,47 @@ def limit_by(select_stmt: Select, by_clauses: Any, limit: int, offset: int | Non |
158 | 163 | return new_stmt |
159 | 164 |
|
160 | 165 |
|
| 166 | +def _validate_cte_options(recursive: bool, materialized: bool) -> None: |
| 167 | + if recursive and materialized: |
| 168 | + raise ValueError("materialized CTEs cannot be recursive") |
| 169 | + |
| 170 | + |
| 171 | +def _apply_materialized(new_cte: CTE, materialized: bool) -> CTE: |
| 172 | + if not materialized: |
| 173 | + return new_cte |
| 174 | + # SQLAlchemy renders CTE prefixes between the name and the body, which is exactly |
| 175 | + # where ClickHouse expects the keyword. Scoping it to this dialect keeps a statement |
| 176 | + # shared with another backend compiling unchanged there. |
| 177 | + return new_cte.prefix_with(_MATERIALIZED_KEYWORD, dialect=driver_name) |
| 178 | + |
| 179 | + |
| 180 | +def cte( |
| 181 | + statement: HasCTE, |
| 182 | + name: str | None = None, |
| 183 | + recursive: bool = False, |
| 184 | + nesting: bool = False, |
| 185 | + materialized: bool = False, |
| 186 | +) -> CTE: |
| 187 | + """Standard SQLAlchemy `cte()` plus `materialized=True` for `WITH <name> AS MATERIALIZED (...)`. |
| 188 | +
|
| 189 | + A materialized CTE body is computed once instead of being inlined at every reference. |
| 190 | + Requires ClickHouse 26.3 or later. The server only honors the keyword when the |
| 191 | + analyzer and the `enable_materialized_cte` setting are enabled for the query: |
| 192 | +
|
| 193 | + stmt = select(...).execution_options( |
| 194 | + settings={"enable_materialized_cte": 1, "enable_analyzer": 1} |
| 195 | + ) |
| 196 | +
|
| 197 | + Use this with the standard `sqlalchemy.select`. Statements built with |
| 198 | + `cc_sqlalchemy.select` have the same options on their own `.cte()` method. |
| 199 | + Raises `ValueError` when `recursive` and `materialized` are both true. |
| 200 | + """ |
| 201 | + if not isinstance(statement, HasCTE): |
| 202 | + raise TypeError(f"cte() expects a SQLAlchemy statement that supports CTEs. Got {type(statement).__name__}") |
| 203 | + _validate_cte_options(recursive, materialized) |
| 204 | + return _apply_materialized(statement.cte(name=name, recursive=recursive, nesting=nesting), materialized) |
| 205 | + |
| 206 | + |
161 | 207 | def _select_ch_join( |
162 | 208 | self: Select, |
163 | 209 | right: Any, |
@@ -278,6 +324,20 @@ def left_array_join(self, *cols: Any, alias: Any = None) -> "ClickHouseSelect": |
278 | 324 | def prewhere(self, whereclause: Any) -> "ClickHouseSelect": |
279 | 325 | return cast("ClickHouseSelect", prewhere(self, whereclause)) |
280 | 326 |
|
| 327 | + def cte( |
| 328 | + self, |
| 329 | + name: str | None = None, |
| 330 | + recursive: bool = False, |
| 331 | + nesting: bool = False, |
| 332 | + materialized: bool = False, |
| 333 | + ) -> CTE: |
| 334 | + """Standard `Select.cte()` plus `materialized=True` for `AS MATERIALIZED`. |
| 335 | +
|
| 336 | + See :func:`cte` for the server requirements. |
| 337 | + """ |
| 338 | + _validate_cte_options(recursive, materialized) |
| 339 | + return _apply_materialized(super().cte(name=name, recursive=recursive, nesting=nesting), materialized) |
| 340 | + |
281 | 341 | def limit_by(self, by_clauses: Any, limit: int, offset: int | None = None) -> "ClickHouseSelect": |
282 | 342 | return cast("ClickHouseSelect", limit_by(self, by_clauses, limit, offset)) |
283 | 343 |
|
|
0 commit comments