Add disable_samples configuration flag

- Add table-level disable_samples flag that overrides PII detection
- Flag is set via meta configuration: meta: {disable_samples: true}
- Works independently of PII tags and takes precedence
- Add comprehensive integration tests for all scenarios
- Addresses Linear issue ELE-4833 additional requirements

Co-Authored-By: Yosef Arbiv <yosef.arbiv@gmail.com>

Author: devin-ai-integration[bot]

integration_tests/tests/test_sampling_pii.py

@@ -108,3 +108,96 @@ def test_sampling_pii_feature_disabled(test_id: str, dbt_project: DbtProject):
         for row in dbt_project.run_query(SAMPLES_QUERY.format(test_id=test_id))
     ]
     assert len(samples) == TEST_SAMPLE_ROW_COUNT
+
+
+@pytest.mark.skip_targets(["clickhouse"])
+def test_sampling_disable_samples_flag(test_id: str, dbt_project: DbtProject):
+    """Test that disable_samples flag prevents sample collection regardless of PII tags"""
+    null_count = 50
+    data = [{COLUMN_NAME: None} for _ in range(null_count)]
+
+    test_result = dbt_project.test(
+        test_id,
+        "not_null",
+        dict(column_name=COLUMN_NAME),
+        data=data,
+        as_model=True,
+        model_config={
+            "config": {"meta": {"disable_samples": True}, "tags": ["normal"]}
+        },
+        test_vars={
+            "enable_elementary_test_materialization": True,
+            "test_sample_row_count": TEST_SAMPLE_ROW_COUNT,
+            "disable_samples_on_pii_tables": False,
+            "pii_table_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) == 0
+
+
+@pytest.mark.skip_targets(["clickhouse"])
+def test_sampling_disable_samples_overrides_pii(test_id: str, dbt_project: DbtProject):
+    """Test that disable_samples flag overrides PII detection when both are present"""
+    null_count = 50
+    data = [{COLUMN_NAME: None} for _ in range(null_count)]
+
+    test_result = dbt_project.test(
+        test_id,
+        "not_null",
+        dict(column_name=COLUMN_NAME),
+        data=data,
+        as_model=True,
+        model_config={"config": {"meta": {"disable_samples": True}, "tags": ["pii"]}},
+        test_vars={
+            "enable_elementary_test_materialization": True,
+            "test_sample_row_count": TEST_SAMPLE_ROW_COUNT,
+            "disable_samples_on_pii_tables": True,
+            "pii_table_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) == 0
+
+
+@pytest.mark.skip_targets(["clickhouse"])
+def test_sampling_disable_samples_false_allows_samples(
+    test_id: str, dbt_project: DbtProject
+):
+    """Test that disable_samples: false allows sample collection normally"""
+    null_count = 50
+    data = [{COLUMN_NAME: None} for _ in range(null_count)]
+
+    test_result = dbt_project.test(
+        test_id,
+        "not_null",
+        dict(column_name=COLUMN_NAME),
+        data=data,
+        as_model=True,
+        model_config={
+            "config": {"meta": {"disable_samples": False}, "tags": ["normal"]}
+        },
+        test_vars={
+            "enable_elementary_test_materialization": True,
+            "test_sample_row_count": TEST_SAMPLE_ROW_COUNT,
+            "disable_samples_on_pii_tables": False,
+            "pii_table_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) == TEST_SAMPLE_ROW_COUNT

macros/edr/materializations/test/test.sql

@@ -51,9 +51,14 @@
 {% macro handle_dbt_test(flattened_test, materialization_macro) %}
   {% set result = materialization_macro() %}
   {% set sample_limit = elementary.get_config_var('test_sample_row_count') %}
-  {% if elementary.is_pii_table(flattened_test) %}
+  
+  {% set disable_samples = elementary.insensitive_get_dict_value(flattened_test, 'meta', {}).get('disable_samples', false) %}
+  {% if disable_samples %}
+    {% set sample_limit = 0 %}
+  {% elif elementary.is_pii_table(flattened_test) %}
     {% set sample_limit = 0 %}
   {% endif %}
+  
   {% set result_rows = elementary.query_test_result_rows(sample_limit=sample_limit,
                                                          ignore_passed_tests=true) %}
   {% set elementary_test_results_row = elementary.get_dbt_test_result_row(flattened_test, result_rows) %}