Coverage for polar/kit/db/models/base.py: 93%

1from datetime import datetime 

2from uuid import UUID 

3 

4from sqlalchemy import TIMESTAMP, MetaData, Uuid, inspect 

5from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column 

6 

7from polar.enums import RateLimitGroup 

8from polar.kit.extensions.sqlalchemy.types import StringEnum 

9from polar.kit.utils import generate_uuid, utc_now 

10 

11my_metadata = MetaData( 

12 naming_convention={ 

13 "ix": "ix_%(column_0_N_label)s", 

14 "uq": "%(table_name)s_%(column_0_N_name)s_key", 

15 "ck": "%(table_name)s_%(constraint_name)s_check", 

16 "fk": "%(table_name)s_%(column_0_N_name)s_fkey", 

17 "pk": "%(table_name)s_pkey", 

18 } 

19) 

20 

21 

22class Model(DeclarativeBase): 

23 __abstract__ = True 

24 

25 metadata = my_metadata 

26 

27 

28class TimestampedModel(Model): 

29 __abstract__ = True 

30 

31 created_at: Mapped[datetime] = mapped_column( 

32 TIMESTAMP(timezone=True), nullable=False, default=utc_now, index=True 

33 ) 

34 modified_at: Mapped[datetime | None] = mapped_column( 

35 TIMESTAMP(timezone=True), onupdate=utc_now, nullable=True, default=None 

36 ) 

37 deleted_at: Mapped[datetime | None] = mapped_column( 

38 TIMESTAMP(timezone=True), nullable=True, default=None, index=True 

39 ) 

40 

41 def set_modified_at(self) -> None: 

42 self.modified_at = utc_now() 

43 

44 def set_deleted_at(self) -> None: 

45 self.deleted_at = utc_now() 

46 

47 

48class IDModel(Model): 

49 __abstract__ = True 

50 

51 id: Mapped[UUID] = mapped_column(Uuid, primary_key=True, default=generate_uuid) 

52 

53 def __eq__(self, __value: object) -> bool: 

54 return isinstance(__value, self.__class__) and self.id == __value.id 

55 

56 def __hash__(self) -> int: 

57 return self.id.int 

58 

59 def __repr__(self) -> str: 

60 # We do this complex thing because we might be outside a session with 

61 # an expired object; typically when Sentry tries to serialize the object for 

62 # error reporting. 

63 # But basically, we want to show the ID if we have it. 

64 insp = inspect(self) 

65 if insp.identity is not None: 

66 id_value = insp.identity[0] 

67 return f"{self.__class__.__name__}(id={id_value!r})" 

68 return f"{self.__class__.__name__}(id=None)" 

69 

70 @classmethod 

71 def generate_id(cls) -> UUID: 

72 return generate_uuid() 

73 

74 

75class RecordModel(IDModel, TimestampedModel): 

76 __abstract__ = True 

77 

78 

79class RateLimitGroupMixin: 

80 __abstract__ = True 

81 

82 rate_limit_group: Mapped[RateLimitGroup] = mapped_column( 

83 StringEnum(RateLimitGroup, length=16), 

84 nullable=False, 

85 default=RateLimitGroup.default, 

86 )