Support databricks_tags for MV/STs

ericj-db · ericj-db · commit a8cd583e01bc · 2025-08-11T10:15:46.000-07:00
diff --git a/dbt/adapters/databricks/relation.py b/dbt/adapters/databricks/relation.py
@@ -43,6 +43,10 @@ class DatabricksRelationType(StrEnum):
     MetricView = "metric_view"
     Unknown = "unknown"
 
+    def render(self) -> str:
+        """Return the type formatted for SQL statements (replace underscores with spaces)"""
+        return self.value.replace("_", " ")
+
 
 class DatabricksTableType(StrEnum):
     External = "external"
diff --git a/dbt/include/databricks/macros/adapters/persist_docs.sql b/dbt/include/databricks/macros/adapters/persist_docs.sql
@@ -28,7 +28,7 @@ COMMENT ON COLUMN {{ column_path }} IS '{{ escaped_comment }}'
 {% endmacro %}
 
 {% macro alter_relation_comment_sql(relation, description) %}
-COMMENT ON {{ relation.type.upper() }} {{ relation.render() }} IS '{{ description | replace("'", "\\'") }}'
+COMMENT ON {{ relation.type.render().upper() }} {{ relation.render() }} IS '{{ description | replace("'", "\\'") }}'
 {% endmacro %}
 
 {% macro alter_column_comments(relation, column_dict) %}
diff --git a/dbt/include/databricks/macros/materializations/materialized_view.sql b/dbt/include/databricks/macros/materializations/materialized_view.sql
@@ -64,9 +64,12 @@
     {{ run_hooks(pre_hooks, inside_transaction=True) }}
 
     {% set grant_config = config.get('grants') %}
+    {% set tags = config.get('databricks_tags') %}
 
     {{ execute_multiple_statements(build_sql) }}
 
+    {%- do apply_tags(target_relation, tags) -%}
+
     {% set column_tags = adapter.get_column_tags_from_model(config.model) %}
     {% if column_tags %}
       {{ apply_column_tags(target_relation, column_tags) }}
diff --git a/dbt/include/databricks/macros/materializations/streaming_table.sql b/dbt/include/databricks/macros/materializations/streaming_table.sql
@@ -63,9 +63,12 @@
     {{ run_hooks(pre_hooks, inside_transaction=True) }}
 
     {% set grant_config = config.get('grants') %}
+    {% set tags = config.get('databricks_tags') %}
 
     {{ execute_multiple_statements(build_sql) }}
 
+    {%- do apply_tags(target_relation, tags) -%}
+
     {% set should_revoke = should_revoke(existing_relation, full_refresh_mode=True) %}
     {% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %}
 
diff --git a/dbt/include/databricks/macros/relations/components/column_mask.sql b/dbt/include/databricks/macros/relations/components/column_mask.sql
@@ -41,13 +41,13 @@
 {%- endmacro -%}
 
 {% macro alter_drop_column_mask(relation, column) -%}
-  ALTER {{ relation.type }} {{ relation.render() }}
+  ALTER {{ relation.type.render() }} {{ relation.render() }}
   ALTER COLUMN `{{ column }}`
   DROP MASK;
 {%- endmacro -%}
 
 {% macro alter_set_column_mask(relation, column, mask) -%}
-  ALTER {{ relation.type }} {{ relation.render() }}
+  ALTER {{ relation.type.render() }} {{ relation.render() }}
   ALTER COLUMN `{{ column }}`
   SET MASK {{ mask.function }}
   {%- if mask.using_columns %}
diff --git a/dbt/include/databricks/macros/relations/components/column_tags.sql b/dbt/include/databricks/macros/relations/components/column_tags.sql
@@ -38,7 +38,7 @@
   {%- if relation.type == 'view' -%}
     ALTER TABLE {{ relation.render() }}
   {%- else -%}
-    ALTER {{ relation.type | replace('_', ' ') }} {{ relation.render() }}
+    ALTER {{ relation.type.render() }} {{ relation.render() }}
   {%- endif -%}
   ALTER COLUMN `{{ column }}`
   SET TAGS (
diff --git a/dbt/include/databricks/macros/relations/components/constraints.sql b/dbt/include/databricks/macros/relations/components/constraints.sql
@@ -119,23 +119,23 @@
 {%- endmacro -%}
 
 {% macro alter_set_non_null_constraint(relation, column) -%}
-  ALTER {{ relation.type }} {{ relation.render() }} ALTER COLUMN {{ column }} SET NOT NULL;
+  ALTER {{ relation.type.render() }} {{ relation.render() }} ALTER COLUMN {{ column }} SET NOT NULL;
 {%- endmacro -%}
 
 {% macro alter_unset_non_null_constraint(relation, column) -%}
-  ALTER {{ relation.type }} {{ relation.render() }} ALTER COLUMN {{ column }} DROP NOT NULL;
+  ALTER {{ relation.type.render() }} {{ relation.render() }} ALTER COLUMN {{ column }} DROP NOT NULL;
 {%- endmacro -%}
 
 {% macro alter_set_constraint(relation, constraint) -%}
-  ALTER {{ relation.type }} {{ relation.render() }} ADD {{ constraint.render() }};
+  ALTER {{ relation.type.render() }} {{ relation.render() }} ADD {{ constraint.render() }};
 {%- endmacro -%}
 
 {% macro alter_unset_constraint(relation, constraint) -%}
   {% set constraint_type = constraint.type %}
   {% if constraint_type == 'primary_key' %}
     {# Need to only add CASCADE to PK constraints because dropping check constraints break when adding CASCADE #}
-    ALTER {{ relation.type }} {{ relation.render() }} DROP CONSTRAINT {{ constraint.name }} CASCADE;
+    ALTER {{ relation.type.render() }} {{ relation.render() }} DROP CONSTRAINT {{ constraint.name }} CASCADE;
   {% else %}
-    ALTER {{ relation.type }} {{ relation.render() }} DROP CONSTRAINT IF EXISTS {{ constraint.name }};
+    ALTER {{ relation.type.render() }} {{ relation.render() }} DROP CONSTRAINT IF EXISTS {{ constraint.name }};
   {% endif %}
 {%- endmacro -%}
diff --git a/dbt/include/databricks/macros/relations/components/query.sql b/dbt/include/databricks/macros/relations/components/query.sql
@@ -8,7 +8,7 @@
 {% endmacro %}
 
 {% macro get_alter_query_sql(target_relation, query) -%}
-  ALTER {{ target_relation.type|upper }} {{ target_relation.render() }} AS (
+  ALTER {{ target_relation.type.render()|upper }} {{ target_relation.render() }} AS (
     {{ query }}
   )
 {%- endmacro %}
diff --git a/dbt/include/databricks/macros/relations/liquid_clustering.sql b/dbt/include/databricks/macros/relations/liquid_clustering.sql
@@ -16,15 +16,15 @@
   {%- set auto_cluster = liquid_clustering.auto_cluster -%}
   {%- if cols and cols != [] %}
     {%- call statement('set_cluster_by_columns') -%}
-      ALTER {{ target_relation.type }} {{ target_relation.render() }} CLUSTER BY ({{ cols | join(', ') }})
+      ALTER {{ target_relation.type.render() }} {{ target_relation.render() }} CLUSTER BY ({{ cols | join(', ') }})
     {%- endcall -%}
   {%- elif auto_cluster -%}
     {%- call statement('set_cluster_by_auto') -%}
-      ALTER {{ target_relation.type }} {{ target_relation.render() }} CLUSTER BY AUTO
+      ALTER {{ target_relation.type.render() }} {{ target_relation.render() }} CLUSTER BY AUTO
     {%- endcall -%}
   {% else %}
     {%- call statement('unset_cluster_by') -%}
-      ALTER {{ target_relation.type }} {{ target_relation.render() }} CLUSTER BY NONE
+      ALTER {{ target_relation.type.render() }} {{ target_relation.render() }} CLUSTER BY NONE
     {%- endcall -%}
   {%- endif %}
 {%- endmacro -%}
diff --git a/dbt/include/databricks/macros/relations/tags.sql b/dbt/include/databricks/macros/relations/tags.sql
@@ -29,7 +29,7 @@
 {%- endmacro -%}
 
 {% macro alter_set_tags(relation, tags) -%}
-  ALTER {{ relation.type }} {{ relation.render() }} SET TAGS (
+  ALTER {{ relation.type.render() }} {{ relation.render() }} SET TAGS (
     {% for tag in tags -%}
       '{{ tag }}' = '{{ tags[tag] }}' {%- if not loop.last %}, {% endif -%}
     {%- endfor %}
diff --git a/dbt/include/databricks/macros/relations/tblproperties.sql b/dbt/include/databricks/macros/relations/tblproperties.sql
@@ -17,7 +17,7 @@
   {% set tblproperty_statment = databricks__tblproperties_clause(tblproperties) %}
   {% if tblproperty_statment %}
     {%- call statement('main') -%}
-      ALTER {{ relation.type }} {{ relation.render() }} SET {{ tblproperty_statment}}
+      ALTER {{ relation.type.render() }} {{ relation.render() }} SET {{ tblproperty_statment}}
     {%- endcall -%}
   {% endif %}
 {%- endmacro -%}
diff --git a/tests/functional/adapter/tags/fixtures.py b/tests/functional/adapter/tags/fixtures.py
@@ -7,6 +7,15 @@
 select cast(1 as bigint) as id, 'hello' as msg, 'blue' as color
 """
 
+streaming_table_tags_sql = """
+{{ config(
+    materialized='streaming_table',
+    databricks_tags = {'a': 'b', 'c': 'd'},
+) }}
+
+select * from stream {{ ref('my_seed') }}
+"""
+
 simple_python_model = """
 import pandas
 
diff --git a/tests/functional/adapter/tags/test_databricks_tags.py b/tests/functional/adapter/tags/test_databricks_tags.py
@@ -1,6 +1,7 @@
 import pytest
 
 from dbt.tests import util
+from dbt.tests.adapter.materialized_view.files import MY_SEED
 from tests.functional.adapter.tags import fixtures
 
 
@@ -37,6 +38,34 @@ def models(self):
         return {"tags.sql": fixtures.tags_sql.replace("table", "incremental")}
 
 
+@pytest.mark.dlt
+@pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster")
+class TestMaterializedViewTags(TestTags):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "tags.sql": fixtures.tags_sql.replace("table", "materialized_view"),
+        }
+
+
+@pytest.mark.dlt
+@pytest.mark.skip_profile("databricks_cluster", "databricks_uc_cluster")
+class TestStreamingTableTags(TestTags):
+    @pytest.fixture(scope="class")
+    def seeds(self):
+        return {"my_seed.csv": MY_SEED}
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "tags.sql": fixtures.streaming_table_tags_sql,
+        }
+
+    def test_tags(self, project):
+        util.run_dbt(["seed"])
+        super().test_tags(project)
+
+
 @pytest.mark.python
 @pytest.mark.skip_profile("databricks_cluster")
 class TestPythonTags(TestTags):

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@`
`8`	`8`	`{% endmacro %}`
`9`	`9`
`10`	`10`	`{% macro get_alter_query_sql(target_relation, query) -%}`
`11`		`- ALTER {{ target_relation.type\|upper }} {{ target_relation.render() }} AS (`
	`11`	`+ ALTER {{ target_relation.type.render()\|upper }} {{ target_relation.render() }} AS (`
`12`	`12`	`{{ query }}`
`13`	`13`	`)`
`14`	`14`	`{%- endmacro %}`