Ele 3941 collect column info (#772)

* replace old flags with new * change condition - meta & tags * pre-commit * fix upload columns
elementary-data · Jan 8, 2025 · 408e679 · 408e679
1 parent 2c599d1
commit 408e679
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 13 deletions.
diff --git a/integration_tests/tests/dbt_project.py b/integration_tests/tests/dbt_project.py
@@ -18,7 +18,7 @@
 _DEFAULT_VARS = {
     "disable_dbt_invocation_autoupload": True,
     "disable_dbt_artifacts_autoupload": True,
-    "disable_dbt_columns_autoupload": True,
+    "columns_upload_strategy": "none",
     "disable_run_results": True,
     "disable_freshness_results": True,
     "debug_logs": True,

diff --git a/integration_tests/tests/test_dbt_artifacts/test_columns.py b/integration_tests/tests/test_dbt_artifacts/test_columns.py
@@ -28,17 +28,17 @@
 @Parametrization.autodetect_parameters()
 @Parametrization.case(
     name="default",
-    only_with_description=None,
+    columns_upload_strategy=None,
     expected_columns=["with_description"],
 )
 @Parametrization.case(
     name="only_with_description",
-    only_with_description=True,
+    columns_upload_strategy="enriched_only",
     expected_columns=["with_description"],
 )
 @Parametrization.case(
     name="all",
-    only_with_description=False,
+    columns_upload_strategy="all",
     expected_columns=[
         "with_description",
         "without_description",
@@ -48,13 +48,11 @@
 )
 def test_flatten_table_columns(
     dbt_project: DbtProject,
-    only_with_description: Optional[bool],
+    columns_upload_strategy: Optional[str],
     expected_columns: List[str],
 ) -> None:
-    if only_with_description is not None:
-        dbt_project.dbt_runner.vars[
-            "upload_only_columns_with_descriptions"
-        ] = only_with_description
+    if columns_upload_strategy is not None:
+        dbt_project.dbt_runner.vars["columns_upload_strategy"] = columns_upload_strategy
     flattened_columns = json.loads(
         dbt_project.dbt_runner.run_operation(
             "elementary.flatten_table_columns", macro_args={"table_node": TABLE_NODE}

diff --git a/macros/edr/dbt_artifacts/upload_dbt_artifacts.sql b/macros/edr/dbt_artifacts/upload_dbt_artifacts.sql
@@ -11,7 +11,7 @@
       }
     %}
 
-    {% if not elementary.get_config_var('disable_dbt_columns_autoupload') %}
+    {% if elementary.get_config_var('columns_upload_strategy') != 'none' %}
       {% do model_upload_func_map.update({"dbt_columns": elementary.upload_dbt_columns}) %}
     {% endif %}
 

diff --git a/macros/edr/dbt_artifacts/upload_dbt_columns.sql b/macros/edr/dbt_artifacts/upload_dbt_columns.sql
@@ -35,7 +35,17 @@
 
     {% set flattened_columns = [] %}
     {% for column_node in column_nodes.values() %}
-        {% if not elementary.get_config_var('upload_only_columns_with_descriptions') or column_node.get('description') %}
+        {% set config_dict = elementary.safe_get_with_default(column_node, 'config', {}) %}
+        {% set config_meta_dict = elementary.safe_get_with_default(config_dict, 'meta') %}
+        {% set meta_dict = elementary.safe_get_with_default(column_node, 'meta', {}) %}
+        {% set has_meta = config_meta_dict or meta_dict | length > 0 %}
+
+        {% set config_tags = elementary.safe_get_with_default(config_dict, 'tags') %}
+        {% set global_tags = elementary.safe_get_with_default(column_node, 'tags') %}
+        {% set meta_tags = elementary.safe_get_with_default(meta_dict, 'tags') %}
+        {% set has_tags = config_tags or global_tags or meta_tags %}
+
+        {% if elementary.get_config_var('columns_upload_strategy') == 'all' or column_node.get('description') or has_meta or has_tags %}
             {% set flat_column = elementary.flatten_column(table_node, column_node) %}
             {% do flattened_columns.append(flat_column) %}
         {% endif %}

diff --git a/macros/edr/system/system_utils/get_config_var.sql b/macros/edr/system/system_utils/get_config_var.sql
@@ -32,8 +32,7 @@
     'disable_freshness_results': false,
     'disable_tests_results': false,
     'disable_dbt_artifacts_autoupload': false,
-    'disable_dbt_columns_autoupload': false,
-    'upload_only_columns_with_descriptions': true,
+    'columns_upload_strategy': 'enriched_only',
     'disable_dbt_invocation_autoupload': false,
     'disable_skipped_model_alerts': true,
     'disable_skipped_test_alerts': true,