Athena support POC

macros/commands/generate_elementary_cli_profile.sql

@@ -76,6 +76,26 @@ elementary:
       threads: {{ target.threads }}
 {% endmacro %}
 
+{% macro athena__generate_elementary_cli_profile(method, elementary_database, elementary_schema) %}
+elementary:
+  outputs:
+    default:
+      type: "{{ target.type }}"
+      s3_staging_dir: "{{ target.s3_staging_dir }}"
+      region_name: "{{ target.region_name }}"
+      database: "{{ target.database }}"
+      aws_profile_name: "{{ target.aws_profile_name }}"
+      work_group: "{{ target.work_group }}"
+      aws_access_key_id: "<AWS_ACCESS_KEY_ID>"
+      aws_secret_access_key: "<AWS_SECRET_ACCESS_KEY>"
+      {%- if elementary_database %}
+      catalog: "{{ elementary_database }}"
+      {%- endif %}
+      schema: "{{ elementary_schema }}"
+      threads: {{ target.threads }}
+{% endmacro %}
+
+
 {% macro default__generate_elementary_cli_profile(method, elementary_database, elementary_schema) %}
 Adapter "{{ target.type }}" is not supported on Elementary.
 {% endmacro %}

macros/edr/alerts/anomaly_detection_description.sql

@@ -9,21 +9,21 @@
 {% endmacro %}
 
 {% macro freshness_description() %}
-    'Last update was at ' || anomalous_value || ', ' || abs(round({{ elementary.edr_cast_as_numeric('metric_value/3600') }}, 2)) || ' hours ago. Usually the table is updated within ' || abs(round({{ elementary.edr_cast_as_numeric('training_avg/3600') }}, 2)) || ' hours.'
+    'Last update was at ' || anomalous_value || ', ' || {{ elementary.edr_cast_as_string('abs(round(' ~ elementary.edr_cast_as_numeric('metric_value/3600') ~ ', 2))') }} || ' hours ago. Usually the table is updated within ' || {{ elementary.edr_cast_as_string('abs(round(' ~ elementary.edr_cast_as_numeric('training_avg/3600') ~ ', 2))') }} || ' hours.'
 {% endmacro %}
 
 {% macro table_metric_description() %}
-    'The last ' || metric_name || ' value is ' || round({{ elementary.edr_cast_as_numeric('metric_value') }}, 3) ||
-    '. The average for this metric is ' || round({{ elementary.edr_cast_as_numeric('training_avg') }}, 3) || '.'
+    'The last ' || metric_name || ' value is ' || {{ elementary.edr_cast_as_string('round(' ~ elementary.edr_cast_as_numeric('metric_value') ~ ', 3)') }} ||
+    '. The average for this metric is ' || {{ elementary.edr_cast_as_string('round(' ~ elementary.edr_cast_as_numeric('training_avg') ~ ', 3)') }} || '.'
 {% endmacro %}
 
 {% macro column_metric_description() %}
-    'In column ' || column_name || ', the last ' || metric_name || ' value is ' || round({{ elementary.edr_cast_as_numeric('metric_value') }}, 3) ||
-    '. The average for this metric is ' || round({{ elementary.edr_cast_as_numeric('training_avg') }}, 3) || '.'
+    'In column ' || column_name || ', the last ' || metric_name || ' value is ' || {{ elementary.edr_cast_as_string('round(' ~ elementary.edr_cast_as_numeric('metric_value') ~ ', 3)') }} ||
+    '. The average for this metric is ' || {{ elementary.edr_cast_as_string('round(' ~ elementary.edr_cast_as_numeric('training_avg') ~ ', 3)') }} || '.'
 {% endmacro %}
 
 {% macro dimension_metric_description() %}
     'The last ' || metric_name || ' value for dimension ' || dimension || ' - ' ||
-    case when dimension_value is null then 'NULL' else dimension_value end || ' is ' || round({{ elementary.edr_cast_as_numeric('metric_value') }}, 3) ||
-    '. The average for this metric is ' || round({{ elementary.edr_cast_as_numeric('training_avg') }}, 3) || '.'
+    case when dimension_value is null then 'NULL' else dimension_value end || ' is ' || {{ elementary.edr_cast_as_string('round(' ~ elementary.edr_cast_as_numeric('metric_value') ~ ', 3)') }} ||
+    '. The average for this metric is ' || {{ elementary.edr_cast_as_string('round(' ~ elementary.edr_cast_as_numeric('training_avg') ~ ', 3)') }} || '.'
 {% endmacro %}

macros/edr/data_monitoring/anomaly_detection/get_anomaly_scores_query.sql

@@ -143,23 +143,23 @@
                 end as anomaly_score,
                 {{ test_configuration.anomaly_sensitivity }} as anomaly_score_threshold,
                 source_value as anomalous_value,
-                bucket_start,
-                bucket_end,
+                {{ elementary.edr_cast_as_timestamp('bucket_start') }} as bucket_start,
+                {{ elementary.edr_cast_as_timestamp('bucket_end') }} as bucket_end,
                 bucket_seasonality,
                 metric_value,
                 case
                     when training_stddev is null then null
                     else (-1) * {{ test_configuration.anomaly_sensitivity }} * training_stddev + training_avg
                 end as min_metric_value,
-                case 
+                case
                     when training_stddev is null then null
                     else {{ test_configuration.anomaly_sensitivity }} * training_stddev + training_avg
                 end as max_metric_value,
                 training_avg,
                 training_stddev,
                 training_set_size,
-                training_start,
-                training_end,
+                {{ elementary.edr_cast_as_timestamp('training_start') }} as training_start,
+                {{ elementary.edr_cast_as_timestamp('training_end') }} as training_end,
                 dimension,
                 dimension_value
             from time_window_aggregation
@@ -172,4 +172,4 @@
 
     {% endset %}
     {{ return(anomaly_scores_query) }}
-{% endmacro %}
+{% endmacro %}

macros/edr/materializations/tests/test.sql

@@ -178,7 +178,7 @@
         'test_execution_id': test_execution_id,
         'test_unique_id': elementary.insensitive_get_dict_value(flattened_test, 'unique_id'),
         'model_unique_id': parent_model_unique_id,
-        'detected_at': elementary.insensitive_get_dict_value(flattened_test, 'generated_at'),
+        'detected_at': elementary.edr_cast_as_timestamp("'" ~ elementary.insensitive_get_dict_value(flattened_test, 'generated_at') ~ "'"),
         'database_name': elementary.insensitive_get_dict_value(flattened_test, 'database_name'),
         'schema_name': elementary.insensitive_get_dict_value(flattened_test, 'schema_name'),
         'table_name': parent_model_name,

macros/edr/metadata_collection/get_columns_from_information_schema.sql

@@ -43,4 +43,17 @@
     from information_schema.columns
     where upper(table_schema) = upper('{{ schema_name }}')
 
-{% endmacro %}
+{% endmacro %}
+
+{% macro athena__get_columns_from_information_schema(database_name, schema_name) %}
+    select
+        upper(table_catalog || '.' || table_schema || '.' || table_name) as full_table_name,
+        upper(table_catalog) as database_name,
+        upper(table_schema) as schema_name,
+        upper(table_name) as table_name,
+        upper(column_name) as column_name,
+        data_type
+    from information_schema.columns
+    where upper(table_schema) = upper('{{ schema_name }}')
+
+{% endmacro %}

macros/edr/metadata_collection/get_tables_from_information_schema.sql

@@ -86,3 +86,22 @@
         table_name
     from information_schema_tables
 {% endmacro %}
+
+{% macro athena__get_tables_from_information_schema(schema_tuple) %}
+    {%- set database_name, schema_name = schema_tuple %}
+
+    select
+        {{ elementary.full_table_name() }} as full_table_name,
+        upper(database_name || '.' || schema_name) as full_schema_name,
+        database_name,
+        schema_name,
+        table_name
+    from (
+        select
+            upper(table_catalog) as database_name,
+            upper(table_schema) as schema_name,
+            upper(table_name) as table_name
+        from information_schema.tables
+        where upper(table_schema) = upper('{{ schema_name }}') and upper(table_catalog) = upper('{{ database_name }}')
+    )
+{% endmacro %}

macros/edr/system/system_utils/buckets_cte.sql

@@ -89,3 +89,14 @@
     {%- endset %}
     {{ return(complete_buckets_cte) }}
 {% endmacro %}
+
+{% macro athena__complete_buckets_cte(time_bucket, bucket_end_expr, min_bucket_start_expr, max_bucket_end_expr) %}
+    {%- set complete_buckets_cte %}
+        select
+          edr_bucket_start,
+          {{ bucket_end_expr }} as edr_bucket_end
+        from unnest(sequence({{ min_bucket_start_expr }}, {{ max_bucket_end_expr }}, interval '{{ time_bucket.count }}' {{ time_bucket.period }})) as t(edr_bucket_start)
+        where {{ bucket_end_expr }} <= {{ max_bucket_end_expr }}
+    {%- endset %}
+    {{ return(complete_buckets_cte) }}
+{% endmacro %}

macros/utils/cross_db_utils/current_timestamp.sql

@@ -45,4 +45,12 @@
 
 {% macro spark__edr_current_timestamp_in_utc() %}
     unix_timestamp()
-{% endmacro %}
+{% endmacro %}
+
+{% macro athena__edr_current_timestamp() -%}
+    CURRENT_TIMESTAMP
+{%- endmacro -%}
+
+{% macro athena__edr_current_timestamp_in_utc() -%}
+    cast(CURRENT_TIMESTAMP AT TIME ZONE 'utc' AS TIMESTAMP)
+{%- endmacro -%}

macros/utils/cross_db_utils/datediff.sql

@@ -24,4 +24,12 @@
     {%- else %}
         {{ exceptions.raise_compiler_error("Unsupported date_part in edr_datediff: ".format(date_part)) }}
     {%- endif %}
-{% endmacro %}
+{% endmacro %}
+
+{% macro athena__edr_datediff(first_date, second_date, date_part) %}
+    {% set macro = dbt.datediff or dbt_utils.datediff %}
+    {% if not macro %}
+        {{ exceptions.raise_compiler_error("Did not find a `datediff` macro.") }}
+    {% endif %}
+    {{ return(macro(elementary.edr_cast_as_date(first_date), elementary.edr_cast_as_date(second_date), date_part)) }}
+{% endmacro %}

macros/utils/cross_db_utils/target_database.sql

@@ -18,3 +18,7 @@
 {% macro bigquery__target_database() %}
     {% do return(target.project) %}
 {% endmacro %}
+
+{% macro athena__target_database() %}
+    {% do return(target.database) %}
+{% endmacro %}

macros/utils/cross_db_utils/timeadd.sql

@@ -25,3 +25,7 @@
 {% macro redshift__edr_timeadd(date_part, number, timestamp_expression) %}
     dateadd({{ date_part }}, {{ elementary.edr_cast_as_int(number) }}, {{ elementary.edr_cast_as_timestamp(timestamp_expression) }})
 {% endmacro %}
+
+{% macro athena__edr_timeadd(date_part, number, timestamp_expression) %}
+    date_add('{{ date_part }}', {{ elementary.edr_cast_as_int(number) }}, {{ elementary.edr_cast_as_timestamp(timestamp_expression) }})
+{% endmacro %}

macros/utils/data_types/cast_column.sql

@@ -1,7 +1,19 @@
 {%- macro edr_cast_as_timestamp(timestamp_field) -%}
+    {{ return(adapter.dispatch('edr_cast_as_timestamp', 'elementary')(timestamp_field)) }}
+{%- endmacro -%}
+
+{%- macro default__edr_cast_as_timestamp(timestamp_field) -%}
     cast({{ timestamp_field }} as {{ elementary.edr_type_timestamp() }})
 {%- endmacro -%}
 
+{# Athena needs explicit conversion for ISO8601 timestamps used in buckets_cte #}
+{%- macro athena__edr_cast_as_timestamp(timestamp_field) -%}
+    coalesce(
+        try_cast({{ timestamp_field }} as {{ elementary.edr_type_timestamp() }}),
+        cast(from_iso8601_timestamp(cast({{ timestamp_field }} AS {{ elementary.edr_type_string() }})) AS {{ elementary.edr_type_timestamp() }})
+    )
+{%- endmacro -%}
+
 {%- macro edr_cast_as_float(column) -%}
     cast({{ column }} as {{ elementary.edr_type_float() }})
 {%- endmacro -%}
@@ -43,6 +55,15 @@
     cast({{ elementary.edr_cast_as_timestamp(timestamp_field) }} as {{ elementary.edr_type_date() }})
 {%- endmacro -%}
 
+{# Athena needs explicit conversion for ISO8601 timestamps used in buckets_cte #}
+{%- macro athena__edr_cast_as_date(timestamp_field) -%}
+    coalesce(
+        try_cast({{ timestamp_field }} as {{ elementary.edr_type_date() }}),
+        cast(from_iso8601_timestamp(cast({{ timestamp_field }} AS {{ elementary.edr_type_string() }})) AS {{ elementary.edr_type_date() }})
+    )
+{%- endmacro -%}
+
+
 {%- macro const_as_text(string) -%}
     {{ return(adapter.dispatch('const_as_text', 'elementary')(string)) }}
 {%- endmacro -%}

macros/utils/data_types/data_type.sql

@@ -43,6 +43,11 @@
     {% do return("string") %}
 {% endmacro %}
 
+{% macro athena__edr_type_string() %}
+    {% do return("varchar") %}
+{% endmacro %}
+
+
 
 
 {%- macro edr_type_long_string() -%}
@@ -93,6 +98,10 @@
 
 
 {% macro edr_type_timestamp() %}
+    {{ return(adapter.dispatch('edr_type_timestamp', 'elementary')()) }}
+{% endmacro %}
+
+{% macro default__edr_type_timestamp() %}
     {% set macro = dbt.type_timestamp or dbt_utils.type_timestamp %}
     {% if not macro %}
         {{ exceptions.raise_compiler_error("Did not find a `type_timestamp` macro.") }}
@@ -121,3 +130,13 @@
 {% macro bigquery__edr_type_date() %}
     date
 {% endmacro %}
+
+{% macro athena__edr_type_timestamp() %}
+  {%- set config = model.get('config', {}) -%}
+  {%- set table_type = config.get('table_type', 'glue') -%}
+  {%- if table_type == 'iceberg' -%}
+    timestamp(6)
+  {%- else -%}
+    timestamp
+  {%- endif -%}
+{% endmacro %}

macros/utils/table_operations/create_or_replace.sql

@@ -25,3 +25,12 @@
     {% do run_query(dbt.create_table_as(temporary, relation, sql_query)) %}
     {% do adapter.commit() %}
 {% endmacro %}
+
+{% macro athena__create_or_replace(temporary, relation, sql_query) %}
+    {% set drop_query %}
+        drop table if exists {{ relation.schema }}.{{ relation.identifier }}
+    {% endset %}
+    {% do elementary.run_query(drop_query) %}
+    {% do run_query(dbt.create_table_as(temporary, relation, sql_query)) %}
+    {% do adapter.commit() %}
+{% endmacro %}

macros/utils/table_operations/delete_and_insert.sql

@@ -76,3 +76,26 @@
 
     {% do return(queries) %}
 {% endmacro %}
+
+{% macro athena__get_delete_and_insert_queries(relation, insert_relation, delete_relation, delete_column_key) %}
+    {% set queries = [] %}
+
+    {% if delete_relation %}
+        {% set delete_query %}
+            delete from {{ relation }}
+            where
+            {{ delete_column_key }} is null
+            or {{ delete_column_key }} in (select {{ delete_column_key }} from {{ delete_relation }});
+        {% endset %}
+        {% do queries.append(delete_query) %}
+    {% endif %}
+
+    {% if insert_relation %}
+        {% set insert_query %}
+            insert into {{ relation }} select * from {{ insert_relation }};
+        {% endset %}
+        {% do queries.append(insert_query) %}
+    {% endif %}
+
+    {% do return(queries) %}
+{% endmacro %}

macros/utils/table_operations/insert_rows.sql

@@ -111,12 +111,22 @@
     {{- return(string_value | replace("'", "''")) -}}
 {%- endmacro -%}
 
+{%- macro athena__escape_special_chars(string_value) -%}
+    {{- return(string_value | replace("'", "''")) -}}
+{%- endmacro -%}
+
 {%- macro render_value(value) -%}
     {%- if value is defined and value is not none -%}
         {%- if value is number -%}
             {{- value -}}
         {%- elif value is string -%}
-            '{{- elementary.escape_special_chars(value) -}}'
+            {%- if value.startswith('cast(')  -%}
+              {{- value -}}
+            {%- elif value.startswith('coalesce(') -%}
+              {{- value -}}
+            {%- else -%}
+              '{{- elementary.escape_special_chars(value) -}}'
+            {%- endif -%}
         {%- elif value is mapping or value is sequence -%}
             '{{- elementary.escape_special_chars(tojson(value)) -}}'
         {%- else -%}

macros/utils/table_operations/merge_sql.sql

@@ -19,3 +19,29 @@
     {% do return(macro(target_relation, tmp_relation, unique_key, dest_columns)) %}
     {{ return(merge_sql) }}
 {% endmacro %}
+
+{% macro athena__merge_sql(target_relation, tmp_relation, unique_key, dest_columns, incremental_predicates) %}
+
+    {% set query %}
+      merge into {{ target_relation }} as target using {{ tmp_relation }} as src
+      ON (target.{{unique_key}} = src.{{ unique_key}})
+      when matched
+      then update set
+      {%- for col in dest_columns %}
+          {{ col.column }} = src.{{ col.column }} {{ ", " if not loop.last }}
+      {%- endfor %}
+      when not matched
+        then insert (
+        {%- for col in dest_columns %}
+            {{ col.column }} {{ ", " if not loop.last }}
+        {%- endfor %}
+
+        )
+        values (
+        {%- for col in dest_columns %}
+            src.{{ col.column }} {{ ", " if not loop.last }}
+        {%- endfor %}
+        )
+    {% endset %}
+    {% do return(query) %}
+{% endmacro %}

macros/utils/table_operations/replace_table_data.sql

@@ -30,3 +30,9 @@
 
     {% do adapter.drop_relation(intermediate_relation) %}
 {% endmacro %}
+
+{% macro athena__replace_table_data(relation, rows) %}
+    {% set intermediate_relation = elementary.create_intermediate_relation(relation, rows, temporary=True) %}
+    {% do elementary.create_or_replace(False, relation, 'select * from {}'.format(intermediate_relation)) %}
+    {% do adapter.drop_relation(intermediate_relation) %}
+{% endmacro %}

models/edr/alerts/alerts_anomaly_detection.sql

@@ -15,7 +15,7 @@ alerts_anomaly_detection as (
            test_execution_id,
            test_unique_id,
            model_unique_id,
-           detected_at,
+           {{ elementary.edr_cast_as_timestamp("detected_at") }} as detected_at,
            database_name,
            schema_name,
            table_name,

models/edr/alerts/alerts_dbt_models.sql

@@ -63,7 +63,7 @@ with error_models as (
 
 select model_execution_id as alert_id,
        unique_id,
-       generated_at as detected_at,
+       {{ elementary.edr_cast_as_timestamp("generated_at") }} as detected_at,
        database_name,
        materialization,
        path,