Skip to content

Commit

Permalink
feat: add output-to-lower
Browse files Browse the repository at this point in the history 
This command allows for the outputs to be created in lower case. Currently this only affects data_types.
  • Loading branch information
@bkimjin
bkimjin committed Nov 21, 2024
1 parent 8fecb0d commit 79202da
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 22 deletions.
82 changes: 62 additions & 20 deletions src/dbt_osmosis/core/osmosis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ def __init__(
use_unrendered_descriptions: bool = False,
profile: Optional[str] = None,
add_inheritance_for_specified_keys: Optional[List[str]] = None,
output_to_lower: bool = False,
):
"""Initializes the DbtYamlManager class."""
super().__init__(target, profiles_dir, project_dir, threads, vars=vars, profile=profile)
Expand All @@ -145,6 +146,7 @@ def __init__(
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 []
self.output_to_lower = output_to_lower

if len(list(self.filtered_models())) == 0:
logger().warning(
Expand Down Expand Up @@ -846,6 +848,7 @@ def _run(
node: ManifestNode,
schema_map: Dict[str, SchemaFileLocation],
force_inheritance: bool = False,
output_to_lower: bool = False,
):
try:
with self.mutex:
Expand Down Expand Up @@ -930,6 +933,7 @@ def _run(
node,
section,
columns_db_meta,
output_to_lower,
)
if (
n_cols_added
Expand Down Expand Up @@ -1002,13 +1006,17 @@ def _sort_columns(column_info: dict) -> int:
logger().error("Error occurred while processing model %s: %s", unique_id, e)
raise e

def propagate_documentation_downstream(self, force_inheritance: bool = False) -> None:
def propagate_documentation_downstream(
self, force_inheritance: bool = False, output_to_lower: bool = False
) -> None:
schema_map = self.build_schema_folder_mapping()
futs = []
with self.adapter.connection_named("dbt-osmosis"):
for unique_id, node in self.filtered_models():
futs.append(
self.tpe.submit(self._run, unique_id, node, schema_map, force_inheritance)
self.tpe.submit(
self._run, unique_id, node, schema_map, force_inheritance, output_to_lower
)
)
wait(futs)

Expand Down Expand Up @@ -1041,6 +1049,7 @@ def update_columns_attribute(
attribute_name: str,
meta_key: str,
skip_attribute_update: Any,
output_to_lower: bool = False,
) -> int:
changes_committed = 0
if (skip_attribute_update is True) or (skip_attribute_update is None):
Expand All @@ -1059,7 +1068,10 @@ def update_columns_attribute(
setattr(node.columns[cased_column_name], attribute_name, column_meta)
for model_column in yaml_file_model_section["columns"]:
if self.column_casing(model_column["name"]) == cased_column_name:
model_column.update({attribute_name: column_meta})
if output_to_lower:
model_column.update({attribute_name: column_meta.lower()})
else:
model_column.update({attribute_name: column_meta})
changes_committed += 1
return changes_committed

Expand All @@ -1069,25 +1081,42 @@ def add_missing_cols_to_node_and_model(
node: ManifestNode,
yaml_file_model_section: Dict[str, Any],
columns_db_meta: Dict[str, ColumnMetadata],
output_to_lower: bool,
) -> int:
"""Add missing columns to node and model simultaneously
THIS MUTATES THE NODE AND MODEL OBJECTS so that state is always accurate"""
changes_committed = 0
for column in missing_columns:
node.columns[column] = ColumnInfo.from_dict(
{
"name": column,
"description": columns_db_meta[column].comment or "",
"data_type": columns_db_meta[column].type,
}
)
yaml_file_model_section.setdefault("columns", []).append(
{
"name": column,
"description": columns_db_meta[column].comment or "",
"data_type": columns_db_meta[column].type,
}
)
if output_to_lower:
node.columns[column] = ColumnInfo.from_dict(
{
"name": column.lower(),
"description": columns_db_meta[column].comment or "",
"data_type": columns_db_meta[column].type.lower(),
}
)
yaml_file_model_section.setdefault("columns", []).append(
{
"name": column.lower(),
"description": columns_db_meta[column].comment or "",
"data_type": columns_db_meta[column].type.lower(),
}
)
else:
node.columns[column] = ColumnInfo.from_dict(
{
"name": column,
"description": columns_db_meta[column].comment or "",
"data_type": columns_db_meta[column].type,
}
)
yaml_file_model_section.setdefault("columns", []).append(
{
"name": column,
"description": columns_db_meta[column].comment or "",
"data_type": columns_db_meta[column].type,
}
)
changes_committed += 1
logger().info(
":syringe: Injecting column %s into dbt schema for model %s", column, node.unique_id
Expand All @@ -1102,13 +1131,14 @@ def update_schema_file_and_node(
node: ManifestNode,
section: Dict[str, Any],
columns_db_meta: Dict[str, ColumnMetadata],
output_to_lower: bool,
) -> Tuple[int, int, int, int, int]:
"""Take action on a schema file mirroring changes in the node."""
logger().info(":microscope: Looking for actions for %s", node.unique_id)
n_cols_added = 0
if not self.skip_add_columns:
n_cols_added = self.add_missing_cols_to_node_and_model(
missing_columns, node, section, columns_db_meta
missing_columns, node, section, columns_db_meta, output_to_lower
)

knowledge = ColumnLevelKnowledgePropagator.get_node_columns_with_inherited_knowledge(
Expand All @@ -1131,10 +1161,22 @@ def update_schema_file_and_node(
)
)
n_cols_data_type_updated = self.update_columns_attribute(
node, section, columns_db_meta, "data_type", "type", self.skip_add_data_types
node,
section,
columns_db_meta,
"data_type",
"type",
self.skip_add_data_types,
self.output_to_lower,
)
n_cols_description_updated = self.update_columns_attribute(
node, section, columns_db_meta, "description", "comment", self.catalog_file
node,
section,
columns_db_meta,
"description",
"comment",
self.catalog_file,
self.output_to_lower,
)
n_cols_removed = self.remove_columns_not_in_database(extra_columns, node, section)
return (
Expand Down
27 changes: 25 additions & 2 deletions src/dbt_osmosis/main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,11 @@ def wrapper(*args, **kwargs):
type=click.STRING,
help="If specified, will add inheritance for the specified keys.",
)
@click.option(
"--output-to-lower",
is_flag=True,
help="If specified, output yaml file in lowercase if possible.",
)
@click.argument("models", nargs=-1)
def refactor(
target: Optional[str] = None,
Expand All @@ -215,6 +220,7 @@ def refactor(
vars: Optional[str] = None,
use_unrendered_descriptions: bool = False,
add_inheritance_for_specified_keys: Optional[List[str]] = None,
output_to_lower: bool = False,
):
"""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
Expand Down Expand Up @@ -249,12 +255,15 @@ def refactor(
vars=vars,
use_unrendered_descriptions=use_unrendered_descriptions,
add_inheritance_for_specified_keys=add_inheritance_for_specified_keys,
output_to_lower=output_to_lower,
)

# Conform project structure & bootstrap undocumented models injecting columns
if runner.commit_project_restructure_to_disk():
runner.safe_parse_project(reinit=True)
runner.propagate_documentation_downstream(force_inheritance=force_inheritance)
runner.propagate_documentation_downstream(
force_inheritance=force_inheritance, output_to_lower=output_to_lower
)
if check and runner.mutations > 0:
exit(1)

Expand Down Expand Up @@ -345,6 +354,11 @@ def refactor(
type=click.STRING,
help="If specified, will add inheritance for the specified keys.",
)
@click.option(
"--output-to-lower",
is_flag=True,
help="If specified, output yaml file in lowercase if possible.",
)
@click.argument("models", nargs=-1)
def organize(
target: Optional[str] = None,
Expand All @@ -364,6 +378,7 @@ def organize(
profile: Optional[str] = None,
vars: Optional[str] = None,
add_inheritance_for_specified_keys: Optional[List[str]] = None,
output_to_lower: bool = False,
):
"""Organizes schema ymls based on config and injects undocumented models
Expand Down Expand Up @@ -395,6 +410,7 @@ def organize(
profile=profile,
vars=vars,
add_inheritance_for_specified_keys=add_inheritance_for_specified_keys,
output_to_lower=output_to_lower,
)

# Conform project structure & bootstrap undocumented models injecting columns
Expand Down Expand Up @@ -514,6 +530,11 @@ def organize(
type=click.STRING,
help="If specified, will add inheritance for the specified keys.",
)
@click.option(
"--output-to-lower",
is_flag=True,
help="If specified, output yaml file in lowercase if possible.",
)
@click.argument("models", nargs=-1)
def document(
target: Optional[str] = None,
Expand All @@ -536,6 +557,7 @@ def document(
vars: Optional[str] = None,
use_unrendered_descriptions: bool = False,
add_inheritance_for_specified_keys: Optional[List[str]] = None,
output_to_lower: bool = False,
):
"""Column level documentation inheritance for existing models
Expand Down Expand Up @@ -569,10 +591,11 @@ def document(
vars=vars,
use_unrendered_descriptions=use_unrendered_descriptions,
add_inheritance_for_specified_keys=add_inheritance_for_specified_keys,
output_to_lower=output_to_lower,
)

# Propagate documentation & inject/remove schema file columns to align with model in database
runner.propagate_documentation_downstream(force_inheritance)
runner.propagate_documentation_downstream(force_inheritance, output_to_lower)
if check and runner.mutations > 0:
exit(1)

Expand Down

0 comments on commit 79202da

Please sign in to comment.