move rows inserted to its own macro

dbt-labs · Apr 29, 2024 · 6c9880d · 6c9880d
1 parent d8935c3
commit 6c9880d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 10 deletions.
diff --git a/insert_by_period/macros/get_rows_inserted.sql b/insert_by_period/macros/get_rows_inserted.sql
@@ -0,0 +1,25 @@
+{% macro get_rows_inserted(result) -%}
+  {{ return(adapter.dispatch('get_rows_inserted', 'insert_by_period')(result)) }}
+{% endmacro %}
+
+{% macro default__get_rows_inserted(result) %}
+
+  {% if 'response' in result.keys() %} {# added in v0.19.0 #}
+    {% set rows_inserted = result['response']['rows_affected'] %}
+  {% else %} {# older versions #}
+    {% set rows_inserted = result['status'].split(" ")[2] | int %}
+  {% endif %}
+
+  {{return(rows_inserted)}}
+
+{% endmacro %}
+
+{% macro databricks__get_rows_inserted(result) %}
+
+  {% if 'data' in result.keys() %}
+    {% set rows_inserted = result['data'][0][0] | int %}
+  {% endif %}
+
+  {{return(rows_inserted)}}
+
+{% endmacro %}
diff --git a/insert_by_period/macros/insert_by_period_materialization.sql b/insert_by_period/macros/insert_by_period_materialization.sql
@@ -98,20 +98,12 @@
     {%- endcall %}
     {% set result = load_result('main-' ~ i) %}
 
-    {% if target.type == 'databricks' %} {# databricks provides rows affected in 'data' and not 'response' #}
-      {% if 'data' in result.keys() %}
-        {% set rows_inserted = result['data'][0][0] | int %}
-      {% endif %}
-    {% elif 'response' in result.keys() %} {# added in v0.19.0 #}
-        {% set rows_inserted = result['response']['rows_affected'] %}
-    {% else %} {# older versions #}
-        {% set rows_inserted = result['status'].split(" ")[2] | int %}
-    {% endif %}
+    {% set rows_inserted = insert_by_period.get_rows_inserted(result) %}
 
     {%- set sum_rows_inserted = loop_vars['sum_rows_inserted'] + rows_inserted -%}
     {%- if loop_vars.update({'sum_rows_inserted': sum_rows_inserted}) %} {% endif -%}
 
-    {%- set msg = "Ran for " ~ period ~ " " ~ (i + 1) ~ " of " ~ (num_periods) ~ "; " ~ rows_inserted ~ " records inserted" -%}
+    {%- set msg = "Ran for " ~ period ~ " " ~ (i + 1) ~ " of " ~ (num_periods) ~ "; " ~ rows_inserted ~ " record(s) inserted" -%}
     {{ print(msg) }}
 
   {%- endfor %}