Coverage for dj/models/column.py: 100%

40 statements  

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

1""" 

2Models for columns. 

3""" 

4from typing import TYPE_CHECKING, List, Optional, Tuple, TypedDict 

5 

6from pydantic import root_validator 

7from sqlalchemy import TypeDecorator 

8from sqlalchemy.sql.schema import Column as SqlaColumn 

9from sqlalchemy.types import Text 

10from sqlmodel import Field, Relationship 

11 

12from dj.models.base import BaseSQLModel 

13from dj.sql.parsing.types import ColumnType 

14 

15if TYPE_CHECKING: 

16 from dj.models.attribute import ColumnAttribute 

17 from dj.models.node import Node 

18 

19 

20class ColumnYAML(TypedDict, total=False): 

21 """ 

22 Schema of a column in the YAML file. 

23 """ 

24 

25 type: str 

26 dimension: str 

27 

28 

29class ColumnTypeDecorator(TypeDecorator): # pylint: disable=abstract-method 

30 """ 

31 Converts a column type from the database to a `ColumnType` class 

32 """ 

33 

34 impl = Text 

35 

36 def process_bind_param(self, value: ColumnType, dialect): 

37 return str(value) 

38 

39 def process_result_value(self, value, dialect): 

40 from dj.sql.parsing.backends.antlr4 import ( # pylint: disable=import-outside-toplevel 

41 parse_rule, 

42 ) 

43 

44 return parse_rule(value, "dataType") 

45 

46 

47class Column(BaseSQLModel, table=True): # type: ignore 

48 """ 

49 A column. 

50 

51 Columns can be physical (associated with ``Table`` objects) or abstract (associated 

52 with ``Node`` objects). 

53 """ 

54 

55 id: Optional[int] = Field(default=None, primary_key=True) 

56 name: str 

57 type: ColumnType = Field(sa_column=SqlaColumn(ColumnTypeDecorator, nullable=False)) 

58 

59 dimension_id: Optional[int] = Field(default=None, foreign_key="node.id") 

60 dimension: "Node" = Relationship( 

61 sa_relationship_kwargs={ 

62 "lazy": "joined", 

63 }, 

64 ) 

65 dimension_column: Optional[str] = None 

66 

67 attributes: List["ColumnAttribute"] = Relationship( 

68 back_populates="column", 

69 sa_relationship_kwargs={ 

70 "lazy": "joined", 

71 }, 

72 ) 

73 

74 def identifier(self) -> Tuple[str, ColumnType]: 

75 """ 

76 Unique identifier for this column. 

77 """ 

78 return self.name, self.type 

79 

80 def __hash__(self) -> int: 

81 return hash(self.id) 

82 

83 class Config: # pylint: disable=missing-class-docstring, too-few-public-methods 

84 arbitrary_types_allowed = True 

85 

86 @root_validator 

87 def type_string(cls, values): # pylint: disable=no-self-argument 

88 """ 

89 Processes the column type 

90 """ 

91 values["type"] = str(values.get("type")) 

92 return values 

93 

94 

95class ColumnAttributeInput(BaseSQLModel): 

96 """ 

97 A column attribute input 

98 """ 

99 

100 attribute_type_namespace: Optional[str] = "system" 

101 attribute_type_name: str 

102 column_name: str