sqlmesh.core.engine_adapter.duckdb

  1from __future__ import annotations
  2
  3import typing as t
  4from sqlglot import exp
  5from pathlib import Path
  6
  7from sqlmesh.core.engine_adapter.mixins import (
  8    GetCurrentCatalogFromFunctionMixin,
  9    LogicalMergeMixin,
 10    RowDiffMixin,
 11)
 12from sqlmesh.core.engine_adapter.shared import (
 13    CatalogSupport,
 14    CommentCreationTable,
 15    CommentCreationView,
 16    DataObject,
 17    DataObjectType,
 18    SourceQuery,
 19    set_catalog,
 20)
 21
 22if t.TYPE_CHECKING:
 23    from sqlmesh.core._typing import SchemaName, TableName
 24    from sqlmesh.core.engine_adapter._typing import DF
 25
 26
 27@set_catalog(override_mapping={"_get_data_objects": CatalogSupport.REQUIRES_SET_CATALOG})
 28class DuckDBEngineAdapter(LogicalMergeMixin, GetCurrentCatalogFromFunctionMixin, RowDiffMixin):
 29    DIALECT = "duckdb"
 30    SUPPORTS_TRANSACTIONS = False
 31    SCHEMA_DIFFER_KWARGS = {
 32        "parameterized_type_defaults": {
 33            exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(18, 3), (0,)],
 34        },
 35    }
 36    COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
 37    COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 38    SUPPORTS_CREATE_DROP_CATALOG = True
 39    SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA", "TABLE", "VIEW"]
 40
 41    @property
 42    def catalog_support(self) -> CatalogSupport:
 43        return CatalogSupport.FULL_SUPPORT
 44
 45    def set_current_catalog(self, catalog: str) -> None:
 46        """Sets the catalog name of the current connection."""
 47        self.execute(exp.Use(this=exp.to_identifier(catalog)))
 48
 49    def _create_catalog(self, catalog_name: exp.Identifier) -> None:
 50        if not self._is_motherduck:
 51            db_filename = f"{catalog_name.output_name}.db"
 52            self.execute(
 53                exp.Attach(
 54                    this=exp.alias_(exp.Literal.string(db_filename), catalog_name), exists=True
 55                )
 56            )
 57        else:
 58            self.execute(
 59                exp.Create(this=exp.Table(this=catalog_name), kind="DATABASE", exists=True)
 60            )
 61
 62    def _drop_catalog(self, catalog_name: exp.Identifier) -> None:
 63        if not self._is_motherduck:
 64            db_file_path = Path(f"{catalog_name.output_name}.db")
 65            self.execute(exp.Detach(this=catalog_name, exists=True))
 66            if db_file_path.exists():
 67                db_file_path.unlink()
 68        else:
 69            self.execute(
 70                exp.Drop(
 71                    this=exp.Table(this=catalog_name), kind="DATABASE", cascade=True, exists=True
 72                )
 73            )
 74
 75    def _df_to_source_queries(
 76        self,
 77        df: DF,
 78        target_columns_to_types: t.Dict[str, exp.DataType],
 79        batch_size: int,
 80        target_table: TableName,
 81        source_columns: t.Optional[t.List[str]] = None,
 82    ) -> t.List[SourceQuery]:
 83        temp_table = self._get_temp_table(target_table)
 84        temp_table_sql = (
 85            exp.select(*self._casted_columns(target_columns_to_types, source_columns))
 86            .from_("df")
 87            .sql(dialect=self.dialect)
 88        )
 89        self.cursor.sql(f"CREATE TABLE {temp_table} AS {temp_table_sql}")
 90        return [
 91            SourceQuery(
 92                query_factory=lambda: self._select_columns(target_columns_to_types).from_(
 93                    temp_table
 94                ),  # type: ignore
 95                cleanup_func=lambda: self.drop_table(temp_table),
 96            )
 97        ]
 98
 99    def _get_data_objects(
100        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
101    ) -> t.List[DataObject]:
102        """
103        Returns all the data objects that exist in the given schema and optionally catalog.
104        """
105        catalog = self.get_current_catalog()
106
107        if isinstance(schema_name, exp.Table):
108            # Ensures we don't generate identifier quotes
109            schema_name = ".".join(part.name for part in schema_name.parts)
110
111        query = (
112            exp.select(
113                exp.column("table_name").as_("name"),
114                exp.column("table_schema").as_("schema"),
115                exp.case(exp.column("table_type"))
116                .when(
117                    exp.Literal.string("BASE TABLE"),
118                    exp.Literal.string("table"),
119                )
120                .when(
121                    exp.Literal.string("VIEW"),
122                    exp.Literal.string("view"),
123                )
124                .when(
125                    exp.Literal.string("LOCAL TEMPORARY"),
126                    exp.Literal.string("table"),
127                )
128                .as_("type"),
129            )
130            .from_(exp.to_table("system.information_schema.tables"))
131            .where(
132                exp.column("table_catalog").eq(catalog), exp.column("table_schema").eq(schema_name)
133            )
134        )
135        if object_names:
136            query = query.where(exp.column("table_name").isin(*object_names))
137        df = self.fetchdf(query)
138        return [
139            DataObject(
140                catalog=catalog,  # type: ignore
141                schema=row.schema,  # type: ignore
142                name=row.name,  # type: ignore
143                type=DataObjectType.from_str(row.type),  # type: ignore
144            )
145            for row in df.itertuples()
146        ]
147
148    def _normalize_decimal_value(self, col: exp.Expr, precision: int) -> exp.Expr:
149        """
150        duckdb truncates instead of rounding when casting to decimal.
151
152        other databases: select cast(3.14159 as decimal(38,3)) -> 3.142
153        duckdb: select cast(3.14159 as decimal(38,3)) -> 3.141
154
155        however, we can get the behaviour of other databases by casting to double first:
156        select cast(cast(3.14159 as double) as decimal(38, 3)) -> 3.142
157        """
158        return exp.cast(
159            exp.cast(col, "DOUBLE"),
160            f"DECIMAL(38, {precision})",
161        )
162
163    def _create_table(
164        self,
165        table_name_or_schema: t.Union[exp.Schema, TableName],
166        expression: t.Optional[exp.Expr],
167        exists: bool = True,
168        replace: bool = False,
169        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
170        table_description: t.Optional[str] = None,
171        column_descriptions: t.Optional[t.Dict[str, str]] = None,
172        table_kind: t.Optional[str] = None,
173        track_rows_processed: bool = True,
174        **kwargs: t.Any,
175    ) -> None:
176        catalog = self.get_current_catalog()
177        catalog_type_tuple = self.fetchone(
178            exp.select("type")
179            .from_("duckdb_databases()")
180            .where(exp.column("database_name").eq(catalog))
181        )
182        catalog_type = catalog_type_tuple[0] if catalog_type_tuple else None
183
184        partitioned_by_exps = None
185        if catalog_type == "ducklake":
186            partitioned_by_exps = kwargs.pop("partitioned_by", None)
187
188        super()._create_table(
189            table_name_or_schema,
190            expression,
191            exists,
192            replace,
193            target_columns_to_types,
194            table_description,
195            column_descriptions,
196            table_kind,
197            track_rows_processed=track_rows_processed,
198            **kwargs,
199        )
200
201        if partitioned_by_exps:
202            # Schema object contains column definitions, so we extract Table
203            table_name = (
204                table_name_or_schema.this
205                if isinstance(table_name_or_schema, exp.Schema)
206                else table_name_or_schema
207            )
208            table_name_str = (
209                table_name.sql(dialect=self.dialect)
210                if isinstance(table_name, exp.Table)
211                else table_name
212            )
213            partitioned_by_str = ", ".join(
214                expr.sql(dialect=self.dialect) for expr in partitioned_by_exps
215            )
216            self.execute(f"ALTER TABLE {table_name_str} SET PARTITIONED BY ({partitioned_by_str});")
217
218    @property
219    def _is_motherduck(self) -> bool:
220        return self._extra_config.get("is_motherduck", False)