Added athena support for upload_tests

brooklyn-data · Oct 2, 2024 · a13230a · a13230a
1 parent fb6207e
commit a13230a
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 3 deletions.
diff --git a/macros/upload_individual_datasets/upload_tests.sql b/macros/upload_individual_datasets/upload_tests.sql
@@ -41,6 +41,34 @@
     {% endif %}
 {% endmacro -%}
 
+{% macro athena__get_tests_dml_sql(tests) -%}
+    {% if tests != [] %}
+        {% set test_values %}
+            values
+            {% for test in tests -%}
+                (
+                    '{{ invocation_id }}', {# command_invocation_id #}
+                    '{{ test.unique_id }}', {# node_id #}
+                    CAST('{{ run_started_at }}' AS timestamp), {# run_started_at #}
+                    '{{ test.name }}', {# name #}
+                    '{{ tojson(test.depends_on.nodes) }}', {# depends_on_nodes #}
+                    '{{ test.package_name }}', {# package_name #}
+                    '{{ test.original_file_path | replace('\\', '\\\\') }}', {# test_path #}
+                    '{{ tojson(test.tags) }}', {# tags #}
+                    {% if var('dbt_artifacts_exclude_all_results', false) %}
+                        null
+                    {% else %}
+                        '{{ adapter.dispatch('parse_json', 'dbt_artifacts')(tojson(test) | replace("\\", "\\\\") | replace("'","''")) }}' {# all_fields #}
+                    {% endif %}
+                )
+                {%- if not loop.last %},{%- endif %}
+            {%- endfor %}
+        {% endset %}
+        {{ test_values }}
+    {% else %} {{ return("") }}
+    {% endif %}
+{%- endmacro %}
+
 {% macro bigquery__get_tests_dml_sql(tests) -%}
     {% if tests != [] %}
         {% set test_values %}

diff --git a/macros/upload_results/insert_into_metadata_table.sql b/macros/upload_results/insert_into_metadata_table.sql
@@ -18,8 +18,6 @@
     {{ content }}
     {% endset %}
 
-    {% do log(insert_into_table_query, true) %}
-
     {% do run_query(insert_into_table_query) %}
 
 {%- endmacro %}

diff --git a/macros/upload_results/upload_results.sql b/macros/upload_results/upload_results.sql
@@ -19,7 +19,7 @@
             {% set objects = dbt_artifacts.get_dataset_content(dataset) %}
 
             {# Upload in chunks to reduce the query size #}
-            {% if dataset == 'models' %}
+            {% if dataset in ('models', 'tests') %}
                 {% set upload_limit = 50 if target.type in ('athena', 'bigquery') else 100 %}
             {% else %}
                 {% set upload_limit = 300 if target.type in ('athena', 'bigquery') else 5000 %}