Coverage for src/dataknobs_data/database.py: 26%

1from abc import ABC, abstractmethod

2from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, Union

4from .query import Query

5from .query_logic import ComplexQuery

6from .records import Record

7from .streaming import StreamConfig, StreamResult

10class AsyncDatabase(ABC):

11 """Abstract base class for async database implementations."""

13 def __init__(self, config: Dict[str, Any] | None = None):

14 """Initialize the database with optional configuration.

16 Args:

17 config: Backend-specific configuration parameters

18 """

19 self.config = config or {}

20 self._initialize()

22 def _initialize(self) -> None:

23 """Initialize the database backend. Override in subclasses."""

24 pass

26 @abstractmethod

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

28 """Create a new record in the database.

30 Args:

31 record: The record to create

33 Returns:

34 The ID of the created record

35 """

36 raise NotImplementedError

38 @abstractmethod

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

40 """Read a record by ID.

42 Args:

43 id: The record ID

45 Returns:

46 The record if found, None otherwise

47 """

48 raise NotImplementedError

50 @abstractmethod

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

52 """Update an existing record.

54 Args:

55 id: The record ID

56 record: The updated record

58 Returns:

59 True if the record was updated, False if not found

60 """

61 raise NotImplementedError

63 @abstractmethod

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

65 """Delete a record by ID.

67 Args:

68 id: The record ID

70 Returns:

71 True if the record was deleted, False if not found

72 """

73 raise NotImplementedError

75 @abstractmethod

76 async def search(self, query: Union[Query, ComplexQuery]) -> List[Record]:

77 """Search for records matching a query.

79 Args:

80 query: The search query (simple or complex)

82 Returns:

83 List of matching records

84 """

85 raise NotImplementedError

87 async def _search_with_complex_query(self, query: ComplexQuery) -> List[Record]:

88 """Default implementation for ComplexQuery using in-memory filtering.

90 Backends can override this for native boolean logic support.

92 Args:

93 query: Complex query with boolean logic

95 Returns:

96 List of matching records

97 """

98 # Try to convert to simple query if possible

99 try:

100 simple_query = query.to_simple_query()

101 return await self.search(simple_query)

102 except ValueError:

103 # Can't convert - need to do in-memory filtering

104 # Get all records (or use a base filter if possible)

105 all_records = await self.search(Query())

106

107 # Apply complex condition filtering

108 results = []

109 for record in all_records:

110 if query.matches(record):

111 results.append(record)

112

113 # Apply sorting

114 if query.sort_specs:

115 for sort_spec in reversed(query.sort_specs):

116 reverse = sort_spec.order.value == "desc"

117 results.sort(

118 key=lambda r: r.get_value(sort_spec.field, ""),

119 reverse=reverse

120 )

121

122 # Apply offset and limit

123 if query.offset_value:

124 results = results[query.offset_value:]

125 if query.limit_value:

126 results = results[:query.limit_value]

127

128 # Apply field projection

129 if query.fields:

130 results = [r.project(query.fields) for r in results]

131

132 return results

133

134 @abstractmethod

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

136 """Check if a record exists.

137

138 Args:

139 id: The record ID

140

141 Returns:

142 True if the record exists, False otherwise

143 """

144 raise NotImplementedError

145

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

147 """Update or insert a record.

148

149 Args:

150 id: The record ID

151 record: The record to upsert

152

153 Returns:

154 The record ID

155 """

156 if await self.exists(id):

157 await self.update(id, record)

158 else:

159 return await self.create(record)

160 return id

161

162 async def create_batch(self, records: List[Record]) -> List[str]:

163 """Create multiple records in batch.

164

165 Args:

166 records: List of records to create

167

168 Returns:

169 List of created record IDs

170 """

171 ids = []

172 for record in records:

173 id = await self.create(record)

174 ids.append(id)

175 return ids

176

177 async def read_batch(self, ids: List[str]) -> List[Record | None]:

178 """Read multiple records by ID.

179

180 Args:

181 ids: List of record IDs

182

183 Returns:

184 List of records (None for not found)

185 """

186 records = []

187 for id in ids:

188 record = await self.read(id)

189 records.append(record)

190 return records

191

192 async def delete_batch(self, ids: List[str]) -> List[bool]:

193 """Delete multiple records by ID.

194

195 Args:

196 ids: List of record IDs

197

198 Returns:

199 List of deletion results

200 """

201 results = []

202 for id in ids:

203 result = await self.delete(id)

204 results.append(result)

205 return results

206

207 async def update_batch(self, updates: List[tuple[str, Record]]) -> List[bool]:

208 """Update multiple records.

209

210 Default implementation calls update() for each ID/record pair.

211 Override for better performance.

212

213 Args:

214 updates: List of (id, record) tuples to update

215

216 Returns:

217 List of success flags for each update

218 """

219 results = []

220 for id, record in updates:

221 result = await self.update(id, record)

222 results.append(result)

223 return results

224

225 async def count(self, query: Query | None = None) -> int:

226 """Count records matching a query.

227

228 Args:

229 query: Optional search query (counts all if None)

230

231 Returns:

232 Number of matching records

233 """

234 if query:

235 results = await self.search(query)

236 return len(results)

237 else:

238 return await self._count_all()

239

240 @abstractmethod

241 async def _count_all(self) -> int:

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

243 raise NotImplementedError

244

245 async def clear(self) -> int:

246 """Clear all records from the database.

247

248 Returns:

249 Number of records deleted

250 """

251 raise NotImplementedError

252

253 async def connect(self) -> None:

254 """Connect to the database."""

255 pass

256

257 async def close(self) -> None:

258 """Close the database connection."""

259 pass

260

261 async def disconnect(self) -> None:

262 """Disconnect from the database (alias for close)."""

263 await self.close()

264

265 async def __aenter__(self):

266 """Async context manager entry."""

267 await self.connect()

268 return self

269

270 async def __aexit__(self, exc_type, exc_val, exc_tb):

271 """Async context manager exit."""

272 await self.close()

273

274 @abstractmethod

275 async def stream_read(

276 self,

277 query: Optional[Query] = None,

278 config: Optional[StreamConfig] = None

279 ) -> AsyncIterator[Record]:

280 """Stream records from database.

281

282 Yields records one at a time, fetching in batches internally.

283

284 Args:

285 query: Optional query to filter records

286 config: Streaming configuration

287

288 Yields:

289 Records matching the query

290 """

291 raise NotImplementedError

292

293 @abstractmethod

294 async def stream_write(

295 self,

296 records: AsyncIterator[Record],

297 config: Optional[StreamConfig] = None

298 ) -> StreamResult:

299 """Stream records into database.

300

301 Accepts an iterator and writes in batches.

302

303 Args:

304 records: Iterator of records to write

305 config: Streaming configuration

306

307 Returns:

308 Result of the streaming operation

309 """

310 raise NotImplementedError

311

312 async def stream_transform(

313 self,

314 query: Optional[Query] = None,

315 transform: Optional[Callable[[Record], Optional[Record]]] = None,

316 config: Optional[StreamConfig] = None

317 ) -> AsyncIterator[Record]:

318 """Stream records through a transformation.

319

320 Default implementation, can be overridden for efficiency.

321

322 Args:

323 query: Optional query to filter records

324 transform: Optional transformation function

325 config: Streaming configuration

326

327 Yields:

328 Transformed records

329 """

330 async for record in self.stream_read(query, config):

331 if transform:

332 transformed = transform(record)

333 if transformed: # None means filter out

334 yield transformed

335 else:

336 yield record

337

338 @classmethod

339 async def create(cls, backend: str, config: Dict[str, Any] | None = None) -> "AsyncDatabase":

340 """Factory method to create and connect a database instance.

341

342 Args:

343 backend: The backend type ("memory", "file", "s3", "postgres", "elasticsearch")

344 config: Backend-specific configuration

345

346 Returns:

347 Connected AsyncDatabase instance

348 """

349 from .backends import BACKEND_REGISTRY

350

351 backend_class = BACKEND_REGISTRY.get(backend)

352 if not backend_class:

353 raise ValueError(

354 f"Unknown backend: {backend}. Available: {list(BACKEND_REGISTRY.keys())}"

355 )

356

357 instance = backend_class(config)

358 await instance.connect()

359 return instance

360

361

362class SyncDatabase(ABC):

363 """Synchronous variant of the Database abstract base class."""

364

365 def __init__(self, config: Dict[str, Any] | None = None):

366 """Initialize the database with optional configuration.

367

368 Args:

369 config: Backend-specific configuration parameters

370 """

371 self.config = config or {}

372 self._initialize()

373

374 def _initialize(self) -> None:

375 """Initialize the database backend. Override in subclasses."""

376 pass

377

378 @abstractmethod

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

380 """Create a new record in the database."""

381 raise NotImplementedError

382

383 @abstractmethod

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

385 """Read a record by ID."""

386 raise NotImplementedError

387

388 @abstractmethod

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

390 """Update an existing record."""

391 raise NotImplementedError

392

393 @abstractmethod

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

395 """Delete a record by ID."""

396 raise NotImplementedError

397

398 @abstractmethod

399 def search(self, query: Union[Query, ComplexQuery]) -> List[Record]:

400 """Search for records matching a query (simple or complex)."""

401 raise NotImplementedError

402

403 def _search_with_complex_query(self, query: ComplexQuery) -> List[Record]:

404 """Default implementation for ComplexQuery using in-memory filtering.

405

406 Backends can override this for native boolean logic support.

407

408 Args:

409 query: Complex query with boolean logic

410

411 Returns:

412 List of matching records

413 """

414 # Try to convert to simple query if possible

415 try:

416 simple_query = query.to_simple_query()

417 return self.search(simple_query)

418 except ValueError:

419 # Can't convert - need to do in-memory filtering

420 # Get all records (or use a base filter if possible)

421 all_records = self.search(Query())

422

423 # Apply complex condition filtering

424 results = []

425 for record in all_records:

426 if query.matches(record):

427 results.append(record)

428

429 # Apply sorting

430 if query.sort_specs:

431 for sort_spec in reversed(query.sort_specs):

432 reverse = sort_spec.order.value == "desc"

433 results.sort(

434 key=lambda r: r.get_value(sort_spec.field, ""),

435 reverse=reverse

436 )

437

438 # Apply offset and limit

439 if query.offset_value:

440 results = results[query.offset_value:]

441 if query.limit_value:

442 results = results[:query.limit_value]

443

444 # Apply field projection

445 if query.fields:

446 results = [r.project(query.fields) for r in results]

447

448 return results

449

450 @abstractmethod

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

452 """Check if a record exists."""

453 raise NotImplementedError

454

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

456 """Update or insert a record."""

457 if self.exists(id):

458 self.update(id, record)

459 else:

460 return self.create(record)

461 return id

462

463 def create_batch(self, records: List[Record]) -> List[str]:

464 """Create multiple records in batch."""

465 ids = []

466 for record in records:

467 id = self.create(record)

468 ids.append(id)

469 return ids

470

471 def read_batch(self, ids: List[str]) -> List[Record | None]:

472 """Read multiple records by ID."""

473 records = []

474 for id in ids:

475 record = self.read(id)

476 records.append(record)

477 return records

478

479 def delete_batch(self, ids: List[str]) -> List[bool]:

480 """Delete multiple records by ID."""

481 results = []

482 for id in ids:

483 result = self.delete(id)

484 results.append(result)

485 return results

486

487 def update_batch(self, updates: List[tuple[str, Record]]) -> List[bool]:

488 """Update multiple records.

489

490 Default implementation calls update() for each ID/record pair.

491 Override for better performance.

492

493 Args:

494 updates: List of (id, record) tuples to update

495

496 Returns:

497 List of success flags for each update

498 """

499 results = []

500 for id, record in updates:

501 result = self.update(id, record)

502 results.append(result)

503 return results

504

505 def count(self, query: Query | None = None) -> int:

506 """Count records matching a query."""

507 if query:

508 results = self.search(query)

509 return len(results)

510 else:

511 return self._count_all()

512

513 @abstractmethod

514 def _count_all(self) -> int:

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

516 raise NotImplementedError

517

518 def clear(self) -> int:

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

520 raise NotImplementedError

521

522 def connect(self) -> None:

523 """Connect to the database."""

524 pass

525

526 def close(self) -> None:

527 """Close the database connection."""

528 pass

529

530 def disconnect(self) -> None:

531 """Disconnect from the database (alias for close)."""

532 self.close()

533

534 def __enter__(self):

535 """Context manager entry."""

536 self.connect()

537 return self

538

539 def __exit__(self, exc_type, exc_val, exc_tb):

540 """Context manager exit."""

541 self.close()

542

543 @abstractmethod

544 def stream_read(

545 self,

546 query: Optional[Query] = None,

547 config: Optional[StreamConfig] = None

548 ) -> Iterator[Record]:

549 """Stream records from database.

550

551 Yields records one at a time, fetching in batches internally.

552

553 Args:

554 query: Optional query to filter records

555 config: Streaming configuration

556

557 Yields:

558 Records matching the query

559 """

560 raise NotImplementedError

561

562 @abstractmethod

563 def stream_write(

564 self,

565 records: Iterator[Record],

566 config: Optional[StreamConfig] = None

567 ) -> StreamResult:

568 """Stream records into database.

569

570 Accepts an iterator and writes in batches.

571

572 Args:

573 records: Iterator of records to write

574 config: Streaming configuration

575

576 Returns:

577 Result of the streaming operation

578 """

579 raise NotImplementedError

580

581 def stream_transform(

582 self,

583 query: Optional[Query] = None,

584 transform: Optional[Callable[[Record], Optional[Record]]] = None,

585 config: Optional[StreamConfig] = None

586 ) -> Iterator[Record]:

587 """Stream records through a transformation.

588

589 Default implementation, can be overridden for efficiency.

590

591 Args:

592 query: Optional query to filter records

593 transform: Optional transformation function

594 config: Streaming configuration

595

596 Yields:

597 Transformed records

598 """

599 for record in self.stream_read(query, config):

600 if transform:

601 transformed = transform(record)

602 if transformed: # None means filter out

603 yield transformed

604 else:

605 yield record

606

607 @classmethod

608 def create(cls, backend: str, config: Dict[str, Any] | None = None) -> "SyncDatabase":

609 """Factory method to create and connect a synchronous database instance.

610

611 Args:

612 backend: The backend type ("memory", "file", "s3", "postgres", "elasticsearch")

613 config: Backend-specific configuration

614

615 Returns:

616 Connected SyncDatabase instance

617 """

618 from .backends import SYNC_BACKEND_REGISTRY

619

620 backend_class = SYNC_BACKEND_REGISTRY.get(backend)

621 if not backend_class:

622 raise ValueError(

623 f"Unknown backend: {backend}. Available: {list(SYNC_BACKEND_REGISTRY.keys())}"

624 )

625

626 instance = backend_class(config)

627 instance.connect()

628 return instance