diff --git a/dbt_loom/__init__.py b/dbt_loom/__init__.py index e44556a..48fea15 100644 --- a/dbt_loom/__init__.py +++ b/dbt_loom/__init__.py @@ -92,12 +92,7 @@ def convert_model_nodes_to_model_node_args( unique_id: LoomModelNodeArgs( schema=node.schema_name, identifier=node.identifier, - **( - # Small bit of logic to support both pydantic 2 and pydantic 1 - node.model_dump(exclude={"schema_name", "depends_on", "node_config"}) # type: ignore - if hasattr(node, "model_dump") - else node.dict(exclude={"schema_name", "depends_on", "node_config"}) - ), + **(node.dump()), ) for unique_id, node in selected_nodes.items() if node is not None diff --git a/dbt_loom/manifests.py b/dbt_loom/manifests.py index 2bfd5bf..e7a3b0a 100644 --- a/dbt_loom/manifests.py +++ b/dbt_loom/manifests.py @@ -33,8 +33,9 @@ class ManifestNode(BaseModel): """A basic ManifestNode that can be referenced across projects.""" name: str - resource_type: NodeType package_name: str + unique_id: str + resource_type: NodeType schema_name: str = Field(alias="schema") database: Optional[str] = None relation_name: Optional[str] = None @@ -58,6 +59,15 @@ def default_depends_on_nodes(cls, v, values): node for node in depends_on.nodes if node.split(".")[0] not in ("source") ] + @validator("resource_type", always=True) + def fix_resource_types(cls, v, values): + """If the resource type does not match the unique_id prefix, then rewrite the resource type.""" + + node_type = values.get("unique_id").split(".")[0] + if v != node_type: + return node_type + return v + @property def identifier(self) -> str: if not self.relation_name: @@ -65,6 +75,14 @@ def identifier(self) -> str: return self.relation_name.split(".")[-1].replace('"', "").replace("`", "") + def dump(self) -> Dict: + """Dump the ManifestNode to a Dict, with support for pydantic 1 and 2""" + exclude_set = {"schema_name", "depends_on", "node_config", "unique_id"} + if hasattr(self, "model_dump"): + return self.model_dump(exclude=exclude_set) # type: ignore + + return self.dict(exclude=exclude_set) + class ManifestLoader: def __init__(self): @@ -81,9 +99,9 @@ def load_from_local_filesystem(config: FileReferenceConfig) -> Dict: """Load a manifest dictionary from a local file""" if not config.path.exists(): raise LoomConfigurationError(f"The path `{config.path}` does not exist.") - - if config.path.suffix == '.gz': - with gzip.open(config.path, 'rt') as file: + + if config.path.suffix == ".gz": + with gzip.open(config.path, "rt") as file: return json.load(file) else: return json.load(open(config.path)) diff --git a/tests/test_mainfest_node.py b/tests/test_mainfest_node.py new file mode 100644 index 0000000..6cf0f10 --- /dev/null +++ b/tests/test_mainfest_node.py @@ -0,0 +1,23 @@ +from dbt_loom.manifests import ManifestNode + + +try: + from dbt.artifacts.resources.types import NodeType +except ModuleNotFoundError: + from dbt.node_types import NodeType # type: ignore + + +def test_rewrite_resource_types(): + """Confirm that resource types are rewritten if they are incorrect due to previous injections.""" + + node = { + "unique_id": "seed.example.foo", + "name": "foo", + "package_name": "example", + "schema": "bar", + "resource_type": "model", + } + + manifest_node = ManifestNode(**(node)) # type: ignore + + assert manifest_node.resource_type == NodeType.Seed