Support data_tests

Signed-off-by: Yu Ishikawa <[email protected]>
ubie-oss · Jul 31, 2024 · aeee514 · aeee514
1 parent 8d3dfbe
commit aeee514
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/integration_tests/macros/tests/codegen/test_generate_secured_model_schema_v2.sql b/integration_tests/macros/tests/codegen/test_generate_secured_model_schema_v2.sql
@@ -17,6 +17,11 @@
           "meta": {
             "data_privacy": {
               "level": "internal",
+              "test_project__test_dataset__test_table": {
+                "data_tests": [
+                  "unique"
+                ],
+              },
             },
           },
         },
@@ -109,13 +114,15 @@ models:
         meta:
           data_privacy:
             level: internal
+        data_tests:
+          - unique
       - name: user_id
         description: |-
           User ID
         meta:
           data_privacy:
             level: internal
-        tests:
+        data_tests:
           - not_null
 {%- endraw -%}
   {%- endset %}

diff --git a/macros/codegen/generate_secured_model_schema_v2.sql b/macros/codegen/generate_secured_model_schema_v2.sql
@@ -76,12 +76,18 @@ models:
             {#- Think of the downgraded data security level #}
             level: {{ column.meta.data_privacy.level }}
         {%- endif %}
+        {#- We support both `data_tests` and `tests` for a while -#}
+        {#- TODO: remove `tests` after dbt no longer supports it -#}
         {%- if 'data_privacy' in column.meta
             and name in column.meta.data_privacy
-            and 'tests' in column.meta.data_privacy[name]
-            and column.meta.data_privacy[name].tests | length > 0 %}
-        tests: {%- for test in column.meta.data_privacy[name].tests %}
-          - {{ test }}
+            and (
+              column.meta.data_privacy[name].get('data_tests', []) | length > 0
+              or column.meta.data_privacy[name].get('tests', []) | length > 0
+            ) %}
+        {%- set data_tests = column.meta.data_privacy[name].get('data_tests', [])
+            + column.meta.data_privacy[name].get('tests', []) %}
+        data_tests: {%- for data_test in data_tests %}
+          - {{ data_test }}
         {%- endfor %}
         {%- endif %}
     {%- endfor %}