Support RECORD (#59)

* Support `RECORD` Signed-off-by: Yu Ishikawa <[email protected]> * Improve the test case Signed-off-by: Yu Ishikawa <[email protected]> --------- Signed-off-by: Yu Ishikawa <[email protected]>
ubie-oss · Nov 8, 2023 · 1116f04 · 1116f04
1 parent 391813d
commit 1116f04
Show file tree

Hide file tree

Showing 7 changed files with 15 additions and 6 deletions.
diff --git a/integration_tests/macros/tests/codegen/test_generate_privacy_protected_model_sql.sql b/integration_tests/macros/tests/codegen/test_generate_privacy_protected_model_sql.sql
@@ -77,6 +77,14 @@
               "level": "restricted",
             },
           },
+        },
+        "dummy_record": {
+          "data_type": "RECORD",
+          "meta": {
+            "data_privacy": {
+              "level": "internal",
+            },
+          },
         }
       },
       where="1 = 1",
@@ -141,6 +149,7 @@ WITH privacy_protected_model AS (
     ) AS `consents`,
     dummy_column AS `dummy_column`,
     dummy_array AS `dummy_array`,
+    dummy_record AS `dummy_record`,
   FROM
     {{ ref('test_restricted_users') }} AS __original_table
   WHERE

diff --git a/macros/pseudonymization/sha256.sql b/macros/pseudonymization/sha256.sql
@@ -5,7 +5,7 @@
 {%- macro bigquery__sha256(expression, data_type=none) -%}
   {% set secured_expression = "SHA256(CAST({} AS STRING))".format(expression)%}
 
-  {% if data_type | upper == "ARRAY" %}
+  {% if data_type | upper in ["ARRAY", "RECORD"] %}
     {% set secured_expression = "ARRAY(SELECT SHA256(CAST(e AS STRING)) FROM UNNEST({}) AS e)".format(expression)%}
   {% endif %}
 

diff --git a/macros/pseudonymization/sha512.sql b/macros/pseudonymization/sha512.sql
@@ -5,7 +5,7 @@
 {%- macro bigquery__sha512(expression, data_type=none) -%}
   {% set secured_expression = "SHA512(CAST({} AS STRING))".format(expression)%}
 
-  {% if data_type | upper == "ARRAY" %}
+  {% if data_type | upper in ["ARRAY", "RECORD"] %}
     {% set secured_expression = "ARRAY(SELECT SHA512(CAST(e AS STRING)) FROM UNNEST({}) AS e)".format(expression)%}
   {% endif %}
 

diff --git a/macros/pseudonymization/udf_hash.sql b/macros/pseudonymization/udf_hash.sql
@@ -9,7 +9,7 @@
 {%- macro bigquery__udf_hash(expression, udf_function, data_type=none) -%}
   {% set secured_expression = "{}(CAST({} AS STRING))".format(udf_function, expression)%}
 
-  {% if data_type | upper == "ARRAY" %}
+  {% if data_type | upper in ["ARRAY", "RECORD"] %}
     {% set secured_expression = "ARRAY(SELECT {}(CAST(e AS STRING)) FROM UNNEST({}) AS e)".format(udf_function, expression)%}
   {% endif %}
 

diff --git a/macros/utils/expressions/convert_restructured_column_to_expression.sql b/macros/utils/expressions/convert_restructured_column_to_expression.sql
@@ -20,7 +20,7 @@
   {% set is_array = false %}
   {% if restructured_column.original_info is defined
       and restructured_column.original_info.data_type is defined
-      and restructured_column.original_info.data_type == "ARRAY" %}
+      and restructured_column.original_info.data_type in ["ARRAY", "RECORD"] %}
     {% set is_array = true %}
   {% endif %}
 

diff --git a/macros/utils/expressions/flatten_restructured_column_for_schema.sql b/macros/utils/expressions/flatten_restructured_column_for_schema.sql
@@ -17,7 +17,7 @@
   {% set is_array = false %}
   {% if restructured_column.original_info is defined
       and restructured_column.original_info.data_type is defined
-      and restructured_column.original_info.data_type == "ARRAY" %}
+      and restructured_column.original_info.data_type in ["ARRAY", "RECORD"] %}
     {% set is_array = true %}
   {% endif %}
 

diff --git a/macros/utils/expressions/get_secured_restructured_column.sql b/macros/utils/expressions/get_secured_restructured_column.sql
@@ -26,7 +26,7 @@
   {% if restructured_column.original_info is defined
       and restructured_column.original_info is mapping
       and restructured_column.original_info.data_type is defined
-      and restructured_column.original_info.data_type == "ARRAY" %}
+      and restructured_column.original_info.data_type in ["ARRAY", "RECORD"] %}
     {% set is_array = true %}
   {% endif %}