Read upstream descriptions from sources (#154)

* Read upstream descriptions from source

* Update README.md

* Add test

* Update test_helper_get_models.sql

CHANGELOG.md

@@ -1,3 +1,6 @@
+## New features
+- `generate_model_yaml` with `upstream_descriptions=True` now reads from upstream sources in addition to models. 
+
 ## Fixes
 - Column `description` fields are now correctly escaped in `generate_model_yaml` ([#142](https://github.com/dbt-labs/dbt-codegen/issues/142))
 

README.md

@@ -220,7 +220,7 @@ schema.yml file.
 
 ### Arguments:
 * `model_names` (required): The model(s) you wish to generate YAML for.
-* `upstream_descriptions` (optional, default=False): Whether you want to include descriptions for identical column names from upstream models.
+* `upstream_descriptions` (optional, default=False): Whether you want to include descriptions for identical column names from upstream models and sources.
 * `include_data_types` (optional, default=True): Whether you want to add data types to your model column definitions.
 
 ### Usage:

integration_tests/models/model_from_source.sql

@@ -0,0 +1,3 @@
+select
+    *
+from {{ source('codegen_integration_tests__data_source_schema', 'codegen_integration_tests__data_source_table') }}

integration_tests/models/source.yml

@@ -4,4 +4,9 @@ sources:
   - name: codegen_integration_tests__data_source_schema
     tables:
       - name: codegen_integration_tests__data_source_table
+        columns:
+          - name: my_integer_col
+            description: My Integer Column
+          - name: my_bool_col
+            description: My Boolean Column
       - name: codegen_integration_tests__data_source_table_case_sensitive

integration_tests/tests/test_generate_model_yaml_upstream_source_descriptions.sql

@@ -0,0 +1,23 @@
+{% set actual_model_yaml = codegen.generate_model_yaml(
+    model_names=['model_from_source'],
+    upstream_descriptions=True,
+    include_data_types=False
+  )
+%}
+
+{% set expected_model_yaml %}
+version: 2
+
+models:
+  - name: model_from_source
+    description: ""
+    columns:
+      - name: my_integer_col
+        description: "My Integer Column"
+
+      - name: my_bool_col
+        description: "My Boolean Column"
+
+{% endset %}
+
+{{ assert_equal (actual_model_yaml | trim, expected_model_yaml | trim) }}

integration_tests/tests/test_helper_get_models.sql

@@ -7,6 +7,6 @@
 {% set actual_list = codegen.get_models(prefix='model_')|sort %}
 {% endif %}
 
-{% set expected_list = ['model_data_a', 'model_struct', 'model_without_any_ctes', 'model_without_import_ctes'] %}
+{% set expected_list = ['model_data_a', 'model_from_source', 'model_struct', 'model_without_any_ctes', 'model_without_import_ctes'] %}
 
 {{ assert_equal (actual_list, expected_list) }}

macros/helpers/helpers.sql

@@ -7,8 +7,16 @@
 
 
 {# add to an input dictionary entries containing all the column descriptions of a given model #}
-{% macro add_model_column_descriptions_to_dict(model_name,dict_with_descriptions={}) %}
-    {% for node in graph.nodes.values() | selectattr('name', "equalto", model_name) %}
+{% macro add_model_column_descriptions_to_dict(resource_type, model_name, dict_with_descriptions={}) %}
+    {% if resource_type == 'source' %}
+        {# sources aren't part of graph.nodes #}
+        {% set nodes = graph.sources %}
+    {% else %}
+        {% set nodes = graph.nodes %}
+    {% endif %}
+    {% for node in nodes.values()
+        | selectattr('resource_type', 'equalto', resource_type)
+        | selectattr('name', 'equalto', model_name) %}
         {% for col_name, col_values in node.columns.items() %}
             {% do dict_with_descriptions.update( {col_name: col_values.description} ) %}
         {% endfor %}
@@ -22,7 +30,9 @@
     {% if execute %}
         {% set glob_dict = {} %}
         {% for full_model in codegen.get_model_dependencies(model_name) %}
-            {% do codegen.add_model_column_descriptions_to_dict(full_model.split('.')[-1],glob_dict) %}
+            {% do codegen.add_model_column_descriptions_to_dict(
+                full_model.split('.')[0], full_model.split('.')[-1], glob_dict
+            ) %}
         {% endfor %}
         {{ return(glob_dict) }}
     {% endif %}