From 6f6aea6339fd498b33e6db0cdbd07eef82e5b53d Mon Sep 17 00:00:00 2001
From: Doug Beatty <doug.beatty@dbtlabs.com>
Date: Mon, 25 Sep 2023 16:51:54 -0600
Subject: [PATCH] Enable custom overrides of data type formatting via multiple
 dispatch

---
 CHANGELOG.md                   | 17 ++++++++++++++++-
 macros/generate_model_yaml.sql |  3 +--
 macros/generate_source.sql     |  3 +--
 macros/helpers/helpers.sql     | 22 +++++++++++++++++++++-
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ac9a82e..368d5fc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,9 +14,24 @@
 
 ## 🚨 Breaking change
 
-- `include_data_types` parameter added to `generate_model_yaml` and behavior changed for `generate_source`. Both default to `true`
+`include_data_types` parameter added to `generate_model_yaml` and behavior changed for `generate_source`. Both default to `true`
 and are lowercase to align with the dbt style guide. Scale & precision are **not** included. Previous logic for `generate_source` defaulted to `false` and the resulting data types were uppercase and included scale & precision ([#122](https://github.com/dbt-labs/dbt-codegen/pull/122)).
 
+[Dispatch](https://docs.getdbt.com/reference/dbt-jinja-functions/dispatch) can be used to utilize the column data type formatting of previous versions. Namely, by adding this macro to your project:
+```sql
+{% macro default__data_type_format_source(column) %}
+    {{ return(column.data_type | upper) }}
+{% endmacro %}
+```
+
+And then adding this within `dbt_project.yml`:
+```yaml
+dispatch:
+  - macro_namespace: codegen
+    search_order: ['my_project', 'codegen']
+```
+
+
 ## New features
 - Addition of the [create_base_models](macros/create_base_models.sql)
 This macro generates a series of terminal commands (appended w) bash script which creates a new file in your dbt project based off the results of the [generate_base_model](macros/generate_base_model.sql) macro. Therefore, instead of outputting in the terminal, it will create the file for you.
diff --git a/macros/generate_model_yaml.sql b/macros/generate_model_yaml.sql
index ae6b325..f226536 100644
--- a/macros/generate_model_yaml.sql
+++ b/macros/generate_model_yaml.sql
@@ -7,8 +7,7 @@
 
     {% do model_yaml.append('      - name: ' ~ column_name  | lower ) %}
     {% if include_data_types %}
-        {% set formatted = codegen.format_column(column) %}
-        {% do model_yaml.append('        data_type: ' ~ formatted['data_type'] | lower) %}
+        {% do model_yaml.append('        data_type: ' ~ codegen.data_type_format_model(column)) %}
     {% endif %}
     {% do model_yaml.append('        description: "' ~ column_desc_dict.get(column.name | lower,'') ~ '"') %}
     {% do model_yaml.append('') %}
diff --git a/macros/generate_source.sql b/macros/generate_source.sql
index f3f780a..31fd3e7 100644
--- a/macros/generate_source.sql
+++ b/macros/generate_source.sql
@@ -62,8 +62,7 @@
         {% for column in columns %}
             {% do sources_yaml.append('          - name: ' ~ column.name | lower ) %}
             {% if include_data_types %}
-                {% set formatted = codegen.format_column(column) %}
-                {% do sources_yaml.append('            data_type: ' ~ formatted['data_type'] | lower) %}
+                {% do sources_yaml.append('            data_type: ' ~ codegen.data_type_format_source(column)) %}
             {% endif %}
             {% if include_descriptions %}
                 {% do sources_yaml.append('            description: ""' ) %}
diff --git a/macros/helpers/helpers.sql b/macros/helpers/helpers.sql
index fd28c87..557858e 100644
--- a/macros/helpers/helpers.sql
+++ b/macros/helpers/helpers.sql
@@ -57,4 +57,24 @@
         {% endfor %}
     {% endif %}
     {{ return(model_names) }}
-{% endmacro %}
\ No newline at end of file
+{% endmacro %}
+
+{% macro data_type_format_source(column) -%}
+  {{ return(adapter.dispatch('data_type_format_source', 'codegen')(column)) }}
+{%- endmacro %}
+
+{# format a column data type for a source #}
+{% macro default__data_type_format_source(column) %}
+    {% set formatted = codegen.format_column(column) %}
+    {{ return(formatted['data_type'] | lower) }}
+{% endmacro %}
+
+{% macro data_type_format_model(column) -%}
+  {{ return(adapter.dispatch('data_type_format_model', 'codegen')(column)) }}
+{%- endmacro %}
+
+{# format a column data type for a model #}
+{% macro default__data_type_format_model(column) %}
+    {% set formatted = codegen.format_column(column) %}
+    {{ return(formatted['data_type'] | lower) }}
+{% endmacro %}