Repository: TobikoData/sqlmesh
Pull request #2632 (2024-05-17 22:47:37)

Feat: Introduce support for 'pre ping' to detect stale or lost connections
For long running plans, a state connection established at the beginning of a plan application might go stale by the end of it.

This update introduces a "pre-ping" mechanism, ensuring that the connection is not stale before initiating a new transaction. The ping implementation itself is dialect-specific but essentially comes down to executing a `SELECT 1` query in all cases except MySQL, which has a first-class API for this purpose.

The new behavior is configurable and is turned off by default for all engines except for those likely to be used to store the SQLMesh state: MySQL, MSSQL, and Postgres. This is also only supported for engines with transaction support.


sungchun12: Really glad to see this PR :) 

Patch:
diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md
index c3f53f6c..6858446d 100644
--- a/docs/concepts/macros/sqlmesh_macros.md
+++ b/docs/concepts/macros/sqlmesh_macros.md
@@ -600,12 +600,14 @@ If the column data types are known, the resulting query `CAST`s columns to their
 
 - `relation`: The relation/table whose columns are being selected
 - `alias` (optional): The alias of the relation (if it has one)
-- `except_` (optional): A list of columns to exclude
+- `exclude` (optional): A list of columns to exclude
 - `prefix` (optional): A string to use as a prefix for all selected column names
 - `suffix` (optional): A string to use as a suffix for all selected column names
 - `quote_identifiers` (optional): Whether to quote the resulting identifiers, defaults to true
 
-Like all SQLMesh macro functions, omitting an argument when calling `@STAR` requires passing all subsequent arguments with their name and the special `:=` keyword operator. For example, we might omit the `alias` argument with `@STAR(foo, except_ := [c])`. Learn more about macro function arguments [below](#positional-and-keyword-arguments).
+**NOTE**: the `exclude` argument used to be named `except_`. The latter is still supported but we discourage its use because it will be deprecated in the future.
+
+Like all SQLMesh macro functions, omitting an argument when calling `@STAR` requires passing all subsequent arguments with their name and the special `:=` keyword operator. For example, we might omit the `alias` argument with `@STAR(foo, exclude := [c])`. Learn more about macro function arguments [below](#positional-and-keyword-arguments).
 
 As a `@STAR` example, consider the following query:
 
@@ -635,15 +637,15 @@ FROM foo AS bar
 Note these aspects of the rendered query:
 - Each column is `CAST` to its data type in the table `foo` (e.g., `a` to `TEXT`)
 - Each column selection uses the alias `bar` (e.g., `"bar"."a"`)
-- Column `c` is not present because it was passed to `@STAR`'s `except_` argument
+- Column `c` is not present because it was passed to `@STAR`'s `exclude` argument
 - Each column alias is prefixed with `baz_` and suffixed with `_qux` (e.g., `"baz_a_qux"`)
 
 Now consider a more complex example that provides different prefixes to `a` and `b` than to `d` and includes an explicit column `my_column`:
 
 ```sql linenums="1"
 SELECT
-  @STAR(foo, bar, except_=[c, d], 'ab_pre_'),
-  @STAR(foo, bar, except_=[a, b, c], 'd_pre_'),
+  @STAR(foo, bar, exclude := [c, d], 'ab_pre_'),
+  @STAR(foo, bar, exclude := [a, b, c], 'd_pre_'),
   my_column
 FROM foo AS bar
 ```
@@ -661,7 +663,7 @@ FROM foo AS bar
 
 Note these aspects of the rendered query:
 - Columns `a` and `b` have the prefix `"ab_pre_"` , while column `d` has the prefix `"d_pre_"`
-- Column `c` is not present because it was passed to the `except_` argument in both `@STAR` calls
+- Column `c` is not present because it was passed to the `exclude` argument in both `@STAR` calls
 - `my_column` is present in the query
 
 ### @GENERATE_SURROGATE_KEY
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index df647c23..3e69fbb2 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -159,11 +159,12 @@ Most parameters are specific to the connection engine `type` - see [below](#engi
 
 #### General
 
-| Option              | Description                                                                                                                 | Type | Required |
-| ------------------- | --------------------------------------------------------------------------------------------------------------------------- | :--: | :------: |
-| `type`              | The engine type name, listed in engine-specific configuration pages below.                                                  | str  |    Y     |
-| `concurrent_tasks`  | The maximum number of concurrent tasks that will be run by SQLMesh. (Default: 4 for engines that support concurrent tasks.) | int  |    N     |
-| `register_comments` | Whether SQLMesh should register model comments with the SQL engine (if the engine supports it). (Default: `true`.)          | bool |    N     |
+| Option              | Description                                                                                                                                                             | Type | Required |
+|---------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:----:|:--------:|
+| `type`              | The engine type name, listed in engine-specific configuration pages below.                                                                                              | str  | Y        |
+| `concurrent_tasks`  | The maximum number of concurrent tasks that will be run by SQLMesh. (Default: 4 for engines that support concurrent tasks.)                                             | int  | N        |
+| `register_comments` | Whether SQLMesh should register model comments with the SQL engine (if the engine supports it). (Default: `true`.)                                                      | bool | N        |
+| `pre_ping`          | Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive. This can only be enabled for engines with transaction support. | bool | N        |
 
 #### Engine-specific
 
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 8bca2a43..ded2e600 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -44,6 +44,7 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     type_: str
     concurrent_tasks: int
     register_comments: bool
+    pre_ping: bool
 
     @property
     @abc.abstractmethod
@@ -105,6 +106,7 @@ class ConnectionConfig(abc.ABC, BaseConfig):
             default_catalog=self.get_catalog(),
             cursor_init=self._cursor_init,
             register_comments=register_comments_override or self.register_comments,
+            pre_ping=self.pre_ping,
             **self._extra_engine_config,
         )
 
@@ -127,6 +129,7 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
         connector_config: A dictionary of configuration to pass into the duckdb connector.
         concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
         register_comments: Whether or not to register model comments with the SQL engine.
+        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
     """
 
     extensions: t.List[str] = []
@@ -134,6 +137,7 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: Literal[1] = 1
     register_comments: bool = True
+    pre_ping: Literal[False] = False
 
     @property
     def _engine_adapter(self) -> t.Type[EngineAdapter]:
@@ -319,6 +323,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
         private_key_path: The optional path to the private key to use for authentication. This would be used instead of `private_key`.
         private_key_passphrase: The optional passphrase to use to decrypt `private_key` or `private_key_path`. Keys can be created without encryption so only provide this if needed.
         register_comments: Whether or not to register model comments with the SQL engine.
+        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
     """
 
     account: str
@@ -337,6 +342,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: bool = False
 
     type_: Literal["snowflake"] = Field(alias="type", default="snowflake")
 
@@ -496,6 +502,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
             Defaults to deriving the cluster id from the `http_path` value.
         force_databricks_connect: Force all queries to run using Databricks Connect instead of the SQL connector.
         disable_databricks_connect: Even if databricks connect is installed, do not use it.
+        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
     """
 
     server_hostname: t.Optional[str] = None
@@ -512,6 +519,7 @@ class DatabricksConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 1
     register_comments: bool = True
+    pre_ping: Literal[False] = False
 
     type_: Literal["databricks"] = Field(alias="type", default="databricks")
 
@@ -680,6 +688,7 @@ class BigQueryConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 1
     register_comments: bool = True
+    pre_ping: Literal[False] = False
 
     type_: Literal["bigquery"] = Field(alias="type", default="bigquery")
 
@@ -766,6 +775,7 @@ class GCPPostgresConnectionConfig(ConnectionConfig):
         password: The postgres user's password. Only needed when the user is a postgres user.
         enable_iam_auth: Set to True when user is an IAM user.
         db: Name of the db to connect to.
+        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
     """
 
     instance_connection_string: str
@@ -779,6 +789,7 @@ class GCPPostgresConnectionConfig(ConnectionConfig):
     type_: Literal["gcp_postgres"] = Field(alias="type", default="gcp_postgres")
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: bool = True
 
     @model_validator(mode="before")
     @model_validator_v1_args
@@ -853,6 +864,7 @@ class RedshiftConnectionConfig(ConnectionConfig):
         is_serverless: Redshift end-point is serverless or provisional. Default value false.
         serverless_acct_id: The account ID of the serverless. Default value None
         serverless_work_group: The name of work group for serverless end point. Default value None.
+        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
     """
 
     user: t.Optional[str] = None
@@ -879,6 +891,7 @@ class RedshiftConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: bool = False
 
     type_: Literal["redshift"] = Field(alias="type", default="redshift")
 
@@ -932,6 +945,7 @@ class PostgresConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: bool = True
 
     type_: Literal["postgres"] = Field(alias="type", default="postgres")
 
@@ -970,6 +984,7 @@ class MySQLConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: bool = True
 
     type_: Literal["mysql"] = Field(alias="type", default="mysql")
 
@@ -1022,6 +1037,7 @@ class MSSQLConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: bool = True
 
     type_: Literal["mssql"] = Field(alias="type", default="mssql")
 
@@ -1064,6 +1080,7 @@ class SparkConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: Literal[False] = False
 
     type_: Literal["spark"] = Field(alias="type", default="spark")
 
@@ -1176,6 +1193,7 @@ class TrinoConnectionConfig(ConnectionConfig):
 
     concurrent_tasks: int = 4
     register_comments: bool = True
+    pre_ping: Literal[False] = False
 
     type_: Literal["trino"] = Field(alias="type", default="trino")
 
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 08be3a45..55df8968 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -109,26 +109,28 @@ class EngineAdapter:
         default_catalog: t.Optional[str] = None,
         execute_log_level: int = logging.DEBUG,
         register_comments: bool = True,
+        pre_ping: bool = False,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
         self._connection_pool = create_connection_pool(
             connection_factory, multithreaded, cursor_kwargs=cursor_kwargs, cursor_init=cursor_init
         )
-        self.sql_gen_kwargs = sql_gen_kwargs or {}
+        self._sql_gen_kwargs = sql_gen_kwargs or {}
         self._default_catalog = default_catalog
         self._execute_log_level = execute_log_level
         self._extra_config = kwargs
-        self.register_comments = register_comments
+        self._register_comments = register_comments
+        self._pre_ping = pre_ping
 
     def with_log_level(self, level: int) -> EngineAdapter:
         adapter = self.__class__(
             lambda: None,
             dialect=self.dialect,
-            sql_gen_kwargs=self.sql_gen_kwargs,
+            sql_gen_kwargs=self._sql_gen_kwargs,
             default_catalog=self._default_catalog,
             execute_log_level=level,
-            register_comments=self.register_comments,
+            register_comments=self._register_comments,
             **self._extra_config,
         )
 
@@ -146,7 +148,7 @@ class EngineAdapter:
 
     @property
     def comments_enabled(self) -> bool:
-        return self.register_comments and self.COMMENT_CREATION_TABLE.is_supported
+        return self._register_comments and self.COMMENT_CREATION_TABLE.is_supported
 
     @classmethod
     def is_pandas_df(cls, value: t.Any) -> bool:
@@ -1741,6 +1743,15 @@ class EngineAdapter:
         ):
             yield
             return
+
+        if self._pre_ping:
+            try:
+                logger.debug("Pinging the database to check the connection")
+                self._ping()
+            except Exception:
+                logger.info("Connection to the database was lost. Reconnecting...")
+                self._connection_pool.close()
+
         self._connection_pool.begin()
         try:
             yield
@@ -1921,7 +1932,7 @@ class EngineAdapter:
             "dialect": self.dialect,
             "pretty": False,
             "comments": False,
-            **self.sql_gen_kwargs,
+            **self._sql_gen_kwargs,
             **kwargs,
         }
 
@@ -2038,6 +2049,12 @@ class EngineAdapter:
     ) -> None:
         self.execute(exp.rename_table(old_table_name, new_table_name))
 
+    def _ping(self) -> None:
+        try:
+            self._execute(exp.select("1").sql(dialect=self.dialect))
+        finally:
+            self._connection_pool.close_cursor()
+
     @classmethod
     def _select_columns(cls, columns: t.Iterable[str]) -> exp.Select:
         return exp.select(*(exp.column(c, quoted=True) for c in columns))
diff --git a/sqlmesh/core/engine_adapter/mysql.py b/sqlmesh/core/engine_adapter/mysql.py
index 3642cfc0..628fd1f1 100644
--- a/sqlmesh/core/engine_adapter/mysql.py
+++ b/sqlmesh/core/engine_adapter/mysql.py
@@ -144,3 +144,6 @@ class MySQLEngineAdapter(
                 f"Column comments for table '{table.alias_or_name}' not registered - this may be due to limited permissions.",
                 exc_info=True,
             )
+
+    def _ping(self) -> None:
+        self._connection_pool.get().ping(reconnect=False)
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 7432613b..430a599d 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -791,10 +791,11 @@ def star(
     evaluator: MacroEvaluator,
     relation: exp.Table,
     alias: exp.Column = t.cast(exp.Column, exp.column("")),
-    except_: t.Union[exp.Array, exp.Tuple] = exp.Tuple(this=[]),
+    exclude: t.Union[exp.Array, exp.Tuple] = exp.Tuple(this=[]),
     prefix: exp.Literal = exp.Literal.string(""),
     suffix: exp.Literal = exp.Literal.string(""),
     quote_identifiers: exp.Boolean = exp.true(),
+    except_: t.Union[exp.Array, exp.Tuple] = exp.Tuple(this=[]),
 ) -> t.List[exp.Alias]:
     """Returns a list of projections for the given relation.
 
@@ -802,10 +803,11 @@ def star(
         evaluator: MacroEvaluator that invoked the macro
         relation: The relation to select star from
         alias: The alias of the relation
-        except_: Columns to exclude
+        exclude: Columns to exclude
         prefix: A prefix to use for all selections
         suffix: A suffix to use for all selections
         quote_identifiers: Whether or not quote the resulting aliases, defaults to true
+        except_: Alias for exclude (TODO: deprecate this, update docs)
 
     Returns:
         An array of columns.
@@ -814,14 +816,20 @@ def star(
         >>> from sqlglot import parse_one, exp
         >>> from sqlglot.schema import MappingSchema
         >>> from sqlmesh.core.macros import MacroEvaluator
-        >>> sql = "SELECT @STAR(foo, bar, except_ := [c], prefix := 'baz_') FROM foo AS bar"
+        >>> sql = "SELECT @STAR(foo, bar, exclude := [c], prefix := 'baz_') FROM foo AS bar"
         >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": exp.DataType.build("string"), "b": exp.DataType.build("string"), "c": exp.DataType.build("string"), "d": exp.DataType.build("int")}})).transform(parse_one(sql)).sql()
         'SELECT CAST("bar"."a" AS TEXT) AS "baz_a", CAST("bar"."b" AS TEXT) AS "baz_b", CAST("bar"."d" AS INT) AS "baz_d" FROM foo AS bar'
     """
     if alias and not isinstance(alias, (exp.Identifier, exp.Column)):
         raise SQLMeshError(f"Invalid alias '{alias}'. Expected an identifier.")
-    if except_ and not isinstance(except_, (exp.Array, exp.Tuple)):
-        raise SQLMeshError(f"Invalid except '{except_}'. Expected an array.")
+    if exclude and not isinstance(exclude, (exp.Array, exp.Tuple)):
+        raise SQLMeshError(f"Invalid exclude '{exclude}'. Expected an array.")
+    if except_:
+        logger.warning(
+            "The 'except_' argument in @STAR will soon be deprecated. Use 'exclude' instead."
+        )
+        if not isinstance(exclude, (exp.Array, exp.Tuple)):
+            raise SQLMeshError(f"Invalid exclude_ '{exclude}'. Expected an array.")
     if prefix and not isinstance(prefix, exp.Literal):
         raise SQLMeshError(f"Invalid prefix '{prefix}'. Expected a literal.")
     if suffix and not isinstance(suffix, exp.Literal):
@@ -829,12 +837,12 @@ def star(
     if not isinstance(quote_identifiers, exp.Boolean):
         raise SQLMeshError(f"Invalid quote_identifiers '{quote_identifiers}'. Expected a boolean.")
 
-    exclude = {e.name for e in except_.expressions}
+    excluded_names = {e.name for e in exclude.expressions or except_.expressions}
     quoted = quote_identifiers.this
     table_identifier = alias.name or relation.name
 
     columns_to_types = {
-        k: v for k, v in evaluator.columns_to_types(relation).items() if k not in exclude
+        k: v for k, v in evaluator.columns_to_types(relation).items() if k not in excluded_names
     }
     if columns_to_types_all_known(columns_to_types):
         return [


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index 027dcff9..577cd19b 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -2661,3 +2661,19 @@ def test_get_data_objects_batching(mocker: MockerFixture, make_mocked_engine_ada
 
     assert len(calls[0][0][1]) == adapter.DATA_OBJECT_FILTER_BATCH_SIZE
     assert len(calls[1][0][1]) == 1
+
+
+def test_pre_ping(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+    adapter._pre_ping = True
+    adapter.cursor.execute.side_effect = RuntimeError("error")
+
+    with pytest.raises(RuntimeError):
+        adapter.execute("SELECT 'test'")
+
+    assert to_sql_calls(adapter) == [
+        "SELECT 1",  # ping
+        "SELECT 'test'",
+    ]
+
+    adapter._connection_pool.get().close.assert_called_once()
diff --git a/tests/core/engine_adapter/test_mysql.py b/tests/core/engine_adapter/test_mysql.py
index 46d23bbc..1feccf79 100644
--- a/tests/core/engine_adapter/test_mysql.py
+++ b/tests/core/engine_adapter/test_mysql.py
@@ -62,3 +62,16 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
         f"ALTER TABLE `test_table` COMMENT = '{truncated_table_comment}'",
         f"ALTER TABLE `test_table` MODIFY `a` INT COMMENT '{truncated_column_comment}'",
     ]
+
+
+def test_pre_ping(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(MySQLEngineAdapter)
+    adapter._pre_ping = True
+
+    adapter.execute("SELECT 'test'")
+
+    assert to_sql_calls(adapter) == [
+        "SELECT 'test'",
+    ]
+
+    adapter._connection_pool.get().ping.assert_called_once_with(reconnect=False)
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 9778ad01..d74efbd4 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -475,6 +475,7 @@ def test_connection_config_serialization():
         "register_comments": True,
         "type": "duckdb",
         "extensions": [],
+        "pre_ping": False,
         "connector_config": {},
         "database": "my_db",
     }
@@ -483,6 +484,7 @@ def test_connection_config_serialization():
         "register_comments": True,
         "type": "duckdb",
         "extensions": [],
+        "pre_ping": False,
         "connector_config": {},
         "database": "my_test_db",
     }