Feat: MySQL support (#1232)

Author: izeigerman

docs/integrations/engines.md

@@ -4,6 +4,11 @@
 ### BigQuery - Local/Built-in Scheduler
 **Engine Adapter Type**: `bigquery`
 
+**Installation**:
+```
+pip install "sqlmesh[bigquery]"
+```
+
 | Option                          | Description                                                                                                    |  Type  | Required |
 |---------------------------------|----------------------------------------------------------------------------------------------------------------|:------:|:--------:|
 | `method`                        | Connection methods. Can be `oath`, `oauth-secrets`, `service-account`, `service-account-json`. Default: `oath` | string |    N     |
@@ -71,6 +76,11 @@ sqlmesh_airflow = SQLMeshAirflow(
 ### Databricks - Local/Built-in Scheduler
 **Engine Adapter Type**: `databricks`
 
+**Installation**:
+```
+pip install "sqlmesh[databricks]"
+```
+
 If you are always running SQLMesh commands directly on a Databricks Cluster (like in a Databricks Notebook) then the only relevant configuration is `catalog` and it is optional.
 The SparkSession provided by Databricks will be used to execute all SQLMesh commands.
 
@@ -170,21 +180,44 @@ sqlmesh_airflow = SQLMeshAirflow(
 ### DuckDB - Airflow
 DuckDB only works when running locally; therefore it does not support Airflow.
 
+## MySQL
+### MySQL - Local/Built-in Scheduler
+**Engine Adapter Type**: `mysql`
+
+**Installation**:
+```
+pip install "sqlmesh[mysql]"
+```
+
+| Option         | Description                                                  | Type   | Required |
+|----------------|--------------------------------------------------------------|:------:|:--------:|
+| `host`         | The hostname of the MysQL server                             | string | Y        |
+| `user`         | The username to use for authentication with the MySQL server | string | Y        |
+| `password`     | The password to use for authentication with the MySQL server | string | Y        |
+| `port`         | The port number of the MySQL server                          | int    | Y        |
+| `charset`      | The character set used for the connection                    | string | N        |
+| `ssl_disabled` | Is SSL disabled                                              | bool   | N        |
+
 ## Postgres
 ### Postgres - Local/Built-in Scheduler
 **Engine Adapter Type**: `postgres`
 
-| Option            | Description                                                                     |  Type  | Required |
+**Installation**:
+```
+pip install "sqlmesh[postgres]"
+```
+
+| Option            | Description                                                                     | Type   | Required |
 |-------------------|---------------------------------------------------------------------------------|:------:|:--------:|
-| `host`            | The hostname of the Postgres server                                             | string |    Y     |
-| `user`            | The username to use for authentication with the Postgres server                 | string |    Y     |
-| `password`        | The password to use for authentication with the Postgres server                 | string |    Y     |
-| `port`            | The port number of the Postgres server                                          |  int   |    Y     |
-| `database`        | The name of the database instance to connect to                                 | string |    Y     |
-| `keepalives_idle` | The number of seconds between each keepalive packet sent to the server.         |  int   |    N     |
-| `connect_timeout` | The number of seconds to wait for the connection to the server. (Default: `10`) |  int   |    N     |
-| `role`            | The role to use for authentication with the Postgres server                     | string |    N     |
-| `sslmode`         | The security of the connection to the Postgres server.                          | string |    N     |
+| `host`            | The hostname of the Postgres server                                             | string | Y        |
+| `user`            | The username to use for authentication with the Postgres server                 | string | Y        |
+| `password`        | The password to use for authentication with the Postgres server                 | string | Y        |
+| `port`            | The port number of the Postgres server                                          | int    | Y        |
+| `database`        | The name of the database instance to connect to                                 | string | Y        |
+| `keepalives_idle` | The number of seconds between each keepalive packet sent to the server.         | int    | N        |
+| `connect_timeout` | The number of seconds to wait for the connection to the server. (Default: `10`) | int    | N        |
+| `role`            | The role to use for authentication with the Postgres server                     | string | N        |
+| `sslmode`         | The security of the connection to the Postgres server                           | string | N        |
 
 ### Postgres - Airflow
 **Engine Name:** `postgres`
@@ -213,6 +246,11 @@ sqlmesh_airflow = SQLMeshAirflow(
 ### Redshift - Local/Built-in Scheduler
 **Engine Adapter Type**: `Redshift`
 
+**Installation**:
+```
+pip install "sqlmesh[redshift]"
+```
+
 | Option                  | Description                                                                                                 |  Type  | Required |
 |-------------------------|-------------------------------------------------------------------------------------------------------------|:------:|:--------:|
 | `user`                  | The username to use for authentication with the Amazon Redshift cluster                                     | string |    N     |
@@ -262,6 +300,11 @@ sqlmesh_airflow = SQLMeshAirflow(
 ### GCP Postgres - Local/Built-in Scheduler
 **Engine Adapter Type**: `postgres`
 
+**Installation**:
+```
+pip install "sqlmesh[gcppostgres]"
+```
+
 | Option                    | Description                                                                         |  Type   | Required |
 |---------------------------|-------------------------------------------------------------------------------------|:-------:|:--------:|
 | `instance_connection_str` | Connection name for the postgres instance                                           | string  |    Y     |
@@ -274,6 +317,11 @@ sqlmesh_airflow = SQLMeshAirflow(
 ### Snowflake - Local/Built-in Scheduler
 **Engine Adapter Type**: `snowflake`
 
+**Installation**:
+```
+pip install "sqlmesh[snowflake]"
+```
+
 | Option          | Description                        |  Type  | Required |
 |-----------------|------------------------------------|:------:|:--------:|
 | `user`          | The Snowflake username             | string |    N     |

setup.cfg

@@ -66,6 +66,9 @@ ignore_missing_imports = True
 [mypy-fsspec]
 ignore_missing_imports = True
 
+[mypy-mysql.*]
+ignore_missing_imports = True
+
 [mypy-psycopg2.*]
 ignore_missing_imports = True
 

setup.py

@@ -46,7 +46,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot~=17.8.0",
+        "sqlglot~=17.8.4",
     ],
     extras_require={
         "bigquery": [
@@ -104,6 +104,9 @@
             "langchain",
             "openai",
         ],
+        "mysql": [
+            "mysql-connector-python",
+        ],
         "postgres": [
             "psycopg2",
         ],

sqlmesh/core/audit/builtin.py

@@ -29,7 +29,7 @@
   SELECT
     @EACH(
       @columns,
-      c -> row_number() OVER (PARTITION BY c ORDER BY 1) AS rank_@c
+      c -> row_number() OVER (PARTITION BY c ORDER BY c) AS rank_@c
     )
   FROM @this_model
 )

sqlmesh/core/config/__init__.py

@@ -5,6 +5,7 @@
     DatabricksConnectionConfig,
     DuckDBConnectionConfig,
     GCPPostgresConnectionConfig,
+    MySQLConnectionConfig,
     PostgresConnectionConfig,
     RedshiftConnectionConfig,
     SnowflakeConnectionConfig,

sqlmesh/core/config/connection.py

@@ -51,6 +51,11 @@ def _extra_engine_config(self) -> t.Dict[str, t.Any]:
         """kwargs that are for execution config only"""
         return {}
 
+    @property
+    def _cursor_kwargs(self) -> t.Optional[t.Dict[str, t.Any]]:
+        """Key-value arguments that will be passed during cursor construction."""
+        return None
+
     def create_engine_adapter(self) -> EngineAdapter:
         """Returns a new instance of the Engine Adapter."""
         return self._engine_adapter(
@@ -61,6 +66,7 @@ def create_engine_adapter(self) -> EngineAdapter:
                 }
             ),
             multithreaded=self.concurrent_tasks > 1,
+            cursor_kwargs=self._cursor_kwargs,
             **self._extra_engine_config,
         )
 
@@ -621,6 +627,48 @@ def _connection_factory(self) -> t.Callable:
         return connect
 
 
+class MySQLConnectionConfig(_ConnectionConfig):
+    host: str
+    user: str
+    password: str
+    port: t.Optional[int] = None
+    charset: t.Optional[str] = None
+    ssl_disabled: t.Optional[bool] = None
+
+    concurrent_tasks: int = 4
+
+    type_: Literal["mysql"] = Field(alias="type", default="mysql")
+
+    _cursor_kwargs = {"buffered": True}
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        connection_keys = {
+            "host",
+            "user",
+            "password",
+            "port",
+            "database",
+        }
+        if self.port is not None:
+            connection_keys.add("port")
+        if self.charset is not None:
+            connection_keys.add("charset")
+        if self.ssl_disabled is not None:
+            connection_keys.add("ssl_disabled")
+        return connection_keys
+
+    @property
+    def _engine_adapter(self) -> t.Type[EngineAdapter]:
+        return engine_adapter.MySQLEngineAdapter
+
+    @property
+    def _connection_factory(self) -> t.Callable:
+        from mysql.connector import connect
+
+        return connect
+
+
 class SparkConnectionConfig(_ConnectionConfig):
     """
     Vanilla Spark Connection Configuration. Use `DatabricksConnectionConfig` for Databricks.
@@ -675,6 +723,7 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
         GCPPostgresConnectionConfig,
         DatabricksConnectionConfig,
         DuckDBConnectionConfig,
+        MySQLConnectionConfig,
         PostgresConnectionConfig,
         RedshiftConnectionConfig,
         SnowflakeConnectionConfig,

sqlmesh/core/engine_adapter/__init__.py

@@ -9,6 +9,7 @@
 from sqlmesh.core.engine_adapter.bigquery import BigQueryEngineAdapter
 from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter
 from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
+from sqlmesh.core.engine_adapter.mysql import MySQLEngineAdapter
 from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
 from sqlmesh.core.engine_adapter.redshift import RedshiftEngineAdapter
 from sqlmesh.core.engine_adapter.shared import TransactionType
@@ -23,7 +24,7 @@
     "databricks": DatabricksEngineAdapter,
     "redshift": RedshiftEngineAdapter,
     "postgres": PostgresEngineAdapter,
-    "mysql": EngineAdapterWithIndexSupport,
+    "mysql": MySQLEngineAdapter,
     "mssql": EngineAdapterWithIndexSupport,
 }
 

sqlmesh/core/engine_adapter/base.py

@@ -95,10 +95,13 @@ def __init__(
         dialect: str = "",
         sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None,
         multithreaded: bool = False,
+        cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
-        self._connection_pool = create_connection_pool(connection_factory, multithreaded)
+        self._connection_pool = create_connection_pool(
+            connection_factory, multithreaded, cursor_kwargs=cursor_kwargs
+        )
         self.sql_gen_kwargs = sql_gen_kwargs or {}
         self._extra_config = kwargs
 
@@ -489,7 +492,8 @@ def columns(
         self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE"))
         describe_output = self.cursor.fetchall()
         return {
-            column_name: exp.DataType.build(column_type, dialect=self.dialect)
+            # Note: MySQL  returns the column type as bytes.
+            column_name: exp.DataType.build(_decoded_str(column_type), dialect=self.dialect)
             for column_name, column_type, *_ in itertools.takewhile(
                 lambda t: not t[0].startswith("#"),
                 describe_output,
@@ -1005,3 +1009,9 @@ def _add_where_to_query(self, query: Query, where: t.Optional[exp.Expression]) -
 
 class EngineAdapterWithIndexSupport(EngineAdapter):
     SUPPORTS_INDEXES = True
+
+
+def _decoded_str(value: t.Union[str, bytes]) -> str:
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    return value

sqlmesh/core/engine_adapter/base_postgres.py

@@ -4,7 +4,7 @@
 
 from sqlglot import exp
 
-from sqlmesh.core.engine_adapter.base import EngineAdapter
+from sqlmesh.core.engine_adapter.mixins import CommitOnExecuteMixin
 from sqlmesh.core.engine_adapter.shared import (
     DataObject,
     DataObjectType,
@@ -17,7 +17,7 @@
     from sqlmesh.core.engine_adapter._typing import QueryOrDF
 
 
-class BasePostgresEngineAdapter(EngineAdapter):
+class BasePostgresEngineAdapter(CommitOnExecuteMixin):
     COLUMNS_TABLE = "information_schema.columns"
     SUPPORTS_MATERIALIZED_VIEWS = True
 
@@ -96,28 +96,6 @@ def create_view(
                 **create_kwargs,
             )
 
-    def execute(
-        self,
-        expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
-        ignore_unsupported_errors: bool = False,
-        quote_identifiers: bool = True,
-        **kwargs: t.Any,
-    ) -> None:
-        """
-        To make sure that inserts and updates take effect we need to commit explicitly unless the
-        statement is executed as part of an active transaction.
-
-        Reference: https://www.psycopg.org/psycopg3/docs/basic/transactions.html
-        """
-        super().execute(
-            expressions,
-            ignore_unsupported_errors=ignore_unsupported_errors,
-            quote_identifiers=quote_identifiers,
-            **kwargs,
-        )
-        if not self._connection_pool.is_transaction_active:
-            self._connection_pool.commit()
-
     def _get_data_objects(
         self, schema_name: str, catalog_name: t.Optional[str] = None
     ) -> t.List[DataObject]:

sqlmesh/core/engine_adapter/databricks.py

@@ -33,13 +33,15 @@ def __init__(
         dialect: str = "",
         sql_gen_kwargs: t.Optional[t.Dict[str, Dialect | bool | str]] = None,
         multithreaded: bool = False,
+        cursor_kwargs: t.Optional[t.Dict[str, t.Any]] = None,
         **kwargs: t.Any,
     ):
         super().__init__(
             connection_factory,
             dialect,
             sql_gen_kwargs,
             multithreaded,
+            cursor_kwargs,
             **kwargs,
         )
         self._spark: t.Optional[PySparkSession] = None