1"""Contains MSSQLEngineAdapter."""
  2
  3from __future__ import annotations
  4
  5from textwrap import dedent
  6import typing as t
  7import logging
  8
  9from sqlglot import exp
 10
 11from sqlmesh.core.dialect import to_schema, add_table
 12from sqlmesh.core.engine_adapter.base import (
 13    EngineAdapterWithIndexSupport,
 14    EngineAdapter,
 15    InsertOverwriteStrategy,
 16    MERGE_SOURCE_ALIAS,
 17    MERGE_TARGET_ALIAS,
 18    _get_data_object_cache_key,
 19)
 20from sqlmesh.core.engine_adapter.mixins import (
 21    GetCurrentCatalogFromFunctionMixin,
 22    PandasNativeFetchDFSupportMixin,
 23    VarcharSizeWorkaroundMixin,
 24    RowDiffMixin,
 25)
 26from sqlmesh.core.engine_adapter.shared import (
 27    CatalogSupport,
 28    CommentCreationTable,
 29    CommentCreationView,
 30    DataObject,
 31    DataObjectType,
 32    SourceQuery,
 33    set_catalog,
 34)
 35from sqlmesh.utils import get_source_columns_to_types
 36
 37if t.TYPE_CHECKING:
 38    from sqlmesh.core._typing import SchemaName, TableName
 39    from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF
 40
 41
 42logger = logging.getLogger(__name__)
 43
 44
 45@set_catalog()
 46class MSSQLEngineAdapter(
 47    EngineAdapterWithIndexSupport,
 48    PandasNativeFetchDFSupportMixin,
 49    GetCurrentCatalogFromFunctionMixin,
 50    VarcharSizeWorkaroundMixin,
 51    RowDiffMixin,
 52):
 53    DIALECT: str = "tsql"
 54    SUPPORTS_TUPLE_IN = False
 55    SUPPORTS_MATERIALIZED_VIEWS = False
 56    CURRENT_CATALOG_EXPRESSION = exp.func("db_name")
 57    COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
 58    COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
 59    SUPPORTS_REPLACE_TABLE = False
 60    MAX_IDENTIFIER_LENGTH = 128
 61    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 62    SCHEMA_DIFFER_KWARGS = {
 63        "parameterized_type_defaults": {
 64            exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(18, 0), (0,)],
 65            exp.DataType.build("BINARY", dialect=DIALECT).this: [(1,)],
 66            exp.DataType.build("VARBINARY", dialect=DIALECT).this: [(1,)],
 67            exp.DataType.build("CHAR", dialect=DIALECT).this: [(1,)],
 68            exp.DataType.build("VARCHAR", dialect=DIALECT).this: [(1,)],
 69            exp.DataType.build("NCHAR", dialect=DIALECT).this: [(1,)],
 70            exp.DataType.build("NVARCHAR", dialect=DIALECT).this: [(1,)],
 71            exp.DataType.build("TIME", dialect=DIALECT).this: [(7,)],
 72            exp.DataType.build("DATETIME2", dialect=DIALECT).this: [(7,)],
 73            exp.DataType.build("DATETIMEOFFSET", dialect=DIALECT).this: [(7,)],
 74        },
 75        "max_parameter_length": {
 76            exp.DataType.build("VARBINARY", dialect=DIALECT).this: 2147483647,  # 2 GB
 77            exp.DataType.build("VARCHAR", dialect=DIALECT).this: 2147483647,
 78            exp.DataType.build("NVARCHAR", dialect=DIALECT).this: 2147483647,
 79        },
 80    }
 81    VARIABLE_LENGTH_DATA_TYPES = {"binary", "varbinary", "char", "varchar", "nchar", "nvarchar"}
 82    INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.MERGE
 83
 84    @property
 85    def catalog_support(self) -> CatalogSupport:
 86        # MSSQL and AzureSQL both use this engine adapter, but they differ in catalog support.
 87        # Therefore, we specify the catalog support in the connection config `_extra_engine_config`
 88        # instead of in the adapter itself.
 89        return self._extra_config["catalog_support"]
 90
 91    def columns(
 92        self,
 93        table_name: TableName,
 94        include_pseudo_columns: bool = True,
 95    ) -> t.Dict[str, exp.DataType]:
 96        """MsSql doesn't support describe so we query information_schema."""
 97
 98        table = exp.to_table(table_name)
 99
100        sql = (
101            exp.select(
102                "COLUMN_NAME",
103                "DATA_TYPE",
104                "CHARACTER_MAXIMUM_LENGTH",
105                "NUMERIC_PRECISION",
106                "NUMERIC_SCALE",
107            )
108            .from_("INFORMATION_SCHEMA.COLUMNS")
109            .where(f"TABLE_NAME = '{table.name}'")
110        )
111        database_name = table.db
112        if database_name:
113            sql = sql.where(f"TABLE_SCHEMA = '{database_name}'")
114
115        columns_raw = self.fetchall(sql, quote_identifiers=True)
116
117        def build_var_length_col(
118            column_name: str,
119            data_type: str,
120            character_maximum_length: t.Optional[int] = None,
121            numeric_precision: t.Optional[int] = None,
122            numeric_scale: t.Optional[int] = None,
123        ) -> tuple:
124            data_type = data_type.lower()
125            if (
126                data_type in self.VARIABLE_LENGTH_DATA_TYPES
127                and character_maximum_length is not None
128                and character_maximum_length > 0
129            ):
130                return (column_name, f"{data_type}({character_maximum_length})")
131            if (
132                data_type in ("varbinary", "varchar", "nvarchar")
133                and character_maximum_length is not None
134                and character_maximum_length == -1
135            ):
136                return (column_name, f"{data_type}(max)")
137            if data_type in ("decimal", "numeric"):
138                return (column_name, f"{data_type}({numeric_precision}, {numeric_scale})")
139            if data_type == "float":
140                return (column_name, f"{data_type}({numeric_precision})")
141
142            return (column_name, data_type)
143
144        columns = [build_var_length_col(*row) for row in columns_raw]
145
146        return {
147            column_name: exp.DataType.build(data_type, dialect=self.dialect)
148            for column_name, data_type in columns
149        }
150
151    def table_exists(self, table_name: TableName) -> bool:
152        """MsSql doesn't support describe so we query information_schema."""
153        table = exp.to_table(table_name)
154        data_object_cache_key = _get_data_object_cache_key(table.catalog, table.db, table.name)
155        if data_object_cache_key in self._data_object_cache:
156            logger.debug("Table existence cache hit: %s", data_object_cache_key)
157            return self._data_object_cache[data_object_cache_key] is not None
158
159        sql = (
160            exp.select("1")
161            .from_("INFORMATION_SCHEMA.TABLES")
162            .where(f"TABLE_NAME = '{table.alias_or_name}'")
163        )
164        database_name = table.db
165        if database_name:
166            sql = sql.where(f"TABLE_SCHEMA = '{database_name}'")
167
168        result = self.fetchone(sql, quote_identifiers=True)
169
170        return result[0] == 1 if result else False
171
172    def set_current_catalog(self, catalog_name: str) -> None:
173        self.execute(exp.Use(this=exp.to_identifier(catalog_name)))
174
175    def drop_schema(
176        self,
177        schema_name: SchemaName,
178        ignore_if_not_exists: bool = True,
179        cascade: bool = False,
180        **drop_args: t.Dict[str, exp.Expr],
181    ) -> None:
182        """
183        MsSql doesn't support CASCADE clause and drops schemas unconditionally.
184        """
185        if cascade:
186            objects = self._get_data_objects(schema_name)
187            for obj in objects:
188                # Build properly quoted table for MSSQL using square brackets when needed
189                object_table = exp.table_(obj.name, obj.schema_name)
190
191                # _get_data_objects is catalog-specific, so these can't accidentally drop view/tables in another catalog
192                if obj.type == DataObjectType.VIEW:
193                    self.drop_view(
194                        object_table,
195                        ignore_if_not_exists=ignore_if_not_exists,
196                    )
197                else:
198                    self.drop_table(
199                        object_table,
200                        exists=ignore_if_not_exists,
201                    )
202        super().drop_schema(schema_name, ignore_if_not_exists=ignore_if_not_exists, cascade=False)
203
204    def merge(
205        self,
206        target_table: TableName,
207        source_table: QueryOrDF,
208        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
209        unique_key: t.Sequence[exp.Expr],
210        when_matched: t.Optional[exp.Whens] = None,
211        merge_filter: t.Optional[exp.Expr] = None,
212        source_columns: t.Optional[t.List[str]] = None,
213        **kwargs: t.Any,
214    ) -> None:
215        mssql_merge_exists = kwargs.get("physical_properties", {}).get("mssql_merge_exists")
216
217        source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types(
218            source_table,
219            target_columns_to_types,
220            target_table=target_table,
221            source_columns=source_columns,
222        )
223        target_columns_to_types = target_columns_to_types or self.columns(target_table)
224        on = exp.and_(
225            *(
226                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
227                for part in unique_key
228            )
229        )
230        if merge_filter:
231            on = exp.and_(merge_filter, on)
232
233        match_expressions = []
234        if not when_matched:
235            unique_key_names = [y.name for y in unique_key]
236            columns_to_types_no_keys = [
237                c for c in target_columns_to_types if c not in unique_key_names
238            ]
239
240            target_columns_no_keys = [
241                exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types_no_keys
242            ]
243            source_columns_no_keys = [
244                exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types_no_keys
245            ]
246
247            match_condition = (
248                exp.Exists(
249                    this=exp.select(*target_columns_no_keys).except_(
250                        exp.select(*source_columns_no_keys)
251                    )
252                )
253                if mssql_merge_exists
254                else None
255            )
256
257            if target_columns_no_keys:
258                match_expressions.append(
259                    exp.When(
260                        matched=True,
261                        source=False,
262                        condition=match_condition,
263                        then=exp.Update(
264                            expressions=[
265                                exp.column(col, MERGE_TARGET_ALIAS).eq(
266                                    exp.column(col, MERGE_SOURCE_ALIAS)
267                                )
268                                for col in columns_to_types_no_keys
269                            ],
270                        ),
271                    )
272                )
273        else:
274            match_expressions.extend(when_matched.copy().expressions)
275
276        match_expressions.append(
277            exp.When(
278                matched=False,
279                source=False,
280                then=exp.Insert(
281                    this=exp.Tuple(
282                        expressions=[exp.column(col) for col in target_columns_to_types]
283                    ),
284                    expression=exp.Tuple(
285                        expressions=[
286                            exp.column(col, MERGE_SOURCE_ALIAS) for col in target_columns_to_types
287                        ]
288                    ),
289                ),
290            )
291        )
292        for source_query in source_queries:
293            with source_query as query:
294                self._merge(
295                    target_table=target_table,
296                    query=query,
297                    on=on,
298                    whens=exp.Whens(expressions=match_expressions),
299                )
300
301    def _convert_df_datetime(self, df: DF, columns_to_types: t.Dict[str, exp.DataType]) -> None:
302        import pandas as pd
303        from pandas.api.types import is_datetime64_any_dtype  # type: ignore
304
305        # pymssql doesn't convert Pandas Timestamp (datetime64) types
306        # - this code is based on snowflake adapter implementation
307        for column, kind in columns_to_types.items():
308            # pymssql errors if the column contains a datetime.date object
309            if kind.is_type("date"):  # type: ignore
310                df[column] = pd.to_datetime(df[column]).dt.strftime("%Y-%m-%d")  # type: ignore
311            elif is_datetime64_any_dtype(df.dtypes[column]):  # type: ignore
312                if getattr(df.dtypes[column], "tz", None) is not None:  # type: ignore
313                    # MSSQL requires a colon in the offset (+00:00) so we use isoformat() instead of strftime()
314                    df[column] = pd.to_datetime(df[column]).map(lambda x: x.isoformat(" "))  # type: ignore
315
316                    # bulk_copy() doesn't work with TZ timestamp, so load into string column and cast to
317                    # timestamp in SELECT statement
318                    columns_to_types[column] = exp.DataType.build("TEXT")
319                else:
320                    df[column] = pd.to_datetime(df[column]).dt.strftime("%Y-%m-%d %H:%M:%S.%f")  # type: ignore
321
322    def _df_to_source_queries(
323        self,
324        df: DF,
325        target_columns_to_types: t.Dict[str, exp.DataType],
326        batch_size: int,
327        target_table: TableName,
328        source_columns: t.Optional[t.List[str]] = None,
329    ) -> t.List[SourceQuery]:
330        import pandas as pd
331        import numpy as np
332
333        assert isinstance(df, pd.DataFrame)
334        temp_table = self._get_temp_table(target_table or "pandas")
335
336        # Return the superclass implementation if the connection pool doesn't support bulk_copy
337        if not hasattr(self._connection_pool.get(), "bulk_copy"):
338            return super()._df_to_source_queries(
339                df, target_columns_to_types, batch_size, target_table, source_columns=source_columns
340            )
341
342        def query_factory() -> Query:
343            # It is possible for the factory to be called multiple times and if so then the temp table will already
344            # be created so we skip creating again. This means we are assuming the first call is the same result
345            # as later calls.
346            if not self.table_exists(temp_table):
347                source_columns_to_types = get_source_columns_to_types(
348                    target_columns_to_types, source_columns
349                )
350                ordered_df = df[
351                    list(source_columns_to_types)
352                ]  # reorder DataFrame so it matches columns_to_types
353                self._convert_df_datetime(ordered_df, source_columns_to_types)
354                self.create_table(temp_table, source_columns_to_types)
355                rows: t.List[t.Tuple[t.Any, ...]] = list(
356                    ordered_df.replace({np.nan: None}).itertuples(index=False, name=None)  # type: ignore
357                )
358                conn = self._connection_pool.get()
359                conn.bulk_copy(temp_table.sql(dialect=self.dialect), rows)
360            return exp.select(
361                *self._casted_columns(target_columns_to_types, source_columns=source_columns)
362            ).from_(temp_table)  # type: ignore
363
364        return [
365            SourceQuery(
366                query_factory=query_factory,
367                cleanup_func=lambda: self.drop_table(temp_table),
368            )
369        ]
370
371    def _get_data_objects(
372        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
373    ) -> t.List[DataObject]:
374        """
375        Returns all the data objects that exist in the given schema and catalog.
376        """
377        import pandas as pd
378
379        catalog = self.get_current_catalog()
380        query = (
381            exp.select(
382                exp.column("TABLE_NAME").as_("name"),
383                exp.column("TABLE_SCHEMA").as_("schema_name"),
384                exp.case()
385                .when(exp.column("TABLE_TYPE").eq("BASE TABLE"), exp.Literal.string("TABLE"))
386                .else_(exp.column("TABLE_TYPE"))
387                .as_("type"),
388            )
389            .from_(exp.table_("TABLES", db="INFORMATION_SCHEMA"))
390            .where(exp.column("TABLE_SCHEMA").eq(to_schema(schema_name).db))
391        )
392        if object_names:
393            query = query.where(exp.column("TABLE_NAME").isin(*object_names))
394        dataframe: pd.DataFrame = self.fetchdf(query)
395        return [
396            DataObject(
397                catalog=catalog,  # type: ignore
398                schema=row.schema_name,  # type: ignore
399                name=row.name,  # type: ignore
400                type=DataObjectType.from_str(row.type),  # type: ignore
401            )
402            for row in dataframe.itertuples()
403        ]
404
405    def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str:
406        sql = super()._to_sql(expression, quote=quote, **kwargs)
407        return f"{sql};"
408
409    def _rename_table(
410        self,
411        old_table_name: TableName,
412        new_table_name: TableName,
413    ) -> None:
414        # The function that renames tables in MSSQL takes string literals as arguments instead of identifiers,
415        # so we shouldn't quote the identifiers.
416        self.execute(exp.rename_table(old_table_name, new_table_name), quote_identifiers=False)
417
418    def _insert_overwrite_by_condition(
419        self,
420        table_name: TableName,
421        source_queries: t.List[SourceQuery],
422        target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
423        where: t.Optional[exp.Condition] = None,
424        insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None,
425        **kwargs: t.Any,
426    ) -> None:
427        # note that this is passed as table_properties here rather than physical_properties
428        use_merge_strategy = kwargs.get("table_properties", {}).get("mssql_merge_exists")
429        if (not where or where == exp.true()) and not use_merge_strategy:
430            # this is a full table replacement, call the base strategy to do DELETE+INSERT
431            # which will result in TRUNCATE+INSERT due to how we have overridden self.delete_from()
432            return EngineAdapter._insert_overwrite_by_condition(
433                self,
434                table_name=table_name,
435                source_queries=source_queries,
436                target_columns_to_types=target_columns_to_types,
437                where=where,
438                insert_overwrite_strategy_override=InsertOverwriteStrategy.DELETE_INSERT,
439                **kwargs,
440            )
441
442        # For conditional overwrites or when mssql_merge_exists is set use MERGE
443        return super()._insert_overwrite_by_condition(
444            table_name=table_name,
445            source_queries=source_queries,
446            target_columns_to_types=target_columns_to_types,
447            where=where,
448            insert_overwrite_strategy_override=insert_overwrite_strategy_override,
449            **kwargs,
450        )
451
452    def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> None:
453        if where == exp.true():
454            # "A TRUNCATE TABLE operation can be rolled back within a transaction."
455            # ref: https://learn.microsoft.com/en-us/sql/t-sql/statements/truncate-table-transact-sql?view=sql-server-ver15#remarks
456            return self.execute(
457                exp.TruncateTable(expressions=[exp.to_table(table_name, dialect=self.dialect)])
458            )
459
460        return super().delete_from(table_name, where)
461
462    def _build_create_comment_table_exp(
463        self, table: exp.Table, table_comment: str, table_kind: str = "TABLE"
464    ) -> exp.Comment | str:
465        template = dedent("""
466            DECLARE @comment sql_variant = {comment};
467            DECLARE @property_name VARCHAR(128) = 'MS_Description';
468            DECLARE @schema_name VARCHAR(128) = {schema_name};
469            DECLARE @object_name VARCHAR(128) = {object_name};
470            DECLARE @object_kind VARCHAR(128) = '{object_kind}';
471            DECLARE @existing sql_variant;
472
473            SELECT TOP 1 @existing = CAST(VALUE AS NVARCHAR) FROM fn_listextendedproperty(@property_name, 'schema', @schema_name, @object_kind, @object_name, DEFAULT, DEFAULT);
474
475            IF @comment IS NULL
476                BEGIN
477                    IF @existing IS NOT NULL
478                        EXEC sp_dropextendedproperty @property_name, 'schema', @schema_name, @object_kind, @object_name;
479                END
480            ELSE
481                BEGIN
482                    IF @existing IS NULL
483                        EXEC sp_addextendedproperty @property_name,@comment, 'schema', @schema_name, @object_kind, @object_name;
484                    ELSE IF @existing != @comment
485                        EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name;
486                END
487        """)
488        tsql_text = template.format(
489            comment=exp.Literal.string(table_comment or "NULL").sql(
490                dialect=self.dialect, identify=False
491            ),
492            schema_name=exp.Literal.string(table.db or "dbo").sql(
493                dialect=self.dialect, identify=False
494            ),
495            object_name=exp.Literal.string(table.name).sql(dialect=self.dialect, identify=False),
496            object_kind=table_kind,
497        )
498        return tsql_text
499
500    def _build_create_comment_column_exp(
501        self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE"
502    ) -> exp.Comment | str:
503        template = dedent("""
504            DECLARE @comment sql_variant = {comment};
505            DECLARE @property_name VARCHAR(128) = 'MS_Description';
506            DECLARE @schema_name VARCHAR(128) = {schema_name};
507            DECLARE @object_name VARCHAR(128) = {object_name};
508            DECLARE @object_kind VARCHAR(128) = '{object_kind}';
509            DECLARE @column_name VARCHAR(128) = {column_name};
510            DECLARE @existing sql_variant;
511
512            SELECT TOP 1 @existing = CAST(VALUE AS NVARCHAR) FROM fn_listextendedproperty(@property_name, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name);
513
514            IF @comment IS NULL
515                BEGIN
516                    IF @existing IS NOT NULL
517                        EXEC sp_dropextendedproperty @property_name, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
518                END
519            ELSE
520                BEGIN
521                    IF @existing IS NULL
522                        EXEC sp_addextendedproperty @property_name,@comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
523                    ELSE IF @existing != @comment
524                        EXEC sp_updateextendedproperty @property_name, @comment, 'schema', @schema_name, @object_kind, @object_name, 'column', @column_name;
525                END
526        """)
527
528        tsql_text = template.format(
529            comment=exp.Literal.string(column_comment or "NULL").sql(
530                dialect=self.dialect, identify=False
531            ),
532            schema_name=exp.Literal.string(table.db or "dbo").sql(
533                dialect=self.dialect, identify=False
534            ),
535            object_name=exp.Literal.string(table.name).sql(dialect=self.dialect, identify=False),
536            object_kind=table_kind,
537            column_name=exp.Literal.string(column_name).sql(dialect=self.dialect, identify=False),
538        )
539
540        return tsql_text