sqlmesh.core.engine_adapter.redshift

  1from __future__ import annotations
  2
  3import logging
  4import typing as t
  5
  6from sqlglot import exp
  7
  8from sqlmesh.core.dialect import to_schema
  9from sqlmesh.core.engine_adapter.base import MERGE_SOURCE_ALIAS, MERGE_TARGET_ALIAS
 10from sqlmesh.core.engine_adapter.base_postgres import BasePostgresEngineAdapter
 11from sqlmesh.core.engine_adapter.mixins import (
 12    GetCurrentCatalogFromFunctionMixin,
 13    NonTransactionalTruncateMixin,
 14    VarcharSizeWorkaroundMixin,
 15    RowDiffMixin,
 16    logical_merge,
 17    GrantsFromInfoSchemaMixin,
 18)
 19from sqlmesh.core.engine_adapter.shared import (
 20    CommentCreationView,
 21    DataObject,
 22    DataObjectType,
 23    SourceQuery,
 24    set_catalog,
 25)
 26from sqlmesh.utils.errors import SQLMeshError
 27
 28if t.TYPE_CHECKING:
 29    import pandas as pd
 30
 31    from sqlmesh.core._typing import SchemaName, TableName
 32    from sqlmesh.core.engine_adapter.base import QueryOrDF, Query
 33
 34logger = logging.getLogger(__name__)
 35
 36
 37@set_catalog()
 38class RedshiftEngineAdapter(
 39    BasePostgresEngineAdapter,
 40    GetCurrentCatalogFromFunctionMixin,
 41    NonTransactionalTruncateMixin,
 42    VarcharSizeWorkaroundMixin,
 43    RowDiffMixin,
 44    GrantsFromInfoSchemaMixin,
 45):
 46    DIALECT = "redshift"
 47    CURRENT_CATALOG_EXPRESSION = exp.func("current_database")
 48    # Redshift doesn't support comments for VIEWs WITH NO SCHEMA BINDING (which we always use)
 49    COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
 50    SUPPORTS_REPLACE_TABLE = False
 51    SUPPORTS_GRANTS = True
 52    SUPPORTS_MULTIPLE_GRANT_PRINCIPALS = True
 53
 54    SCHEMA_DIFFER_KWARGS = {
 55        "parameterized_type_defaults": {
 56            exp.DataType.build("VARBYTE", dialect=DIALECT).this: [(64000,)],
 57            exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(18, 0), (0,)],
 58            exp.DataType.build("CHAR", dialect=DIALECT).this: [(1,)],
 59            exp.DataType.build("VARCHAR", dialect=DIALECT).this: [(256,)],
 60            exp.DataType.build("NCHAR", dialect=DIALECT).this: [(1,)],
 61            exp.DataType.build("NVARCHAR", dialect=DIALECT).this: [(256,)],
 62        },
 63        "max_parameter_length": {
 64            exp.DataType.build("CHAR", dialect=DIALECT).this: 4096,
 65            exp.DataType.build("VARCHAR", dialect=DIALECT).this: 65535,
 66        },
 67        "precision_increase_allowed_types": {exp.DataType.build("VARCHAR", dialect=DIALECT).this},
 68        "drop_cascade": True,
 69    }
 70    VARIABLE_LENGTH_DATA_TYPES = {
 71        "char",
 72        "character",
 73        "nchar",
 74        "varchar",
 75        "character varying",
 76        "nvarchar",
 77        "varbyte",
 78        "varbinary",
 79        "binary varying",
 80    }
 81
 82    def columns(
 83        self,
 84        table_name: TableName,
 85        include_pseudo_columns: bool = True,
 86    ) -> t.Dict[str, exp.DataType]:
 87        table = exp.to_table(table_name)
 88
 89        sql = (
 90            exp.select(
 91                "column_name",
 92                "data_type",
 93                "character_maximum_length",
 94                "numeric_precision",
 95                "numeric_scale",
 96            )
 97            .from_("svv_columns")  # Includes late-binding views
 98            .where(exp.column("table_name").eq(table.alias_or_name))
 99        )
100        if table.args.get("db"):
101            sql = sql.where(exp.column("table_schema").eq(table.args["db"].name))
102
103        columns_raw = self.fetchall(sql, quote_identifiers=True)
104
105        def build_var_length_col(
106            column_name: str,
107            data_type: str,
108            character_maximum_length: t.Optional[int] = None,
109            numeric_precision: t.Optional[int] = None,
110            numeric_scale: t.Optional[int] = None,
111        ) -> tuple:
112            data_type = data_type.lower()
113            if (
114                data_type in self.VARIABLE_LENGTH_DATA_TYPES
115                and character_maximum_length is not None
116            ):
117                return (column_name, f"{data_type}({character_maximum_length})")
118            if data_type in ("decimal", "numeric"):
119                return (column_name, f"{data_type}({numeric_precision}, {numeric_scale})")
120
121            return (column_name, data_type)
122
123        columns = [build_var_length_col(*row) for row in columns_raw]
124
125        return {
126            column_name: exp.DataType.build(data_type, dialect=self.dialect)
127            for column_name, data_type in columns
128        }
129
130    @property
131    def enable_merge(self) -> bool:
132        # Redshift supports the MERGE operation but we use the logical merge
133        # unless the user has opted in by setting enable_merge in the connection.
134        return bool(self._extra_config.get("enable_merge"))
135
136    @property
137    def cursor(self) -> t.Any:
138        # Redshift by default uses a `format` paramstyle that has issues when we try to write our snapshot
139        # data to snapshot table. There doesn't seem to be a way to disable parameter overriding so we just
140        # set it to `qmark` since that doesn't cause issues.
141        cursor = self._connection_pool.get_cursor()
142        cursor.paramstyle = "qmark"
143        return cursor
144
145    def _fetch_native_df(
146        self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False
147    ) -> pd.DataFrame:
148        """Fetches a Pandas DataFrame from the cursor"""
149        import pandas as pd
150
151        self.execute(query, quote_identifiers=quote_identifiers)
152
153        # We manually build the `DataFrame` here because the driver's `fetch_dataframe`
154        # method does not respect the active case-sensitivity configuration.
155        #
156        # Context: https://github.com/aws/amazon-redshift-python-driver/issues/238
157        fetcheddata = self.cursor.fetchall()
158
159        try:
160            columns = [column[0] for column in self.cursor.description]
161        except Exception:
162            columns = None
163            logging.warning(
164                "No row description was found, pandas dataframe will be missing column labels."
165            )
166
167        result = [tuple(row) for row in fetcheddata]
168        return pd.DataFrame(result, columns=columns)
169
170    def _create_table_from_source_queries(
171        self,
172        table_name: TableName,
173        source_queries: t.List[SourceQuery],
174        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
175        exists: bool = True,
176        replace: bool = False,
177        table_description: t.Optional[str] = None,
178        column_descriptions: t.Optional[t.Dict[str, str]] = None,
179        table_kind: t.Optional[str] = None,
180        track_rows_processed: bool = True,
181        **kwargs: t.Any,
182    ) -> None:
183        """
184        Redshift doesn't support `CREATE TABLE IF NOT EXISTS AS...` but does support `CREATE TABLE AS...` so
185        we check if the exists check exists and if not then we can use the base implementation. Otherwise we
186        manually check if it exists and if it does then this is a no-op anyways so we return and if it doesn't
187        then we run the query with exists set to False since we just confirmed it doesn't exist.
188        """
189        if not exists:
190            return super()._create_table_from_source_queries(
191                table_name,
192                source_queries,
193                target_columns_to_types,
194                exists,
195                table_description=table_description,
196                column_descriptions=column_descriptions,
197                **kwargs,
198            )
199        if self.table_exists(table_name):
200            return
201        super()._create_table_from_source_queries(
202            table_name,
203            source_queries,
204            exists=False,
205            table_description=table_description,
206            column_descriptions=column_descriptions,
207            **kwargs,
208        )
209
210    def create_view(
211        self,
212        view_name: TableName,
213        query_or_df: QueryOrDF,
214        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
215        replace: bool = True,
216        materialized: bool = False,
217        materialized_properties: t.Optional[t.Dict[str, t.Any]] = None,
218        table_description: t.Optional[str] = None,
219        column_descriptions: t.Optional[t.Dict[str, str]] = None,
220        view_properties: t.Optional[t.Dict[str, exp.Expr]] = None,
221        source_columns: t.Optional[t.List[str]] = None,
222        **create_kwargs: t.Any,
223    ) -> None:
224        """
225        Redshift views are "binding" by default to their underlying table which means you can't drop that
226        underlying table without dropping the view first. This is a problem for us since we want to be able to
227        swap tables out from under views. Therefore, we create the view as non-binding.
228        """
229        no_schema_binding = True
230        if isinstance(query_or_df, exp.Expr):
231            # We can't include NO SCHEMA BINDING if the query has a recursive CTE
232            has_recursive_cte = any(
233                w.args.get("recursive", False) for w in query_or_df.find_all(exp.With)
234            )
235            no_schema_binding = not has_recursive_cte
236
237        return super().create_view(
238            view_name,
239            query_or_df,
240            target_columns_to_types,
241            replace,
242            materialized,
243            materialized_properties,
244            table_description=table_description,
245            column_descriptions=column_descriptions,
246            no_schema_binding=no_schema_binding,
247            view_properties=view_properties,
248            source_columns=source_columns,
249            **create_kwargs,
250        )
251
252    def replace_query(
253        self,
254        table_name: TableName,
255        query_or_df: QueryOrDF,
256        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
257        table_description: t.Optional[str] = None,
258        column_descriptions: t.Optional[t.Dict[str, str]] = None,
259        source_columns: t.Optional[t.List[str]] = None,
260        supports_replace_table_override: t.Optional[bool] = None,
261        **kwargs: t.Any,
262    ) -> None:
263        """
264        Redshift doesn't support `CREATE OR REPLACE TABLE...` and it also doesn't support `VALUES` expression so we need to specially
265        handle DataFrame replacements.
266
267        If the table doesn't exist then we just create it and load it with insert statements
268        If it does exist then we need to do the:
269            `CREATE TABLE...`, `INSERT INTO...`, `RENAME TABLE...`, `RENAME TABLE...`, DROP TABLE...`  dance.
270        """
271        import pandas as pd
272
273        target_data_object = self.get_data_object(table_name)
274        table_exists = target_data_object is not None
275        if self.drop_data_object_on_type_mismatch(target_data_object, DataObjectType.TABLE):
276            table_exists = False
277
278        if not isinstance(query_or_df, pd.DataFrame) or not table_exists:
279            return super().replace_query(
280                table_name,
281                query_or_df,
282                target_columns_to_types,
283                table_description,
284                column_descriptions,
285                source_columns=source_columns,
286                **kwargs,
287            )
288        source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types(
289            query_or_df,
290            target_columns_to_types,
291            target_table=table_name,
292            source_columns=source_columns,
293        )
294        target_columns_to_types = target_columns_to_types or self.columns(table_name)
295        target_table = exp.to_table(table_name)
296        with self.transaction():
297            temp_table = self._get_temp_table(target_table)
298            old_table = self._get_temp_table(target_table)
299            self.create_table(
300                temp_table,
301                target_columns_to_types,
302                exists=False,
303                table_description=table_description,
304                column_descriptions=column_descriptions,
305                **kwargs,
306            )
307            self._insert_append_source_queries(temp_table, source_queries, target_columns_to_types)
308            self.rename_table(target_table, old_table)
309            self.rename_table(temp_table, target_table)
310            self.drop_table(old_table)
311
312    def _get_data_objects(
313        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
314    ) -> t.List[DataObject]:
315        """
316        Returns all the data objects that exist in the given schema and optionally catalog.
317        """
318        catalog = self.get_current_catalog()
319        table_query = exp.select(
320            exp.column("schemaname").as_("schema_name"),
321            exp.column("tablename").as_("name"),
322            exp.Literal.string("TABLE").as_("type"),
323        ).from_("pg_tables")
324        view_query = (
325            exp.select(
326                exp.column("schemaname").as_("schema_name"),
327                exp.column("viewname").as_("name"),
328                exp.Literal.string("VIEW").as_("type"),
329            )
330            .from_("pg_views")
331            .where(exp.column("definition").ilike("%create materialized view%").not_())
332        )
333        materialized_view_query = (
334            exp.select(
335                exp.column("schemaname").as_("schema_name"),
336                exp.column("viewname").as_("name"),
337                exp.Literal.string("MATERIALIZED_VIEW").as_("type"),
338            )
339            .from_("pg_views")
340            .where(exp.column("definition").ilike("%create materialized view%"))
341        )
342        subquery = exp.union(
343            table_query,
344            exp.union(view_query, materialized_view_query, distinct=False),
345            distinct=False,
346        )
347        query = (
348            exp.select("*")
349            .from_(subquery.subquery(alias="objs"))
350            .where(exp.column("schema_name").eq(to_schema(schema_name).db))
351        )
352        if object_names:
353            query = query.where(exp.column("name").isin(*object_names))
354        df = self.fetchdf(query)
355        return [
356            DataObject(
357                catalog=catalog,
358                schema=row.schema_name,
359                name=row.name,
360                type=DataObjectType.from_str(row.type),  # type: ignore
361            )
362            for row in df.itertuples()
363        ]
364
365    def merge(
366        self,
367        target_table: TableName,
368        source_table: QueryOrDF,
369        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
370        unique_key: t.Sequence[exp.Expr],
371        when_matched: t.Optional[exp.Whens] = None,
372        merge_filter: t.Optional[exp.Expr] = None,
373        source_columns: t.Optional[t.List[str]] = None,
374        **kwargs: t.Any,
375    ) -> None:
376        if self.enable_merge:
377            # By default we use the logical merge unless the user has opted in
378            super().merge(
379                target_table=target_table,
380                source_table=source_table,
381                target_columns_to_types=target_columns_to_types,
382                unique_key=unique_key,
383                when_matched=when_matched,
384                merge_filter=merge_filter,
385                source_columns=source_columns,
386            )
387        else:
388            logical_merge(
389                self,
390                target_table,
391                source_table,
392                target_columns_to_types,
393                unique_key,
394                when_matched=when_matched,
395                merge_filter=merge_filter,
396                source_columns=source_columns,
397            )
398
399    def _merge(
400        self,
401        target_table: TableName,
402        query: Query,
403        on: exp.Expr,
404        whens: exp.Whens,
405    ) -> None:
406        # Redshift does not support table aliases in the target table of a MERGE statement.
407        # So we must use the actual table name instead of an alias, as we do with the source table.
408        def resolve_target_table(expression: exp.Expr) -> exp.Expr:
409            if (
410                isinstance(expression, exp.Column)
411                and expression.table.upper() == MERGE_TARGET_ALIAS
412            ):
413                expression.set("table", exp.to_table(target_table))
414            return expression
415
416        # Ensure that there is exactly one "WHEN MATCHED" and one "WHEN NOT MATCHED" clause.
417        # Since Redshift does not support multiple "WHEN MATCHED" clauses.
418        if (
419            len(whens.expressions) != 2
420            or whens.expressions[0].args["matched"] == whens.expressions[1].args["matched"]
421        ):
422            raise SQLMeshError(
423                "Redshift only supports a single WHEN MATCHED and WHEN NOT MATCHED clause"
424            )
425
426        using = exp.alias_(
427            exp.Subquery(this=query), alias=MERGE_SOURCE_ALIAS, copy=False, table=True
428        )
429        self.execute(
430            exp.Merge(
431                this=target_table,
432                using=using,
433                on=on.transform(resolve_target_table),
434                whens=whens.transform(resolve_target_table),
435            ),
436            track_rows_processed=True,
437        )
438
439    def _normalize_decimal_value(self, expr: exp.Expr, precision: int) -> exp.Expr:
440        # Redshift is finicky. It truncates when the data is already in a table, but rounds when the data is generated as part of a SELECT.
441        #
442        # The following works:
443        #  > select cast(cast(3.14159 as decimal(6, 5)) as decimal(6, 3)); --produces '3.142', the value we want / what every other database produces
444        #
445        # However, if you write that to a table, and then cast it to a less precise decimal, you get _truncation_.
446        #  > create table foo (val decimal(6, 5)); insert into foo(val) values (3.14159);
447        #  > select cast(val as decimal(6, 3)) from foo; --produces '3.141'
448        #
449        # So to make up for this, we force it to round by injecting a round() expression
450        rounded = exp.func("ROUND", expr, precision)
451
452        return super()._normalize_decimal_value(rounded, precision)