Add column mask support

Author: ericj-db

dbt/adapters/databricks/impl.py

@@ -984,6 +984,7 @@ def _describe_relation(
             results["foreign_key_constraints"] = adapter.execute_macro(
                 "fetch_foreign_key_constraints", kwargs=kwargs
             )
+            results["column_masks"] = adapter.execute_macro("fetch_column_masks", kwargs=kwargs)
         results["show_tblproperties"] = adapter.execute_macro("fetch_tbl_properties", kwargs=kwargs)
 
         kwargs = {"table_name": relation}

dbt/adapters/databricks/relation_configs/column_mask.py

@@ -0,0 +1,66 @@
+from dataclasses import asdict
+from typing import ClassVar, Optional
+
+from dbt.adapters.contracts.relation import RelationConfig
+from dbt.adapters.databricks.relation_configs.base import (
+    DatabricksComponentConfig,
+    DatabricksComponentProcessor,
+)
+from dbt.adapters.relation_configs.config_base import RelationResults
+
+
+class ColumnMaskConfig(DatabricksComponentConfig):
+    # column name -> mask
+    set_column_masks: dict[str, str]
+    unset_column_masks: list[str] = []
+
+    def get_diff(self, other: "ColumnMaskConfig") -> Optional["ColumnMaskConfig"]:
+        # Find column masks that need to be unset
+        unset_column_mask = [
+            col for col in other.set_column_masks if col not in self.set_column_masks
+        ]
+
+        # Find column masks that need to be set or updated
+        set_column_mask = {
+            col: mask
+            for col, mask in self.set_column_masks.items()
+            if col not in other.set_column_masks or other.set_column_masks[col] != mask
+        }
+
+        if set_column_mask or unset_column_mask:
+            return ColumnMaskConfig(
+                set_column_masks=set_column_mask,
+                unset_column_masks=unset_column_mask,
+            )
+        return None
+
+
+class ColumnMaskProcessor(DatabricksComponentProcessor[ColumnMaskConfig]):
+    name: ClassVar[str] = "column_masks"
+
+    @classmethod
+    def from_relation_results(cls, results: RelationResults) -> ColumnMaskConfig:
+        column_masks = results.get("column_masks")
+        set_column_masks = {}
+
+        if column_masks:
+            for row in column_masks.rows:
+                set_column_masks[row[0]] = row[1]
+
+        return ColumnMaskConfig(set_column_masks=set_column_masks)
+
+    @classmethod
+    def from_relation_config(cls, relation_config: RelationConfig) -> ColumnMaskConfig:
+        # Extract config from model node
+        columns = getattr(relation_config, "columns", {})
+        columns = [
+            {"name": name, **(col if isinstance(col, dict) else asdict(col))}
+            for name, col in columns.items()
+        ]
+
+        set_column_masks = {}
+        for col in columns:
+            extra = col.get("_extra", {})
+            if extra and "column_mask" in extra:
+                set_column_masks[col["name"]] = extra["column_mask"]
+        return ColumnMaskConfig(set_column_masks=set_column_masks)

dbt/adapters/databricks/relation_configs/incremental.py

@@ -2,6 +2,7 @@
     DatabricksRelationConfigBase,
 )
 from dbt.adapters.databricks.relation_configs.column_comments import ColumnCommentsProcessor
+from dbt.adapters.databricks.relation_configs.column_mask import ColumnMaskProcessor
 from dbt.adapters.databricks.relation_configs.comment import CommentProcessor
 from dbt.adapters.databricks.relation_configs.constraints import ConstraintsProcessor
 from dbt.adapters.databricks.relation_configs.liquid_clustering import LiquidClusteringProcessor
@@ -13,6 +14,7 @@ class IncrementalTableConfig(DatabricksRelationConfigBase):
     config_components = [
         CommentProcessor,
         ColumnCommentsProcessor,
+        ColumnMaskProcessor,
         ConstraintsProcessor,
         TagsProcessor,
         TblPropertiesProcessor,

dbt/include/databricks/macros/relations/components/column_mask.sql

@@ -0,0 +1,69 @@
+{% macro fetch_column_masks(relation) -%}
+  {% if relation.is_hive_metastore() %}
+    {{ exceptions.raise_compiler_error("Column masks are not supported for Hive Metastore") }}
+  {%- endif %}
+  {% call statement('list_column_masks', fetch_result=True) -%}
+    {{ fetch_column_masks_sql(relation) }}
+  {% endcall %}
+  {% do return(load_result('list_column_masks').table) %}
+{%- endmacro -%}
+
+{% macro fetch_column_masks_sql(relation) -%}
+  SELECT 
+    column_name,
+    mask_name
+  FROM `{{ relation.database|lower }}`.information_schema.column_masks
+  WHERE table_catalog = '{{ relation.database|lower }}'
+    AND table_schema = '{{ relation.schema|lower }}'
+    AND table_name = '{{ relation.identifier|lower }}';
+{%- endmacro -%}
+
+{% macro apply_column_masks_from_model_columns(relation) -%}
+  {% if relation.is_hive_metastore() %}
+    {{ exceptions.raise_compiler_error("Column masks are not supported for Hive Metastore") }}
+  {%- endif %}
+  {{ log("Applying column masks from model to relation " ~ relation) }}
+  {% set columns = model.get('columns', {}) %}
+  {% for column_name, column_def in columns.items() %}
+    {% if column_def is mapping and column_def.get('column_mask') %}
+      {%- call statement('main') -%}
+        {{ alter_set_column_mask(relation, column_name, column_def.column_mask) }}
+      {%- endcall -%}
+    {% endif %}
+  {% endfor %}
+{%- endmacro -%}
+
+{% macro apply_column_masks(relation, column_masks) -%}
+  {% if relation.is_hive_metastore() %}
+    {{ exceptions.raise_compiler_error("Column masks are not supported for Hive Metastore") }}
+  {%- endif %}
+  {{ log("Applying column masks to relation " ~ relation) }}
+  {%- if column_masks.unset_column_mask %}
+    {%- for column in column_masks.unset_column_mask -%}
+      {%- call statement('main') -%}
+        {{ alter_drop_column_mask(relation, column) }}
+      {%- endcall -%}
+    {%- endfor -%}
+  {%- endif %}
+  {%- if column_masks.set_column_mask %}
+    {%- for column, mask in column_masks.set_column_mask.items() -%}
+      {%- call statement('main') -%}
+        {{ alter_set_column_mask(relation, column, mask) }}
+      {%- endcall -%}
+    {%- endfor -%}
+  {%- endif %}
+{%- endmacro -%}
+
+{% macro alter_drop_column_mask(relation, column) -%}
+  ALTER TABLE {{ relation.render() }}
+  ALTER COLUMN {{ column }}
+  DROP MASK;
+{%- endmacro -%}
+
+{% macro alter_set_column_mask(relation, column, mask) -%}
+  ALTER TABLE {{ relation.render() }}
+  ALTER COLUMN {{ column }}
+  SET MASK {{ mask }};
+{%- endmacro -%}
+
+

dbt/include/databricks/macros/relations/table/alter.sql

@@ -7,6 +7,7 @@
       {% set tblproperties = configuration_changes.changes.get("tblproperties") %}
       {% set liquid_clustering = configuration_changes.changes.get("liquid_clustering")%}
       {% set constraints = configuration_changes.changes.get("constraints") %}
+      {% set column_masks = configuration_changes.changes.get("column_masks") %}
       {% if tags is not none %}
         {% do apply_tags(target_relation, tags.set_tags, tags.unset_tags) %}
       {%- endif -%}
@@ -25,5 +26,8 @@
       {% if constraints %}
         {{ apply_constraints(target_relation, constraints) }}
       {% endif %}
+      {% if column_masks %}
+        {{ apply_column_masks(target_relation, column_masks) }}
+      {% endif %}
     {%- endif -%}
 {% endmacro %}

dbt/include/databricks/macros/relations/table/create.sql

@@ -12,6 +12,7 @@
 
   {{ apply_alter_constraints(target_relation) }}
   {{ apply_tags(target_relation, tags) }}
+  {{ apply_column_masks_from_model_columns(target_relation) }}
 
   {% call statement('merge into target') %}
     insert into {{ target_relation }} select * from {{ intermediate_relation }}

tests/functional/adapter/column_masks/fixtures.py

@@ -0,0 +1,19 @@
+base_model_sql = """
+{{ config(
+    materialized = 'table'
+) }}
+SELECT 'abc-123' as id, 'password123' as password;
+"""
+
+
+model = """
+version: 2
+models:
+  - name: base_model
+    columns:
+      - name: id
+        data_type: string
+      - name: password
+        column_mask: password_mask
+        data_type: string
+"""

tests/functional/adapter/column_masks/test_column_mask.py

@@ -0,0 +1,53 @@
+import pytest
+
+from dbt.tests.util import run_dbt
+from tests.functional.adapter.column_masks.fixtures import (
+    base_model_sql,
+    model,
+)
+from tests.functional.adapter.fixtures import MaterializationV2Mixin
+
+
+class TestColumnMask(MaterializationV2Mixin):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "base_model.sql": base_model_sql,
+            "schema.yml": model,
+        }
+
+    def test_column_mask(self, project):
+        # Create the mask function
+        project.run_sql(
+            f"CREATE OR REPLACE FUNCTION {project.database}.{project.test_schema}."
+            "password_mask(password STRING) RETURNS STRING RETURN '*****';"
+        )
+
+        run_dbt(["run"])
+
+        # Verify column mask was created
+        masks = project.run_sql(
+            f"""
+            SELECT column_name, mask_name
+            FROM {project.database}.information_schema.column_masks
+            """,
+            fetch="all",
+        )
+
+        assert len(masks) == 1
+        assert masks[0][0] == "password"  # column_name
+        assert masks[0][1] == f"{project.database}.{project.test_schema}.password_mask"  # mask_name
+
+        # Verify masked value
+        result = project.run_sql("SELECT id, password FROM base_model", fetch="one")
+        assert result[0] == "abc-123"
+        assert result[1] == "*****"  # Masked value should be 5 asterisks
+
+
+class TestIncrementalColumnMask(TestColumnMask):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "base_model.sql": base_model_sql.replace("table", "incremental"),
+            "schema.yml": model,
+        }

tests/functional/adapter/incremental/fixtures.py

@@ -850,3 +850,43 @@ def model(dbt, spark):
         to: ref('fk_referenced_to_table')
         to_columns: [id, version]
 """
+
+
+column_mask_sql = """
+{{ config(
+    materialized = 'incremental',
+    incremental_strategy = 'merge',
+    unique_key = 'id',
+) }}
+
+select cast(1 as bigint) as id, 'hello' as name, 'john.doe@example.com' as email,
+'password123' as password
+"""
+
+column_mask_name = """
+version: 2
+
+models:
+  - name: column_mask_sql
+    columns:
+        - name: id
+        - name: name
+          column_mask: full_mask
+        - name: email
+          column_mask: full_mask
+        - name: password
+"""
+
+column_mask_password = """
+version: 2
+
+models:
+  - name: column_mask_sql
+    columns:
+        - name: id
+        - name: name
+        - name: email
+          column_mask: email_mask
+        - name: password
+          column_mask: full_mask
+"""

tests/functional/adapter/incremental/test_incremental_column_masks.py

@@ -0,0 +1,63 @@
+import pytest
+
+from dbt.tests import util
+from tests.functional.adapter.fixtures import MaterializationV2Mixin
+from tests.functional.adapter.incremental import fixtures
+
+
+@pytest.mark.skip_profile("databricks_cluster")
+class TestIncrementalColumnMasks(MaterializationV2Mixin):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "column_mask_sql.sql": fixtures.column_mask_sql,
+            "schema.yml": fixtures.column_mask_name,
+        }
+
+    def test_changing_column_masks(self, project):
+        # Create the mask functions
+        project.run_sql(
+            f"""
+            CREATE OR REPLACE FUNCTION
+                {project.database}.{project.test_schema}.full_mask(value STRING)
+            RETURNS STRING
+            RETURN '*****';
+            """
+        )
+        # Masks all characters before the @ symbol
+        project.run_sql(
+            f"""
+            CREATE OR REPLACE FUNCTION
+                {project.database}.{project.test_schema}.email_mask(value STRING)
+            RETURNS STRING
+            RETURN CONCAT(
+                REPEAT('*', POSITION('@' IN value) - 1),
+                SUBSTR(value, POSITION('@' IN value))
+            );
+            """
+        )
+
+        # First run with name masked
+        util.run_dbt(["run"])
+        masks = project.run_sql(
+            "SELECT id, name, email, password FROM column_mask_sql",
+            fetch="all",
+        )
+        assert len(masks) == 1
+        assert masks[0][0] == 1
+        assert masks[0][1] == "*****"  # name (masked)
+        assert masks[0][2] == "*****"  # email (masked)
+        assert masks[0][3] == "password123"  # password (unmasked)
+
+        # Update masks and verify changes
+        util.write_file(fixtures.column_mask_password, "models", "schema.yml")
+        util.run_dbt(["run"])
+
+        result = project.run_sql(
+            "SELECT id, name, email, password FROM column_mask_sql", fetch="all"
+        )
+        assert len(result) == 1
+        assert result[0][0] == 1
+        assert result[0][1] == "hello"  # name (unmasked)
+        assert result[0][2] == "********@example.com"  # email (partially masked)
+        assert result[0][3] == "*****"  # password (masked)