Feat: allow setting register_comments for dbt connection (#2212)

* Allow setting register_comments parameters for dbt config

* Add docs

* remove unnecessary kwargs pop

Author: treysp

docs/integrations/dbt.md

@@ -22,8 +22,13 @@ Models **require** a start date for backfilling data through use of the `start`
 >   +start: Jan 1 2000
 ```
 
+### Configuration
 
-### Runtime vars
+SQLMesh determines a project's configuration settings from its dbt configuration files.
+
+This section describes using runtime variables to create multiple configurations and how to disable SQLMesh's automatic model description and comment registration.
+
+#### Runtime vars
 
 dbt supports passing variable values at runtime with its [CLI `vars` option](https://docs.getdbt.com/docs/build/project-variables#defining-variables-on-the-command-line).
 
@@ -70,6 +75,21 @@ sqlmesh --config marketing_config plan
 
 Note that the `--config` option is specified between the word `sqlmesh` and the command being executed (e.g., `plan`, `run`).
 
+#### 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).
+
+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.
+
+SQLMesh's project-level comment registration defaults are overridden with the `sqlmesh_config()` `register_comments` argument. For example, this configuration turns comment registration off:
+
+```python
+config = sqlmesh_config(
+    Path(__file__).parent,
+    register_comments=False,
+    )
+```
+
 ### Running SQLMesh
 
 Run SQLMesh as with a SQLMesh project, generating and applying [plans](../concepts/overview.md#make-a-plan), running [tests](../concepts/overview.md#tests) or [audits](../concepts/overview.md#audits), and executing models with a [scheduler](../guides/scheduling.md) if desired.

sqlmesh/dbt/loader.py

@@ -36,6 +36,7 @@ def sqlmesh_config(
     state_connection: t.Optional[ConnectionConfig] = None,
     dbt_target_name: t.Optional[str] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
+    register_comments: t.Optional[bool] = None,
     **kwargs: t.Any,
 ) -> Config:
     project_root = project_root or Path()
@@ -48,9 +49,13 @@ def sqlmesh_config(
     if variables is not None:
         loader_kwargs["variables"] = variables
 
+    target_to_sqlmesh_args = {}
+    if register_comments is not None:
+        target_to_sqlmesh_args["register_comments"] = register_comments
+
     return Config(
         default_gateway=profile.target_name,
-        gateways={profile.target_name: GatewayConfig(connection=profile.target.to_sqlmesh(), state_connection=state_connection)},  # type: ignore
+        gateways={profile.target_name: GatewayConfig(connection=profile.target.to_sqlmesh(**target_to_sqlmesh_args), state_connection=state_connection)},  # type: ignore
         loader=DbtLoader,
         loader_kwargs=loader_kwargs,
         model_defaults=model_defaults,

sqlmesh/dbt/target.py

@@ -112,7 +112,7 @@ def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         """The default incremental strategy for the db"""
         raise NotImplementedError
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         """Converts target config to SQLMesh connection config"""
         raise NotImplementedError
 
@@ -187,8 +187,7 @@ def relation_class(cls) -> t.Type[BaseRelation]:
 
         return DuckDBRelation
 
-    def to_sqlmesh(self) -> ConnectionConfig:
-        kwargs: t.Dict[str, t.Any] = {}
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         if self.extensions is not None:
             kwargs["extensions"] = self.extensions
         if self.settings is not None:
@@ -277,7 +276,7 @@ def column_class(cls) -> t.Type[Column]:
 
         return SnowflakeColumn
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         return SnowflakeConnectionConfig(
             user=self.user,
             password=self.password,
@@ -291,6 +290,7 @@ def to_sqlmesh(self) -> ConnectionConfig:
             private_key=self.private_key,
             private_key_path=self.private_key_path,
             private_key_passphrase=self.private_key_passphrase,
+            **kwargs,
         )
 
     @classproperty
@@ -347,7 +347,7 @@ def _validate_port(cls, v: t.Union[int, str]) -> int:
     def default_incremental_strategy(self, kind: IncrementalKind) -> str:
         return "delete+insert" if kind is IncrementalByUniqueKeyKind else "append"
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         return PostgresConnectionConfig(
             host=self.host,
             user=self.user,
@@ -359,6 +359,7 @@ def to_sqlmesh(self) -> ConnectionConfig:
             connect_timeout=self.connect_timeout,
             role=self.role,
             sslmode=self.sslmode,
+            **kwargs,
         )
 
 
@@ -419,7 +420,7 @@ def column_class(cls) -> t.Type[Column]:
         else:
             return super(RedshiftConfig, cls).column_class
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         return RedshiftConnectionConfig(
             user=self.user,
             password=self.password,
@@ -429,6 +430,7 @@ def to_sqlmesh(self) -> ConnectionConfig:
             sslmode=self.sslmode,
             timeout=self.connect_timeout,
             concurrent_tasks=self.threads,
+            **kwargs,
         )
 
 
@@ -465,13 +467,14 @@ def column_class(cls) -> t.Type[Column]:
 
         return DatabricksColumn
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         return DatabricksConnectionConfig(
             server_hostname=self.host,
             http_path=self.http_path,
             access_token=self.token,
             concurrent_tasks=self.threads,
             catalog=self.database,
+            **kwargs,
         )
 
 
@@ -548,7 +551,7 @@ def column_class(cls) -> t.Type[Column]:
 
         return BigQueryColumn
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         job_retries = self.job_retries if self.job_retries is not None else self.retries
         job_execution_timeout_seconds = (
             self.job_execution_timeout_seconds
@@ -574,6 +577,7 @@ def to_sqlmesh(self) -> ConnectionConfig:
             job_retry_deadline_seconds=self.job_retry_deadline_seconds,
             priority=self.priority,
             maximum_bytes_billed=self.maximum_bytes_billed,
+            **kwargs,
         )
 
 
@@ -664,7 +668,7 @@ def column_class(cls) -> t.Type[Column]:
     def dialect(self) -> str:
         return "tsql"
 
-    def to_sqlmesh(self) -> ConnectionConfig:
+    def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig:
         return MSSQLConnectionConfig(
             host=self.host,
             user=self.user,
@@ -674,6 +678,7 @@ def to_sqlmesh(self) -> ConnectionConfig:
             timeout=self.query_timeout,
             login_timeout=self.login_timeout,
             concurrent_tasks=self.threads,
+            **kwargs,
         )
 
 

tests/dbt/test_config.py

@@ -670,6 +670,16 @@ def test_sqlserver_config():
     )
 
 
+def test_connection_args(tmp_path):
+    dbt_project_dir = "tests/fixtures/dbt/sushi_test"
+
+    config = sqlmesh_config(dbt_project_dir)
+    assert config.gateways["in_memory"].connection.register_comments
+
+    config = sqlmesh_config(dbt_project_dir, register_comments=False)
+    assert not config.gateways["in_memory"].connection.register_comments
+
+
 @pytest.mark.cicdonly
 def test_db_type_to_relation_class():
     from dbt.adapters.bigquery.relation import BigQueryRelation