Coverage for src/dataknobs_data/backends/file.py: 15%

1"""File-based database backend implementation.""" 

2 

3import asyncio 

4import csv 

5import gzip 

6import json 

7import os 

8import platform 

9import tempfile 

10import threading 

11import time 

12import uuid 

13from pathlib import Path 

14from typing import Any, AsyncIterator, Iterator, Optional 

15 

16from dataknobs_config import ConfigurableBase 

17 

18from ..database import AsyncDatabase, SyncDatabase 

19from ..query import Query 

20from ..records import Record 

21from ..streaming import AsyncStreamingMixin, StreamConfig, StreamResult, StreamingMixin 

22 

23 

24class FileLock: 

25 """Cross-platform file locking.""" 

26 

27 def __init__(self, filepath: str): 

28 self.filepath = filepath 

29 self.lockfile = filepath + ".lock" 

30 self.lock_handle = None 

31 

32 def acquire(self): 

33 """Acquire the file lock.""" 

34 if platform.system() == "Windows": 

35 import msvcrt 

36 

37 while True: 

38 try: 

39 self.lock_handle = open(self.lockfile, "wb") 

40 msvcrt.locking(self.lock_handle.fileno(), msvcrt.LK_NBLCK, 1) 

41 break 

42 except OSError: 

43 if self.lock_handle: 

44 self.lock_handle.close() 

45 import time 

46 

47 time.sleep(0.01) 

48 else: 

49 import fcntl 

50 

51 self.lock_handle = open(self.lockfile, "wb") 

52 fcntl.lockf(self.lock_handle, fcntl.LOCK_EX) 

53 

54 def release(self): 

55 """Release the file lock.""" 

56 if self.lock_handle: 

57 if platform.system() == "Windows": 

58 import msvcrt 

59 

60 try: 

61 msvcrt.locking(self.lock_handle.fileno(), msvcrt.LK_UNLCK, 1) 

62 except (OSError, IOError): 

63 pass 

64 self.lock_handle.close() 

65 try: 

66 os.remove(self.lockfile) 

67 except (OSError, FileNotFoundError): 

68 pass 

69 

70 def __enter__(self): 

71 self.acquire() 

72 return self 

73 

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

75 self.release() 

76 

77 

78class FileFormat: 

79 """Base class for file format handlers.""" 

80 

81 @staticmethod 

82 def load(filepath: str) -> dict[str, dict[str, Any]]: 

83 """Load data from file.""" 

84 raise NotImplementedError 

85 

86 @staticmethod 

87 def save(filepath: str, data: dict[str, dict[str, Any]]): 

88 """Save data to file.""" 

89 raise NotImplementedError 

90 

91 

92class JSONFormat(FileFormat): 

93 """JSON file format handler.""" 

94 

95 @staticmethod 

96 def load(filepath: str) -> dict[str, dict[str, Any]]: 

97 """Load data from JSON file.""" 

98 if not os.path.exists(filepath): 

99 return {} 

100 

101 # Check if file is empty 

102 if os.path.getsize(filepath) == 0: 

103 return {} 

104 

105 try: 

106 if filepath.endswith(".gz"): 

107 try: 

108 with gzip.open(filepath, "rt", encoding="utf-8") as f: 

109 content = f.read() 

110 if not content.strip(): 

111 return {} 

112 data = json.loads(content) 

113 except (gzip.BadGzipFile, OSError): 

114 # File has .gz extension but isn't gzipped, treat as regular file 

115 with open(filepath, encoding="utf-8") as f: 

116 content = f.read() 

117 if not content.strip(): 

118 return {} 

119 data = json.loads(content) 

120 else: 

121 with open(filepath, encoding="utf-8") as f: 

122 content = f.read() 

123 if not content.strip(): 

124 return {} 

125 data = json.loads(content) 

126 

127 return data 

128 except json.JSONDecodeError: 

129 return {} 

130 

131 @staticmethod 

132 def save(filepath: str, data: dict[str, dict[str, Any]]): 

133 """Save data to JSON file.""" 

134 if filepath.endswith(".gz"): 

135 with gzip.open(filepath, "wt", encoding="utf-8") as f: 

136 json.dump(data, f, indent=2, ensure_ascii=False) 

137 else: 

138 with open(filepath, "w", encoding="utf-8") as f: 

139 json.dump(data, f, indent=2, ensure_ascii=False) 

140 

141 

142class CSVFormat(FileFormat): 

143 """CSV file format handler.""" 

144 

145 @staticmethod 

146 def load(filepath: str) -> dict[str, dict[str, Any]]: 

147 """Load data from CSV file.""" 

148 if not os.path.exists(filepath): 

149 return {} 

150 

151 # Check if file is empty 

152 if os.path.getsize(filepath) == 0: 

153 return {} 

154 

155 data = {} 

156 try: 

157 if filepath.endswith(".gz"): 

158 with gzip.open(filepath, "rt", encoding="utf-8") as f: 

159 reader = csv.DictReader(f) 

160 for row in reader: 

161 if "__id__" in row: 

162 record_id = row.pop("__id__") 

163 # Try to deserialize JSON strings back to objects 

164 fields = {} 

165 for key, value in row.items(): 

166 if value and isinstance(value, str): 

167 # Try to parse as JSON if it looks like JSON 

168 if (value.startswith('{') and value.endswith('}')) or \ 

169 (value.startswith('[') and value.endswith(']')): 

170 try: 

171 fields[key] = json.loads(value) 

172 except json.JSONDecodeError: 

173 fields[key] = value 

174 else: 

175 fields[key] = value 

176 else: 

177 fields[key] = value 

178 data[record_id] = {"fields": fields} 

179 else: 

180 with open(filepath, encoding="utf-8") as f: 

181 reader = csv.DictReader(f) 

182 for row in reader: 

183 if "__id__" in row: 

184 record_id = row.pop("__id__") 

185 # Try to deserialize JSON strings back to objects 

186 fields = {} 

187 for key, value in row.items(): 

188 if value and isinstance(value, str): 

189 # Try to parse as JSON if it looks like JSON 

190 if (value.startswith('{') and value.endswith('}')) or \ 

191 (value.startswith('[') and value.endswith(']')): 

192 try: 

193 fields[key] = json.loads(value) 

194 except json.JSONDecodeError: 

195 fields[key] = value 

196 else: 

197 fields[key] = value 

198 else: 

199 fields[key] = value 

200 data[record_id] = {"fields": fields} 

201 except (OSError, csv.Error): 

202 return {} 

203 

204 return data 

205 

206 @staticmethod 

207 def save(filepath: str, data: dict[str, dict[str, Any]]): 

208 """Save data to CSV file.""" 

209 if not data: 

210 if filepath.endswith(".gz"): 

211 with gzip.open(filepath, "wt", encoding="utf-8") as f: 

212 f.write("") 

213 else: 

214 with open(filepath, "w", encoding="utf-8") as f: 

215 f.write("") 

216 return 

217 

218 # Extract all field names and prepare flattened data 

219 all_fields = set() 

220 flattened_data = {} 

221 for record_id, record_data in data.items(): 

222 if "fields" in record_data: 

223 # Flatten field values for CSV format 

224 flat_fields = {} 

225 for field_name, field_data in record_data["fields"].items(): 

226 # Handle both full field dicts and simple values 

227 if isinstance(field_data, dict) and "value" in field_data: 

228 value = field_data["value"] 

229 else: 

230 value = field_data 

231 

232 # Serialize complex types as JSON strings 

233 if isinstance(value, (dict, list)): 

234 flat_fields[field_name] = json.dumps(value) 

235 else: 

236 flat_fields[field_name] = value 

237 all_fields.add(field_name) 

238 flattened_data[record_id] = flat_fields 

239 

240 fieldnames = ["__id__"] + sorted(list(all_fields)) 

241 

242 if filepath.endswith(".gz"): 

243 with gzip.open(filepath, "wt", encoding="utf-8", newline="") as f: 

244 writer = csv.DictWriter(f, fieldnames=fieldnames) 

245 writer.writeheader() 

246 for record_id, fields in flattened_data.items(): 

247 row = {"__id__": record_id} 

248 row.update(fields) 

249 writer.writerow(row) 

250 else: 

251 with open(filepath, "w", encoding="utf-8", newline="") as f: 

252 writer = csv.DictWriter(f, fieldnames=fieldnames) 

253 writer.writeheader() 

254 for record_id, fields in flattened_data.items(): 

255 row = {"__id__": record_id} 

256 row.update(fields) 

257 writer.writerow(row) 

258 

259 

260class ParquetFormat(FileFormat): 

261 """Parquet file format handler.""" 

262 

263 @staticmethod 

264 def load(filepath: str) -> dict[str, dict[str, Any]]: 

265 """Load data from Parquet file.""" 

266 if not os.path.exists(filepath): 

267 return {} 

268 

269 try: 

270 import pandas as pd 

271 import pyarrow.parquet as pq 

272 

273 df = pd.read_parquet(filepath) 

274 data = {} 

275 

276 for idx, row in df.iterrows(): 

277 row_dict = row.to_dict() 

278 if "__id__" in row_dict: 

279 record_id = row_dict.pop("__id__") 

280 else: 

281 record_id = str(idx) 

282 

283 # Remove NaN values 

284 fields = {k: v for k, v in row_dict.items() if pd.notna(v)} 

285 data[record_id] = {"fields": fields} 

286 

287 return data 

288 except ImportError: 

289 raise ImportError("Parquet support requires pandas and pyarrow packages") 

290 

291 @staticmethod 

292 def save(filepath: str, data: dict[str, dict[str, Any]]): 

293 """Save data to Parquet file.""" 

294 try: 

295 import pandas as pd 

296 

297 if not data: 

298 # Create empty DataFrame 

299 df = pd.DataFrame() 

300 else: 

301 rows = [] 

302 for record_id, record_data in data.items(): 

303 row = {"__id__": record_id} 

304 if "fields" in record_data: 

305 # Flatten field values for Parquet format 

306 for field_name, field_data in record_data["fields"].items(): 

307 # Handle both full field dicts and simple values 

308 if isinstance(field_data, dict) and "value" in field_data: 

309 row[field_name] = field_data["value"] 

310 else: 

311 row[field_name] = field_data 

312 rows.append(row) 

313 

314 df = pd.DataFrame(rows) 

315 

316 df.to_parquet(filepath, index=False, compression="snappy") 

317 except ImportError: 

318 raise ImportError("Parquet support requires pandas and pyarrow packages") 

319 

320 

321class AsyncFileDatabase(AsyncDatabase, AsyncStreamingMixin, ConfigurableBase): 

322 """Async file-based database implementation.""" 

323 

324 FORMAT_HANDLERS = { 

325 ".json": JSONFormat, 

326 ".csv": CSVFormat, 

327 ".tsv": CSVFormat, 

328 ".parquet": ParquetFormat, 

329 ".pq": ParquetFormat, 

330 } 

331 

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

333 super().__init__(config) 

334 self.filepath = self.config.get("path", "data.json") 

335 self.format = self.config.get("format") 

336 self.compression = self.config.get("compression", None) 

337 self._lock = asyncio.Lock() 

338 self._file_lock = FileLock(self.filepath) 

339 

340 # Detect format from file extension if not specified 

341 if not self.format: 

342 path = Path(self.filepath) 

343 # Check for compression 

344 if path.suffix == ".gz": 

345 self.compression = "gzip" 

346 path = Path(path.stem) 

347 

348 ext = path.suffix.lower() 

349 if ext in self.FORMAT_HANDLERS: 

350 self.format = ext.lstrip(".") 

351 else: 

352 self.format = "json" # Default to JSON 

353 

354 # Apply compression to filepath if specified 

355 if self.compression == "gzip" and not self.filepath.endswith(".gz"): 

356 self.filepath += ".gz" 

357 

358 # Get the appropriate format handler 

359 ext = f".{self.format}" 

360 self.handler = self.FORMAT_HANDLERS.get(ext, JSONFormat) 

361 

362 @classmethod 

363 def from_config(cls, config: dict) -> "AsyncFileDatabase": 

364 """Create from config dictionary.""" 

365 return cls(config) 

366 

367 async def connect(self) -> None: 

368 """Connect to the database (no-op for file backend).""" 

369 pass 

370 

371 def _generate_id(self) -> str: 

372 """Generate a unique ID for a record.""" 

373 return str(uuid.uuid4()) 

374 

375 async def _load_data(self) -> dict[str, Record]: 

376 """Load all data from file.""" 

377 with self._file_lock: 

378 raw_data = self.handler.load(self.filepath) 

379 data = {} 

380 for record_id, record_dict in raw_data.items(): 

381 data[record_id] = Record.from_dict(record_dict) 

382 return data 

383 

384 async def _save_data(self, data: dict[str, Record]): 

385 """Save all data to file atomically.""" 

386 # Convert records to dictionaries 

387 raw_data = {} 

388 for record_id, record in data.items(): 

389 raw_data[record_id] = record.to_dict(include_metadata=True, flatten=False) 

390 

391 # Write to temporary file first 

392 temp_fd, temp_path = tempfile.mkstemp(dir=os.path.dirname(self.filepath) or ".") 

393 os.close(temp_fd) 

394 

395 try: 

396 with self._file_lock: 

397 self.handler.save(temp_path, raw_data) 

398 # Atomic rename 

399 os.replace(temp_path, self.filepath) 

400 except Exception: 

401 # Clean up temp file on error 

402 if os.path.exists(temp_path): 

403 os.remove(temp_path) 

404 raise 

405 

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

407 """Create a new record in the file.""" 

408 async with self._lock: 

409 data = await self._load_data() 

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

411 record_id = record.id if record.id else self._generate_id() 

412 data[record_id] = record.copy(deep=True) 

413 await self._save_data(data) 

414 return record_id 

415 

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

417 """Read a record from the file.""" 

418 async with self._lock: 

419 data = await self._load_data() 

420 record = data.get(id) 

421 return record.copy(deep=True) if record else None 

422 

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

424 """Update a record in the file.""" 

425 async with self._lock: 

426 data = await self._load_data() 

427 if id in data: 

428 data[id] = record.copy(deep=True) 

429 await self._save_data(data) 

430 return True 

431 return False 

432 

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

434 """Delete a record from the file.""" 

435 async with self._lock: 

436 data = await self._load_data() 

437 if id in data: 

438 del data[id] 

439 await self._save_data(data) 

440 return True 

441 return False 

442 

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

444 """Check if a record exists in the file.""" 

445 async with self._lock: 

446 data = await self._load_data() 

447 return id in data 

448 

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

450 """Update or insert a record with the specified ID.""" 

451 async with self._lock: 

452 data = await self._load_data() 

453 data[id] = record.copy(deep=True) 

454 await self._save_data(data) 

455 return id 

456 

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

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

459 async with self._lock: 

460 data = await self._load_data() 

461 results = [] 

462 

463 for record_id, record in data.items(): 

464 # Apply filters 

465 matches = True 

466 for filter in query.filters: 

467 field_value = record.get_value(filter.field) 

468 if not filter.matches(field_value): 

469 matches = False 

470 break 

471 

472 if matches: 

473 results.append((record_id, record)) 

474 

475 # Apply sorting 

476 if query.sort_specs: 

477 for sort_spec in reversed(query.sort_specs): 

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

479 results.sort(key=lambda x: x[1].get_value(sort_spec.field, ""), reverse=reverse) 

480 

481 # Extract records 

482 records = [record for _, record in results] 

483 

484 # Apply offset and limit 

485 if query.offset_value: 

486 records = records[query.offset_value :] 

487 if query.limit_value: 

488 records = records[: query.limit_value] 

489 

490 # Apply field projection 

491 if query.fields: 

492 projected_records = [] 

493 for record in records: 

494 projected_records.append(record.project(query.fields)) 

495 records = projected_records 

496 

497 # Return deep copies 

498 return [record.copy(deep=True) for record in records] 

499 

500 async def _count_all(self) -> int: 

501 """Count all records in the file.""" 

502 async with self._lock: 

503 data = await self._load_data() 

504 return len(data) 

505 

506 async def clear(self) -> int: 

507 """Clear all records from the file.""" 

508 async with self._lock: 

509 data = await self._load_data() 

510 count = len(data) 

511 await self._save_data({}) 

512 return count 

513 

514 async def create_batch(self, records: list[Record]) -> list[str]: 

515 """Create multiple records efficiently.""" 

516 async with self._lock: 

517 data = await self._load_data() 

518 ids = [] 

519 for record in records: 

520 record_id = self._generate_id() 

521 data[record_id] = record.copy(deep=True) 

522 ids.append(record_id) 

523 await self._save_data(data) 

524 return ids 

525 

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

527 """Read multiple records efficiently.""" 

528 async with self._lock: 

529 data = await self._load_data() 

530 results = [] 

531 for record_id in ids: 

532 record = data.get(record_id) 

533 results.append(record.copy(deep=True) if record else None) 

534 return results 

535 

536 async def delete_batch(self, ids: list[str]) -> list[bool]: 

537 """Delete multiple records efficiently.""" 

538 async with self._lock: 

539 data = await self._load_data() 

540 results = [] 

541 modified = False 

542 for record_id in ids: 

543 if record_id in data: 

544 del data[record_id] 

545 results.append(True) 

546 modified = True 

547 else: 

548 results.append(False) 

549 

550 if modified: 

551 await self._save_data(data) 

552 

553 return results 

554 

555 async def stream_read( 

556 self, 

557 query: Optional[Query] = None, 

558 config: Optional[StreamConfig] = None 

559 ) -> AsyncIterator[Record]: 

560 """Stream records from file.""" 

561 # For file backend, we can use the default implementation 

562 # since we need to load all data anyway 

563 config = config or StreamConfig() 

564 

565 # Use search to get all matching records 

566 if query: 

567 records = await self.search(query) 

568 else: 

569 records = await self.search(Query()) 

570 

571 # Yield records in batches for consistency 

572 for i in range(0, len(records), config.batch_size): 

573 batch = records[i:i + config.batch_size] 

574 for record in batch: 

575 yield record 

576 

577 async def stream_write( 

578 self, 

579 records: AsyncIterator[Record], 

580 config: Optional[StreamConfig] = None 

581 ) -> StreamResult: 

582 """Stream records into file.""" 

583 # Use the default implementation from mixin 

584 return await self._default_stream_write(records, config) 

585 

586 

587class SyncFileDatabase(SyncDatabase, StreamingMixin, ConfigurableBase): 

588 """Synchronous file-based database implementation.""" 

589 

590 FORMAT_HANDLERS = { 

591 ".json": JSONFormat, 

592 ".csv": CSVFormat, 

593 ".tsv": CSVFormat, 

594 ".parquet": ParquetFormat, 

595 ".pq": ParquetFormat, 

596 } 

597 

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

599 super().__init__(config) 

600 self.filepath = self.config.get("path", "data.json") 

601 self.format = self.config.get("format") 

602 self.compression = self.config.get("compression", None) 

603 self._lock = threading.RLock() 

604 self._file_lock = FileLock(self.filepath) 

605 

606 # Detect format from file extension if not specified 

607 if not self.format: 

608 path = Path(self.filepath) 

609 # Check for compression 

610 if path.suffix == ".gz": 

611 self.compression = "gzip" 

612 path = Path(path.stem) 

613 

614 ext = path.suffix.lower() 

615 if ext in self.FORMAT_HANDLERS: 

616 self.format = ext.lstrip(".") 

617 else: 

618 self.format = "json" # Default to JSON 

619 

620 # Apply compression to filepath if specified 

621 if self.compression == "gzip" and not self.filepath.endswith(".gz"): 

622 self.filepath += ".gz" 

623 

624 # Get the appropriate format handler 

625 ext = f".{self.format}" 

626 self.handler = self.FORMAT_HANDLERS.get(ext, JSONFormat) 

627 

628 @classmethod 

629 def from_config(cls, config: dict) -> "SyncFileDatabase": 

630 """Create from config dictionary.""" 

631 return cls(config) 

632 

633 def connect(self) -> None: 

634 """Connect to the database (no-op for file backend).""" 

635 pass 

636 

637 def _generate_id(self) -> str: 

638 """Generate a unique ID for a record.""" 

639 return str(uuid.uuid4()) 

640 

641 def _load_data(self) -> dict[str, Record]: 

642 """Load all data from file.""" 

643 with self._file_lock: 

644 raw_data = self.handler.load(self.filepath) 

645 data = {} 

646 for record_id, record_dict in raw_data.items(): 

647 data[record_id] = Record.from_dict(record_dict) 

648 return data 

649 

650 def _save_data(self, data: dict[str, Record]): 

651 """Save all data to file atomically.""" 

652 # Convert records to dictionaries 

653 raw_data = {} 

654 for record_id, record in data.items(): 

655 raw_data[record_id] = record.to_dict(include_metadata=True, flatten=False) 

656 

657 # Write to temporary file first 

658 temp_fd, temp_path = tempfile.mkstemp(dir=os.path.dirname(self.filepath) or ".") 

659 os.close(temp_fd) 

660 

661 try: 

662 with self._file_lock: 

663 self.handler.save(temp_path, raw_data) 

664 # Atomic rename 

665 os.replace(temp_path, self.filepath) 

666 except Exception: 

667 # Clean up temp file on error 

668 if os.path.exists(temp_path): 

669 os.remove(temp_path) 

670 raise 

671 

672 def _do_set_data(self, data: dict[str, Record], record: Record) -> str: 

673 """Ensure record has an ID, set data[id]=record.copy() and return the ID""" 

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

675 if not record.id: 

676 record.id = self._generate_id() 

677 data[record.id] = record.copy(deep=True) 

678 return record.id 

679 

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

681 """Create a new record in the file.""" 

682 with self._lock: 

683 data = self._load_data() 

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

685 record_id = self._do_set_data(data, record) 

686 self._save_data(data) 

687 return record_id 

688 

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

690 """Read a record from the file.""" 

691 with self._lock: 

692 data = self._load_data() 

693 record = data.get(id) 

694 return record.copy(deep=True) if record else None 

695 

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

697 """Update a record in the file.""" 

698 with self._lock: 

699 data = self._load_data() 

700 if id in data: 

701 data[id] = record.copy(deep=True) 

702 self._save_data(data) 

703 return True 

704 return False 

705 

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

707 """Delete a record from the file.""" 

708 with self._lock: 

709 data = self._load_data() 

710 if id in data: 

711 del data[id] 

712 self._save_data(data) 

713 return True 

714 return False 

715 

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

717 """Check if a record exists in the file.""" 

718 with self._lock: 

719 data = self._load_data() 

720 return id in data 

721 

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

723 """Update or insert a record with the specified ID.""" 

724 with self._lock: 

725 data = self._load_data() 

726 data[id] = record.copy(deep=True) 

727 self._save_data(data) 

728 return id 

729 

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

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

732 with self._lock: 

733 data = self._load_data() 

734 results = [] 

735 

736 for record_id, record in data.items(): 

737 # Apply filters 

738 matches = True 

739 for filter in query.filters: 

740 field_value = record.get_value(filter.field) 

741 if not filter.matches(field_value): 

742 matches = False 

743 break 

744 

745 if matches: 

746 results.append((record_id, record)) 

747 

748 # Apply sorting 

749 if query.sort_specs: 

750 for sort_spec in reversed(query.sort_specs): 

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

752 results.sort(key=lambda x: x[1].get_value(sort_spec.field, ""), reverse=reverse) 

753 

754 # Extract records 

755 records = [record for _, record in results] 

756 

757 # Apply offset and limit 

758 if query.offset_value: 

759 records = records[query.offset_value :] 

760 if query.limit_value: 

761 records = records[: query.limit_value] 

762 

763 # Apply field projection 

764 if query.fields: 

765 projected_records = []