[feature] add inheritance for specified keys (#147)

* add inheritance for specified keys

* update unittest

* remove unused variable

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:

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)

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

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):