fix: add Athena-specific edr_drop_schema and always-run schema cleanup step (#960)

Author: devin-ai-integration[bot]

.github/workflows/test-warehouse.yml

@@ -202,3 +202,14 @@ jobs:
         with:
           name: detailed_report_${{ inputs.warehouse-type }}_dbt_${{ inputs.dbt-version }}
           path: ${{ env.TESTS_DIR }}/tests/detailed_report_${{ inputs.warehouse-type }}_dbt_${{ inputs.dbt-version }}.html
+
+      - name: Drop test schemas
+        if: >-
+          always() &&
+          contains(fromJSON('["snowflake","bigquery","redshift","databricks_catalog","athena"]'), inputs.warehouse-type)
+        working-directory: ${{ env.TESTS_DIR }}
+        continue-on-error: true
+        run: |
+          ${{ (inputs.dbt-version == 'fusion' && '~/.local/bin/dbt') || 'dbt' }} run-operation elementary_tests.drop_test_schemas \
+            --project-dir dbt_project \
+            -t "${{ inputs.warehouse-type }}"

integration_tests/dbt_project/macros/clear_env.sql

@@ -8,6 +8,32 @@
     ) %}
 {% endmacro %}
 
+{% macro drop_test_schemas(num_workers=8) %}
+    {#
+      Drop every schema that a CI test run may have created.
+      This covers the base schema (no xdist suffix) as well as
+      each pytest-xdist worker schema (_gw0 … _gw<N-1>).
+      Called from the workflow with `if: always()` so that schemas
+      are cleaned up even when the pytest process is cancelled or
+      crashes before its own teardown runs.
+    #}
+    {% set database = elementary.target_database() %}
+    {% set base_schema = target.schema %}
+    {% set suffixes = [""] %}
+    {% for i in range(num_workers) %} {% do suffixes.append("_gw" ~ i) %} {% endfor %}
+
+    {% for suffix in suffixes %}
+        {% set test_schema = base_schema ~ suffix %}
+        {% set elementary_schema = base_schema ~ "_elementary" ~ suffix %}
+        {% do log(
+            "Dropping schemas: " ~ test_schema ~ ", " ~ elementary_schema,
+            info=true,
+        ) %}
+        {% do elementary_tests.edr_drop_schema(database, elementary_schema) %}
+        {% do elementary_tests.edr_drop_schema(database, test_schema) %}
+    {% endfor %}
+{% endmacro %}
+
 {% macro edr_drop_schema(database_name, schema_name) %}
     {% do return(
         adapter.dispatch("edr_drop_schema", "elementary_tests")(
@@ -34,6 +60,24 @@
     {% do run_query("DROP DATABASE IF EXISTS `" ~ safe_schema ~ "` CASCADE") %}
 {% endmacro %}
 
+{% macro athena__edr_drop_schema(database_name, schema_name) %}
+    {#
+      Athena's SQL `DROP SCHEMA … CASCADE` can fail when the schema
+      contains Iceberg tables.  Work around this by first dropping every
+      relation individually (the adapter handles Iceberg vs Hive
+      differences in its drop_relation implementation) and then removing
+      the now-empty schema.
+    #}
+    {% set schema_relation = api.Relation.create(
+        database=database_name, schema=schema_name
+    ) %}
+    {% set relations = adapter.list_relations_without_caching(schema_relation) %}
+    {% for relation in relations %}
+        {% do adapter.drop_relation(relation) %}
+    {% endfor %}
+    {% do dbt.drop_schema(schema_relation) %}
+{% endmacro %}
+
 {% macro duckdb__edr_drop_schema(database_name, schema_name) %}
     {% do run_query("DROP SCHEMA IF EXISTS " ~ schema_name ~ " CASCADE") %}
     {% do adapter.commit() %}