Coverage for polar/models/custom_field.py: 85%

1from datetime import datetime 

2from enum import StrEnum 

3from typing import TYPE_CHECKING, Annotated, Any, Literal, NotRequired, TypedDict 

4from uuid import UUID 

5 

6from annotated_types import Ge, Len, MinLen 

7from pydantic import AfterValidator, Field, ValidationInfo 

8from sqlalchemy import ForeignKey, String, UniqueConstraint, Uuid 

9from sqlalchemy.dialects.postgresql import CITEXT, JSONB 

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

11 

12from polar.kit.db.models import RecordModel 

13from polar.kit.metadata import MetadataMixin 

14 

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

16 from polar.models import Organization 

17 

18 

19class CustomFieldType(StrEnum): 

20 text = "text" 

21 number = "number" 

22 date = "date" 

23 checkbox = "checkbox" 

24 select = "select" 

25 

26 def get_model(self) -> type["CustomField"]: 

27 return { 

28 CustomFieldType.text: CustomFieldText, 

29 CustomFieldType.number: CustomFieldNumber, 

30 CustomFieldType.date: CustomFieldDate, 

31 CustomFieldType.checkbox: CustomFieldCheckbox, 

32 CustomFieldType.select: CustomFieldSelect, 

33 }[self] 

34 

35 

36PositiveInt = Annotated[int, Ge(0)] 

37NonEmptyString = Annotated[str, Len(min_length=1)] 

38 

39 

40def validate_ge_le(v: int, info: ValidationInfo) -> int: 

41 """Validate that le is greater than or equal to ge when both are provided.""" 

42 ge = info.data.get("ge") 

43 if ge is not None and v is not None and ge > v: 

44 raise ValueError( 

45 "Greater than or equal (ge) must be less than or equal to " 

46 "Less than or equal (le)" 

47 ) 

48 return v 

49 

50 

51class CustomFieldProperties(TypedDict): 

52 form_label: NotRequired[NonEmptyString] 

53 form_help_text: NotRequired[NonEmptyString] 

54 form_placeholder: NotRequired[NonEmptyString] 

55 

56 

57class CustomFieldTextProperties(CustomFieldProperties): 

58 textarea: NotRequired[bool] 

59 min_length: NotRequired[PositiveInt] 

60 max_length: NotRequired[PositiveInt] 

61 

62 

63class ComparableProperties(TypedDict): 

64 ge: NotRequired[int] 

65 le: NotRequired[Annotated[int, AfterValidator(validate_ge_le)]] 

66 

67 

68class CustomFieldNumberProperties(CustomFieldProperties, ComparableProperties): 

69 pass 

70 

71 

72class CustomFieldDateProperties(CustomFieldProperties, ComparableProperties): 

73 pass 

74 

75 

76class CustomFieldCheckboxProperties(CustomFieldProperties): 

77 pass 

78 

79 

80class CustomFieldSelectOption(TypedDict): 

81 value: NonEmptyString 

82 label: NonEmptyString 

83 

84 

85class CustomFieldSelectProperties(CustomFieldProperties): 

86 options: Annotated[list[CustomFieldSelectOption], MinLen(1)] 

87 

88 

89class CustomField(MetadataMixin, RecordModel): 

90 __tablename__ = "custom_fields" 

91 __table_args__ = (UniqueConstraint("slug", "organization_id"),) 

92 

93 type: Mapped[CustomFieldType] = mapped_column(String, nullable=False, index=True) 

94 slug: Mapped[str] = mapped_column( 

95 CITEXT, 

96 nullable=False, 

97 # Don't create an index for slug 

98 # as it's covered by the unique constraint, being the leading column of it 

99 index=False, 

100 ) 

101 name: Mapped[str] = mapped_column(String, nullable=False) 

102 properties: Mapped[CustomFieldProperties] = mapped_column( 

103 JSONB, nullable=False, default=dict 

104 ) 

105 

106 organization_id: Mapped[UUID] = mapped_column( 

107 Uuid, 

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

109 nullable=False, 

110 index=True, 

111 ) 

112 

113 @declared_attr 

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

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

116 

117 def get_field_definition(self, required: bool) -> tuple[Any, Any]: 

118 raise NotImplementedError() 

119 

120 __mapper_args__ = { 

121 "polymorphic_on": "type", 

122 } 

123 

124 

125class CustomFieldText(CustomField): 

126 type: Mapped[Literal[CustomFieldType.text]] = mapped_column( 

127 use_existing_column=True 

128 ) 

129 properties: Mapped[CustomFieldTextProperties] = mapped_column( 

130 use_existing_column=True 

131 ) 

132 

133 __mapper_args__ = { 

134 "polymorphic_identity": CustomFieldType.text, 

135 "polymorphic_load": "inline", 

136 } 

137 

138 def get_field_definition(self, required: bool) -> tuple[Any, Any]: 

139 return ( 

140 str if required else str | None, 

141 Field( 

142 default=None if not required else ..., 

143 min_length=self.properties.get("min_length"), 

144 max_length=self.properties.get("max_length"), 

145 ), 

146 ) 

147 

148 

149class CustomFieldNumber(CustomField): 

150 type: Mapped[Literal[CustomFieldType.number]] = mapped_column( 

151 use_existing_column=True 

152 ) 

153 properties: Mapped[CustomFieldNumberProperties] = mapped_column( 

154 use_existing_column=True 

155 ) 

156 

157 __mapper_args__ = { 

158 "polymorphic_identity": CustomFieldType.number, 

159 "polymorphic_load": "inline", 

160 } 

161 

162 def get_field_definition(self, required: bool) -> tuple[Any, Any]: 

163 return ( 

164 int if required else int | None, 

165 Field( 

166 default=None if not required else ..., 

167 ge=self.properties.get("ge"), 

168 le=self.properties.get("le"), 

169 ), 

170 ) 

171 

172 

173class CustomFieldDate(CustomField): 

174 type: Mapped[Literal[CustomFieldType.date]] = mapped_column( 

175 use_existing_column=True 

176 ) 

177 properties: Mapped[CustomFieldDateProperties] = mapped_column( 

178 use_existing_column=True 

179 ) 

180 

181 __mapper_args__ = { 

182 "polymorphic_identity": CustomFieldType.date, 

183 "polymorphic_load": "inline", 

184 } 

185 

186 def get_field_definition(self, required: bool) -> tuple[Any, Any]: 

187 ge = self.properties.get("ge") 

188 ge_date = datetime.fromtimestamp(ge).date() if ge else None 

189 le = self.properties.get("le") 

190 le_date = datetime.fromtimestamp(le).date() if le else None 

191 return ( 

192 datetime if required else datetime | None, 

193 Field(default=None if not required else ..., ge=ge_date, le=le_date), 

194 ) 

195 

196 

197class CustomFieldCheckbox(CustomField): 

198 type: Mapped[Literal[CustomFieldType.checkbox]] = mapped_column( 

199 use_existing_column=True 

200 ) 

201 properties: Mapped[CustomFieldCheckboxProperties] = mapped_column( 

202 use_existing_column=True 

203 ) 

204 

205 __mapper_args__ = { 

206 "polymorphic_identity": CustomFieldType.checkbox, 

207 "polymorphic_load": "inline", 

208 } 

209 

210 def get_field_definition(self, required: bool) -> tuple[Any, Any]: 

211 return ( 

212 Literal[True] if required else bool, 

213 Field(default=False if not required else ...), 

214 ) 

215 

216 

217class CustomFieldSelect(CustomField): 

218 type: Mapped[Literal[CustomFieldType.select]] = mapped_column( 

219 use_existing_column=True 

220 ) 

221 properties: Mapped[CustomFieldSelectProperties] = mapped_column( 

222 use_existing_column=True 

223 ) 

224 

225 __mapper_args__ = { 

226 "polymorphic_identity": CustomFieldType.select, 

227 "polymorphic_load": "inline", 

228 } 

229 

230 def get_field_definition(self, required: bool) -> tuple[Any, Any]: 

231 literal_type = Literal[ # type: ignore 

232 tuple(option["value"] for option in self.properties["options"]) 

233 ] 

234 return ( 

235 literal_type if required else literal_type | None, # pyright: ignore 

236 Field( 

237 default=None if not required else ..., 

238 ), 

239 )