diff --git a/src/dbt_osmosis/core/column_level_knowledge_propagator.py b/src/dbt_osmosis/core/column_level_knowledge_propagator.py index a332873..e1efd63 100644 --- a/src/dbt_osmosis/core/column_level_knowledge_propagator.py +++ b/src/dbt_osmosis/core/column_level_knowledge_propagator.py @@ -194,6 +194,7 @@ def update_undocumented_columns_with_prior_knowledge( skip_add_tags: bool, skip_merge_meta: bool, add_progenitor_to_meta: bool, + add_inheritance_for_specified_keys: Iterable[str] = [], ) -> int: """Update undocumented columns with prior knowledge in node and model simultaneously THIS MUTATES THE NODE AND MODEL OBJECTS so that state is always accurate""" @@ -202,6 +203,9 @@ def update_undocumented_columns_with_prior_knowledge( inheritables.append("tags") if not skip_merge_meta: inheritables.append("meta") + for key in add_inheritance_for_specified_keys: + if key not in inheritables: + inheritables.append(key) changes_committed = 0 for column in undocumented_columns: diff --git a/src/dbt_osmosis/core/osmosis.py b/src/dbt_osmosis/core/osmosis.py index b61b770..6bb51dc 100644 --- a/src/dbt_osmosis/core/osmosis.py +++ b/src/dbt_osmosis/core/osmosis.py @@ -109,6 +109,7 @@ def __init__( vars: Optional[str] = None, use_unrendered_descriptions: bool = False, profile: Optional[str] = None, + add_inheritance_for_specified_keys: Optional[List[str]] = None, ): """Initializes the DbtYamlManager class.""" super().__init__(target, profiles_dir, project_dir, threads, vars=vars, profile=profile) @@ -123,6 +124,7 @@ def __init__( self.skip_merge_meta = skip_merge_meta self.add_progenitor_to_meta = add_progenitor_to_meta self.use_unrendered_descriptions = use_unrendered_descriptions + self.add_inheritance_for_specified_keys = add_inheritance_for_specified_keys or [] if len(list(self.filtered_models())) == 0: logger().warning( @@ -1081,6 +1083,7 @@ def update_schema_file_and_node( self.skip_add_tags, self.skip_merge_meta, self.add_progenitor_to_meta, + self.add_inheritance_for_specified_keys, ) ) n_cols_data_type_updated = self.update_columns_data_type(node, section, columns_db_meta) diff --git a/src/dbt_osmosis/main.py b/src/dbt_osmosis/main.py index 87ef64e..e9cd3e4 100644 --- a/src/dbt_osmosis/main.py +++ b/src/dbt_osmosis/main.py @@ -177,6 +177,14 @@ def wrapper(*args, **kwargs): "This is useful for propogating docs blocks" ), ) +@click.option( + "--add-inheritance-for-specified-keys", + multiple=True, + type=click.STRING, + help=( + "If specified, will add inheritance for the specified keys." + ) +) @click.argument("models", nargs=-1) def refactor( target: Optional[str] = None, @@ -196,6 +204,7 @@ def refactor( profile: Optional[str] = None, vars: Optional[str] = None, use_unrendered_descriptions: bool = False, + add_inheritance_for_specified_keys: Optional[List[str]] = None, ): """Executes organize which syncs yaml files with database schema and organizes the dbt models directory, reparses the project, then executes document passing down inheritable documentation @@ -227,6 +236,7 @@ def refactor( profile=profile, vars=vars, use_unrendered_descriptions=use_unrendered_descriptions, + add_inheritance_for_specified_keys=add_inheritance_for_specified_keys, ) # Conform project structure & bootstrap undocumented models injecting columns @@ -307,6 +317,14 @@ def refactor( " my_value}'" ), ) +@click.option( + "--add-inheritance-for-specified-keys", + multiple=True, + type=click.STRING, + help=( + "If specified, will add inheritance for the specified keys." + ) +) @click.argument("models", nargs=-1) def organize( target: Optional[str] = None, @@ -323,6 +341,7 @@ def organize( add_progenitor_to_meta: bool = False, profile: Optional[str] = None, vars: Optional[str] = None, + add_inheritance_for_specified_keys: Optional[List[str]] = None, ): """Organizes schema ymls based on config and injects undocumented models @@ -351,6 +370,7 @@ def organize( add_progenitor_to_meta=add_progenitor_to_meta, profile=profile, vars=vars, + add_inheritance_for_specified_keys=add_inheritance_for_specified_keys, ) # Conform project structure & bootstrap undocumented models injecting columns @@ -454,6 +474,14 @@ def organize( "This is useful for propogating docs blocks" ), ) +@click.option( + "--add-inheritance-for-specified-keys", + multiple=True, + type=click.STRING, + help=( + "If specified, will add inheritance for the specified keys." + ) +) @click.argument("models", nargs=-1) def document( target: Optional[str] = None, @@ -473,6 +501,7 @@ def document( profile: Optional[str] = None, vars: Optional[str] = None, use_unrendered_descriptions: bool = False, + add_inheritance_for_specified_keys: Optional[List[str]] = None, ): """Column level documentation inheritance for existing models @@ -503,6 +532,7 @@ def document( profile=profile, vars=vars, use_unrendered_descriptions=use_unrendered_descriptions, + add_inheritance_for_specified_keys=add_inheritance_for_specified_keys, ) # Propagate documentation & inject/remove schema file columns to align with model in database diff --git a/tests/test_column_level_knowledge_propagator.py b/tests/test_column_level_knowledge_propagator.py index 7cda8d1..69755e8 100644 --- a/tests/test_column_level_knowledge_propagator.py +++ b/tests/test_column_level_knowledge_propagator.py @@ -486,6 +486,78 @@ def test_update_undocumented_columns_with_prior_knowledge_add_progenitor_to_meta } assert set(target_node.columns["customer_id"].tags) == set(["my_tag1", "my_tag2"]) +def test_update_undocumented_columns_with_prior_knowledge_with_add_inheritance_for_specified_keys(): + manifest = load_manifest() + manifest.nodes["model.jaffle_shop_duckdb.stg_customers"].columns[ + "customer_id" + ].description = "THIS COLUMN IS UPDATED FOR TESTING" + manifest.nodes["model.jaffle_shop_duckdb.stg_customers"].columns["customer_id"].meta = { + "my_key": "my_value" + } + manifest.nodes["model.jaffle_shop_duckdb.stg_customers"].columns["customer_id"].tags = [ + "my_tag1", + "my_tag2", + ] + manifest.nodes["model.jaffle_shop_duckdb.stg_customers"].columns["customer_id"]._extra = { + "policy_tags": ["my_policy_tag1"], + } + + target_node_name = "model.jaffle_shop_duckdb.customers" + manifest.nodes[target_node_name].columns["customer_id"].tags = set( + [ + "my_tag3", + "my_tag4", + ] + ) + manifest.nodes[target_node_name].columns["customer_id"].meta = { + "my_key": "my_old_value", + "my_new_key": "my_new_value", + } + target_node = manifest.nodes[target_node_name] + knowledge = ColumnLevelKnowledgePropagator.get_node_columns_with_inherited_knowledge( + manifest, target_node, placeholders=[""] + ) + yaml_file_model_section = { + "columns": [ + { + "name": "customer_id", + } + ] + } + undocumented_columns = target_node.columns.keys() + ColumnLevelKnowledgePropagator.update_undocumented_columns_with_prior_knowledge( + undocumented_columns, + target_node, + yaml_file_model_section, + knowledge, + skip_add_tags=False, + skip_merge_meta=False, + add_progenitor_to_meta=False, + add_inheritance_for_specified_keys=["policy_tags"], + ) + + assert yaml_file_model_section["columns"][0]["name"] == "customer_id" + assert ( + yaml_file_model_section["columns"][0]["description"] == "THIS COLUMN IS UPDATED FOR TESTING" + ) + assert yaml_file_model_section["columns"][0]["meta"] == { + "my_key": "my_value", + "my_new_key": "my_new_value", + } + assert set(yaml_file_model_section["columns"][0]["tags"]) == set( + ["my_tag1", "my_tag2", "my_tag3", "my_tag4"] + ) + assert set(yaml_file_model_section["columns"][0]["policy_tags"]) == set(["my_policy_tag1"]) + + assert target_node.columns["customer_id"].description == "THIS COLUMN IS UPDATED FOR TESTING" + assert target_node.columns["customer_id"].meta == { + "my_key": "my_value", + "my_new_key": "my_new_value", + } + assert set(target_node.columns["customer_id"].tags) == set( + ["my_tag1", "my_tag2", "my_tag3", "my_tag4"] + ) + assert set(target_node.columns["customer_id"]._extra["policy_tags"]) == set(["my_policy_tag1"]) @pytest.mark.parametrize("use_unrendered_descriptions", [True, False]) def test_use_unrendered_descriptions(use_unrendered_descriptions):