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

1from datetime import datetime 

2from enum import StrEnum 

3from typing import TYPE_CHECKING 

4from uuid import UUID 

5 

6from sqlalchemy import ( 

7 TIMESTAMP, 

8 ForeignKey, 

9 Integer, 

10 String, 

11 Uuid, 

12) 

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

14 

15from polar.kit.db.models import RecordModel 

16from polar.kit.utils import utc_now 

17 

18from .benefit import Benefit 

19from .customer import Customer 

20 

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

22 from .license_key_activation import LicenseKeyActivation 

23 from .organization import Organization 

24 

25 

26class LicenseKeyStatus(StrEnum): 

27 granted = "granted" 

28 revoked = "revoked" 

29 disabled = "disabled" 

30 

31 

32class LicenseKey(RecordModel): 

33 __tablename__ = "license_keys" 

34 

35 organization_id: Mapped[UUID] = mapped_column( 

36 Uuid, 

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

38 nullable=False, 

39 index=True, 

40 ) 

41 

42 @declared_attr 

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

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

45 

46 customer_id: Mapped[UUID] = mapped_column( 

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

48 ) 

49 

50 @declared_attr 

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

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

53 

54 benefit_id: Mapped[UUID] = mapped_column( 

55 Uuid, 

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

57 nullable=False, 

58 index=True, 

59 ) 

60 

61 @declared_attr 

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

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

64 

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

66 

67 status: Mapped[LicenseKeyStatus] = mapped_column( 

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

69 ) 

70 

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

72 

73 @declared_attr 

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

75 return relationship( 

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

77 ) 

78 

79 @declared_attr 

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

81 return relationship( 

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) 

94 

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

96 

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

98 

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

100 TIMESTAMP(timezone=True), 

101 nullable=True, 

102 ) 

103 

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

105 TIMESTAMP(timezone=True), 

106 nullable=True, 

107 ) 

108 

109 @property 

110 def display_key(self) -> str: 

111 prefix = "****" 

112 last_six_digits = self.key[-6:] 

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

114 

115 @property 

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

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

118 

119 @activation.setter 

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

121 self._activation = value 

122 

123 def mark_revoked(self) -> None: 

124 self.status = LicenseKeyStatus.revoked 

125 

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

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: 

133 return self.status == LicenseKeyStatus.granted