From d7d6dc3a2e21812902fadd8281d719ec44171f33 Mon Sep 17 00:00:00 2001
From: Doug Guthrie <douglas.p.guthrie@gmail.com>
Date: Mon, 1 May 2023 10:09:27 -0600
Subject: [PATCH] Add macros to get query stats

---
 macros/snowflake/create_query_stats_table.sql | 30 +++++++++++++
 macros/snowflake/get_query_stats.sql          | 44 +++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 macros/snowflake/create_query_stats_table.sql
 create mode 100644 macros/snowflake/get_query_stats.sql

diff --git a/macros/snowflake/create_query_stats_table.sql b/macros/snowflake/create_query_stats_table.sql
new file mode 100644
index 0000000..f99aa4d
--- /dev/null
+++ b/macros/snowflake/create_query_stats_table.sql
@@ -0,0 +1,30 @@
+{% macro create_query_stats_table() %}
+
+    {% if execute %}
+
+        -- Create the table if it doesn't exist
+        {% set create_table_sql %}
+
+            create table if not exists {{ target.database }}.{{ target.schema }}.query_operator_stats as
+            select
+                1::integer as dbt_cloud_project_id,
+                1::integer as dbt_cloud_run_id,
+                1::integer as dbt_cloud_job_id,
+                ''::string as dbt_cloud_run_reason,
+                ''::string as dbt_version,
+                ''::string as invocation_id,
+                ''::string as status,
+                ''::string as thread_id,
+                1.0::float as execution_time,
+                ''::string as unique_id,
+                *
+            from table(get_query_operator_stats(last_query_id()))
+            limit 0;
+
+        {% endset %}
+
+        {% do run_query(create_table_sql) %}
+
+    {% endif %}
+
+{% endmacro %}
\ No newline at end of file
diff --git a/macros/snowflake/get_query_stats.sql b/macros/snowflake/get_query_stats.sql
new file mode 100644
index 0000000..5ba8bee
--- /dev/null
+++ b/macros/snowflake/get_query_stats.sql
@@ -0,0 +1,44 @@
+{% macro get_query_stats(results) %}
+
+    {% if execute %}
+
+        {% if results != [] %}
+
+            {% for result in results %}
+
+                {% if result.adapter_response %}
+
+                    {% set query_id = result.adapter_response.query_id %}
+
+                    {% do log('Uploading stats for query ' ~ query_id, true) %}
+
+                    {% set insert_sql %}
+
+                        insert into {{ target.database }}.{{ target.schema }}.query_operator_stats
+                        select
+                            {{ env_var('DBT_CLOUD_PROJECT_ID', 'null') }},
+                            {{ env_var('DBT_CLOUD_RUN_ID', 'null') }},
+                            {{ env_var('DBT_CLOUD_JOB_ID', 'null') }},
+                            '{{ env_var('DBT_CLOUD_RUN_REASON', '') }}',
+                            '{{ dbt_version }}',
+                            '{{ invocation_id }}',
+                            '{{ result.status }}',
+                            '{{ result.thread_id }}',
+                            {{ result.execution_time }},
+                            '{{ result.unique_id }}',
+                            *
+                        from table(get_query_operator_stats('{{ query_id }}'));
+
+                    {% endset %}
+
+                    {% do run_query(insert_sql) %}
+
+                {% endif %}
+
+            {% endfor %}
+
+        {% endif %}
+
+    {% endif %}
+
+{% endmacro %}
\ No newline at end of file