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

29 statements  

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

1""" 

2Models for tables. 

3""" 

4 

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

6 

7from sqlmodel import Field, Relationship 

8 

9from dj.models.base import BaseSQLModel 

10 

11if TYPE_CHECKING: 

12 from dj.models.catalog import Catalog 

13 from dj.models.column import Column 

14 from dj.models.database import Database 

15 from dj.models.node import NodeRevision 

16 

17 

18class TableYAML(TypedDict, total=False): 

19 """ 

20 Schema of a table in the YAML file. 

21 """ 

22 

23 catalog: Optional[str] 

24 schema: Optional[str] 

25 table: str 

26 cost: float 

27 

28 

29class TableColumns(BaseSQLModel, table=True): # type: ignore 

30 """ 

31 Join table for table columns. 

32 """ 

33 

34 table_id: Optional[int] = Field( 

35 default=None, 

36 foreign_key="table.id", 

37 primary_key=True, 

38 ) 

39 column_id: Optional[int] = Field( 

40 default=None, 

41 foreign_key="column.id", 

42 primary_key=True, 

43 ) 

44 

45 

46class TableBase(BaseSQLModel): 

47 """ 

48 A base table. 

49 """ 

50 

51 schema_: Optional[str] = Field(default=None, alias="schema") 

52 table: str 

53 cost: float = 1.0 

54 

55 

56class Table(TableBase, table=True): # type: ignore 

57 """ 

58 A table with data. 

59 

60 Nodes can have data in multiple tables, in different databases. 

61 """ 

62 

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

64 

65 database_id: int = Field(foreign_key="database.id") 

66 database: "Database" = Relationship(back_populates="tables") 

67 

68 columns: List["Column"] = Relationship( 

69 link_model=TableColumns, 

70 sa_relationship_kwargs={ 

71 "primaryjoin": "Table.id==TableColumns.table_id", 

72 "secondaryjoin": "Column.id==TableColumns.column_id", 

73 "cascade": "all, delete", 

74 }, 

75 ) 

76 

77 def identifier( 

78 self, 

79 ) -> Tuple[Optional[str], Optional[str], str]: # pragma: no cover 

80 """ 

81 Unique identifier for this table. 

82 """ 

83 # Catalogs will soon be required and this return can be simplified 

84 return ( 

85 self.catalog.name if self.catalog else None, # pylint: disable=no-member 

86 self.schema_, 

87 self.table, 

88 ) 

89 

90 def __hash__(self): 

91 return hash(self.id) 

92 

93 

94class CreateColumn(BaseSQLModel): 

95 """ 

96 A column creation request 

97 """ 

98 

99 name: str 

100 type: str 

101 

102 

103class CreateTable(TableBase): 

104 """ 

105 Create table input 

106 """ 

107 

108 database_name: str 

109 catalog_name: str 

110 columns: List[CreateColumn]