Fix disable_samples to exclude columns instead of skipping sampling

- Remove early return that skipped sampling entirely for disable_samples
- Create unified get_columns_to_exclude_from_sampling helper macro
- Integrate disable_samples with existing PII column exclusion logic
- Update integration tests to verify column exclusion behavior
- Add test for interaction between PII and disable_samples columns
- Maintain backward compatibility with existing functionality

Addresses GitHub feedback on PR #833

Co-Authored-By: Yosef Arbiv <[email protected]>

Author: devin-ai-integration[bot]

integration_tests/tests/test_column_pii_sampling.py

@@ -126,7 +126,7 @@ def test_column_pii_sampling_all_columns_pii(test_id: str, dbt_project: DbtProje
 
     assert len(samples) == TEST_SAMPLE_ROW_COUNT
     for sample in samples:
-        assert "_no_non_pii_columns" in sample
-        assert sample["_no_non_pii_columns"] == 1
+        assert "_no_non_excluded_columns" in sample
+        assert sample["_no_non_excluded_columns"] == 1
         assert SENSITIVE_COLUMN not in sample
         assert SAFE_COLUMN not in sample

integration_tests/tests/test_disable_samples_config.py

@@ -51,7 +51,11 @@ def test_disable_samples_config_prevents_sampling(
         json.loads(row["result_row"])
         for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
     ]
-    assert len(samples) == 0
+    assert len(samples) == 5
+    for sample in samples:
+        assert "_no_non_excluded_columns" in sample
+        assert sample["_no_non_excluded_columns"] == 1
+        assert COLUMN_NAME not in sample
 
 
 @pytest.mark.skip_targets(["clickhouse"])
@@ -84,7 +88,9 @@ def test_disable_samples_false_allows_sampling(test_id: str, dbt_project: DbtPro
         for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
     ]
     assert len(samples) == 5
-    assert all([row == {COLUMN_NAME: None} for row in samples])
+    for sample in samples:
+        assert COLUMN_NAME in sample
+        assert sample[COLUMN_NAME] is None
 
 
 @pytest.mark.skip_targets(["clickhouse"])
@@ -119,4 +125,84 @@ def test_disable_samples_config_overrides_pii_tags(
         json.loads(row["result_row"])
         for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
     ]
-    assert len(samples) == 0
+    assert len(samples) == 5
+    for sample in samples:
+        assert "_no_non_excluded_columns" in sample
+        assert sample["_no_non_excluded_columns"] == 1
+        assert COLUMN_NAME not in sample
+
+
+@pytest.mark.skip_targets(["clickhouse"])
+def test_disable_samples_and_pii_interaction(test_id: str, dbt_project: DbtProject):
+    """Test that disable_samples and PII columns both get excluded"""
+    data = [
+        {"col1": None, "col2": f"pii{i}", "col3": f"disabled{i}"} for i in range(10)
+    ]
+
+    columns = [
+        {"name": "col1", "tests": [{"not_null": {}}]},
+        {"name": "col2", "config": {"tags": ["pii"]}},
+        {"name": "col3", "config": {"disable_samples": True}},
+    ]
+
+    test_result = dbt_project.test(
+        test_id,
+        "not_null",
+        columns=columns,
+        data=data,
+        test_vars={
+            "enable_elementary_test_materialization": True,
+            "test_sample_row_count": 5,
+            "disable_samples_on_pii_columns": True,
+            "pii_column_tags": ["pii"],
+        },
+    )
+    assert test_result["status"] == "fail"
+
+    samples = [
+        json.loads(row["result_row"])
+        for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
+    ]
+
+    assert len(samples) == 5
+    for sample in samples:
+        assert "col1" in sample
+        assert "col2" not in sample
+        assert "col3" not in sample
+
+
+@pytest.mark.skip_targets(["clickhouse"])
+def test_disable_samples_with_multiple_columns(test_id: str, dbt_project: DbtProject):
+    """Test that disable_samples excludes only the disabled column"""
+    data = [{"col1": None, "col2": f"value{i}"} for i in range(10)]
+
+    columns = [
+        {
+            "name": "col1",
+            "config": {"disable_samples": True},
+            "tests": [{"not_null": {}}],
+        },
+        {"name": "col2"},
+    ]
+
+    test_result = dbt_project.test(
+        test_id,
+        "not_null",
+        columns=columns,
+        data=data,
+        test_vars={
+            "enable_elementary_test_materialization": True,
+            "test_sample_row_count": 5,
+        },
+    )
+    assert test_result["status"] == "fail"
+
+    samples = [
+        json.loads(row["result_row"])
+        for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
+    ]
+
+    assert len(samples) == 5
+    for sample in samples:
+        assert "col1" not in sample
+        assert "col2" in sample

macros/edr/materializations/test/test.sql

@@ -113,18 +113,10 @@
     {% do return([]) %}
   {% endif %}
 
-  {% if flattened_test and elementary.is_sampling_disabled_for_column(flattened_test) %}
-    {% do elementary.debug_log("Skipping sample query because disable_samples is true for this column.") %}
-    {% do return([]) %}
-  {% endif %}
-  
-  {% set pii_columns = [] %}
-  {% if flattened_test %}
-    {% set pii_columns = elementary.get_pii_columns_from_parent_model(flattened_test) %}
-  {% endif %}
+  {% set columns_to_exclude = elementary.get_columns_to_exclude_from_sampling(flattened_test) %}
 
   {% set select_clause = "*" %}
-  {% if pii_columns %}
+  {% if columns_to_exclude %}
     {% set query_to_get_columns %}
       with test_results as (
         {{ sql }}
@@ -133,11 +125,11 @@
     {% endset %}
     {% set columns_result = elementary.run_query(query_to_get_columns) %}
     {% set all_columns = columns_result.column_names %}
-    {% set safe_columns = all_columns | reject("in", pii_columns) | list %}
+    {% set safe_columns = all_columns | reject("in", columns_to_exclude) | list %}
     {% if safe_columns %}
       {% set select_clause = safe_columns | join(", ") %}
     {% else %}
-      {% set select_clause = "1 as _no_non_pii_columns" %}
+      {% set select_clause = "1 as _no_non_excluded_columns" %}
     {% endif %}
   {% endif %}
 
@@ -150,6 +142,26 @@
   {% do return(elementary.agate_to_dicts(elementary.run_query(query))) %}
 {% endmacro %}
 
+{% macro get_columns_to_exclude_from_sampling(flattened_test) %}
+  {% set columns_to_exclude = [] %}
+  
+  {% if not flattened_test %}
+    {% do return(columns_to_exclude) %}
+  {% endif %}
+  
+  {% set pii_columns = elementary.get_pii_columns_from_parent_model(flattened_test) %}
+  {% set columns_to_exclude = columns_to_exclude + pii_columns %}
+  
+  {% if elementary.is_sampling_disabled_for_column(flattened_test) %}
+    {% set test_column_name = elementary.insensitive_get_dict_value(flattened_test, 'test_column_name') %}
+    {% if test_column_name and test_column_name not in columns_to_exclude %}
+      {% do columns_to_exclude.append(test_column_name) %}
+    {% endif %}
+  {% endif %}
+  
+  {% do return(columns_to_exclude) %}
+{% endmacro %}
+
 {% macro is_sampling_disabled_for_column(flattened_test) %}
   {% set test_column_name = elementary.insensitive_get_dict_value(flattened_test, 'test_column_name') %}
   {% set parent_model_unique_id = elementary.insensitive_get_dict_value(flattened_test, 'parent_model_unique_id') %}