sqlmesh.core.engine_adapter.mysql

  1from __future__ import annotations
  2
  3import logging
  4import typing as t
  5
  6from sqlglot import exp, parse_one
  7
  8from sqlmesh.core.dialect import to_schema
  9from sqlmesh.core.engine_adapter.mixins import (
 10    LogicalMergeMixin,
 11    NonTransactionalTruncateMixin,
 12    PandasNativeFetchDFSupportMixin,
 13    RowDiffMixin,
 14)
 15from sqlmesh.core.engine_adapter.shared import (
 16    CommentCreationTable,
 17    CommentCreationView,
 18    DataObject,
 19    DataObjectType,
 20    set_catalog,
 21)
 22
 23if t.TYPE_CHECKING:
 24    from sqlmesh.core._typing import SchemaName, TableName
 25
 26logger = logging.getLogger(__name__)
 27
 28
 29@set_catalog()
 30class MySQLEngineAdapter(
 31    LogicalMergeMixin, PandasNativeFetchDFSupportMixin, NonTransactionalTruncateMixin, RowDiffMixin
 32):
 33    DEFAULT_BATCH_SIZE = 200
 34    DIALECT = "mysql"
 35    SUPPORTS_INDEXES = True
 36    COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_NO_CTAS
 37    COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
 38    MAX_TABLE_COMMENT_LENGTH = 2048
 39    MAX_COLUMN_COMMENT_LENGTH = 1024
 40    SUPPORTS_REPLACE_TABLE = False
 41    MAX_IDENTIFIER_LENGTH = 64
 42    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 43    SCHEMA_DIFFER_KWARGS = {
 44        "parameterized_type_defaults": {
 45            exp.DataType.build("BIT", dialect=DIALECT).this: [(1,)],
 46            exp.DataType.build("BINARY", dialect=DIALECT).this: [(1,)],
 47            exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(10, 0), (0,)],
 48            exp.DataType.build("CHAR", dialect=DIALECT).this: [(1,)],
 49            exp.DataType.build("NCHAR", dialect=DIALECT).this: [(1,)],
 50            exp.DataType.build("TEXT", dialect=DIALECT).this: [(65535,)],
 51            exp.DataType.build("TIME", dialect=DIALECT).this: [(0,)],
 52            exp.DataType.build("DATETIME", dialect=DIALECT).this: [(0,)],
 53            exp.DataType.build("TIMESTAMP", dialect=DIALECT).this: [(0,)],
 54        },
 55    }
 56
 57    def get_current_catalog(self) -> t.Optional[str]:
 58        """Returns the catalog name of the current connection."""
 59        return None
 60
 61    def create_index(
 62        self,
 63        table_name: TableName,
 64        index_name: str,
 65        columns: t.Tuple[str, ...],
 66        exists: bool = True,
 67    ) -> None:
 68        # MySQL doesn't support IF EXISTS clause for indexes.
 69        super().create_index(table_name, index_name, columns, exists=False)
 70
 71    def drop_schema(
 72        self,
 73        schema_name: SchemaName,
 74        ignore_if_not_exists: bool = True,
 75        cascade: bool = False,
 76        **drop_args: t.Dict[str, exp.Expr],
 77    ) -> None:
 78        # MySQL doesn't support CASCADE clause and drops schemas unconditionally.
 79        super().drop_schema(schema_name, ignore_if_not_exists=ignore_if_not_exists, cascade=False)
 80
 81    def _get_data_objects(
 82        self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
 83    ) -> t.List[DataObject]:
 84        """
 85        Returns all the data objects that exist in the given schema and optionally catalog.
 86        """
 87        query = (
 88            exp.select(
 89                exp.column("table_name").as_("name"),
 90                exp.column("table_schema").as_("schema_name"),
 91                exp.case()
 92                .when(
 93                    exp.column("table_type").eq("BASE TABLE"),
 94                    exp.Literal.string("table"),
 95                )
 96                .when(
 97                    exp.column("table_type").eq("VIEW"),
 98                    exp.Literal.string("view"),
 99                )
100                .else_("table_type")
101                .as_("type"),
102            )
103            .from_(exp.table_("tables", db="information_schema"))
104            .where(exp.column("table_schema").eq(to_schema(schema_name).db))
105        )
106        if object_names:
107            query = query.where(exp.column("table_name").isin(*object_names))
108        df = self.fetchdf(query)
109        return [
110            DataObject(
111                schema=row.schema_name,
112                name=row.name,
113                type=DataObjectType.from_str(row.type),  # type: ignore
114            )
115            for row in df.itertuples()
116        ]
117
118    def _build_create_comment_table_exp(
119        self, table: exp.Table, table_comment: str, table_kind: str
120    ) -> exp.Comment | str:
121        table_sql = table.sql(dialect=self.dialect, identify=True)
122
123        truncated_comment = self._truncate_table_comment(table_comment)
124        comment_sql = exp.Literal.string(truncated_comment).sql(dialect=self.dialect)
125
126        return f"ALTER TABLE {table_sql} COMMENT = {comment_sql}"
127
128    def _create_column_comments(
129        self,
130        table_name: TableName,
131        column_comments: t.Dict[str, str],
132        table_kind: str = "TABLE",
133        materialized_view: bool = False,
134    ) -> None:
135        table = exp.to_table(table_name)
136        table_sql = table.sql(dialect=self.dialect, identify=True)
137
138        # MySQL ALTER TABLE MODIFY completely replaces the column (overwriting options and constraints).
139        # self.columns() only returns the column types so doesn't allow us to fully/correctly replace a column definition.
140        # To get the full column definition we retrieve and parse the table's CREATE TABLE statement.
141        create_table_exp = parse_one(
142            self.fetchone(f"SHOW CREATE TABLE {table_sql}")[1],  # type: ignore
143            dialect=self.dialect,
144        )
145        col_def_exps = {
146            col_def.name: col_def.copy()
147            for col_def in create_table_exp.find(exp.Schema).find_all(exp.ColumnDef)  # type: ignore
148        }
149
150        for col in column_comments:
151            col_def = col_def_exps.get(col)
152            if col_def:
153                col_def.args["constraints"].extend(
154                    self._build_col_comment_exp(col_def.alias_or_name, column_comments)
155                )
156
157                try:
158                    self.execute(
159                        f"ALTER TABLE {table_sql} MODIFY {col_def.sql(dialect=self.dialect, identify=True)}",
160                    )
161                except Exception:
162                    logger.warning(
163                        f"Column comments for column '{col_def.alias_or_name}' in table '{table.alias_or_name}' not registered - this may be due to limited permissions.",
164                        exc_info=True,
165                    )
166
167    def _create_table_like(
168        self,
169        target_table_name: TableName,
170        source_table_name: TableName,
171        exists: bool,
172        **kwargs: t.Any,
173    ) -> None:
174        self.execute(
175            exp.Create(
176                this=exp.to_table(target_table_name),
177                kind="TABLE",
178                exists=exists,
179                properties=exp.Properties(
180                    expressions=[
181                        exp.LikeProperty(
182                            this=exp.to_table(source_table_name),
183                        ),
184                    ],
185                ),
186            )
187        )
188
189    def ping(self) -> None:
190        self._connection_pool.get().ping(reconnect=False)