Fixed Athena temp tables cleaning (#714)

* - Added has_temp_table_support redefinition for Athena adapter - Added clean_up_tables to handle_tests_results and clean_elementary_test_tables to delete files from S3 * - Fix handle_tests_resuls data_monitoring_metrics cleaning up operation * Typo fix * Added adapter specific implementation for clean_up_tables and fully_drop_relation * Fixed handle_tests_results * Bugfix --------- Co-authored-by: Itamar Hartstein <[email protected]>
elementary-data · May 28, 2024 · b9de5f8 · b9de5f8
1 parent a59d552
commit b9de5f8
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 2 deletions.
diff --git a/macros/edr/tests/on_run_end/handle_tests_results.sql b/macros/edr/tests/on_run_end/handle_tests_results.sql
@@ -111,7 +111,7 @@
     {% do elementary.run_query(insert_query) %}
 
     {% if not elementary.has_temp_table_support() %}
-        {% do adapter.drop_relation(temp_relation) %}
+        {% do elementary.fully_drop_relation(temp_relation) %}
     {% endif %}
 {% endmacro %}
 
@@ -159,7 +159,7 @@
     {% do elementary.run_query(insert_query) %}
 
     {% if not elementary.has_temp_table_support() %}
-        {% do adapter.drop_relation(temp_relation) %}
+        {% do elementary.fully_drop_relation(temp_relation) %}
     {% endif %}
 {% endmacro %}
 

diff --git a/macros/edr/tests/test_utils/clean_elementary_test_tables.sql b/macros/edr/tests/test_utils/clean_elementary_test_tables.sql
@@ -8,6 +8,9 @@
             {% endfor %}
         {% endfor %}
 
+        {# Extra entry-point to clean up tables before dropping the relation #}
+        {% do elementary.clean_up_tables(test_table_relations) %}
+
         {% do elementary.file_log("Deleting temporary Elementary test tables: {}".format(test_table_relations)) %}
         {% set queries = elementary.get_clean_elementary_test_tables_queries(test_table_relations) %}
         {% for query in queries %}

diff --git a/macros/edr/tests/test_utils/clean_up_tables.sql b/macros/edr/tests/test_utils/clean_up_tables.sql
@@ -0,0 +1,13 @@
+{% macro clean_up_tables(relations) %}
+    {{ return(adapter.dispatch('clean_up_tables', 'elementary')(relations)) }}
+{% endmacro %}
+
+{% macro default__clean_up_tables(relations) %}
+    {# Default implementation does nothing #}
+{% endmacro %}
+
+{% macro athena__clean_up_tables(relations) %}
+    {% for relation in relations %}
+        {% do adapter.clean_up_table(relation) %}
+    {% endfor %}
+{% endmacro %}
diff --git a/macros/utils/table_operations/fully_drop_relation.sql b/macros/utils/table_operations/fully_drop_relation.sql
@@ -0,0 +1,12 @@
+{% macro fully_drop_relation(relation) %}
+    {{ return(adapter.dispatch('fully_drop_relation', 'elementary')(relation)) }}
+{% endmacro %}
+
+{% macro default__fully_drop_relation(relation) %}
+    {% do adapter.drop_relation(relation) %}
+{% endmacro %}
+
+{% macro athena__fully_drop_relation(relation) %}
+    {% do adapter.clean_up_table(relation) %}
+    {% do adapter.drop_relation(relation) %}
+{% endmacro %}
diff --git a/macros/utils/table_operations/has_temp_table_support.sql b/macros/utils/table_operations/has_temp_table_support.sql
@@ -14,3 +14,7 @@
     {% do return(false) %}
 {% endmacro %}
 
+{% macro athena__has_temp_table_support() %}
+    {% do return(false) %}
+{% endmacro %}
+