Coverage for polar/models/license_key.py: 77%

1from datetime import datetime 2 ctx1ab

2from enum import StrEnum 2 ctx1ab

3from typing import TYPE_CHECKING 2 ctx1ab

4from uuid import UUID 2 ctx1ab

5 

6from sqlalchemy import ( 2 ctx1ab

7 TIMESTAMP, 

8 ForeignKey, 

9 Integer, 

10 String, 

11 Uuid, 

12) 

13from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship 2 ctx1ab

14 

15from polar.kit.db.models import RecordModel 2 ctx1ab

16from polar.kit.utils import utc_now 2 ctx1ab

17 

18from .benefit import Benefit 2 ctx1ab

19from .customer import Customer 2 ctx1ab

20 

21if TYPE_CHECKING: 21 ↛ 22line 21 didn't jump to line 22 because the condition on line 21 was never true2 ctx1ab

22 from .license_key_activation import LicenseKeyActivation 

23 from .organization import Organization 

24 

25 

26class LicenseKeyStatus(StrEnum): 2 ctx1ab

27 granted = "granted" 2 ctx1ab

28 revoked = "revoked" 2 ctx1ab

29 disabled = "disabled" 2 ctx1ab

30 

31 

32class LicenseKey(RecordModel): 2 ctx1ab

33 __tablename__ = "license_keys" 2 ctx1ab

34 

35 organization_id: Mapped[UUID] = mapped_column( 2 ctx1ab

36 Uuid, 

37 ForeignKey("organizations.id", ondelete="cascade"), 

38 nullable=False, 

39 index=True, 

40 ) 

41 

42 @declared_attr 2 ctx1ab

43 def organization(cls) -> Mapped["Organization"]: 2 ctx1ab

44 return relationship("Organization", lazy="raise") 2 ctx1ab

45 

46 customer_id: Mapped[UUID] = mapped_column( 2 ctx1ab

47 Uuid, ForeignKey("customers.id", ondelete="cascade"), nullable=False, index=True 

48 ) 

49 

50 @declared_attr 2 ctx1ab

51 def customer(cls) -> Mapped[Customer]: 2 ctx1ab

52 return relationship("Customer", lazy="raise") 2 ctx1ab

53 

54 benefit_id: Mapped[UUID] = mapped_column( 2 ctx1ab

55 Uuid, 

56 ForeignKey("benefits.id", ondelete="cascade"), 

57 nullable=False, 

58 index=True, 

59 ) 

60 

61 @declared_attr 2 ctx1ab

62 def benefit(cls) -> Mapped[Benefit]: 2 ctx1ab

63 return relationship("Benefit", lazy="raise") 2 ctx1ab

64 

65 key: Mapped[str] = mapped_column(String, nullable=False) 2 ctx1ab

66 

67 status: Mapped[LicenseKeyStatus] = mapped_column( 2 ctx1ab

68 String, nullable=False, default=LicenseKeyStatus.granted 

69 ) 

70 

71 limit_activations: Mapped[int | None] = mapped_column(Integer, nullable=True) 2 ctx1ab

72 

73 @declared_attr 2 ctx1ab

74 def all_activations(cls) -> Mapped[list["LicenseKeyActivation"]]: 2 ctx1ab

75 return relationship( 2 ctx1ab

76 "LicenseKeyActivation", lazy="raise", back_populates="license_key" 

77 ) 

78 

79 @declared_attr 2 ctx1ab

80 def activations(cls) -> Mapped[list["LicenseKeyActivation"]]: 2 ctx1ab

81 return relationship( 2 ctx1ab

82 "LicenseKeyActivation", 

83 lazy="raise", 

84 primaryjoin=( 

85 "and_(" 

86 "LicenseKeyActivation.license_key_id == LicenseKey.id, " 

87 "LicenseKeyActivation.deleted_at.is_(None)" 

88 ")" 

89 ), 

90 viewonly=True, 

91 ) 

92 

93 usage: Mapped[int] = mapped_column(Integer, nullable=False, default=0) 2 ctx1ab

94 

95 limit_usage: Mapped[int | None] = mapped_column(Integer, nullable=True) 2 ctx1ab

96 

97 validations: Mapped[int] = mapped_column(Integer, nullable=False, default=0) 2 ctx1ab

98 

99 last_validated_at: Mapped[datetime | None] = mapped_column( 2 ctx1ab

100 TIMESTAMP(timezone=True), 

101 nullable=True, 

102 ) 

103 

104 expires_at: Mapped[datetime | None] = mapped_column( 2 ctx1ab

105 TIMESTAMP(timezone=True), 

106 nullable=True, 

107 ) 

108 

109 @property 2 ctx1ab

110 def display_key(self) -> str: 2 ctx1ab

111 prefix = "****" 

112 last_six_digits = self.key[-6:] 

113 return f"{prefix}-{last_six_digits}" 

114 

115 @property 2 ctx1ab

116 def activation(self) -> "LicenseKeyActivation | None": 2 ctx1ab

117 return getattr(self, "_activation", None) 

118 

119 @activation.setter 2 ctx1ab

120 def activation(self, value: "LicenseKeyActivation") -> None: 2 ctx1ab

121 self._activation = value 

122 

123 def mark_revoked(self) -> None: 2 ctx1ab

124 self.status = LicenseKeyStatus.revoked 

125 

126 def mark_validated(self, increment_usage: int | None = None) -> None: 2 ctx1ab

127 self.validations += 1 

128 self.last_validated_at = utc_now() 

129 if increment_usage: 

130 self.usage += increment_usage 

131 

132 def is_active(self) -> bool: 2 ctx1ab

133 return self.status == LicenseKeyStatus.granted