|
5 | 5 | import re |
6 | 6 | from functools import cached_property |
7 | 7 | from sqlglot import exp, maybe_parse |
| 8 | +from sqlglot.optimizer.scope import Scope, find_all_in_scope, traverse_scope |
8 | 9 | from sqlmesh.core.dialect import to_schema |
9 | 10 | from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin |
10 | 11 | from sqlmesh.core.engine_adapter.base import EngineAdapterWithIndexSupport |
|
20 | 21 | ) |
21 | 22 | from sqlmesh.core.schema_diff import TableAlterOperation |
22 | 23 | from sqlmesh.utils import get_source_columns_to_types |
| 24 | +from sqlmesh.utils.errors import SQLMeshError |
23 | 25 |
|
24 | 26 | if t.TYPE_CHECKING: |
25 | 27 | import pandas as pd |
@@ -71,17 +73,164 @@ def inject_virtual_catalog(self, gateway: str) -> None: |
71 | 73 |
|
72 | 74 | def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: |
73 | 75 | """Render SQL without the virtual catalog, which ClickHouse does not support.""" |
74 | | - if self._default_catalog and any(self._virtual_catalog_references(expression)): |
| 76 | + virtual_catalog = self._default_catalog or self._extra_config.get("virtual_catalog") |
| 77 | + if virtual_catalog and any(self._virtual_catalog_references(expression)): |
75 | 78 | expression = expression.copy() |
76 | 79 | for reference in list(self._virtual_catalog_references(expression)): |
77 | 80 | reference.set("catalog", None) |
| 81 | + if isinstance(expression, exp.Query): |
| 82 | + expression = self._map_catalog_references(expression, preserve_unmapped=True) |
| 83 | + elif isinstance(expression, exp.Delete): |
| 84 | + where = expression.args.get("where") |
| 85 | + if where and any(where.find_all(exp.Query)): |
| 86 | + expression = expression.copy() |
| 87 | + where = expression.args["where"] |
| 88 | + expression.set( |
| 89 | + "where", |
| 90 | + where.transform( |
| 91 | + lambda node: self._map_catalog_references(node) |
| 92 | + if isinstance(node, exp.Query) |
| 93 | + else node, |
| 94 | + copy=False, |
| 95 | + ), |
| 96 | + ) |
| 97 | + elif isinstance(expression, exp.Insert) and isinstance(expression.expression, exp.Query): |
| 98 | + expression = expression.copy() |
| 99 | + expression.set("expression", self._map_catalog_references(expression.expression)) |
78 | 100 | return super()._to_sql(expression, quote=quote, **kwargs) |
79 | 101 |
|
80 | | - def _virtual_catalog_references(self, expression: exp.Expr) -> t.Iterator[exp.Expr]: |
| 102 | + def _virtual_catalog_references( |
| 103 | + self, expression: exp.Expr |
| 104 | + ) -> t.Iterator[t.Union[exp.Table, exp.Column]]: |
| 105 | + virtual_catalog = self._default_catalog or self._extra_config.get("virtual_catalog") |
81 | 106 | return ( |
82 | 107 | reference |
83 | | - for reference in expression.find_all(exp.Table, exp.Column) |
84 | | - if reference.text("catalog") == self._default_catalog |
| 108 | + for reference in self._catalog_references(expression) |
| 109 | + if virtual_catalog and reference.text("catalog") == virtual_catalog |
| 110 | + ) |
| 111 | + |
| 112 | + def _get_source_queries( |
| 113 | + self, |
| 114 | + query_or_df: QueryOrDF, |
| 115 | + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], |
| 116 | + target_table: TableName, |
| 117 | + *, |
| 118 | + batch_size: t.Optional[int] = None, |
| 119 | + source_columns: t.Optional[t.List[str]] = None, |
| 120 | + ) -> t.List[SourceQuery]: |
| 121 | + if isinstance(query_or_df, exp.Query): |
| 122 | + query_or_df = self._map_catalog_references(query_or_df) |
| 123 | + return super()._get_source_queries( |
| 124 | + query_or_df, |
| 125 | + target_columns_to_types, |
| 126 | + target_table, |
| 127 | + batch_size=batch_size, |
| 128 | + source_columns=source_columns, |
| 129 | + ) |
| 130 | + |
| 131 | + def _map_catalog_references(self, query: Query, *, preserve_unmapped: bool = False) -> Query: |
| 132 | + mapping = self._extra_config.get("datalake_catalog_mapping") or {} |
| 133 | + if not self._default_catalog and not mapping: |
| 134 | + return query |
| 135 | + |
| 136 | + mapped_catalogs = {self._default_catalog, *mapping} |
| 137 | + query = query.copy() |
| 138 | + for reference in self._catalog_references(query): |
| 139 | + if not reference.catalog or ( |
| 140 | + preserve_unmapped and reference.catalog not in mapped_catalogs |
| 141 | + ): |
| 142 | + continue |
| 143 | + if isinstance(reference, exp.Table): |
| 144 | + self._map_catalog_reference(reference) |
| 145 | + continue |
| 146 | + |
| 147 | + table = self._map_catalog_reference( |
| 148 | + exp.Table( |
| 149 | + this=reference.args["table"].copy(), |
| 150 | + db=reference.args["db"].copy(), |
| 151 | + catalog=reference.args["catalog"].copy(), |
| 152 | + ) |
| 153 | + ) |
| 154 | + for column_key, table_key in (("table", "this"), ("db", "db"), ("catalog", "catalog")): |
| 155 | + value = table.args.get(table_key) |
| 156 | + reference.set(column_key, value.copy() if value is not None else None) |
| 157 | + return query |
| 158 | + |
| 159 | + def _catalog_references(self, expression: exp.Expr) -> t.List[t.Union[exp.Table, exp.Column]]: |
| 160 | + """Use query scopes to distinguish qualified columns from nested-field paths.""" |
| 161 | + tables = list(expression.find_all(exp.Table)) |
| 162 | + references: t.List[t.Union[exp.Table, exp.Column]] = list(tables) |
| 163 | + root_tables = list(find_all_in_scope(expression, exp.Table)) |
| 164 | + scoped_columns = set() |
| 165 | + virtual_catalog = self._default_catalog or self._extra_config.get("virtual_catalog") |
| 166 | + |
| 167 | + def matches_source_table(column: exp.Column, table: exp.Table) -> bool: |
| 168 | + if (column.catalog, column.db, column.table) == (table.catalog, table.db, table.name): |
| 169 | + return True |
| 170 | + # ClickHouse strips the virtual catalog from DML targets before SQL rendering. |
| 171 | + return ( |
| 172 | + column.catalog == virtual_catalog |
| 173 | + and not table.catalog |
| 174 | + and (column.db, column.table) == (table.db, table.name) |
| 175 | + ) |
| 176 | + |
| 177 | + def add_column_reference(column: exp.Column, scope: t.Optional[Scope]) -> None: |
| 178 | + if not column.catalog: |
| 179 | + return |
| 180 | + current_scope = scope |
| 181 | + while current_scope is not None: |
| 182 | + # The nearest alias or CTE name takes precedence as a nested-field path. |
| 183 | + if column.catalog in current_scope.sources: |
| 184 | + return |
| 185 | + if any( |
| 186 | + isinstance(source, exp.Table) and matches_source_table(column, source) |
| 187 | + for source in current_scope.sources.values() |
| 188 | + ): |
| 189 | + references.append(column) |
| 190 | + return |
| 191 | + current_scope = current_scope.parent |
| 192 | + |
| 193 | + if column.catalog in {table.alias_or_name for table in root_tables}: |
| 194 | + return |
| 195 | + if any(matches_source_table(column, table) for table in root_tables): |
| 196 | + references.append(column) |
| 197 | + |
| 198 | + for scope in traverse_scope(expression): |
| 199 | + for column in scope.find_all(exp.Column): |
| 200 | + scoped_columns.add(id(column)) |
| 201 | + add_column_reference(column, scope) |
| 202 | + |
| 203 | + for column in find_all_in_scope(expression, exp.Column): |
| 204 | + if id(column) not in scoped_columns: |
| 205 | + add_column_reference(column, None) |
| 206 | + return references |
| 207 | + |
| 208 | + def _map_catalog_reference(self, table: exp.Table) -> exp.Table: |
| 209 | + """Map catalog.schema.table to a ClickHouse DataLakeCatalog database.table reference.""" |
| 210 | + catalog = table.catalog |
| 211 | + if not catalog: |
| 212 | + return table |
| 213 | + if catalog == self._default_catalog: |
| 214 | + table.set("catalog", None) |
| 215 | + return table |
| 216 | + |
| 217 | + target_database = (self._extra_config.get("datalake_catalog_mapping") or {}).get(catalog) |
| 218 | + if target_database: |
| 219 | + if not table.db: |
| 220 | + raise SQLMeshError( |
| 221 | + f"ClickHouse DataLakeCatalog mapping for '{catalog}' requires a " |
| 222 | + "three-level catalog.schema.table reference" |
| 223 | + ) |
| 224 | + source_table = f"{table.db}.{table.name}" |
| 225 | + table.set("catalog", None) |
| 226 | + table.set("db", exp.to_identifier(target_database)) |
| 227 | + table.set("this", exp.to_identifier(source_table, quoted=True)) |
| 228 | + return table |
| 229 | + |
| 230 | + raise SQLMeshError( |
| 231 | + f"{self.dialect} requires that all catalog operations be against the virtual " |
| 232 | + f"catalog {self._default_catalog!r} or a configured DataLakeCatalog mapping. " |
| 233 | + f"Provided catalog: {catalog}" |
85 | 234 | ) |
86 | 235 |
|
87 | 236 | @property |
|
0 commit comments