Coverage for src/dataknobs_data/query.py: 41%

1from dataclasses import dataclass, field 

2from enum import Enum 

3from typing import Any, Dict, List, Union 

4 

5 

6class Operator(Enum): 

7 """Query operators for filtering.""" 

8 

9 EQ = "=" # Equal 

10 NEQ = "!=" # Not equal 

11 GT = ">" # Greater than 

12 GTE = ">=" # Greater than or equal 

13 LT = "<" # Less than 

14 LTE = "<=" # Less than or equal 

15 IN = "in" # In list 

16 NOT_IN = "not_in" # Not in list 

17 LIKE = "like" # String pattern matching (SQL LIKE) 

18 REGEX = "regex" # Regular expression matching 

19 EXISTS = "exists" # Field exists 

20 NOT_EXISTS = "not_exists" # Field does not exist 

21 BETWEEN = "between" # Value between two bounds (inclusive) 

22 NOT_BETWEEN = "not_between" # Value not between two bounds 

23 

24 

25class SortOrder(Enum): 

26 """Sort order for query results.""" 

27 

28 ASC = "asc" 

29 DESC = "desc" 

30 

31 

32@dataclass 

33class Filter: 

34 """Represents a filter condition.""" 

35 

36 field: str 

37 operator: Operator 

38 value: Any = None 

39 

40 def matches(self, record_value: Any) -> bool: 

41 """Check if a record value matches this filter. 

42  

43 Supports type-aware comparisons for ranges and special handling 

44 for datetime/date objects. 

45 """ 

46 if self.operator == Operator.EXISTS: 

47 return record_value is not None 

48 elif self.operator == Operator.NOT_EXISTS: 

49 return record_value is None 

50 elif record_value is None: 

51 return False 

52 

53 if self.operator == Operator.EQ: 

54 return record_value == self.value 

55 elif self.operator == Operator.NEQ: 

56 return record_value != self.value 

57 elif self.operator == Operator.GT: 

58 return self._compare_values(record_value, self.value, lambda a, b: a > b) 

59 elif self.operator == Operator.GTE: 

60 return self._compare_values(record_value, self.value, lambda a, b: a >= b) 

61 elif self.operator == Operator.LT: 

62 return self._compare_values(record_value, self.value, lambda a, b: a < b) 

63 elif self.operator == Operator.LTE: 

64 return self._compare_values(record_value, self.value, lambda a, b: a <= b) 

65 elif self.operator == Operator.IN: 

66 return record_value in self.value 

67 elif self.operator == Operator.NOT_IN: 

68 return record_value not in self.value 

69 elif self.operator == Operator.BETWEEN: 

70 if not isinstance(self.value, (list, tuple)) or len(self.value) != 2: 

71 return False 

72 lower, upper = self.value 

73 return self._compare_values(record_value, lower, lambda a, b: a >= b) and \ 

74 self._compare_values(record_value, upper, lambda a, b: a <= b) 

75 elif self.operator == Operator.NOT_BETWEEN: 

76 if not isinstance(self.value, (list, tuple)) or len(self.value) != 2: 

77 return True 

78 lower, upper = self.value 

79 return not (self._compare_values(record_value, lower, lambda a, b: a >= b) and \ 

80 self._compare_values(record_value, upper, lambda a, b: a <= b)) 

81 elif self.operator == Operator.LIKE: 

82 if not isinstance(record_value, str): 

83 return False 

84 import re 

85 

86 pattern = self.value.replace("%", ".*").replace("_", ".") 

87 return bool(re.match(f"^{pattern}$", record_value)) 

88 elif self.operator == Operator.REGEX: 

89 if not isinstance(record_value, str): 

90 return False 

91 import re 

92 

93 return bool(re.search(self.value, record_value)) 

94 

95 return False 

96 

97 def _compare_values(self, a: Any, b: Any, comparator) -> bool: 

98 """Compare two values with type awareness. 

99  

100 Handles special cases: 

101 - Datetime strings are parsed for comparison 

102 - Mixed numeric types are converted appropriately 

103 - String comparisons are case-sensitive 

104 """ 

105 from datetime import datetime, date 

106 

107 # Handle datetime/date comparisons 

108 if isinstance(a, str) and isinstance(b, (datetime, date)): 

109 try: 

110 a = datetime.fromisoformat(a.replace("Z", "+00:00")) 

111 except (ValueError, AttributeError): 

112 return False 

113 elif isinstance(b, str) and isinstance(a, (datetime, date)): 

114 try: 

115 b = datetime.fromisoformat(b.replace("Z", "+00:00")) 

116 except (ValueError, AttributeError): 

117 return False 

118 elif isinstance(a, str) and isinstance(b, str): 

119 # Check if both look like dates 

120 if "T" in a or "-" in a: 

121 try: 

122 a = datetime.fromisoformat(a.replace("Z", "+00:00")) 

123 b = datetime.fromisoformat(b.replace("Z", "+00:00")) 

124 except (ValueError, AttributeError): 

125 pass # Keep as strings 

126 

127 # Handle numeric comparisons 

128 if isinstance(a, (int, float)) and isinstance(b, (int, float)): 

129 return comparator(a, b) 

130 

131 # Try direct comparison 

132 try: 

133 return comparator(a, b) 

134 except TypeError: 

135 # Types not comparable 

136 return False 

137 

138 def to_dict(self) -> Dict[str, Any]: 

139 """Convert filter to dictionary representation.""" 

140 return {"field": self.field, "operator": self.operator.value, "value": self.value} 

141 

142 @classmethod 

143 def from_dict(cls, data: Dict[str, Any]) -> "Filter": 

144 """Create filter from dictionary representation.""" 

145 return cls( 

146 field=data["field"], operator=Operator(data["operator"]), value=data.get("value") 

147 ) 

148 

149 

150@dataclass 

151class SortSpec: 

152 """Represents a sort specification.""" 

153 

154 field: str 

155 order: SortOrder = SortOrder.ASC 

156 

157 def to_dict(self) -> Dict[str, str]: 

158 """Convert sort spec to dictionary representation.""" 

159 return {"field": self.field, "order": self.order.value} 

160 

161 @classmethod 

162 def from_dict(cls, data: Dict[str, str]) -> "SortSpec": 

163 """Create sort spec from dictionary representation.""" 

164 return cls(field=data["field"], order=SortOrder(data.get("order", "asc"))) 

165 

166 

167@dataclass 

168class Query: 

169 """Represents a database query with filters, sorting, and pagination.""" 

170 

171 filters: List[Filter] = field(default_factory=list) 

172 sort_specs: List[SortSpec] = field(default_factory=list) 

173 limit_value: int | None = None 

174 offset_value: int | None = None 

175 fields: List[str] | None = None # Field projection 

176 

177 @property 

178 def sort_property(self) -> List[SortSpec]: 

179 """Get sort specifications (backward compatibility).""" 

180 return self.sort_specs 

181 

182 @property 

183 def limit_property(self) -> int | None: 

184 """Get limit value (backward compatibility).""" 

185 return self.limit_value 

186 

187 @property 

188 def offset_property(self) -> int | None: 

189 """Get offset value (backward compatibility).""" 

190 return self.offset_value 

191 

192 def filter(self, field: str, operator: Union[str, Operator], value: Any = None) -> "Query": 

193 """Add a filter to the query (fluent interface). 

194 

195 Args: 

196 field: The field name to filter on 

197 operator: The operator (string or Operator enum) 

198 value: The value to compare against 

199 

200 Returns: 

201 Self for method chaining 

202 """ 

203 if isinstance(operator, str): 

204 op_map = { 

205 "=": Operator.EQ, 

206 "==": Operator.EQ, 

207 "!=": Operator.NEQ, 

208 ">": Operator.GT, 

209 ">=": Operator.GTE, 

210 "<": Operator.LT, 

211 "<=": Operator.LTE, 

212 "in": Operator.IN, 

213 "IN": Operator.IN, 

214 "not_in": Operator.NOT_IN, 

215 "NOT IN": Operator.NOT_IN, 

216 "like": Operator.LIKE, 

217 "LIKE": Operator.LIKE, 

218 "regex": Operator.REGEX, 

219 "exists": Operator.EXISTS, 

220 "not_exists": Operator.NOT_EXISTS, 

221 "between": Operator.BETWEEN, 

222 "BETWEEN": Operator.BETWEEN, 

223 "not_between": Operator.NOT_BETWEEN, 

224 "NOT BETWEEN": Operator.NOT_BETWEEN, 

225 } 

226 operator = op_map.get(operator, Operator.EQ) 

227 

228 self.filters.append(Filter(field=field, operator=operator, value=value)) 

229 return self 

230 

231 def sort_by(self, field: str, order: Union[str, SortOrder] = "asc") -> "Query": 

232 """Add a sort specification to the query (fluent interface). 

233 

234 Args: 

235 field: The field name to sort by 

236 order: The sort order ("asc", "desc", or SortOrder enum) 

237 

238 Returns: 

239 Self for method chaining 

240 """ 

241 if isinstance(order, str): 

242 order = SortOrder.ASC if order.lower() == "asc" else SortOrder.DESC 

243 

244 self.sort_specs.append(SortSpec(field=field, order=order)) 

245 return self 

246 

247 def sort(self, field: str, order: Union[str, SortOrder] = "asc") -> "Query": 

248 """Add sorting (fluent interface).""" 

249 return self.sort_by(field, order) 

250 

251 def set_limit(self, limit: int) -> "Query": 

252 """Set the result limit (fluent interface). 

253 

254 Args: 

255 limit: Maximum number of results 

256 

257 Returns: 

258 Self for method chaining 

259 """ 

260 self.limit_value = limit 

261 return self 

262 

263 def limit(self, value: int) -> "Query": 

264 """Set limit (fluent interface).""" 

265 return self.set_limit(value) 

266 

267 def set_offset(self, offset: int) -> "Query": 

268 """Set the result offset (fluent interface). 

269 

270 Args: 

271 offset: Number of results to skip 

272 

273 Returns: 

274 Self for method chaining 

275 """ 

276 self.offset_value = offset 

277 return self 

278 

279 def offset(self, value: int) -> "Query": 

280 """Set offset (fluent interface).""" 

281 return self.set_offset(value) 

282 

283 def select(self, *fields: str) -> "Query": 

284 """Set field projection (fluent interface). 

285 

286 Args: 

287 fields: Field names to include in results 

288 

289 Returns: 

290 Self for method chaining 

291 """ 

292 self.fields = list(fields) if fields else None 

293 return self 

294 

295 def clear_filters(self) -> "Query": 

296 """Clear all filters (fluent interface).""" 

297 self.filters = [] 

298 return self 

299 

300 def clear_sort(self) -> "Query": 

301 """Clear all sort specifications (fluent interface).""" 

302 self.sort_specs = [] 

303 return self 

304 

305 def to_dict(self) -> Dict[str, Any]: 

306 """Convert query to dictionary representation.""" 

307 result = { 

308 "filters": [f.to_dict() for f in self.filters], 

309 "sort": [s.to_dict() for s in self.sort_specs], 

310 } 

311 if self.limit_value is not None: 

312 result["limit"] = self.limit_value 

313 if self.offset_value is not None: 

314 result["offset"] = self.offset_value 

315 if self.fields is not None: 

316 result["fields"] = self.fields 

317 return result 

318 

319 @classmethod 

320 def from_dict(cls, data: Dict[str, Any]) -> "Query": 

321 """Create query from dictionary representation.""" 

322 query = cls() 

323 

324 for filter_data in data.get("filters", []): 

325 query.filters.append(Filter.from_dict(filter_data)) 

326 

327 for sort_data in data.get("sort", []): 

328 query.sort_specs.append(SortSpec.from_dict(sort_data)) 

329 

330 query.limit_value = data.get("limit") 

331 query.offset_value = data.get("offset") 

332 query.fields = data.get("fields") 

333 

334 return query 

335 

336 def copy(self) -> "Query": 

337 """Create a copy of the query.""" 

338 import copy 

339 

340 return Query( 

341 filters=copy.deepcopy(self.filters), 

342 sort_specs=copy.deepcopy(self.sort_specs), 

343 limit_value=self.limit_value, 

344 offset_value=self.offset_value, 

345 fields=self.fields.copy() if self.fields else None, 

346 ) 

347 

348 def or_(self, *filters: Union[Filter, "Query"]) -> "ComplexQuery": 

349 """Create a ComplexQuery with OR logic. 

350  

351 The current query's filters become an AND group, combined with OR conditions. 

352 Example: Query with filters [A, B] calling or_(C, D) creates: (A AND B) AND (C OR D) 

353  

354 Args: 

355 filters: Filter objects or Query objects to OR together 

356  

357 Returns: 

358 ComplexQuery with OR logic 

359 """ 

360 from .query_logic import ComplexQuery, FilterCondition, LogicCondition, LogicOperator 

361 

362 # Build OR conditions from the arguments 

363 or_conditions = [] 

364 for item in filters: 

365 if isinstance(item, Filter): 

366 or_conditions.append(FilterCondition(item)) 

367 elif isinstance(item, Query): 

368 if len(item.filters) == 1: 

369 or_conditions.append(FilterCondition(item.filters[0])) 

370 elif item.filters: 

371 and_cond = LogicCondition(operator=LogicOperator.AND) 

372 for f in item.filters: 

373 and_cond.conditions.append(FilterCondition(f)) 

374 or_conditions.append(and_cond) 

375 

376 # Create the OR condition group 

377 or_group = None 

378 if or_conditions: 

379 if len(or_conditions) == 1: 

380 or_group = or_conditions[0] 

381 else: 

382 or_group = LogicCondition( 

383 operator=LogicOperator.OR, 

384 conditions=or_conditions 

385 ) 

386 

387 # Combine with existing filters (if any) using AND 

388 if self.filters: 

389 # Create AND condition for existing filters 

390 if len(self.filters) == 1: 

391 existing = FilterCondition(self.filters[0]) 

392 else: 

393 existing = LogicCondition(operator=LogicOperator.AND) 

394 for f in self.filters: 

395 existing.conditions.append(FilterCondition(f)) 

396 

397 # Combine existing AND new OR group with AND 

398 if or_group: 

399 root_condition = LogicCondition( 

400 operator=LogicOperator.AND, 

401 conditions=[existing, or_group] 

402 ) 

403 else: 

404 root_condition = existing 

405 else: 

406 # No existing filters, just use OR group 

407 root_condition = or_group 

408 

409 return ComplexQuery( 

410 condition=root_condition, 

411 sort_specs=self.sort_specs.copy(), 

412 limit_value=self.limit_value, 

413 offset_value=self.offset_value, 

414 fields=self.fields.copy() if self.fields else None 

415 ) 

416 

417 def and_(self, *filters: Union[Filter, "Query"]) -> "Query": 

418 """Add more filters with AND logic (convenience method). 

419  

420 Args: 

421 filters: Filter objects or Query objects to AND together 

422  

423 Returns: 

424 Self for chaining 

425 """ 

426 for item in filters: 

427 if isinstance(item, Filter): 

428 self.filters.append(item) 

429 elif isinstance(item, Query): 

430 self.filters.extend(item.filters) 

431 return self 

432 

433 def not_(self, filter: Filter) -> "ComplexQuery": 

434 """Create a ComplexQuery with NOT logic. 

435  

436 Args: 

437 filter: Filter to negate 

438  

439 Returns: 

440 ComplexQuery with NOT logic 

441 """ 

442 from .query_logic import ComplexQuery, FilterCondition, LogicCondition, LogicOperator 

443 

444 # Current filters as AND 

445 conditions = [] 

446 if self.filters: 

447 if len(self.filters) == 1: 

448 conditions.append(FilterCondition(self.filters[0])) 

449 else: 

450 and_cond = LogicCondition(operator=LogicOperator.AND) 

451 for f in self.filters: 

452 and_cond.conditions.append(FilterCondition(f)) 

453 conditions.append(and_cond) 

454 

455 # Add NOT condition 

456 not_cond = LogicCondition( 

457 operator=LogicOperator.NOT, 

458 conditions=[FilterCondition(filter)] 

459 ) 

460 conditions.append(not_cond) 

461 

462 # Create root condition 

463 if len(conditions) == 1: 

464 root_condition = conditions[0] 

465 else: 

466 root_condition = LogicCondition( 

467 operator=LogicOperator.AND, 

468 conditions=conditions 

469 ) 

470 

471 return ComplexQuery( 

472 condition=root_condition, 

473 sort_specs=self.sort_specs.copy(), 

474 limit_value=self.limit_value, 

475 offset_value=self.offset_value, 

476 fields=self.fields.copy() if self.fields else None 

477 )