Repository: TobikoData/sqlmesh
Pull request #2970 (2024-07-30 19:09:28)

Feat: Add support for dbt adapter method quote_as_configured


Patch:
diff --git a/sqlmesh/dbt/adapter.py b/sqlmesh/dbt/adapter.py
index 1820e72b..a8349c9a 100644
--- a/sqlmesh/dbt/adapter.py
+++ b/sqlmesh/dbt/adapter.py
@@ -30,11 +30,15 @@ class BaseAdapter(abc.ABC):
         jinja_macros: JinjaMacroRegistry,
         jinja_globals: t.Optional[t.Dict[str, t.Any]] = None,
         project_dialect: t.Optional[str] = None,
+        quote_policy: t.Optional[Policy] = None,
     ):
+        from dbt.adapters.base.relation import Policy
+
         self.jinja_macros = jinja_macros
         self.jinja_globals = jinja_globals.copy() if jinja_globals else {}
         self.jinja_globals["adapter"] = self
         self.project_dialect = project_dialect
+        self.quote_policy = quote_policy or Policy()
 
     @abc.abstractmethod
     def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
@@ -95,6 +99,10 @@ class BaseAdapter(abc.ABC):
         """Returns a quoted identifier."""
         return exp.to_column(identifier).sql(dialect=self.project_dialect, identify=True)
 
+    def quote_as_configured(self, value: str, component_type: str) -> str:
+        """Returns the value quoted according to the quote policy."""
+        return self.quote(value) if getattr(self.quote_policy, component_type, False) else value
+
     def dispatch(self, name: str, package: t.Optional[str] = None) -> t.Callable:
         """Returns a dialect-specific version of a macro with the given name."""
         target_type = self.jinja_globals["target"]["type"]
@@ -216,12 +224,12 @@ class RuntimeAdapter(BaseAdapter):
     ):
         from dbt.adapters.base import BaseRelation
         from dbt.adapters.base.column import Column
-        from dbt.adapters.base.relation import Policy
 
         super().__init__(
             jinja_macros,
             jinja_globals=jinja_globals,
             project_dialect=project_dialect or engine_adapter.dialect,
+            quote_policy=quote_policy,
         )
 
         table_mapping = table_mapping or {}
@@ -229,7 +237,6 @@ class RuntimeAdapter(BaseAdapter):
         self.engine_adapter = engine_adapter
         self.relation_type = relation_type or BaseRelation
         self.column_type = column_type or Column
-        self.quote_policy = quote_policy or Policy()
         self.table_mapping = {
             **to_table_mapping((snapshots or {}).values(), deployability_index),
             **table_mapping,
diff --git a/sqlmesh/dbt/builtin.py b/sqlmesh/dbt/builtin.py
index b3d55461..d7ebcfd3 100644
--- a/sqlmesh/dbt/builtin.py
+++ b/sqlmesh/dbt/builtin.py
@@ -382,6 +382,7 @@ def create_builtin_globals(
             jinja_macros,
             jinja_globals={**builtin_globals, **jinja_globals},
             project_dialect=project_dialect,
+            quote_policy=api.quote_policy,
         )
 
     sql_execution = SQLExecution(adapter)


Test patch:
diff --git a/tests/dbt/test_adapter.py b/tests/dbt/test_adapter.py
index 320a64be..bfd71545 100644
--- a/tests/dbt/test_adapter.py
+++ b/tests/dbt/test_adapter.py
@@ -14,10 +14,12 @@ from sqlglot import exp, parse_one
 from sqlmesh import Context
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.snapshot import SnapshotId
+from sqlmesh.dbt.adapter import ParsetimeAdapter
 from sqlmesh.dbt.project import Project
 from sqlmesh.dbt.relation import Policy
 from sqlmesh.dbt.target import SnowflakeConfig
 from sqlmesh.utils.errors import ConfigError
+from sqlmesh.utils.jinja import JinjaMacroRegistry
 
 pytestmark = pytest.mark.dbt
 
@@ -257,3 +259,14 @@ def test_feature_flag_scd_type_2(copy_to_temp_path, caplog):
             "Skipping loading Snapshot (SCD Type 2) models due to the feature flag disabling this feature"
             in caplog.text
         )
+
+
+def test_quote_as_configured():
+    adapter = ParsetimeAdapter(
+        JinjaMacroRegistry(),
+        project_dialect="duckdb",
+        quote_policy=Policy(schema=False, identifier=True),
+    )
+    adapter.quote_as_configured("foo", "identifier") == '"foo"'
+    adapter.quote_as_configured("foo", "schema") == "foo"
+    adapter.quote_as_configured("foo", "database") == "foo"