Repository: TobikoData/sqlmesh
Pull request #2369 (2024-03-29 21:18:27)

Fix!: expand snowflake VIEW column comments and make comment default true
Snowflake's documentation does not describe or demonstrate a way to alter/add column comments to VIEWs after creation, but `ALTER VIEW ... ALTER COLUMN ...` command does this. 

This PR:
- Adds the ability for snowflake to alter/add column comments to VIEWs after creation.
- Removes the `CommentCreationView.IN_SCHEMA_DEF_NO_COLUMN_COMMAND` enum value because it was only used by snowflake
- Makes comment registration `True` by default for Snowflake


Patch:
diff --git a/docs/concepts/models/overview.md b/docs/concepts/models/overview.md
index e64ca89d..733f1ded 100644
--- a/docs/concepts/models/overview.md
+++ b/docs/concepts/models/overview.md
@@ -73,7 +73,7 @@ Model files may contain SQL comments in a format supported in the model's SQL di
 
 Some SQL engines support registering comments as metadata associated with a table or view. They may support table-level comments (e.g., "Revenue data for each customer") and/or column-level comments (e.g., "Customer's unique ID").
 
-SQLMesh will automatically register comments if the engine supports it and the [gateway's connection `register_comments` configuration](../../reference/configuration.md#connection) is `true` (`true` by default for all engines other than Snowflake). Engines vary in their support for comments - see [tables below](#engine-comment-support).
+SQLMesh will automatically register comments if the engine supports it and the [gateway's connection `register_comments` configuration](../../reference/configuration.md#connection) is `true` (`true` by default). Engines vary in their support for comments - see [tables below](#engine-comment-support).
 
 #### Model comments
 
diff --git a/docs/integrations/dbt.md b/docs/integrations/dbt.md
index 0c8fcdef..ac6ee71c 100644
--- a/docs/integrations/dbt.md
+++ b/docs/integrations/dbt.md
@@ -77,7 +77,7 @@ Note that the `--config` option is specified between the word `sqlmesh` and the
 
 #### Registering comments
 
-SQLMesh automatically registers model descriptions and column comments with the target SQL engine, as described in the [Models Overview documentation](../concepts/models/overview#model-description-and-comments). Comment registration is on by default for all engines that support it (but off by default for Snowflake).
+SQLMesh automatically registers model descriptions and column comments with the target SQL engine, as described in the [Models Overview documentation](../concepts/models/overview#model-description-and-comments). Comment registration is on by default for all engines that support it.
 
 dbt offers similar comment registration functionality via its [`persist_docs` model configuration parameter](https://docs.getdbt.com/reference/resource-configs/persist_docs), specified by model. SQLMesh comment registration is configured at the project level, so it does not use dbt's model-specific `persist_docs` configuration.
 
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index f29d44bc..012a45ba 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -148,7 +148,7 @@ Most parameters are specific to the connection engine `type` - see [below](#engi
 | ------------------- | --------------------------------------------------------------------------------------------------------------------------- | :--: | :------: |
 | `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`; `false` for Snowflake engine.)                            | bool |    N     |
+| `register_comments` | Whether SQLMesh should register model comments with the SQL engine (if the engine supports it). (Default: `true`.)                            | bool |    N     |
 
 #### Engine-specific
 
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 26f36aca..66756e8f 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -304,7 +304,7 @@ class SnowflakeConnectionConfig(ConnectionConfig):
     private_key_passphrase: t.Optional[str] = None
 
     concurrent_tasks: int = 4
-    register_comments: bool = False
+    register_comments: bool = True
 
     type_: Literal["snowflake"] = Field(alias="type", default="snowflake")
 
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 3eae8b5d..428bc341 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -110,7 +110,7 @@ class EngineAdapter:
         cursor_init: t.Optional[t.Callable[[t.Any], None]] = None,
         default_catalog: t.Optional[str] = None,
         execute_log_level: int = logging.DEBUG,
-        register_comments: bool = False,
+        register_comments: bool = True,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
diff --git a/sqlmesh/core/engine_adapter/shared.py b/sqlmesh/core/engine_adapter/shared.py
index b6b0e328..1480be2f 100644
--- a/sqlmesh/core/engine_adapter/shared.py
+++ b/sqlmesh/core/engine_adapter/shared.py
@@ -108,17 +108,13 @@ class CommentCreationView(Enum):
                                    and in post-creation commands
     IN_SCHEMA_DEF_NO_COMMANDS = all comments can be registered in CREATE VIEW schema definitions,
                                   but not in post-creation commands
-    IN_SCHEMA_DEF_NO_COLUMN_COMMAND = all comments can be registered in CREATE VIEW schema definitions,
-                                        view comments can be registered in post-creation commands,
-                                        column comments cannot be registered in post-creation commands
     COMMENT_COMMAND_ONLY = comments can only be registered via a post-creation command like `COMMENT` or `ALTER`
     """
 
     UNSUPPORTED = 1
     IN_SCHEMA_DEF_AND_COMMANDS = 2
     IN_SCHEMA_DEF_NO_COMMANDS = 3
-    IN_SCHEMA_DEF_NO_COLUMN_COMMAND = 4
-    COMMENT_COMMAND_ONLY = 5
+    COMMENT_COMMAND_ONLY = 4
 
     @property
     def is_unsupported(self) -> bool:
@@ -132,10 +128,6 @@ class CommentCreationView(Enum):
     def is_in_schema_def_no_commands(self) -> bool:
         return self == CommentCreationView.IN_SCHEMA_DEF_NO_COMMANDS
 
-    @property
-    def is_in_schema_def_no_column_command(self) -> bool:
-        return self == CommentCreationView.IN_SCHEMA_DEF_NO_COLUMN_COMMAND
-
     @property
     def is_comment_command_only(self) -> bool:
         return self == CommentCreationView.COMMENT_COMMAND_ONLY
@@ -149,7 +141,6 @@ class CommentCreationView(Enum):
         return self in (
             CommentCreationView.IN_SCHEMA_DEF_AND_COMMANDS,
             CommentCreationView.IN_SCHEMA_DEF_NO_COMMANDS,
-            CommentCreationView.IN_SCHEMA_DEF_NO_COLUMN_COMMAND,
         )
 
     @property
diff --git a/sqlmesh/core/engine_adapter/snowflake.py b/sqlmesh/core/engine_adapter/snowflake.py
index e423b68d..aad5ac77 100644
--- a/sqlmesh/core/engine_adapter/snowflake.py
+++ b/sqlmesh/core/engine_adapter/snowflake.py
@@ -13,7 +13,6 @@ from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.mixins import GetCurrentCatalogFromFunctionMixin
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
-    CommentCreationView,
     DataObject,
     DataObjectType,
     SourceQuery,
@@ -41,7 +40,6 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin):
     SUPPORTS_CLONING = True
     CATALOG_SUPPORT = CatalogSupport.FULL_SUPPORT
     CURRENT_CATALOG_EXPRESSION = exp.func("current_database")
-    COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_NO_COLUMN_COMMAND
 
     @contextlib.contextmanager
     def session(self, properties: SessionProperties) -> t.Iterator[None]:
@@ -189,3 +187,11 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin):
 
     def set_current_catalog(self, catalog: str) -> None:
         self.execute(exp.Use(this=exp.to_identifier(catalog)))
+
+    def _build_create_comment_column_exp(
+        self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE"
+    ) -> exp.Comment | str:
+        table_sql = table.sql(dialect=self.dialect, identify=True)
+        column_sql = exp.column(column_name).sql(dialect=self.dialect, identify=True)
+
+        return f"ALTER {table_kind} {table_sql} ALTER COLUMN {column_sql} COMMENT '{self._truncate_column_comment(column_comment)}'"


Test patch:
diff --git a/tests/core/engine_adapter/config.yaml b/tests/core/engine_adapter/config.yaml
index fd29007e..66be470e 100644
--- a/tests/core/engine_adapter/config.yaml
+++ b/tests/core/engine_adapter/config.yaml
@@ -5,7 +5,6 @@ gateways:
       catalogs:
         memory: ':memory:'
         testing: 'testing.duckdb'
-      register_comments: true
   inttest_trino:
     connection:
       type: trino
@@ -15,7 +14,6 @@ gateways:
       catalog: datalake
       http_scheme: http
       retries: 20
-      register_comments: true
     state_connection:
       type: duckdb
   inttest_trino_iceberg:
@@ -27,7 +25,6 @@ gateways:
       catalog: datalake_iceberg
       http_scheme: http
       retries: 20
-      register_comments: true
     state_connection:
       type: duckdb
   inttest_trino_delta:
@@ -39,7 +36,6 @@ gateways:
       catalog: datalake_delta
       http_scheme: http
       retries: 20
-      register_comments: true
     state_connection:
       type: duckdb
   inttest_spark:
@@ -47,7 +43,6 @@ gateways:
       type: spark
       config:
         spark.remote: sc://localhost
-      register_comments: true
     state_connection:
       type: duckdb
   inttest_mssql:
@@ -56,7 +51,6 @@ gateways:
       host: localhost
       user: sa
       password: 1StrongPwd@@
-      register_comments: true
   inttest_postgres:
     connection:
       type: postgres
@@ -66,7 +60,6 @@ gateways:
       host: localhost
       port: 5432
       concurrent_tasks: 1
-      register_comments: true
   inttest_mysql:
     connection:
       type: mysql
@@ -75,7 +68,6 @@ gateways:
       password: mysql
       port: 3306
       charset: utf8
-      register_comments: true
 
 model_defaults:
     dialect: duckdb
diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py
index b5029a4a..c2425b44 100644
--- a/tests/core/engine_adapter/test_snowflake.py
+++ b/tests/core/engine_adapter/test_snowflake.py
@@ -2,7 +2,7 @@ import typing as t
 
 import pytest
 from pytest_mock.plugin import MockerFixture
-from sqlglot import exp
+from sqlglot import exp, parse_one
 
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.engine_adapter import SnowflakeEngineAdapter
@@ -97,3 +97,49 @@ def test_session(
         )
 
     assert to_sql_calls(adapter) == expected_calls
+
+
+def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
+
+    adapter.create_table(
+        "test_table",
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
+        table_description="table description",
+        column_descriptions={"a": "a column description"},
+    )
+
+    adapter.ctas(
+        "test_table",
+        parse_one("SELECT a, b FROM source_table"),
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")},
+        table_description="table description",
+        column_descriptions={"a": "a column description"},
+    )
+
+    adapter.create_view(
+        "test_view",
+        parse_one("SELECT a, b FROM source_table"),
+        table_description="table description",
+        column_descriptions={"a": "a column description"},
+    )
+
+    adapter._create_table_comment(
+        "test_table",
+        "table description",
+    )
+
+    adapter._create_column_comments(
+        "test_table",
+        {"a": "a column description"},
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        """CREATE TABLE IF NOT EXISTS "test_table" ("a" INT COMMENT 'a column description', "b" INT) COMMENT='table description'""",
+        """CREATE TABLE IF NOT EXISTS "test_table" ("a" INT COMMENT 'a column description', "b" INT) COMMENT='table description' AS SELECT "a", "b" FROM "source_table\"""",
+        """CREATE OR REPLACE VIEW "test_view" COMMENT='table description' AS SELECT "a", "b" FROM "source_table\"""",
+        """ALTER VIEW "test_view" ALTER COLUMN "a" COMMENT 'a column description'""",
+        """COMMENT ON TABLE "test_table" IS 'table description'""",
+        """ALTER TABLE "test_table" ALTER COLUMN "a" COMMENT 'a column description'""",
+    ]