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:
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>",
     ]