Skip to content

feat: add an option to pretty sql #3455

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
izeigerman merged 2 commits into from
Dec 6, 2024
Merged

feat: add an option to pretty sql #3455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2181714
feat: add an option to pretty sql
MikeWallis42 Sep 6, 2024
e1907b1
fix magics info test
MikeWallis42 Dec 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions docs/reference/configuration.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 2 additions & 0 deletions sqlmesh/core/config/connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand Down
5 changes: 3 additions & 2 deletions sqlmesh/core/engine_adapter/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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__(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions tests/core/test_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -562,6 +562,7 @@ def test_connection_config_serialization():
"type": "duckdb",
"extensions": [],
"pre_ping": False,
"pretty_sql": False,
"connector_config": {},
"database": "my_db",
}
Expand All @@ -571,6 +572,7 @@ def test_connection_config_serialization():
"type": "duckdb",
"extensions": [],
"pre_ping": False,
"pretty_sql": False,
"connector_config": {},
"database": "my_test_db",
}
Expand Down
28 changes: 28 additions & 0 deletions tests/core/test_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 4 additions & 4 deletions tests/integrations/jupyter/test_magics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>",
]

Expand Down