|
| 1 | +from typing import List, Tuple, Union |
| 2 | + |
| 3 | +from dbterd.adapters.algos import base |
| 4 | +from dbterd.adapters.meta import Ref, SemanticEntity, Table |
| 5 | +from dbterd.constants import TEST_META_RELATIONSHIP_TYPE |
| 6 | +from dbterd.helpers.log import logger |
| 7 | +from dbterd.types import Catalog, Manifest |
| 8 | + |
| 9 | + |
| 10 | +def parse_metadata(data, **kwargs) -> Tuple[List[Table], List[Ref]]: |
| 11 | + raise NotImplementedError() # pragma: no cover |
| 12 | + |
| 13 | + |
| 14 | +def parse( |
| 15 | + manifest: Manifest, catalog: Union[str, Catalog], **kwargs |
| 16 | +) -> Tuple[List[Table], List[Ref]]: |
| 17 | + # Parse metadata |
| 18 | + if catalog == "metadata": # pragma: no cover |
| 19 | + return parse_metadata(data=manifest, **kwargs) |
| 20 | + |
| 21 | + # Parse Table |
| 22 | + tables = base.get_tables(manifest=manifest, catalog=catalog, **kwargs) |
| 23 | + tables = base.filter_tables_based_on_selection(tables=tables, **kwargs) |
| 24 | + |
| 25 | + # Parse Ref |
| 26 | + relationships = _get_relationships(manifest=manifest, **kwargs) |
| 27 | + relationships = base.make_up_relationships( |
| 28 | + relationships=relationships, tables=tables |
| 29 | + ) |
| 30 | + |
| 31 | + # Fulfill columns in Tables (due to `select *`) |
| 32 | + tables = base.enrich_tables_from_relationships( |
| 33 | + tables=tables, relationships=relationships |
| 34 | + ) |
| 35 | + |
| 36 | + logger.info( |
| 37 | + f"Collected {len(tables)} table(s) and {len(relationships)} relationship(s)" |
| 38 | + ) |
| 39 | + return ( |
| 40 | + sorted(tables, key=lambda tbl: tbl.node_name), |
| 41 | + sorted(relationships, key=lambda rel: rel.name), |
| 42 | + ) |
| 43 | + |
| 44 | + |
| 45 | +def find_related_nodes_by_id( |
| 46 | + manifest: Union[Manifest, dict], node_unique_id: str, type: str = None, **kwargs |
| 47 | +) -> List[str]: |
| 48 | + """Find FK/PK nodes which are linked to the given node |
| 49 | +
|
| 50 | + Args: |
| 51 | + manifest (Union[Manifest, dict]): Manifest data |
| 52 | + node_unique_id (str): Manifest model node unique id |
| 53 | + type (str, optional): Manifest type (local file or metadata). Defaults to None. |
| 54 | +
|
| 55 | + Returns: |
| 56 | + List[str]: Manifest nodes' unique ID |
| 57 | + """ |
| 58 | + found_nodes = [node_unique_id] |
| 59 | + if type == "metadata": # pragma: no cover |
| 60 | + return found_nodes # not supported yet, returned input only |
| 61 | + |
| 62 | + entities = _get_linked_semantic_entities(manifest=manifest) |
| 63 | + for foreign, primary in entities: |
| 64 | + if primary.model == node_unique_id: |
| 65 | + found_nodes.append(foreign.model) |
| 66 | + if foreign.model == node_unique_id: |
| 67 | + found_nodes.append(primary.model) |
| 68 | + |
| 69 | + return list(set(found_nodes)) |
| 70 | + |
| 71 | + |
| 72 | +def _get_relationships(manifest: Manifest, **kwargs) -> List[Ref]: |
| 73 | + """_summary_ |
| 74 | +
|
| 75 | + Args: |
| 76 | + manifest (Manifest): Extract relationships from dbt artifacts based on Semantic Entities |
| 77 | +
|
| 78 | + Returns: |
| 79 | + List[Ref]: List of parsed relationship |
| 80 | + """ |
| 81 | + entities = _get_linked_semantic_entities(manifest=manifest) |
| 82 | + return base.get_unique_refs( |
| 83 | + refs=[ |
| 84 | + Ref( |
| 85 | + name=primary_entity.semantic_model, |
| 86 | + table_map=(primary_entity.model, foreign_entity.model), |
| 87 | + column_map=( |
| 88 | + primary_entity.column_name, |
| 89 | + foreign_entity.column_name, |
| 90 | + ), |
| 91 | + type=primary_entity.relationship_type, |
| 92 | + ) |
| 93 | + for foreign_entity, primary_entity in entities |
| 94 | + ] |
| 95 | + ) |
| 96 | + |
| 97 | + |
| 98 | +def _get_linked_semantic_entities( |
| 99 | + manifest: Manifest, |
| 100 | +) -> List[Tuple[SemanticEntity, SemanticEntity]]: |
| 101 | + """Get filtered list of Semantic Entities which are linked |
| 102 | +
|
| 103 | + Args: |
| 104 | + manifest (Manifest): Manifest data |
| 105 | +
|
| 106 | + Returns: |
| 107 | + List[Tuple[SemanticEntity, SemanticEntity]]: List of (FK, PK) objects |
| 108 | + """ |
| 109 | + foreigns, primaries = _get_semantic_entities(manifest=manifest) |
| 110 | + linked_entities = [] |
| 111 | + for foreign_entity in foreigns: |
| 112 | + for primary_entity in primaries: |
| 113 | + if foreign_entity.entity_name == primary_entity.entity_name: |
| 114 | + linked_entities.append((foreign_entity, primary_entity)) |
| 115 | + return linked_entities |
| 116 | + |
| 117 | + |
| 118 | +def _get_semantic_entities( |
| 119 | + manifest: Manifest, |
| 120 | +) -> Tuple[List[SemanticEntity], List[SemanticEntity]]: |
| 121 | + """Get all Semantic Entities |
| 122 | +
|
| 123 | + Args: |
| 124 | + manifest (Manifest): Manifest data |
| 125 | +
|
| 126 | + Returns: |
| 127 | + Tuple[List[SemanticEntity], List[SemanticEntity]]: FK list and PK list |
| 128 | + """ |
| 129 | + FK = "foreign" |
| 130 | + PK = "primary" |
| 131 | + |
| 132 | + semantic_entities = [] |
| 133 | + for x in _get_semantic_nodes(manifest=manifest): |
| 134 | + semantic_node = manifest.semantic_models[x] |
| 135 | + for e in semantic_node.entities: |
| 136 | + if e.type.value in [PK, FK]: |
| 137 | + semantic_entities.append( |
| 138 | + SemanticEntity( |
| 139 | + semantic_model=x, |
| 140 | + model=semantic_node.depends_on.nodes[0], |
| 141 | + entity_name=e.name, |
| 142 | + entity_type=e.type.value, |
| 143 | + column_name=e.expr or e.name, |
| 144 | + relationship_type=semantic_node.config.meta.get( |
| 145 | + TEST_META_RELATIONSHIP_TYPE, "" |
| 146 | + ), |
| 147 | + ) |
| 148 | + ) |
| 149 | + if semantic_node.primary_entity: |
| 150 | + semantic_entities.append( |
| 151 | + SemanticEntity( |
| 152 | + semantic_model=x, |
| 153 | + model=semantic_node.depends_on.nodes[0], |
| 154 | + entity_name=semantic_node.primary_entity, |
| 155 | + entity_type=PK, |
| 156 | + column_name=semantic_node.primary_entity, |
| 157 | + relationship_type=semantic_node.config.meta.get( |
| 158 | + TEST_META_RELATIONSHIP_TYPE, "" |
| 159 | + ), |
| 160 | + ) |
| 161 | + ) |
| 162 | + |
| 163 | + return ( |
| 164 | + [x for x in semantic_entities if x.entity_type == FK], |
| 165 | + [x for x in semantic_entities if x.entity_type == PK], |
| 166 | + ) |
| 167 | + |
| 168 | + |
| 169 | +def _get_semantic_nodes(manifest: Manifest) -> List: |
| 170 | + """Extract the Semantic Models |
| 171 | +
|
| 172 | + Args: |
| 173 | + manifest (Manifest): Manifest data |
| 174 | +
|
| 175 | + Returns: |
| 176 | + List: List of Semantic Models |
| 177 | + """ |
| 178 | + if not hasattr(manifest, "semantic_models"): |
| 179 | + logger.warning( |
| 180 | + "No relationships will be captured" |
| 181 | + "since dbt version is NOT supported for the Semantic Models" |
| 182 | + ) |
| 183 | + return [] |
| 184 | + |
| 185 | + return [ |
| 186 | + x |
| 187 | + for x in manifest.semantic_models |
| 188 | + if len(manifest.semantic_models[x].depends_on.nodes) |
| 189 | + ] |
0 commit comments