Skip to content

Commit

Permalink
Escape backslashes in comments
Browse files Browse the repository at this point in the history
  • Loading branch information
treysp committed Mar 19, 2024
1 parent d16f0dc commit 5db5fbc
Show file tree
Hide file tree
Showing 8 changed files with 44 additions and 7 deletions.
2 changes: 1 addition & 1 deletion examples/sushi/models/marketing.sql
Expand Up @@ -8,7 +8,7 @@ MODEL (
);

SELECT
customer_id::INT AS customer_id, -- customer_id uniquely identifies customers
customer_id::INT AS customer_id, -- customer_id uniquely identifies customers \
status::TEXT AS status,
updated_at::TIMESTAMP AS updated_at
FROM
Expand Down
8 changes: 7 additions & 1 deletion sqlmesh/core/engine_adapter/base.py
Expand Up @@ -87,6 +87,7 @@ class EngineAdapter:
COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_AND_COMMANDS
MAX_TABLE_COMMENT_LENGTH: t.Optional[int] = None
MAX_COLUMN_COMMENT_LENGTH: t.Optional[int] = None
ESCAPE_COMMENT_BACKSLASH = True
INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT
SUPPORTS_MATERIALIZED_VIEWS = False
SUPPORTS_MATERIALIZED_VIEW_SCHEMA = False
Expand Down Expand Up @@ -1919,7 +1920,12 @@ def _build_view_properties_exp(
return exp.Properties(expressions=properties)
return None

def _truncate_comment(self, comment: str, length: t.Optional[int]) -> str:
def _truncate_comment(
self, comment: str, length: t.Optional[int], escape_backslash: t.Optional[bool] = None
) -> str:
escape = escape_backslash if escape_backslash is not None else self.ESCAPE_COMMENT_BACKSLASH
if escape:
comment = comment.replace("\\", "\\\\")
return comment[:length] if length else comment

def _truncate_table_comment(self, comment: str) -> str:
Expand Down
4 changes: 2 additions & 2 deletions sqlmesh/core/engine_adapter/bigquery.py
Expand Up @@ -444,8 +444,8 @@ def _create_column_comments(
for i in range(len(table_def["schema"]["fields"])):
comment = column_comments.get(table_def["schema"]["fields"][i]["name"], None)
if comment:
table_def["schema"]["fields"][i]["description"] = self._truncate_column_comment(
comment
table_def["schema"]["fields"][i]["description"] = self._truncate_comment(
comment, self.MAX_COLUMN_COMMENT_LENGTH, escape_backslash=False
)

# An "etag" is BQ versioning metadata that changes when an object is updated/modified. `update_table`
Expand Down
6 changes: 4 additions & 2 deletions sqlmesh/core/engine_adapter/mixins.py
Expand Up @@ -221,11 +221,13 @@ def _build_view_properties_exp(
return exp.Properties(expressions=properties)
return None

def _truncate_comment(self, comment: str, length: t.Optional[int]) -> str:
def _truncate_comment(
self, comment: str, length: t.Optional[int], escape_backslash: t.Optional[bool] = None
) -> str:
# iceberg does not have a comment length limit
if self.current_catalog_type == "iceberg":
return comment
return super()._truncate_comment(comment, length)
return super()._truncate_comment(comment, length, escape_backslash)


class GetCurrentCatalogFromFunctionMixin(EngineAdapter):
Expand Down
1 change: 1 addition & 0 deletions sqlmesh/core/engine_adapter/postgres.py
Expand Up @@ -29,6 +29,7 @@ class PostgresEngineAdapter(
HAS_VIEW_BINDING = True
CURRENT_CATALOG_EXPRESSION = exp.column("current_catalog")
SUPPORTS_REPLACE_TABLE = False
ESCAPE_COMMENT_BACKSLASH = False

def _fetch_native_df(
self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
Expand Down
8 changes: 8 additions & 0 deletions tests/core/engine_adapter/test_bigquery.py
Expand Up @@ -606,6 +606,13 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
column_descriptions={"a": long_column_comment},
)

adapter.create_table(
"test_table",
{"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
table_description="\\",
column_descriptions={"a": "\\"},
)

adapter.ctas(
"test_table",
parse_one("SELECT a, b FROM source_table"),
Expand All @@ -628,6 +635,7 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
sql_calls = _to_sql_calls(execute_mock)
assert sql_calls == [
f"CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64 OPTIONS (description='{truncated_column_comment}'), `b` INT64) OPTIONS (description='{truncated_table_comment}')",
"CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64 OPTIONS (description='\\\\'), `b` INT64) OPTIONS (description='\\\\')",
f"CREATE TABLE IF NOT EXISTS `test_table` (`a` INT64 OPTIONS (description='{truncated_column_comment}'), `b` INT64) OPTIONS (description='{truncated_table_comment}') AS SELECT `a`, `b` FROM `source_table`",
f"CREATE OR REPLACE VIEW `test_table` OPTIONS (description='{truncated_table_comment}') AS SELECT `a`, `b` FROM `source_table`",
f"ALTER TABLE `test_table` SET OPTIONS(description = '{truncated_table_comment}')",
Expand Down
2 changes: 1 addition & 1 deletion tests/core/engine_adapter/test_integration.py
Expand Up @@ -1683,7 +1683,7 @@ def test_sushi(mark_gateway: t.Tuple[str, str], ctx: TestContext):
},
"marketing": {
"table": "Sushi marketing data",
"column": {"customer_id": "customer_id uniquely identifies customers"},
"column": {"customer_id": "customer_id uniquely identifies customers \\"},
},
"orders": {
"table": "Table of sushi orders.",
Expand Down
20 changes: 20 additions & 0 deletions tests/core/engine_adapter/test_postgres.py
Expand Up @@ -2,6 +2,8 @@

import pytest
from pytest_mock import MockFixture
from pytest_mock.plugin import MockerFixture
from sqlglot import exp
from sqlglot.helper import ensure_list

from sqlmesh.core.engine_adapter import PostgresEngineAdapter
Expand Down Expand Up @@ -60,3 +62,21 @@ def test_drop_schema_with_catalog(

adapter.drop_schema("test_catalog.test_schema")
assert "requires that all catalog operations be against a single catalog" in caplog.text


def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
adapter = make_mocked_engine_adapter(PostgresEngineAdapter)

adapter.create_table(
"test_table",
{"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
table_description="\\",
column_descriptions={"a": "\\"},
)

sql_calls = to_sql_calls(adapter)
assert sql_calls == [
'CREATE TABLE IF NOT EXISTS "test_table" ("a" INT, "b" INT)',
"""COMMENT ON TABLE "test_table" IS '\\'""",
"""COMMENT ON COLUMN "test_table"."a" IS '\\'""",
]

0 comments on commit 5db5fbc

Please sign in to comment.