Coverage for dj/sql/dag.py: 100%

18 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-04-17 20:05 -0700

1""" 

2DAG related functions. 

3""" 

4import collections 

5from typing import List 

6 

7from dj.models.node import Node, NodeType 

8from dj.utils import get_settings 

9 

10settings = get_settings() 

11 

12 

13def get_dimensions(node: Node) -> List[str]: 

14 """ 

15 Return all available dimensions for a given node. 

16 """ 

17 dimensions = [] 

18 

19 # Start with the node itself or the node's immediate parent if it's a metric node 

20 to_process = collections.deque( 

21 [node, *(node.current.parents if node.type == NodeType.METRIC else [])], 

22 ) 

23 processed = set() 

24 

25 while to_process: 

26 current_node = to_process.popleft() 

27 processed.add(current_node) 

28 

29 for column in current_node.current.columns: 

30 # Include the dimension if it's a column belonging to a dimension node 

31 # or if it's tagged with the dimension column attribute 

32 if current_node.type == NodeType.DIMENSION or any( 

33 attr.attribute_type.name == "dimension" for attr in column.attributes 

34 ): 

35 dimensions.append(f"{current_node.name}.{column.name}") 

36 if column.dimension and column.dimension not in processed: 

37 to_process.append(column.dimension) 

38 return sorted(dimensions)