Support USING COLUMNS syntax

Author: ericj-db

dbt/adapters/databricks/relation_configs/column_mask.py

@@ -1,6 +1,8 @@
 from dataclasses import asdict
 from typing import ClassVar, Optional
 
+from dbt_common.exceptions import DbtRuntimeError
+
 from dbt.adapters.contracts.relation import RelationConfig
 from dbt.adapters.databricks.relation_configs.base import (
     DatabricksComponentConfig,
@@ -10,11 +12,27 @@
 
 
 class ColumnMaskConfig(DatabricksComponentConfig):
-    # column name -> mask
-    set_column_masks: dict[str, str]
+    # column name -> mask config (function name and optional using_columns)
+    set_column_masks: dict[str, dict[str, str]]
     unset_column_masks: list[str] = []
 
     def get_diff(self, other: "ColumnMaskConfig") -> Optional["ColumnMaskConfig"]:
+        # Create a mapping of function names to their using_columns values
+        function_using_columns = {}
+        for mask in self.set_column_masks.values():
+            function_using_columns[mask["function"]] = mask.get("using_columns")
+
+        # Check if any function's using_columns has changed
+        for mask in other.set_column_masks.values():
+            function = mask["function"]
+            if function in function_using_columns and function_using_columns[function] != mask.get(
+                "using_columns"
+            ):
+                raise DbtRuntimeError(
+                    f"The value of using_columns for existing function {function} was updated. "
+                    f"This is not supported. Please use a new function with a different name."
+                )
+
         # 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
@@ -45,7 +63,11 @@ def from_relation_results(cls, results: RelationResults) -> ColumnMaskConfig:
 
         if column_masks:
             for row in column_masks.rows:
-                set_column_masks[row[0]] = row[1]
+                # row contains [column_name, mask_name, using_columns]
+                mask_config = {"function": row[1]}
+                if row[2]:
+                    mask_config["using_columns"] = row[2]
+                set_column_masks[row[0]] = mask_config
 
         return ColumnMaskConfig(set_column_masks=set_column_masks)
 
@@ -61,6 +83,17 @@ def from_relation_config(cls, relation_config: RelationConfig) -> ColumnMaskConf
         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"]
+            column_mask = extra.get("column_mask") if extra else None
+            if column_mask:
+                fully_qualified_function_name = (
+                    column_mask["function"]
+                    if "." in column_mask["function"]
+                    else (
+                        f"{relation_config.database}."
+                        f"{relation_config.schema}."
+                        f"{column_mask['function']}"
+                    )
+                )
+                column_mask["function"] = fully_qualified_function_name
+                set_column_masks[col["name"]] = column_mask
         return ColumnMaskConfig(set_column_masks=set_column_masks)

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

@@ -11,7 +11,8 @@
 {% macro fetch_column_masks_sql(relation) -%}
   SELECT 
     column_name,
-    mask_name
+    mask_name,
+    using_columns
   FROM `{{ relation.database|lower }}`.information_schema.column_masks
   WHERE table_catalog = '{{ relation.database|lower }}'
     AND table_schema = '{{ relation.schema|lower }}'
@@ -63,7 +64,10 @@
 {% macro alter_set_column_mask(relation, column, mask) -%}
   ALTER TABLE {{ relation.render() }}
   ALTER COLUMN {{ column }}
-  SET MASK {{ mask }};
+  SET MASK {{ mask.function }}
+  {%- if mask.using_columns %}
+  USING COLUMNS ({{ mask.using_columns }})
+  {%- endif %};
 {%- endmacro -%}
 
 {% macro column_mask_exists() %}

tests/functional/adapter/column_masks/fixtures.py

@@ -2,7 +2,7 @@
 {{ config(
     materialized = 'table'
 ) }}
-SELECT 'abc-123' as id, 'password123' as password;
+SELECT 'abc123' as id, 'password123' as password;
 """
 
 
@@ -14,6 +14,21 @@
       - name: id
         data_type: string
       - name: password
-        column_mask: password_mask
+        column_mask:
+          function: password_mask
         data_type: string
 """
+
+model_with_extra_args = """
+version: 2
+models:
+  - name: base_model
+    columns:
+      - name: id
+        data_type: string
+      - name: password
+        data_type: string
+        column_mask:
+          function: weird_mask
+          using_columns: "id, 'literal_string', 333, true, null, INTERVAL 2 DAYS"
+"""

tests/functional/adapter/column_masks/test_column_mask.py

@@ -1,9 +1,10 @@
 import pytest
 
-from dbt.tests.util import run_dbt
+from dbt.tests.util import run_dbt, write_file
 from tests.functional.adapter.column_masks.fixtures import (
     base_model_sql,
     model,
+    model_with_extra_args,
 )
 from tests.functional.adapter.fixtures import MaterializationV2Mixin
 
@@ -17,7 +18,7 @@ def models(self):
             "schema.yml": model,
         }
 
-    def test_column_mask(self, project):
+    def test_column_mask_no_extra_args(self, project):
         # Create the mask function
         project.run_sql(
             f"CREATE OR REPLACE FUNCTION {project.database}.{project.test_schema}."
@@ -31,6 +32,7 @@ def test_column_mask(self, project):
             f"""
             SELECT column_name, mask_name
             FROM {project.database}.information_schema.column_masks
+            WHERE table_schema = '{project.test_schema}'
             """,
             fetch="all",
         )
@@ -41,9 +43,45 @@ def test_column_mask(self, project):
 
         # Verify masked value
         result = project.run_sql("SELECT id, password FROM base_model", fetch="one")
-        assert result[0] == "abc-123"
+        assert result[0] == "abc123"
         assert result[1] == "*****"  # Masked value should be 5 asterisks
 
+    def test_column_mask_with_extra_args(self, project):
+        write_file(model_with_extra_args, "models", "schema.yml")
+
+        # Create a mask function that concatenates all possible arg types: original column, other
+        # columns, and every supported literal type from
+        # https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-ddl-column-mask
+        project.run_sql(
+            f"""
+            CREATE OR REPLACE FUNCTION {project.database}.{project.test_schema}.weird_mask(
+                password STRING,
+                id STRING,
+                literal STRING,
+                num INT,
+                bool_val BOOLEAN,
+                null_val STRING,
+                interval INTERVAL DAY
+            )
+            RETURNS STRING
+            RETURN CONCAT(
+                password, '-',
+                id, '-',
+                literal, '-',
+                CAST(num AS STRING), '-',
+                CAST(bool_val AS STRING), '-',
+                COALESCE(null_val, 'NULL'), '-',
+                CAST(interval AS STRING)
+            );
+            """
+        )
+        run_dbt(["run"])
+
+        # Not meant to resemble a real life example. Just for the sake of testing different types
+        result = project.run_sql("SELECT id, password FROM base_model", fetch="one")
+        assert result[0] == "abc123"
+        assert result[1] == "password123-abc123-literal_string-333-true-NULL-INTERVAL '2' DAY"
+
 
 @pytest.mark.skip_profile("databricks_cluster")
 class TestIncrementalColumnMask(TestColumnMask):

tests/functional/adapter/incremental/fixtures.py

@@ -891,21 +891,23 @@ def model(dbt, spark):
 'password123' as password
 """
 
-column_mask_name = """
+column_mask_base = """
 version: 2
 
 models:
   - name: column_mask_sql
     columns:
         - name: id
         - name: name
-          column_mask: full_mask
+          column_mask:
+            function: full_mask
         - name: email
-          column_mask: full_mask
+          column_mask:
+            function: full_mask
         - name: password
 """
 
-column_mask_password = """
+column_mask_valid_mask_updates = """
 version: 2
 
 models:
@@ -914,7 +916,25 @@ def model(dbt, spark):
         - name: id
         - name: name
         - name: email
-          column_mask: email_mask
+          column_mask:
+            function: email_mask
         - name: password
-          column_mask: full_mask
+          column_mask:
+            function: full_mask
+"""
+
+column_mask_invalid_update = """
+version: 2
+
+models:
+  - name: column_mask_sql
+    columns:
+        - name: id
+        - name: name
+          column_mask:
+        - name: email
+        - name: password
+          column_mask:
+            function: full_mask
+            using_columns: "id"
 """

tests/functional/adapter/incremental/test_incremental_column_masks.py

@@ -11,15 +11,15 @@ class TestIncrementalColumnMasks(MaterializationV2Mixin):
     def models(self):
         return {
             "column_mask_sql.sql": fixtures.column_mask_sql,
-            "schema.yml": fixtures.column_mask_name,
+            "schema.yml": fixtures.column_mask_base,
         }
 
     def test_changing_column_masks(self, project):
-        # Create the mask functions
+        # Create the full mask function
         project.run_sql(
             f"""
             CREATE OR REPLACE FUNCTION
-                {project.database}.{project.test_schema}.full_mask(value STRING)
+                {project.database}.{project.test_schema}.full_mask(password STRING)
             RETURNS STRING
             RETURN '*****';
             """
@@ -50,7 +50,7 @@ def test_changing_column_masks(self, project):
         assert masks[0][3] == "password123"  # password (unmasked)
 
         # Update masks and verify changes
-        util.write_file(fixtures.column_mask_password, "models", "schema.yml")
+        util.write_file(fixtures.column_mask_valid_mask_updates, "models", "schema.yml")
         util.run_dbt(["run"])
 
         result = project.run_sql(
@@ -61,3 +61,35 @@ def test_changing_column_masks(self, project):
         assert result[0][1] == "hello"  # name (unmasked)
         assert result[0][2] == "********@example.com"  # email (partially masked)
         assert result[0][3] == "*****"  # password (masked)
+
+
+class TestInvalidColumnMaskUpdate(MaterializationV2Mixin):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "column_mask_sql.sql": fixtures.column_mask_sql,
+            "schema.yml": fixtures.column_mask_base,
+        }
+
+    def test_invalid_using_columns_update(self, project):
+        # Create the full mask function
+        project.run_sql(
+            f"""
+            CREATE OR REPLACE FUNCTION
+                {project.database}.{project.test_schema}.full_mask(password STRING)
+            RETURNS STRING
+            RETURN '*****';
+            """
+        )
+
+        # First run with name masked
+        util.run_dbt(["run"])
+
+        # Trying to feed new arguments to using_columns on existing function should fail
+        util.write_file(fixtures.column_mask_invalid_update, "models", "schema.yml")
+        results = util.run_dbt(["run"], expect_pass=False)
+        assert len(results.results) == 1
+        assert (
+            "This is not supported. Please use a new function with a different name."
+            in results.results[0].message
+        )