feat: Add table_type param to ddl_data_type macro (#599)

Co-authored-by: Oleksandr Kot <[email protected]> Co-authored-by: Serhii Dimchenko <[email protected]>
dbt-labs · Mar 14, 2024 · 59c2c05 · 59c2c05
1 parent 9c179a8
commit 59c2c05
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 17 deletions.
diff --git a/dbt/include/athena/macros/materializations/models/incremental/column_helpers.sql b/dbt/include/athena/macros/materializations/models/incremental/column_helpers.sql
@@ -1,4 +1,4 @@
-{% macro alter_relation_add_columns(relation, add_columns = none) -%}
+{% macro alter_relation_add_columns(relation, add_columns = none, table_type = 'hive') -%}
   {% if add_columns is none %}
     {% set add_columns = [] %}
   {% endif %}
@@ -7,7 +7,7 @@
       alter {{ relation.type }} {{ relation.render_hive() }}
           add columns (
             {%- for column in add_columns -%}
-                {{ column.name }} {{ ddl_data_type(column.data_type) }}{{ ', ' if not loop.last }}
+                {{ column.name }} {{ ddl_data_type(column.data_type, table_type) }}{{ ', ' if not loop.last }}
             {%- endfor -%}
           )
   {%- endset -%}
@@ -30,7 +30,7 @@
   {%- endfor -%}
 {% endmacro %}
 
-{% macro alter_relation_replace_columns(relation, replace_columns = none) -%}
+{% macro alter_relation_replace_columns(relation, replace_columns = none, table_type = 'hive') -%}
   {% if replace_columns is none %}
     {% set replace_columns = [] %}
   {% endif %}
@@ -39,7 +39,7 @@
       alter {{ relation.type }} {{ relation.render_hive() }}
           replace columns (
             {%- for column in replace_columns -%}
-                {{ column.name }} {{ ddl_data_type(column.data_type) }}{{ ', ' if not loop.last }}
+                {{ column.name }} {{ ddl_data_type(column.data_type, table_type) }}{{ ', ' if not loop.last }}
             {%- endfor -%}
           )
   {%- endset -%}
@@ -49,10 +49,10 @@
   {% endif %}
 {% endmacro %}
 
-{% macro alter_relation_rename_column(relation, source_column, target_column, target_column_type) -%}
+{% macro alter_relation_rename_column(relation, source_column, target_column, target_column_type, table_type = 'hive') -%}
   {% set sql -%}
       alter {{ relation.type }} {{ relation.render_pure() }}
-          change column {{ source_column }} {{ target_column }} {{ ddl_data_type(target_column_type) }}
+          change column {{ source_column }} {{ target_column }} {{ ddl_data_type(target_column_type, table_type) }}
   {%- endset -%}
   {{ return(run_query(sql)) }}
 {% endmacro %}
diff --git a/dbt/include/athena/macros/materializations/models/incremental/on_schema_change.sql b/dbt/include/athena/macros/materializations/models/incremental/on_schema_change.sql
@@ -7,7 +7,7 @@
   {%- set add_to_target_arr = schema_changes_dict['source_not_in_target'] -%}
   {%- if on_schema_change == 'append_new_columns'-%}
      {%- if add_to_target_arr | length > 0 -%}
-       {%- do alter_relation_add_columns(target_relation, add_to_target_arr) -%}
+       {%- do alter_relation_add_columns(target_relation, add_to_target_arr, table_type) -%}
      {%- endif -%}
   {% elif on_schema_change == 'sync_all_columns' %}
      {%- set remove_from_target_arr = schema_changes_dict['target_not_in_source'] -%}
@@ -20,14 +20,14 @@
        {% for remove_col in remove_from_target_arr if remove_col.column.endswith('__dbt_alter') %}
          {%- set origin_col_name = remove_col.column | replace('__dbt_alter', '') -%}
          {% for add_col in add_to_target_arr if add_col.column == origin_col_name %}
-             {%- do alter_relation_rename_column(target_relation, remove_col.name, add_col.name, add_col.data_type) -%}
+             {%- do alter_relation_rename_column(target_relation, remove_col.name, add_col.name, add_col.data_type, table_type) -%}
              {%- do remove_from_target_arr.remove(remove_col) -%}
              {%- do add_to_target_arr.remove(add_col) -%}
          {% endfor %}
        {% endfor %}
 
        {% if add_to_target_arr | length > 0 %}
-         {%- do alter_relation_add_columns(target_relation, add_to_target_arr) -%}
+         {%- do alter_relation_add_columns(target_relation, add_to_target_arr, table_type) -%}
        {% endif %}
        {% if remove_from_target_arr | length > 0 %}
          {%- do alter_relation_drop_columns(target_relation, remove_from_target_arr) -%}
@@ -42,7 +42,7 @@
      {% else %}
        {%- set replace_with_target_arr = remove_partitions_from_columns(schema_changes_dict['source_columns'], partitioned_by) -%}
        {% if add_to_target_arr | length > 0 or remove_from_target_arr | length > 0 or new_target_types | length > 0 %}
-         {%- do alter_relation_replace_columns(target_relation, replace_with_target_arr) -%}
+         {%- do alter_relation_replace_columns(target_relation, replace_with_target_arr, table_type) -%}
        {% endif %}
      {% endif %}
   {% endif %}
@@ -63,10 +63,11 @@
     3. Drop the existing column
     4. Rename the new column to existing column
   #}
+  {%- set table_type = config.get('table_type', 'hive') -%}
   {%- set tmp_column = column_name + '__dbt_alter' -%}
-  {%- set new_ddl_data_type = ddl_data_type(new_column_type) -%}
+  {%- set new_ddl_data_type = ddl_data_type(new_column_type, table_type) -%}
 
-  {#- do alter_relation_add_columns(relation, [ tmp_column ]) -#}
+  {#- do alter_relation_add_columns(relation, [ tmp_column ], table_type) -#}
   {%- set add_column_query -%}
     alter {{ relation.type }} {{ relation.render_pure() }} add columns({{ tmp_column }} {{ new_ddl_data_type }});
   {%- endset -%}
@@ -83,6 +84,6 @@
   {%- endset -%}
   {%- do run_query(drop_column_query) -%}
 
-  {%- do alter_relation_rename_column(relation, tmp_column, column_name, new_column_type) -%}
+  {%- do alter_relation_rename_column(relation, tmp_column, column_name, new_column_type, table_type) -%}
 
 {% endmacro %}
diff --git a/dbt/include/athena/macros/materializations/snapshots/snapshot.sql b/dbt/include/athena/macros/materializations/snapshots/snapshot.sql
@@ -174,7 +174,7 @@
 
 
       {% if missing_columns %}
-        {% do alter_relation_add_columns(target_relation, missing_columns) %}
+        {% do alter_relation_add_columns(target_relation, missing_columns, table_type) %}
       {% endif %}
 
 

diff --git a/dbt/include/athena/macros/utils/ddl_dml_data_type.sql b/dbt/include/athena/macros/utils/ddl_dml_data_type.sql
@@ -1,8 +1,6 @@
 {# Athena has different types between DML and DDL #}
 {# ref: https://docs.aws.amazon.com/athena/latest/ug/data-types.html #}
-{% macro ddl_data_type(col_type) -%}
-  {%- set table_type = config.get('table_type', 'hive') -%}
-
+{% macro ddl_data_type(col_type, table_type = 'hive') -%}
   -- transform varchar
   {% set re = modules.re %}
   {% set data_type = re.sub('(?:varchar|character varying)(?:\(\d+\))?', 'string', col_type) %}