Repository: TobikoData/sqlmesh (python)

Issue:
Pretty SQL before sending to execution engine
A small request I hope.
Would it be possible to use SQLGlot's pretty functionality before executing SQL statements?
It would be much easier to find syntax issues if it wasn't one long line.

Pull request description:
feat: add an option to pretty sql
Fixes #2935

Patch:
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index d3b1303c..aa5c4819 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -164,10 +164,11 @@ Most parameters are specific to the connection engine `type` - see [below](#engi
 
 | 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        |
+| `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     |
+| `pretty_sql`        | If SQL should be formatted before being executed, not recommended in a production setting. (Default: `false`.)                                                          | bool |    N     |
 
 #### Engine-specific
 
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index b6bc9d87..33cd04e1 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -48,6 +48,7 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     concurrent_tasks: int
     register_comments: bool
     pre_ping: bool
+    pretty_sql: bool = False
 
     @property
     @abc.abstractmethod
@@ -119,6 +120,7 @@ class ConnectionConfig(abc.ABC, BaseConfig):
             cursor_init=self._cursor_init,
             register_comments=register_comments_override or self.register_comments,
             pre_ping=self.pre_ping,
+            pretty_sql=self.pretty_sql,
             **self._extra_engine_config,
         )
 
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 472dee00..5882e947 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1747,7 +1747,6 @@ class GenericContext(BaseContext, t.Generic[C]):
                 start=start,
                 end=end,
                 snapshots=self.snapshots,
-                raise_exception=False,
             ):
                 audit_id = f"{audit_result.audit.name}"
                 if audit_result.model:
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 3e75db0a..27de53a0 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -115,6 +115,7 @@ class EngineAdapter:
         execute_log_level: int = logging.DEBUG,
         register_comments: bool = True,
         pre_ping: bool = False,
+        pretty_sql: bool = False,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
@@ -127,6 +128,7 @@ class EngineAdapter:
         self._extra_config = kwargs
         self._register_comments = register_comments
         self._pre_ping = pre_ping
+        self._pretty_sql = pretty_sql
 
     def with_log_level(self, level: int) -> EngineAdapter:
         adapter = self.__class__(
@@ -2044,7 +2046,6 @@ class EngineAdapter:
         to_sql_kwargs = (
             {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {}
         )
-
         with self.transaction():
             for e in ensure_list(expressions):
                 sql = t.cast(
@@ -2209,7 +2210,7 @@ class EngineAdapter:
         """
         sql_gen_kwargs = {
             "dialect": self.dialect,
-            "pretty": False,
+            "pretty": self._pretty_sql,
             "comments": False,
             **self._sql_gen_kwargs,
             **kwargs,
diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py
index 2c80bb2c..92fb9e3b 100644
--- a/sqlmesh/core/scheduler.py
+++ b/sqlmesh/core/scheduler.py
@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import abc
 import logging
-import traceback
 import typing as t
 
 from sqlglot import exp
@@ -245,7 +244,6 @@ class Scheduler:
             start=start,
             end=end,
             execution_time=execution_time,
-            raise_exception=False,
             snapshots=snapshots,
             deployability_index=deployability_index,
             wap_id=wap_id,
@@ -268,9 +266,10 @@ class Scheduler:
                 )
             if audit_result.blocking:
                 audit_error_to_raise = error
+            else:
+                logger.warning(f"{error}\nAudit is non-blocking so proceeding with execution.")
 
         if audit_error_to_raise:
-            logger.error(f"Audit Failure: {traceback.format_exc()}")
             raise audit_error_to_raise
 
         self.state_sync.add_interval(snapshot, start, end, is_dev=not is_deployable)
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index e3eb504f..02c79855 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -66,7 +66,7 @@ from sqlmesh.utils.concurrency import (
     concurrent_apply_to_values,
 )
 from sqlmesh.utils.date import TimeLike, now
-from sqlmesh.utils.errors import AuditError, ConfigError, SQLMeshError
+from sqlmesh.utils.errors import ConfigError, SQLMeshError
 
 if sys.version_info >= (3, 12):
     from importlib import metadata
@@ -430,7 +430,6 @@ class SnapshotEvaluator:
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
         execution_time: t.Optional[TimeLike] = None,
-        raise_exception: bool = True,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         wap_id: t.Optional[str] = None,
         **kwargs: t.Any,
@@ -443,8 +442,6 @@ class SnapshotEvaluator:
             start: The start datetime to audit. Defaults to epoch start.
             end: The end datetime to audit. Defaults to epoch start.
             execution_time: The date/time time reference to use for execution time.
-            raise_exception: Whether to raise an exception if the audit fails. Blocking rules determine if an
-                AuditError is thrown or if we just warn with logger
             deployability_index: Determines snapshots that are deployable in the context of this evaluation.
             wap_id: The WAP ID if applicable, None otherwise.
             kwargs: Additional kwargs to pass to the renderer.
@@ -494,7 +491,6 @@ class SnapshotEvaluator:
                     start=start,
                     end=end,
                     execution_time=execution_time,
-                    raise_exception=raise_exception,
                     deployability_index=deployability_index,
                     **kwargs,
                 )
@@ -914,7 +910,6 @@ class SnapshotEvaluator:
         start: t.Optional[TimeLike],
         end: t.Optional[TimeLike],
         execution_time: t.Optional[TimeLike],
-        raise_exception: bool,
         deployability_index: t.Optional[DeployabilityIndex],
         **kwargs: t.Any,
     ) -> AuditResult:
@@ -953,18 +948,6 @@ class SnapshotEvaluator:
             select("COUNT(*)").from_(query.subquery("audit")),
             quote_identifiers=True,
         )  # type: ignore
-        if count:
-            audit_error = AuditError(
-                audit_name=audit.name,
-                model=snapshot.model_or_none,
-                count=count,
-                query=query,
-                adapter_dialect=adapter.dialect,
-            )
-            if raise_exception and blocking:
-                raise audit_error
-            else:
-                logger.warning(f"{audit_error}\nAudit is warn only so proceeding with execution.")
 
         return AuditResult(
             audit=audit,
diff --git a/sqlmesh/engines/commands.py b/sqlmesh/engines/commands.py
index 1eaa5c45..602eec47 100644
--- a/sqlmesh/engines/commands.py
+++ b/sqlmesh/engines/commands.py
@@ -1,6 +1,7 @@
 import typing as t
 from enum import Enum
 
+from sqlglot import exp
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
@@ -12,6 +13,7 @@ from sqlmesh.core.snapshot import (
 )
 from sqlmesh.core.state_sync import cleanup_expired_views
 from sqlmesh.utils.date import TimeLike
+from sqlmesh.utils.errors import AuditError
 from sqlmesh.utils.pydantic import PydanticModel
 
 COMMAND_PAYLOAD_FILE_NAME = "payload.json"
@@ -87,7 +89,7 @@ def evaluate(
         deployability_index=command_payload.deployability_index,
         batch_index=command_payload.batch_index,
     )
-    evaluator.audit(
+    audit_results = evaluator.audit(
         snapshot=command_payload.snapshot,
         start=command_payload.start,
         end=command_payload.end,
@@ -97,6 +99,16 @@ def evaluate(
         wap_id=wap_id,
     )
 
+    failed_audit_result = next((r for r in audit_results if r.count and r.blocking), None)
+    if failed_audit_result:
+        raise AuditError(
+            audit_name=failed_audit_result.audit.name,
+            model=command_payload.snapshot.model_or_none,
+            count=t.cast(int, failed_audit_result.count),
+            query=t.cast(exp.Query, failed_audit_result.query),
+            adapter_dialect=evaluator.adapter.dialect,
+        )
+
 
 def promote(
     evaluator: SnapshotEvaluator, command_payload: t.Union[str, PromoteCommandPayload]


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 8b00cea6..b68f746b 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -562,6 +562,7 @@ def test_connection_config_serialization():
         "type": "duckdb",
         "extensions": [],
         "pre_ping": False,
+        "pretty_sql": False,
         "connector_config": {},
         "database": "my_db",
     }
@@ -571,6 +572,7 @@ def test_connection_config_serialization():
         "type": "duckdb",
         "extensions": [],
         "pre_ping": False,
+        "pretty_sql": False,
         "connector_config": {},
         "database": "my_test_db",
     }
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index fe4ba70e..2f697837 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -456,7 +456,7 @@ def test_override_builtin_audit_blocking_mode():
         )
     )
 
-    logger = logging.getLogger("sqlmesh.core.snapshot.evaluator")
+    logger = logging.getLogger("sqlmesh.core.scheduler")
     with patch.object(logger, "warning") as mock_logger:
         plan = context.plan(auto_apply=True, no_prompts=True)
         new_snapshot = next(iter(plan.context_diff.new_snapshots.values()))
@@ -467,7 +467,7 @@ def test_override_builtin_audit_blocking_mode():
                 "Audit 'not_null' for model 'db.x' failed.\n"
                 "Got 1 results, expected 0.\n"
                 f'SELECT * FROM (SELECT * FROM "sqlmesh__db"."db__x__{version}" AS "db__x__{version}") AS "_q_0" WHERE "c" IS NULL AND TRUE\n'
-                "Audit is warn only so proceeding with execution."
+                "Audit is non-blocking so proceeding with execution."
             )
         ]
 
diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py
index 3d3f2310..4b7de2aa 100644
--- a/tests/core/test_scheduler.py
+++ b/tests/core/test_scheduler.py
@@ -610,6 +610,7 @@ def test_audit_failure_notifications(
     mocker.patch("sqlmesh.core.notification_target.NotificationTargetManager.notify", notify_mock)
 
     audit = first(waiter_names.model.audit_definitions.values())
+    query = waiter_names.model.render_query()
 
     def _evaluate():
         scheduler.evaluate(
@@ -620,25 +621,30 @@ def test_audit_failure_notifications(
             DeployabilityIndex.all_deployable(),
             0,
         )
-        _, kwargs = evaluator_audit_mock.call_args_list[0]
-        assert not kwargs["raise_exception"]
 
     evaluator_audit_mock.return_value = [
-        AuditResult(audit=audit, model=waiter_names.model, count=0, skipped=False)
+        AuditResult(audit=audit, model=waiter_names.model, query=query, count=0, skipped=False)
     ]
     _evaluate()
     assert notify_user_mock.call_count == 0
     assert notify_mock.call_count == 0
 
     evaluator_audit_mock.return_value = [
-        AuditResult(audit=audit, model=waiter_names.model, count=None, skipped=True)
+        AuditResult(audit=audit, model=waiter_names.model, query=query, count=None, skipped=True)
     ]
     _evaluate()
     assert notify_user_mock.call_count == 0
     assert notify_mock.call_count == 0
 
     evaluator_audit_mock.return_value = [
-        AuditResult(audit=audit, model=waiter_names.model, count=1, skipped=False, blocking=False)
+        AuditResult(
+            audit=audit,
+            model=waiter_names.model,
+            query=query,
+            count=1,
+            skipped=False,
+            blocking=False,
+        )
     ]
     _evaluate()
     assert notify_user_mock.call_count == 1
@@ -647,7 +653,7 @@ def test_audit_failure_notifications(
     notify_mock.reset_mock()
 
     evaluator_audit_mock.return_value = [
-        AuditResult(audit=audit, model=waiter_names.model, count=1, skipped=False)
+        AuditResult(audit=audit, model=waiter_names.model, query=query, count=1, skipped=False)
     ]
     with pytest.raises(AuditError):
         _evaluate()
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index 4efd26ad..1f33a072 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -51,7 +51,7 @@ from sqlmesh.core.snapshot import (
 from sqlmesh.core.snapshot.evaluator import CustomMaterialization
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.date import to_timestamp
-from sqlmesh.utils.errors import AuditError, ConfigError, SQLMeshError
+from sqlmesh.utils.errors import ConfigError, SQLMeshError
 from sqlmesh.utils.metaprogramming import Executable
 
 
@@ -2468,10 +2468,10 @@ def test_audit_set_blocking_at_use_site(adapter_mock, make_snapshot):
     # Return a non-zero count to indicate audit failure
     adapter_mock.fetchone.return_value = (1,)
 
-    logger = logging.getLogger("sqlmesh.core.snapshot.evaluator")
-    with patch.object(logger, "warning") as mock_logger:
-        evaluator.audit(snapshot, snapshots={})
-        assert "Audit is warn only so proceeding with execution." in mock_logger.call_args[0][0]
+    results = evaluator.audit(snapshot, snapshots={})
+    assert len(results) == 1
+    assert results[0].count == 1
+    assert not results[0].blocking
 
     model = SqlModel(
         name="test_schema.test_table",
@@ -2486,11 +2486,10 @@ def test_audit_set_blocking_at_use_site(adapter_mock, make_snapshot):
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
     adapter_mock.fetchone.return_value = (1,)
 
-    with pytest.raises(
-        AuditError,
-        match="Audit 'always_fail' for model 'test_schema.test_table' failed.",
-    ):
-        evaluator.audit(snapshot, snapshots={})
+    results = evaluator.audit(snapshot, snapshots={})
+    assert len(results) == 1
+    assert results[0].count == 1
+    assert results[0].blocking
 
 
 def test_create_post_statements_use_deployable_table(
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 418da64b..7e443e5e 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -1614,6 +1614,34 @@ test_foo:
     )
 
 
+def test_pretty_query(mocker: MockerFixture) -> None:
+    test = _create_test(
+        body=load_yaml(
+            """
+test_foo:
+  model: xyz
+  schema: my_schema
+  outputs:
+    query:
+      - a: 1
+            """
+        ),
+        test_name="test_foo",
+        model=_create_model("SELECT 1 AS a"),
+        context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
+    )
+    test.engine_adapter._pretty_sql = True
+    spy_execute = mocker.spy(test.engine_adapter, "_execute")
+    _check_successful_or_raise(test.run())
+    spy_execute.assert_has_calls(
+        [
+            call('CREATE SCHEMA IF NOT EXISTS "memory"."my_schema"'),
+            call('SELECT\n  1 AS "a"'),
+            call('DROP SCHEMA IF EXISTS "memory"."my_schema" CASCADE'),
+        ]
+    )
+
+
 def test_complicated_recursive_cte() -> None:
     model_sql = """
 WITH
diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py
index 7888b936..25b766e9 100644
--- a/tests/integrations/jupyter/test_magics.py
+++ b/tests/integrations/jupyter/test_magics.py
@@ -552,16 +552,16 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
         "Models: 17",
         "Macros: 7",
         "",
-        "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  extensions: []\n  connector_config: {}",
-        "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  extensions: []\n  connector_config: {}",
+        "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}",
+        "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}",
         "Data warehouse connection succeeded",
     ]
     assert get_all_html_output(output) == [
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Models: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">17</span></pre>",
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Macros: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span></pre>",
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>",
-        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span></pre>',
-        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span></pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span></pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span></pre>',
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Data warehouse connection <span style=\"color: #008000; text-decoration-color: #008000\">succeeded</span></pre>",
     ]