Coverage for src/dataknobs_data/backends/postgres.py: 12%

1"""PostgreSQL backend implementation with proper connection management.""" 

2 

3import asyncio 

4import asyncpg 

5import json 

6import logging 

7import time 

8import uuid 

9from datetime import datetime 

10from typing import Any, AsyncIterator, Iterator, Optional 

11 

12from dataknobs_config import ConfigurableBase 

13from dataknobs_utils.sql_utils import DotenvPostgresConnector, PostgresDB 

14 

15from ..database import AsyncDatabase, SyncDatabase 

16from ..query import Operator, Query, SortOrder 

17from ..records import Record 

18from ..streaming import StreamConfig, StreamResult 

19from ..pooling import ConnectionPoolManager 

20from ..streaming import async_process_batch_with_fallback, process_batch_with_fallback 

21from ..pooling.postgres import ( 

22 PostgresPoolConfig, 

23 create_asyncpg_pool, 

24 validate_asyncpg_pool 

25) 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30class SyncPostgresDatabase(SyncDatabase, ConfigurableBase): 

31 """Synchronous PostgreSQL database backend with proper connection management.""" 

32 

33 def __init__(self, config: dict[str, Any] | None = None): 

34 """Initialize PostgreSQL database configuration. 

35 

36 Args: 

37 config: Configuration with the following optional keys: 

38 - host: PostgreSQL host (default: from env/localhost) 

39 - port: PostgreSQL port (default: 5432) 

40 - database: Database name (default: from env/postgres) 

41 - user: Username (default: from env/postgres) 

42 - password: Password (default: from env) 

43 - table: Table name (default: "records") 

44 - schema: Schema name (default: "public") 

45 """ 

46 super().__init__(config) 

47 self.db = None # Will be initialized in connect() 

48 self._connected = False 

49 

50 @classmethod 

51 def from_config(cls, config: dict) -> "SyncPostgresDatabase": 

52 """Create from config dictionary.""" 

53 return cls(config) 

54 

55 def connect(self) -> None: 

56 """Connect to the PostgreSQL database.""" 

57 if self._connected: 

58 return # Already connected 

59 

60 config = self.config.copy() 

61 

62 # Extract table configuration 

63 self.table_name = config.pop("table", "records") 

64 self.schema_name = config.pop("schema", "public") 

65 

66 # Create connection using existing utilities 

67 if not any(key in config for key in ["host", "database", "user"]): 

68 # Use dotenv connector for environment-based config 

69 connector = DotenvPostgresConnector() 

70 self.db = PostgresDB(connector) 

71 else: 

72 # Direct configuration - map 'database' to 'db' for PostgresDB 

73 self.db = PostgresDB( 

74 host=config.get("host", "localhost"), 

75 db=config.get("database", "postgres"), # Note: PostgresDB expects 'db' not 'database' 

76 user=config.get("user", "postgres"), 

77 pwd=config.get("password"), # Note: PostgresDB expects 'pwd' not 'password' 

78 port=config.get("port", 5432), 

79 ) 

80 

81 # Create table if it doesn't exist 

82 self._ensure_table() 

83 self._connected = True 

84 

85 def close(self) -> None: 

86 """Close the database connection.""" 

87 if self.db: 

88 # PostgresDB manages its own connections via context managers 

89 # but we can mark as disconnected 

90 self._connected = False 

91 

92 def _initialize(self) -> None: 

93 """Initialize method - connection setup moved to connect().""" 

94 # Configuration parsing stays here, actual connection in connect() 

95 pass 

96 

97 def _ensure_table(self) -> None: 

98 """Ensure the records table exists.""" 

99 if not self.db: 

100 raise RuntimeError("Database not connected. Call connect() first.") 

101 

102 create_table_sql = f""" 

103 CREATE TABLE IF NOT EXISTS {self.schema_name}.{self.table_name} ( 

104 id VARCHAR(255) PRIMARY KEY, 

105 data JSONB NOT NULL, 

106 metadata JSONB, 

107 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, 

108 updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP 

109 ); 

110  

111 CREATE INDEX IF NOT EXISTS idx_{self.table_name}_data  

112 ON {self.schema_name}.{self.table_name} USING GIN (data); 

113  

114 CREATE INDEX IF NOT EXISTS idx_{self.table_name}_metadata 

115 ON {self.schema_name}.{self.table_name} USING GIN (metadata); 

116 """ 

117 self.db.execute(create_table_sql) 

118 

119 def _check_connection(self) -> None: 

120 """Check if database is connected.""" 

121 if not self._connected or not self.db: 

122 raise RuntimeError("Database not connected. Call connect() first.") 

123 

124 def _record_to_row(self, record: Record, id: str | None = None) -> dict[str, Any]: 

125 """Convert a Record to a database row.""" 

126 data = {} 

127 for field_name, field_obj in record.fields.items(): 

128 data[field_name] = field_obj.value 

129 

130 return { 

131 "id": id or str(uuid.uuid4()), 

132 "data": json.dumps(data), 

133 "metadata": json.dumps(record.metadata) if record.metadata else None, 

134 } 

135 

136 def _row_to_record(self, row: dict[str, Any]) -> Record: 

137 """Convert a database row to a Record.""" 

138 data = row.get("data", {}) 

139 if isinstance(data, str): 

140 data = json.loads(data) 

141 

142 metadata = row.get("metadata", {}) 

143 if isinstance(metadata, str) and metadata: 

144 metadata = json.loads(metadata) 

145 elif not metadata: 

146 metadata = {} 

147 

148 return Record(data=data, metadata=metadata) 

149 

150 def create(self, record: Record) -> str: 

151 """Create a new record.""" 

152 self._check_connection() 

153 # Use record's ID if it has one, otherwise generate a new one 

154 id = record.id if record.id else str(uuid.uuid4()) 

155 row = self._record_to_row(record, id) 

156 

157 sql = f""" 

158 INSERT INTO {self.schema_name}.{self.table_name} (id, data, metadata) 

159 VALUES (%(id)s, %(data)s, %(metadata)s) 

160 """ 

161 self.db.execute(sql, row) 

162 return id 

163 

164 def read(self, id: str) -> Record | None: 

165 """Read a record by ID.""" 

166 self._check_connection() 

167 sql = f""" 

168 SELECT id, data, metadata 

169 FROM {self.schema_name}.{self.table_name} 

170 WHERE id = %(id)s 

171 """ 

172 df = self.db.query(sql, {"id": id}) 

173 

174 if df.empty: 

175 return None 

176 

177 row = df.iloc[0].to_dict() 

178 return self._row_to_record(row) 

179 

180 def update(self, id: str, record: Record) -> bool: 

181 """Update an existing record.""" 

182 self._check_connection() 

183 row = self._record_to_row(record, id) 

184 

185 sql = f""" 

186 UPDATE {self.schema_name}.{self.table_name} 

187 SET data = %(data)s, metadata = %(metadata)s, updated_at = CURRENT_TIMESTAMP 

188 WHERE id = %(id)s 

189 """ 

190 result = self.db.execute(sql, row) 

191 # PostgresDB.execute returns number of affected rows 

192 return result > 0 if isinstance(result, int) else False 

193 

194 def delete(self, id: str) -> bool: 

195 """Delete a record by ID.""" 

196 self._check_connection() 

197 sql = f""" 

198 DELETE FROM {self.schema_name}.{self.table_name} 

199 WHERE id = %(id)s 

200 """ 

201 result = self.db.execute(sql, {"id": id}) 

202 return result > 0 if isinstance(result, int) else False 

203 

204 def exists(self, id: str) -> bool: 

205 """Check if a record exists.""" 

206 self._check_connection() 

207 sql = f""" 

208 SELECT 1 FROM {self.schema_name}.{self.table_name} 

209 WHERE id = %(id)s 

210 LIMIT 1 

211 """ 

212 df = self.db.query(sql, {"id": id}) 

213 return not df.empty 

214 

215 def upsert(self, id: str, record: Record) -> str: 

216 """Update or insert a record with a specific ID.""" 

217 self._check_connection() 

218 if self.exists(id): 

219 self.update(id, record) 

220 else: 

221 # Insert with specific ID 

222 row = self._record_to_row(record, id) 

223 sql = f""" 

224 INSERT INTO {self.schema_name}.{self.table_name} (id, data, metadata) 

225 VALUES (%(id)s, %(data)s, %(metadata)s) 

226 """ 

227 self.db.execute(sql, row) 

228 return id 

229 

230 def search(self, query: Query) -> list[Record]: 

231 """Search for records matching the query.""" 

232 self._check_connection() 

233 # Build SQL query from Query object 

234 where_clauses = [] 

235 params = {} 

236 

237 # Build WHERE clauses for filters 

238 for i, filter in enumerate(query.filters): 

239 field_path = f"data->>'{filter.field}'" 

240 param_name = f"param_{i}" 

241 

242 if filter.operator == Operator.EQ: 

243 # Handle different types appropriately 

244 if isinstance(filter.value, bool): 

245 where_clauses.append(f"({field_path})::boolean = %({param_name})s") 

246 params[param_name] = filter.value 

247 elif isinstance(filter.value, (int, float)): 

248 where_clauses.append(f"({field_path})::numeric = %({param_name})s") 

249 params[param_name] = filter.value 

250 else: 

251 where_clauses.append(f"{field_path} = %({param_name})s") 

252 params[param_name] = str(filter.value) 

253 elif filter.operator == Operator.NEQ: 

254 if isinstance(filter.value, bool): 

255 where_clauses.append(f"({field_path})::boolean != %({param_name})s") 

256 params[param_name] = filter.value 

257 elif isinstance(filter.value, (int, float)): 

258 where_clauses.append(f"({field_path})::numeric != %({param_name})s") 

259 params[param_name] = filter.value 

260 else: 

261 where_clauses.append(f"{field_path} != %({param_name})s") 

262 params[param_name] = str(filter.value) 

263 elif filter.operator == Operator.GT: 

264 where_clauses.append(f"({field_path})::numeric > %({param_name})s") 

265 params[param_name] = filter.value 

266 elif filter.operator == Operator.LT: 

267 where_clauses.append(f"({field_path})::numeric < %({param_name})s") 

268 params[param_name] = filter.value 

269 elif filter.operator == Operator.GTE: 

270 where_clauses.append(f"({field_path})::numeric >= %({param_name})s") 

271 params[param_name] = filter.value 

272 elif filter.operator == Operator.LTE: 

273 where_clauses.append(f"({field_path})::numeric <= %({param_name})s") 

274 params[param_name] = filter.value 

275 elif filter.operator == Operator.LIKE: 

276 where_clauses.append(f"{field_path} LIKE %({param_name})s") 

277 params[param_name] = f"%{filter.value}%" 

278 elif filter.operator == Operator.IN: 

279 # Convert values to strings for comparison with JSONB text fields 

280 values = [str(v) for v in filter.value] 

281 where_clauses.append(f"{field_path} = ANY(%({param_name})s)") 

282 params[param_name] = values 

283 elif filter.operator == Operator.NOT_IN: 

284 # Convert values to strings for comparison with JSONB text fields 

285 values = [str(v) for v in filter.value] 

286 where_clauses.append(f"{field_path} != ALL(%({param_name})s)") 

287 params[param_name] = values 

288 elif filter.operator == Operator.BETWEEN: 

289 # Optimize BETWEEN for different data types 

290 if isinstance(filter.value, (list, tuple)) and len(filter.value) == 2: 

291 lower, upper = filter.value 

292 param_lower = f"{param_name}_lower" 

293 param_upper = f"{param_name}_upper" 

294 

295 # Try to determine the type for proper casting 

296 if isinstance(lower, (int, float)) and isinstance(upper, (int, float)): 

297 where_clauses.append( 

298 f"({field_path})::numeric BETWEEN %({param_lower})s AND %({param_upper})s" 

299 ) 

300 elif isinstance(lower, datetime) or isinstance(upper, datetime): 

301 where_clauses.append( 

302 f"({field_path})::timestamp BETWEEN %({param_lower})s AND %({param_upper})s" 

303 ) 

304 else: 

305 # String or unknown type 

306 where_clauses.append( 

307 f"{field_path} BETWEEN %({param_lower})s AND %({param_upper})s" 

308 ) 

309 

310 params[param_lower] = lower 

311 params[param_upper] = upper 

312 elif filter.operator == Operator.NOT_BETWEEN: 

313 # Optimize NOT BETWEEN 

314 if isinstance(filter.value, (list, tuple)) and len(filter.value) == 2: 

315 lower, upper = filter.value 

316 param_lower = f"{param_name}_lower" 

317 param_upper = f"{param_name}_upper" 

318 

319 # Try to determine the type for proper casting 

320 if isinstance(lower, (int, float)) and isinstance(upper, (int, float)): 

321 where_clauses.append( 

322 f"({field_path})::numeric NOT BETWEEN %({param_lower})s AND %({param_upper})s" 

323 ) 

324 elif isinstance(lower, datetime) or isinstance(upper, datetime): 

325 where_clauses.append( 

326 f"({field_path})::timestamp NOT BETWEEN %({param_lower})s AND %({param_upper})s" 

327 ) 

328 else: 

329 # String or unknown type 

330 where_clauses.append( 

331 f"{field_path} NOT BETWEEN %({param_lower})s AND %({param_upper})s" 

332 ) 

333 

334 params[param_lower] = lower 

335 params[param_upper] = upper 

336 

337 # Build SQL 

338 sql = f"SELECT id, data, metadata FROM {self.schema_name}.{self.table_name}" 

339 if where_clauses: 

340 sql += " WHERE " + " AND ".join(where_clauses) 

341 

342 # Add ORDER BY 

343 if query.sort_specs: 

344 order_clauses = [] 

345 for sort_spec in query.sort_specs: 

346 # Try to cast to numeric for proper sorting 

347 # This will sort numbers correctly while still working for strings 

348 field_path = f"data->>'{sort_spec.field}'" 

349 direction = "DESC" if sort_spec.order == SortOrder.DESC else "ASC" 

350 # Use a CASE statement to handle both numeric and string sorting 

351 order_clause = f""" 

352 CASE  

353 WHEN {field_path} ~ '^[0-9]+(\\.[0-9]+)?$'  

354 THEN ({field_path})::numeric  

355 ELSE NULL  

356 END {direction} NULLS LAST, 

357 {field_path} {direction} 

358 """ 

359 order_clauses.append(order_clause) 

360 sql += " ORDER BY " + ", ".join(order_clauses) 

361 

362 # Add LIMIT and OFFSET 

363 if query.limit_value: 

364 sql += f" LIMIT {query.limit_value}" 

365 if query.offset_value: 

366 sql += f" OFFSET {query.offset_value}" 

367 

368 # Execute query 

369 df = self.db.query(sql, params) 

370 

371 # Convert to records 

372 records = [] 

373 for _, row in df.iterrows(): 

374 record = self._row_to_record(row.to_dict()) 

375 

376 # Apply field projection if specified 

377 if query.fields: 

378 record = record.project(query.fields) 

379 

380 records.append(record) 

381 

382 return records 

383 

384 def _count_all(self) -> int: 

385 """Count all records in the database.""" 

386 self._check_connection() 

387 sql = f"SELECT COUNT(*) as count FROM {self.schema_name}.{self.table_name}" 

388 df = self.db.query(sql) 

389 return int(df.iloc[0]["count"]) if not df.empty else 0 

390 

391 def clear(self) -> int: 

392 """Clear all records from the database.""" 

393 self._check_connection() 

394 # Get count first 

395 count = self._count_all() 

396 

397 # Delete all records 

398 sql = f"TRUNCATE TABLE {self.schema_name}.{self.table_name}" 

399 self.db.execute(sql) 

400 

401 return count 

402 

403 def stream_read( 

404 self, 

405 query: Optional[Query] = None, 

406 config: Optional[StreamConfig] = None 

407 ) -> Iterator[Record]: 

408 """Stream records from PostgreSQL.""" 

409 self._check_connection() 

410 config = config or StreamConfig() 

411 

412 # Build SQL query 

413 sql = f"SELECT id, data, metadata FROM {self.schema_name}.{self.table_name}" 

414 params = {} 

415 

416 if query and query.filters: 

417 # Add WHERE clause (simplified for now) 

418 where_clauses = [] 

419 for i, filter in enumerate(query.filters): 

420 field_path = f"data->>'{filter.field}'" 

421 param_name = f"param_{i}" 

422 

423 if filter.operator == Operator.EQ: 

424 where_clauses.append(f"{field_path} = %({param_name})s") 

425 params[param_name] = str(filter.value) 

426 

427 if where_clauses: 

428 sql += " WHERE " + " AND ".join(where_clauses) 

429 

430 # Use cursor for streaming 

431 # Note: PostgresDB may need modification to support cursors 

432 # For now, we'll fetch in batches 

433 sql += f" LIMIT {config.batch_size} OFFSET %(offset)s" 

434 

435 offset = 0 

436 while True: 

437 params["offset"] = offset 

438 df = self.db.query(sql, params) 

439 

440 if df.empty: 

441 break 

442 

443 for _, row in df.iterrows(): 

444 record = self._row_to_record(row.to_dict()) 

445 if query and query.fields: 

446 record = record.project(query.fields) 

447 yield record 

448 

449 offset += config.batch_size 

450 

451 # If we got less than batch_size, we're done 

452 if len(df) < config.batch_size: 

453 break 

454 

455 def stream_write( 

456 self, 

457 records: Iterator[Record], 

458 config: Optional[StreamConfig] = None 

459 ) -> StreamResult: 

460 """Stream records into PostgreSQL.""" 

461 self._check_connection() 

462 config = config or StreamConfig() 

463 result = StreamResult() 

464 start_time = time.time() 

465 quitting = False 

466 

467 batch = [] 

468 for record in records: 

469 batch.append(record) 

470 

471 if len(batch) >= config.batch_size: 

472 # Write batch with graceful fallback 

473 # Use lambda wrapper for _write_batch 

474 continue_processing = process_batch_with_fallback( 

475 batch, 

476 lambda b: self._write_batch(b) or [r.id for r in b], # _write_batch returns None, we need IDs 

477 self.create, 

478 result, 

479 config 

480 ) 

481 

482 if not continue_processing: 

483 quitting = True 

484 break 

485 

486 batch = [] 

487 

488 # Write remaining batch 

489 if batch and not quitting: 

490 process_batch_with_fallback( 

491 batch, 

492 lambda b: self._write_batch(b) or [r.id for r in b], 

493 self.create, 

494 result, 

495 config 

496 ) 

497 

498 result.duration = time.time() - start_time 

499 return result 

500 

501 def _write_batch(self, records: list[Record]) -> None: 

502 """Write a batch of records to the database.""" 

503 # Build batch insert SQL 

504 values = [] 

505 params = {} 

506 

507 for i, record in enumerate(records): 

508 id = str(uuid.uuid4()) 

509 row = self._record_to_row(record, id) 

510 values.append(f"(%(id_{i})s, %(data_{i})s, %(metadata_{i})s)") 

511 params[f"id_{i}"] = row["id"] 

512 params[f"data_{i}"] = row["data"] 

513 params[f"metadata_{i}"] = row["metadata"] 

514 

515 sql = f""" 

516 INSERT INTO {self.schema_name}.{self.table_name} (id, data, metadata) 

517 VALUES {', '.join(values)} 

518 """ 

519 self.db.execute(sql, params) 

520 

521 

522# Global pool manager instance for async PostgreSQL connections 

523_pool_manager = ConnectionPoolManager[asyncpg.Pool]() 

524 

525 

526class AsyncPostgresDatabase(AsyncDatabase, ConfigurableBase): 

527 """Native async PostgreSQL database backend with event loop-aware connection pooling.""" 

528 

529 def __init__(self, config: dict[str, Any] | None = None): 

530 """Initialize async PostgreSQL database.""" 

531 super().__init__(config) 

532 config = config or {} 

533 self._pool_config = PostgresPoolConfig.from_dict(config) 

534 # Add table and schema to pool config from regular config 

535 self.table_name = config.get("table", "records") 

536 self.schema_name = config.get("schema", "public") 

537 self._pool: Optional[asyncpg.Pool] = None 

538 self._connected = False 

539 

540 @classmethod 

541 def from_config(cls, config: dict) -> "AsyncPostgresDatabase": 

542 """Create from config dictionary.""" 

543 return cls(config) 

544 

545 async def connect(self) -> None: 

546 """Connect to the database.""" 

547 if self._connected: 

548 return 

549 

550 # Get or create pool for current event loop 

551 self._pool = await _pool_manager.get_pool( 

552 self._pool_config, 

553 create_asyncpg_pool, 

554 validate_asyncpg_pool 

555 ) 

556 

557 # Ensure table exists 

558 await self._ensure_table() 

559 self._connected = True 

560 

561 async def close(self) -> None: 

562 """Close the database connection.""" 

563 if self._connected: 

564 # Pool manager handles cleanup 

565 self._pool = None 

566 self._connected = False 

567 

568 def _initialize(self) -> None: 

569 """Initialize is handled in connect.""" 

570 pass 

571 

572 async def _ensure_table(self) -> None: 

573 """Ensure the records table exists.""" 

574 if not self._pool: 

575 raise RuntimeError("Database not connected. Call connect() first.") 

576 

577 create_table_sql = f""" 

578 CREATE TABLE IF NOT EXISTS {self.schema_name}.{self.table_name} ( 

579 id VARCHAR(255) PRIMARY KEY, 

580 data JSONB NOT NULL, 

581 metadata JSONB, 

582 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, 

583 updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP 

584 ); 

585  

586 CREATE INDEX IF NOT EXISTS idx_{self.table_name}_data  

587 ON {self.schema_name}.{self.table_name} USING GIN (data); 

588  

589 CREATE INDEX IF NOT EXISTS idx_{self.table_name}_metadata 

590 ON {self.schema_name}.{self.table_name} USING GIN (metadata); 

591 """ 

592 

593 async with self._pool.acquire() as conn: 

594 await conn.execute(create_table_sql) 

595 

596 def _check_connection(self) -> None: 

597 """Check if database is connected.""" 

598 if not self._connected or not self._pool: 

599 raise RuntimeError("Database not connected. Call connect() first.") 

600 

601 def _record_to_row(self, record: Record, id: str | None = None) -> dict[str, Any]: 

602 """Convert a Record to a database row.""" 

603 data = {} 

604 for field_name, field_obj in record.fields.items(): 

605 data[field_name] = field_obj.value 

606 

607 return { 

608 "id": id or str(uuid.uuid4()), 

609 "data": json.dumps(data), 

610 "metadata": json.dumps(record.metadata) if record.metadata else None, 

611 } 

612 

613 def _row_to_record(self, row: asyncpg.Record) -> Record: 

614 """Convert a database row to a Record.""" 

615 data = row.get("data", {}) 

616 if isinstance(data, str): 

617 data = json.loads(data) 

618 

619 metadata = row.get("metadata", {}) 

620 if isinstance(metadata, str) and metadata: 

621 metadata = json.loads(metadata) 

622 elif not metadata: 

623 metadata = {} 

624 

625 return Record(data=data, metadata=metadata) 

626 

627 async def create(self, record: Record) -> str: 

628 """Create a new record.""" 

629 self._check_connection() 

630 # Use record's ID if it has one, otherwise generate a new one 

631 id = record.id if record.id else str(uuid.uuid4()) 

632 row = self._record_to_row(record, id) 

633 

634 sql = f""" 

635 INSERT INTO {self.schema_name}.{self.table_name} (id, data, metadata) 

636 VALUES ($1, $2, $3) 

637 """ 

638 

639 async with self._pool.acquire() as conn: 

640 await conn.execute(sql, row["id"], row["data"], row["metadata"]) 

641 

642 return id 

643 

644 async def read(self, id: str) -> Record | None: 

645 """Read a record by ID.""" 

646 self._check_connection() 

647 sql = f""" 

648 SELECT id, data, metadata 

649 FROM {self.schema_name}.{self.table_name} 

650 WHERE id = $1 

651 """ 

652 

653 async with self._pool.acquire() as conn: 

654 row = await conn.fetchrow(sql, id) 

655 

656 if not row: 

657 return None 

658 

659 return self._row_to_record(row) 

660 

661 async def update(self, id: str, record: Record) -> bool: 

662 """Update an existing record.""" 

663 self._check_connection() 

664 row = self._record_to_row(record, id) 

665 

666 sql = f""" 

667 UPDATE {self.schema_name}.{self.table_name} 

668 SET data = $2, metadata = $3, updated_at = CURRENT_TIMESTAMP 

669 WHERE id = $1 

670 """ 

671 

672 async with self._pool.acquire() as conn: 

673 result = await conn.execute(sql, row["id"], row["data"], row["metadata"]) 

674 

675 # Returns UPDATE n where n is rows affected 

676 return result.split()[-1] != "0" 

677 

678 async def delete(self, id: str) -> bool: 

679 """Delete a record by ID.""" 

680 self._check_connection() 

681 sql = f""" 

682 DELETE FROM {self.schema_name}.{self.table_name} 

683 WHERE id = $1 

684 """ 

685 

686 async with self._pool.acquire() as conn: 

687 result = await conn.execute(sql, id) 

688 

689 # Returns DELETE n where n is rows affected 

690 return result.split()[-1] != "0" 

691 

692 async def exists(self, id: str) -> bool: 

693 """Check if a record exists.""" 

694 self._check_connection() 

695 sql = f""" 

696 SELECT 1 FROM {self.schema_name}.{self.table_name} 

697 WHERE id = $1 

698 LIMIT 1 

699 """ 

700 

701 async with self._pool.acquire() as conn: 

702 row = await conn.fetchrow(sql, id) 

703 

704 return row is not None 

705 

706 async def upsert(self, id: str, record: Record) -> str: 

707 """Update or insert a record with a specific ID.""" 

708 self._check_connection() 

709 row = self._record_to_row(record, id) 

710 

711 sql = f""" 

712 INSERT INTO {self.schema_name}.{self.table_name} (id, data, metadata) 

713 VALUES ($1, $2, $3) 

714 ON CONFLICT (id) DO UPDATE 

715 SET data = EXCLUDED.data, metadata = EXCLUDED.metadata, updated_at = CURRENT_TIMESTAMP 

716 """ 

717 

718 async with self._pool.acquire() as conn: 

719 await conn.execute(sql, row["id"], row["data"], row["metadata"]) 

720 

721 return id 

722 

723 async def search(self, query: Query) -> list[Record]: 

724 """Search for records matching the query.""" 

725 self._check_connection() 

726 

727 # Build SQL query from Query object 

728 where_clauses = [] 

729 params = [] 

730 param_count = 0 

731 

732 # Build WHERE clauses for filters 

733 for filter in query.filters: 

734 param_count += 1 

735 field_path = f"data->>'{filter.field}'" 

736 

737 if filter.operator == Operator.EQ: 

738 if isinstance(filter.value, bool): 

739 where_clauses.append(f"({field_path})::boolean = ${param_count}") 

740 params.append(filter.value) 

741 elif isinstance(filter.value, (int, float)): 

742 where_clauses.append(f"({field_path})::numeric = ${param_count}") 

743 params.append(filter.value) 

744 else: 

745 where_clauses.append(f"{field_path} = ${param_count}") 

746 params.append(str(filter.value)) 

747 elif filter.operator == Operator.NEQ: 

748 if isinstance(filter.value, bool): 

749 where_clauses.append(f"({field_path})::boolean != ${param_count}") 

750 params.append(filter.value) 

751 elif isinstance(filter.value, (int, float)): 

752 where_clauses.append(f"({field_path})::numeric != ${param_count}") 

753 params.append(filter.value) 

754 else: 

755 where_clauses.append(f"{field_path} != ${param_count}") 

756 params.append(str(filter.value)) 

757 elif filter.operator == Operator.GT: 

758 where_clauses.append(f"({field_path})::numeric > ${param_count}") 

759 params.append(filter.value) 

760 elif filter.operator == Operator.LT: 

761 where_clauses.append(f"({field_path})::numeric < ${param_count}") 

762 params.append(filter.value) 

763 elif filter.operator == Operator.GTE: 

764 where_clauses.append(f"({field_path})::numeric >= ${param_count}") 

765 params.append(filter.value) 

766 elif filter.operator == Operator.LTE: 

767 where_clauses.append(f"({field_path})::numeric <= ${param_count}") 

768 params.append(filter.value) 

769 elif filter.operator == Operator.LIKE: 

770 where_clauses.append(f"{field_path} LIKE ${param_count}") 

771 params.append(f"%{filter.value}%") 

772 elif filter.operator == Operator.IN: 

773 values = [str(v) for v in filter.value] 

774 where_clauses.append(f"{field_path} = ANY(${param_count})") 

775 params.append(values) 

776 elif filter.operator == Operator.NOT_IN: 

777 values = [str(v) for v in filter.value] 

778 where_clauses.append(f"{field_path} != ALL(${param_count})") 

779 params.append(values) 

780 elif filter.operator == Operator.BETWEEN: 

781 # Optimize BETWEEN for different data types 

782 if isinstance(filter.value, (list, tuple)) and len(filter.value) == 2: 

783 lower, upper = filter.value 

784 

785 # Try to determine the type for proper casting 

786 if isinstance(lower, (int, float)) and isinstance(upper, (int, float)): 

787 where_clauses.append( 

788 f"({field_path})::numeric BETWEEN ${param_count} AND ${param_count + 1}" 

789 ) 

790 elif isinstance(lower, datetime) or isinstance(upper, datetime): 

791 where_clauses.append( 

792 f"({field_path})::timestamp BETWEEN ${param_count} AND ${param_count + 1}" 

793 ) 

794 else: 

795 # String or unknown type 

796 where_clauses.append( 

797 f"{field_path} BETWEEN ${param_count} AND ${param_count + 1}" 

798 ) 

799 

800 params.append(lower) 

801 params.append(upper) 

802 param_count += 1 # We used two parameters 

803 elif filter.operator == Operator.NOT_BETWEEN: 

804 # Optimize NOT BETWEEN 

805 if isinstance(filter.value, (list, tuple)) and len(filter.value) == 2: 

806 lower, upper = filter.value 

807 

808 # Try to determine the type for proper casting 

809 if isinstance(lower, (int, float)) and isinstance(upper, (int, float)): 

810 where_clauses.append( 

811 f"({field_path})::numeric NOT BETWEEN ${param_count} AND ${param_count + 1}" 

812 ) 

813 elif isinstance(lower, datetime) or isinstance(upper, datetime): 

814 where_clauses.append( 

815 f"({field_path})::timestamp NOT BETWEEN ${param_count} AND ${param_count + 1}" 

816 ) 

817 else: 

818 # String or unknown type 

819 where_clauses.append( 

820 f"{field_path} NOT BETWEEN ${param_count} AND ${param_count + 1}" 

821 ) 

822 

823 params.append(lower)