Feat!: Use pre / post SQL statements instead of hooks (#910)

* feat: improve pre and post hook impl

* Drop support for hooks. Rely on pre/post SQL statements instead

---------

Co-authored-by: z3z1ma <butler.alex2010@gmail.com>

Author: izeigerman

docs/concepts/hooks.md

@@ -132,6 +132,22 @@ SELECT UDF(x)::int AS x
 FROM y
 ```
 
+Additional statements can also be provided **after** the main query, in which case they will run after each evaluation of the SELECT query.
+
+```sql linenums="1" hl_lines="5-7"
+MODEL (
+...
+);
+
+...
+
+SELECT UDF(x)::int AS x
+FROM y;
+
+-- Cleanup statements
+DROP TABLE temp_table;
+```
+
 ## Time column
 Models that are loaded incrementally require a time column to partition data.
 

docs/concepts/models/overview.md

@@ -3,8 +3,6 @@ MODEL (
   kind FULL,
   owner jen,
   cron '@daily',
-  pre NOOP(x=1),
-  post (noop(), noop(y=['a', 2])),
   tags (pii, fact)
 );
 

examples/sushi/hooks/hooks.py

@@ -16,7 +16,6 @@
 from sqlmesh.core.config import Config
 from sqlmesh.core.context import Context, ExecutionContext
 from sqlmesh.core.engine_adapter import EngineAdapter
-from sqlmesh.core.hooks import hook
 from sqlmesh.core.macros import macro
 from sqlmesh.core.model import Model, model
 from sqlmesh.core.snapshot import Snapshot

examples/sushi/models/customers.sql

@@ -32,7 +32,6 @@
 """The default maximum row limit that is used when evaluating a model."""
 
 AUDITS = "audits"
-HOOKS = "hooks"
 MACROS = "macros"
 MODELS = "models"
 TESTS = "tests"

sqlmesh/__init__.py

@@ -52,7 +52,6 @@
 from sqlmesh.core.dialect import format_model_expressions, pandas_to_sql, parse
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import Environment
-from sqlmesh.core.hooks import hook
 from sqlmesh.core.loader import Loader, SqlMeshLoader, update_model_schemas
 from sqlmesh.core.macros import ExecutableOrMacro
 from sqlmesh.core.model import Model
@@ -216,7 +215,6 @@ def __init__(
         self._audits: UniqueKeyDict[str, Audit] = UniqueKeyDict("audits")
         self._macros: UniqueKeyDict[str, ExecutableOrMacro] = UniqueKeyDict("macros")
         self._jinja_macros = JinjaMacroRegistry()
-        self._hooks: UniqueKeyDict[str, hook] = UniqueKeyDict("hooks")
 
         self.path, self.config = t.cast(t.Tuple[Path, Config], next(iter(self.configs.items())))
         self.gateway = gateway
@@ -342,7 +340,6 @@ def load(self, update_schemas: bool = True) -> Context:
         """Load all files in the context's path."""
         with sys_path(*self.configs):
             project = self._loader.load(self, update_schemas)
-            self._hooks = project.hooks
             self._macros = project.macros
             self._models = project.models
             self._audits = project.audits

sqlmesh/core/constants.py

@@ -17,6 +17,7 @@
 import pandas as pd
 from sqlglot import Dialect, exp
 from sqlglot.errors import ErrorLevel
+from sqlglot.helper import ensure_list
 
 from sqlmesh.core.dialect import pandas_to_sql
 from sqlmesh.core.engine_adapter.shared import DataObject, TransactionType
@@ -766,17 +767,19 @@ def supports_transactions(self, transaction_type: TransactionType) -> bool:
 
     def execute(
         self,
-        sql: t.Union[str, exp.Expression],
+        expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
         ignore_unsupported_errors: bool = False,
         **kwargs: t.Any,
     ) -> None:
         """Execute a sql query."""
         to_sql_kwargs = (
             {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {}
         )
-        sql = self._to_sql(sql, **to_sql_kwargs) if isinstance(sql, exp.Expression) else sql
-        logger.debug(f"Executing SQL:\n{sql}")
-        self.cursor.execute(sql, **kwargs)
+
+        for e in ensure_list(expressions):
+            sql = self._to_sql(e, **to_sql_kwargs) if isinstance(e, exp.Expression) else e
+            logger.debug(f"Executing SQL:\n{sql}")
+            self.cursor.execute(sql, **kwargs)
 
     def _create_table_properties(
         self,

sqlmesh/core/context.py

@@ -88,7 +88,7 @@ def create_view(
 
     def execute(
         self,
-        sql: t.Union[str, exp.Expression],
+        expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
         ignore_unsupported_errors: bool = False,
         **kwargs: t.Any,
     ) -> None:
@@ -98,7 +98,7 @@ def execute(
 
         Reference: https://www.psycopg.org/psycopg3/docs/basic/transactions.html
         """
-        super().execute(sql, ignore_unsupported_errors=ignore_unsupported_errors, **kwargs)
+        super().execute(expressions, ignore_unsupported_errors=ignore_unsupported_errors, **kwargs)
         if not self._connection_pool.is_transaction_active:
             self._connection_pool.commit()
 

sqlmesh/core/engine_adapter/base.py

@@ -7,6 +7,7 @@
 import pandas as pd
 from sqlglot import exp
 from sqlglot.errors import ErrorLevel
+from sqlglot.helper import ensure_list
 from sqlglot.transforms import remove_precision_parameterized_types
 
 from sqlmesh.core.engine_adapter.base import EngineAdapter
@@ -383,7 +384,7 @@ def _retryable_execute(
 
     def execute(
         self,
-        sql: t.Union[str, exp.Expression],
+        expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
         ignore_unsupported_errors: bool = False,
         **kwargs: t.Any,
     ) -> None:
@@ -393,14 +394,16 @@ def execute(
         to_sql_kwargs = (
             {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {}
         )
-        sql = self._to_sql(sql, **to_sql_kwargs) if isinstance(sql, exp.Expression) else sql
-        logger.debug(f"Executing SQL:\n{sql}")
-        retry.retry_target(
-            target=functools.partial(self._retryable_execute, sql=sql),
-            predicate=_ErrorCounter(self._extra_config["job_retries"]).should_retry,
-            sleep_generator=retry.exponential_sleep_generator(initial=1.0, maximum=3.0),
-            deadline=self._extra_config.get("job_retry_deadline_seconds"),
-        )
+
+        for e in ensure_list(expressions):
+            sql = self._to_sql(e, **to_sql_kwargs) if isinstance(e, exp.Expression) else e
+            logger.debug(f"Executing SQL:\n{sql}")
+            retry.retry_target(
+                target=functools.partial(self._retryable_execute, sql=sql),
+                predicate=_ErrorCounter(self._extra_config["job_retries"]).should_retry,
+                sleep_generator=retry.exponential_sleep_generator(initial=1.0, maximum=3.0),
+                deadline=self._extra_config.get("job_retry_deadline_seconds"),
+            )
 
     def _get_data_objects(
         self, schema_name: str, catalog_name: t.Optional[str] = None