diff --git a/integration_tests/dbt_project/macros/test_normalize_rows_affected.sql b/integration_tests/dbt_project/macros/test_normalize_rows_affected.sql new file mode 100644 index 000000000..33c0a4a29 --- /dev/null +++ b/integration_tests/dbt_project/macros/test_normalize_rows_affected.sql @@ -0,0 +1,4 @@ +{% macro test_normalize_rows_affected(rows_affected) %} + {% set result = elementary.normalize_rows_affected(rows_affected) %} + {{ return(result) }} +{% endmacro %} diff --git a/integration_tests/tests/test_dbt_artifacts/test_normalize_rows_affected.py b/integration_tests/tests/test_dbt_artifacts/test_normalize_rows_affected.py new file mode 100644 index 000000000..9c7e2bd10 --- /dev/null +++ b/integration_tests/tests/test_dbt_artifacts/test_normalize_rows_affected.py @@ -0,0 +1,34 @@ +import json + +import pytest +from dbt_project import DbtProject + + +@pytest.mark.parametrize( + "input_value,expected_output", + [ + (None, None), + (123, 123), + (0, 0), + (-1, -1), + ("123", 123), + ("0", 0), + ("-1", None), + ("456", 456), + ], +) +def test_normalize_rows_affected(dbt_project: DbtProject, input_value, expected_output): + result = dbt_project.dbt_runner.run_operation( + "elementary_tests.test_normalize_rows_affected", + macro_args={"rows_affected": input_value}, + ) + # When the macro returns None, log_macro_results doesn't log anything, + # so run_operation returns an empty list + if not result: + actual_output = None + else: + actual_output = json.loads(result[0]) + assert actual_output == expected_output, ( + f"normalize_rows_affected({input_value!r}) returned {actual_output!r}, " + f"expected {expected_output!r}" + ) diff --git a/macros/edr/dbt_artifacts/upload_run_results.sql b/macros/edr/dbt_artifacts/upload_run_results.sql index 0a6fe8e70..08fa6eb58 100644 --- a/macros/edr/dbt_artifacts/upload_run_results.sql +++ b/macros/edr/dbt_artifacts/upload_run_results.sql @@ -37,10 +37,25 @@ {{ return(dbt_run_results_empty_table_query) }} {% endmacro %} +{% macro normalize_rows_affected(rows_affected) %} + {% if rows_affected is none %} + {{ return(none) }} + {% elif rows_affected is string %} + {% if rows_affected == '-1' %} + {{ return(none) }} + {% else %} + {{ return(rows_affected | int) }} + {% endif %} + {% else %} + {{ return(rows_affected) }} + {% endif %} +{% endmacro %} + {% macro flatten_run_result(run_result) %} {% set run_result_dict = elementary.get_run_result_dict(run_result) %} {% set node = elementary.safe_get_with_default(run_result_dict, 'node', {}) %} {% set config_dict = elementary.safe_get_with_default(node, 'config', {}) %} + {% set raw_rows_affected = run_result_dict.get('adapter_response', {}).get('rows_affected') %} {% set flatten_run_result_dict = { 'model_execution_id': elementary.get_node_execution_id(node), 'invocation_id': invocation_id, @@ -48,7 +63,7 @@ 'name': node.get('name'), 'message': run_result_dict.get('message'), 'generated_at': elementary.datetime_now_utc_as_string(), - 'rows_affected': run_result_dict.get('adapter_response', {}).get('rows_affected'), + 'rows_affected': elementary.normalize_rows_affected(raw_rows_affected), 'execution_time': run_result_dict.get('execution_time'), 'status': run_result_dict.get('status'), 'resource_type': node.get('resource_type'),