From 21817144d9ccbfab578b539a6648cc1fce82dcf9 Mon Sep 17 00:00:00 2001 From: "michael.wallis" Date: Fri, 6 Sep 2024 08:29:19 +0100 Subject: [PATCH 1/2] feat: add an option to pretty sql --- docs/reference/configuration.md | 9 +++++---- sqlmesh/core/config/connection.py | 2 ++ sqlmesh/core/engine_adapter/base.py | 5 +++-- tests/core/test_config.py | 2 ++ tests/core/test_test.py | 28 ++++++++++++++++++++++++++++ 5 files changed, 40 insertions(+), 6 deletions(-) diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index d3b1303c77..aa5c481970 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 b6bc9d8730..33cd04e18e 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 @@ def create_engine_adapter(self, register_comments_override: bool = False) -> Eng 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/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 3e75db0aa0..27de53a069 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -115,6 +115,7 @@ def __init__( 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 @@ def __init__( 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 @@ def execute( 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 @@ def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.An """ sql_gen_kwargs = { "dialect": self.dialect, - "pretty": False, + "pretty": self._pretty_sql, "comments": False, **self._sql_gen_kwargs, **kwargs, diff --git a/tests/core/test_config.py b/tests/core/test_config.py index 8b00cea622..b68f746b13 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_test.py b/tests/core/test_test.py index 418da64b45..7e443e5e83 100644 --- a/tests/core/test_test.py +++ b/tests/core/test_test.py @@ -1614,6 +1614,34 @@ def test_custom_testing_schema(mocker: MockerFixture) -> None: ) +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 From e1907b12b86c428626f6abb95543a9d546fa646c Mon Sep 17 00:00:00 2001 From: "michael.wallis" Date: Wed, 4 Dec 2024 09:46:50 +0000 Subject: [PATCH 2/2] fix magics info test --- tests/integrations/jupyter/test_magics.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py index 7888b93625..25b766e97c 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) == [ "
Models: 17
", "
Macros: 7
", "
",
-        '
Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  extensions: []  connector_config: {}
',
-        '
Test Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  extensions: []  connector_config: {}
',
+        '
Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  pretty_sql: false  extensions: []  connector_config: {}
',
+        '
Test Connection:  type: duckdb  concurrent_tasks: 1  register_comments: true  pre_ping: false  pretty_sql: false  extensions: []  connector_config: {}
',
         "
Data warehouse connection succeeded
",
     ]