Coverage for polar/kit/metadata.py: 78%

1import inspect 

2import re 

3from typing import Annotated, Any, TypeAlias 

4 

5from fastapi import Depends, Request 

6from pydantic import AliasChoices, BaseModel, Field, StringConstraints 

7from sqlalchemy import ColumnExpressionArgument, Select, and_, or_, true 

8from sqlalchemy.dialects.postgresql import JSONB 

9from sqlalchemy.orm import Mapped, mapped_column 

10 

11MetadataColumn = Annotated[ 

12 dict[str, Any], mapped_column(JSONB, nullable=False, default=dict) 

13] 

14 

15 

16class MetadataMixin: 

17 user_metadata: Mapped[MetadataColumn] 

18 

19 

20MAXIMUM_KEYS = 50 

21_MINIMUM_KEY_LENGTH = 1 

22_MAXIMUM_KEY_LENGTH = 40 

23_MINIMUM_VALUE_LENGTH = 1 

24_MAXIMUM_VALUE_LENGTH = 500 

25MetadataKey = Annotated[ 

26 str, 

27 StringConstraints(min_length=_MINIMUM_KEY_LENGTH, max_length=_MAXIMUM_KEY_LENGTH), 

28] 

29_MetadataValueString = Annotated[ 

30 str, 

31 StringConstraints( 

32 min_length=_MINIMUM_VALUE_LENGTH, max_length=_MAXIMUM_VALUE_LENGTH 

33 ), 

34] 

35MetadataValue = _MetadataValueString | int | float | bool 

36 

37METADATA_DESCRIPTION = inspect.cleandoc( 

38 f""" 

39 {{heading}} 

40 

41 The key must be a string with a maximum length of **{_MAXIMUM_KEY_LENGTH} characters**. 

42 The value must be either: 

43 

44 * A string with a maximum length of **{_MAXIMUM_VALUE_LENGTH} characters** 

45 * An integer 

46 * A floating-point number 

47 * A boolean 

48 

49 You can store up to **{MAXIMUM_KEYS} key-value pairs**. 

50 """ 

51) 

52_description = METADATA_DESCRIPTION.format( 

53 heading="Key-value object allowing you to store additional information." 

54) 

55 

56 

57MetadataField = Annotated[ 

58 dict[MetadataKey, MetadataValue], 

59 Field(max_length=MAXIMUM_KEYS, description=_description), 

60] 

61 

62 

63class MetadataInputMixin(BaseModel): 

64 metadata: MetadataField = Field( 

65 default_factory=dict, serialization_alias="user_metadata" 

66 ) 

67 

68 

69MetadataOutputType: TypeAlias = dict[str, str | int | float | bool] 

70 

71 

72class MetadataOutputMixin(BaseModel): 

73 metadata: MetadataOutputType = Field( 

74 validation_alias=AliasChoices("user_metadata", "metadata") 

75 ) 

76 

77 

78def add_metadata_query_schema(openapi_schema: dict[str, Any]) -> dict[str, Any]: 

79 openapi_schema["components"]["schemas"]["MetadataQuery"] = { 

80 "anyOf": [ 

81 { 

82 "type": "object", 

83 "additionalProperties": { 

84 "anyOf": [ 

85 {"type": "string"}, 

86 {"type": "integer"}, 

87 {"type": "boolean"}, 

88 {"type": "array", "items": {"type": "string"}}, 

89 {"type": "array", "items": {"type": "integer"}}, 

90 {"type": "array", "items": {"type": "boolean"}}, 

91 ] 

92 }, 

93 }, 

94 {"type": "null"}, 

95 ], 

96 "title": "MetadataQuery", 

97 } 

98 return openapi_schema 

99 

100 

101def get_metadata_query_openapi_schema() -> dict[str, Any]: 

102 return { 

103 "name": "metadata", 

104 "in": "query", 

105 "required": False, 

106 "style": "deepObject", 

107 "schema": { 

108 "$ref": "#/components/schemas/MetadataQuery", 

109 }, 

110 "description": ( 

111 "Filter by metadata key-value pairs. " 

112 "It uses the `deepObject` style, e.g. `?metadata[key]=value`." 

113 ), 

114 } 

115 

116 

117_metadata_pattern = r"metadata\[([^\]]+)\]" 

118 

119 

120def _get_metadata_query(request: Request) -> dict[str, list[str]] | None: 

121 query_params = request.query_params 

122 metadata: dict[str, list[str]] = {} 

123 

124 for key, value in query_params.multi_items(): 

125 if match := re.match(_metadata_pattern, key): 

126 metadata_key = match.group(1) 

127 try: 

128 metadata[metadata_key] = [*metadata[metadata_key], value] 

129 except KeyError: 

130 metadata[metadata_key] = [value] 

131 

132 return metadata 

133 

134 

135MetadataQuery = Annotated[dict[str, list[str]], Depends(_get_metadata_query)] 

136 

137 

138def get_metadata_clause[M: MetadataMixin]( 

139 model: type[M], query: MetadataQuery 

140) -> ColumnExpressionArgument[bool]: 

141 clauses: list[ColumnExpressionArgument[bool]] = [] 

142 for key, values in query.items(): 

143 sub_clauses: list[ColumnExpressionArgument[bool]] = [] 

144 for value in values: 

145 sub_clauses.append(model.user_metadata[key].as_string() == value) 

146 clauses.append(or_(*sub_clauses)) 

147 

148 if not clauses: 

149 return true() 

150 

151 return and_(*clauses) 

152 

153 

154def apply_metadata_clause[M: MetadataMixin]( 

155 model: type[M], statement: Select[tuple[M]], query: MetadataQuery 

156) -> Select[tuple[M]]: 

157 clause = get_metadata_clause(model, query) 

158 return statement.where(clause) 

159 

160 

161def extract_metadata_value( 

162 metadata: dict[str, Any], property_selector: str 

163) -> str | None: 

164 """ 

165 Extract a value from metadata using a property selector. 

166 

167 Supports: 

168 - Simple keys: "subject" -> metadata["subject"] 

169 - Nested keys: "metadata.subject" -> metadata["metadata"]["subject"] 

170 - Dot-separated paths of any depth 

171 

172 Returns the value as a string if found, None otherwise. 

173 """ 

174 if not property_selector: 

175 return None 

176 

177 keys = property_selector.split(".") 

178 current: Any = metadata 

179 

180 for key in keys: 

181 if not isinstance(current, dict): 

182 return None 

183 current = current.get(key) 

184 if current is None: 

185 return None 

186 

187 return str(current) if current is not None else None