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

30 statements  

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

1""" 

2Models for columns. 

3""" 

4from datetime import datetime, timezone 

5from functools import partial 

6from typing import TYPE_CHECKING, Dict, List, Optional 

7from uuid import UUID, uuid4 

8 

9from sqlalchemy import DateTime 

10from sqlalchemy.sql.schema import Column as SqlaColumn 

11from sqlalchemy_utils import UUIDType 

12from sqlmodel import JSON, Field, Relationship, SQLModel 

13 

14from dj.models.base import BaseSQLModel 

15from dj.models.engine import Engine, EngineInfo 

16from dj.typing import UTCDatetime 

17 

18if TYPE_CHECKING: 

19 from dj.models import NodeRevision, Table 

20 

21 

22class CatalogEngines(BaseSQLModel, table=True): # type: ignore 

23 """ 

24 Join table for catalogs and engines. 

25 """ 

26 

27 catalog_id: Optional[int] = Field( 

28 default=None, 

29 foreign_key="catalog.id", 

30 primary_key=True, 

31 ) 

32 engine_id: Optional[int] = Field( 

33 default=None, 

34 foreign_key="engine.id", 

35 primary_key=True, 

36 ) 

37 

38 

39class Catalog(BaseSQLModel, table=True): # type: ignore 

40 """ 

41 A catalog. 

42 """ 

43 

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

45 uuid: UUID = Field(default_factory=uuid4, sa_column=SqlaColumn(UUIDType())) 

46 name: str 

47 engines: List[Engine] = Relationship( 

48 link_model=CatalogEngines, 

49 sa_relationship_kwargs={ 

50 "primaryjoin": "Catalog.id==CatalogEngines.catalog_id", 

51 "secondaryjoin": "Engine.id==CatalogEngines.engine_id", 

52 }, 

53 ) 

54 node_revisions: List["NodeRevision"] = Relationship(back_populates="catalog") 

55 created_at: UTCDatetime = Field( 

56 sa_column=SqlaColumn(DateTime(timezone=True)), 

57 default_factory=partial(datetime.now, timezone.utc), 

58 ) 

59 updated_at: UTCDatetime = Field( 

60 sa_column=SqlaColumn(DateTime(timezone=True)), 

61 default_factory=partial(datetime.now, timezone.utc), 

62 ) 

63 extra_params: Dict = Field(default={}, sa_column=SqlaColumn(JSON)) 

64 

65 def __str__(self) -> str: 

66 return self.name 

67 

68 def __hash__(self) -> int: 

69 return hash(self.id) 

70 

71 

72class CatalogInfo(SQLModel): 

73 """ 

74 Class for catalog creation 

75 """ 

76 

77 name: str 

78 engines: List[EngineInfo] = []