Coverage for src/dataknobs_data/streaming.py: 42%
218 statements
« prev ^ index » next coverage.py v7.10.3, created at 2025-08-17 19:59 -0500
1"""Streaming support for database operations."""
3import time
4from dataclasses import dataclass, field
5from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional, TYPE_CHECKING
7from .records import Record
9if TYPE_CHECKING:
10 from .query import Query
13@dataclass
14class StreamConfig:
15 """Configuration for streaming operations."""
17 batch_size: int = 1000
18 prefetch: int = 2 # Number of batches to prefetch
19 timeout: Optional[float] = None
20 on_error: Optional[Callable[[Exception, Record], bool]] = None # Return True to continue
22 def __post_init__(self):
23 """Validate configuration."""
24 if self.batch_size <= 0:
25 raise ValueError("batch_size must be positive")
26 if self.prefetch < 0:
27 raise ValueError("prefetch must be non-negative")
28 if self.timeout is not None and self.timeout <= 0:
29 raise ValueError("timeout must be positive if specified")
32@dataclass
33class StreamResult:
34 """Result of streaming operation."""
36 total_processed: int = 0
37 successful: int = 0
38 failed: int = 0
39 errors: List[Dict[str, Any]] = field(default_factory=list)
40 duration: float = 0.0
41 total_batches: int = 0 # Number of batches processed
42 failed_indices: List[int] = field(default_factory=list) # Indices of failed records
44 @property
45 def success_rate(self) -> float:
46 """Calculate success rate as percentage."""
47 if self.total_processed == 0:
48 return 0.0
49 return (self.successful / self.total_processed) * 100
51 def add_error(self, record_id: Optional[str], error: Exception, index: Optional[int] = None) -> None:
52 """Add an error to the result.
54 Args:
55 record_id: ID of the record that failed
56 error: The exception that occurred
57 index: Optional index of the failed record in the original batch
58 """
59 self.errors.append({
60 "record_id": record_id,
61 "error": str(error),
62 "type": type(error).__name__,
63 "index": index
64 })
65 if index is not None:
66 self.failed_indices.append(index)
68 def merge(self, other: "StreamResult") -> None:
69 """Merge another result into this one."""
70 self.total_processed += other.total_processed
71 self.successful += other.successful
72 self.failed += other.failed
73 self.errors.extend(other.errors)
74 self.duration += other.duration
75 self.total_batches += other.total_batches
76 self.failed_indices.extend(other.failed_indices)
78 def __str__(self) -> str:
79 """Human-readable representation."""
80 return (
81 f"StreamResult(processed={self.total_processed}, "
82 f"successful={self.successful}, failed={self.failed}, "
83 f"success_rate={self.success_rate:.1f}%, "
84 f"duration={self.duration:.2f}s)"
85 )
88def process_batch_with_fallback(
89 batch: List[Record],
90 batch_create_func: Callable[[List[Record]], List[str]],
91 single_create_func: Callable[[Record], str],
92 result: StreamResult,
93 config: StreamConfig,
94 on_quit_signal: Optional[Callable[[], None]] = None,
95 batch_index: int = 0
96) -> bool:
97 """Process a batch with graceful fallback to individual record creation.
99 When a batch operation fails, this function will retry each record individually
100 to identify which specific records are causing the failure, allowing successful
101 records to be processed while only failing the problematic ones.
103 Args:
104 batch: List of records to process
105 batch_create_func: Function to create a batch of records
106 single_create_func: Function to create a single record
107 result: StreamResult to update with statistics
108 config: Stream configuration
109 on_quit_signal: Optional callback when quitting is signaled
111 Returns:
112 True to continue processing, False to quit streaming
113 """
114 try:
115 # Try batch creation first
116 ids = batch_create_func(batch)
117 result.successful += len(ids)
118 result.total_processed += len(batch)
119 result.total_batches += 1
120 return True
121 except Exception as batch_error:
122 # Batch failed, try individual records to identify failures
123 result.total_batches += 1
124 for i, record in enumerate(batch):
125 result.total_processed += 1
126 record_index = batch_index * config.batch_size + i
127 try:
128 single_create_func(record)
129 result.successful += 1
130 except Exception as record_error:
131 # This specific record failed
132 result.failed += 1
133 # Safely get record ID if available
134 record_id = record.id if record and hasattr(record, 'id') else None
135 result.add_error(record_id, record_error, record_index)
137 if config.on_error:
138 # Call error handler
139 if not config.on_error(record_error, record):
140 # Handler returned False, quit streaming
141 if on_quit_signal:
142 on_quit_signal()
143 return False
144 else:
145 # No error handler, quit on first error
146 if on_quit_signal:
147 on_quit_signal()
148 return False
150 return True
153async def async_process_batch_with_fallback(
154 batch: List[Record],
155 batch_create_func: Callable, # Async callable
156 single_create_func: Callable, # Async callable
157 result: StreamResult,
158 config: StreamConfig,
159 on_quit_signal: Optional[Callable[[], None]] = None,
160 batch_index: int = 0
161) -> bool:
162 """Async version of process_batch_with_fallback.
164 When a batch operation fails, this function will retry each record individually
165 to identify which specific records are causing the failure, allowing successful
166 records to be processed while only failing the problematic ones.
168 Args:
169 batch: List of records to process
170 batch_create_func: Async function to create a batch of records
171 single_create_func: Async function to create a single record
172 result: StreamResult to update with statistics
173 config: Stream configuration
174 on_quit_signal: Optional callback when quitting is signaled
176 Returns:
177 True to continue processing, False to quit streaming
178 """
179 try:
180 # Try batch creation first
181 ids = await batch_create_func(batch)
182 result.successful += len(ids)
183 result.total_processed += len(batch)
184 result.total_batches += 1
185 return True
186 except Exception as batch_error:
187 # Batch failed, try individual records to identify failures
188 result.total_batches += 1
189 for i, record in enumerate(batch):
190 result.total_processed += 1
191 record_index = batch_index * config.batch_size + i
192 try:
193 await single_create_func(record)
194 result.successful += 1
195 except Exception as record_error:
196 # This specific record failed
197 result.failed += 1
198 # Safely get record ID if available
199 record_id = record.id if record and hasattr(record, 'id') else None
200 result.add_error(record_id, record_error, record_index)
202 if config.on_error:
203 # Call error handler
204 if not config.on_error(record_error, record):
205 # Handler returned False, quit streaming
206 if on_quit_signal:
207 on_quit_signal()
208 return False
209 else:
210 # No error handler, quit on first error
211 if on_quit_signal:
212 on_quit_signal()
213 return False
215 return True
218class StreamProcessor:
219 """Base class for stream processing utilities."""
221 @staticmethod
222 def batch_iterator(
223 iterator: Iterator[Record],
224 batch_size: int
225 ) -> Iterator[List[Record]]:
226 """Convert a record iterator into batches."""
227 batch = []
228 for record in iterator:
229 batch.append(record)
230 if len(batch) >= batch_size:
231 yield batch
232 batch = []
233 if batch:
234 yield batch
236 @staticmethod
237 def list_to_iterator(records: List[Record]) -> Iterator[Record]:
238 """Convert a list of records to an iterator.
240 Args:
241 records: List of records
243 Yields:
244 Individual records from the list
245 """
246 for record in records:
247 yield record
249 @staticmethod
250 async def list_to_async_iterator(records: List[Record]) -> AsyncIterator[Record]:
251 """Convert a list of records to an async iterator.
253 This adapter allows synchronous lists to be used with async streaming APIs.
255 Args:
256 records: List of records
258 Yields:
259 Individual records from the list
260 """
261 for record in records:
262 yield record
264 @staticmethod
265 async def iterator_to_async_iterator(iterator: Iterator[Record]) -> AsyncIterator[Record]:
266 """Convert a synchronous iterator to an async iterator.
268 This adapter allows synchronous iterators to be used with async streaming APIs.
270 Args:
271 iterator: Synchronous iterator of records
273 Yields:
274 Individual records from the iterator
275 """
276 for record in iterator:
277 yield record
279 @staticmethod
280 async def async_batch_iterator(
281 iterator: AsyncIterator[Record],
282 batch_size: int
283 ) -> AsyncIterator[List[Record]]:
284 """Convert an async record iterator into batches."""
285 batch = []
286 async for record in iterator:
287 batch.append(record)
288 if len(batch) >= batch_size:
289 yield batch
290 batch = []
291 if batch:
292 yield batch
294 @staticmethod
295 def filter_stream(
296 iterator: Iterator[Record],
297 predicate: Callable[[Record], bool]
298 ) -> Iterator[Record]:
299 """Filter records in a stream."""
300 for record in iterator:
301 if predicate(record):
302 yield record
304 @staticmethod
305 async def async_filter_stream(
306 iterator: AsyncIterator[Record],
307 predicate: Callable[[Record], bool]
308 ) -> AsyncIterator[Record]:
309 """Filter records in an async stream."""
310 async for record in iterator:
311 if predicate(record):
312 yield record
314 @staticmethod
315 def transform_stream(
316 iterator: Iterator[Record],
317 transform: Callable[[Record], Optional[Record]]
318 ) -> Iterator[Record]:
319 """Transform records in a stream, filtering out None results."""
320 for record in iterator:
321 result = transform(record)
322 if result is not None:
323 yield result
325 @staticmethod
326 async def async_transform_stream(
327 iterator: AsyncIterator[Record],
328 transform: Callable[[Record], Optional[Record]]
329 ) -> AsyncIterator[Record]:
330 """Transform records in an async stream, filtering out None results."""
331 async for record in iterator:
332 result = transform(record)
333 if result is not None:
334 yield result
337class StreamingMixin:
338 """Mixin class providing common streaming functionality for sync databases."""
340 def _default_stream_read(
341 self,
342 query: Optional["Query"] = None,
343 config: Optional[StreamConfig] = None
344 ) -> Iterator[Record]:
345 """
346 Default implementation of stream_read using search method.
348 This provides a simple streaming wrapper around the search method
349 that most backends can use without modification.
350 """
351 config = config or StreamConfig()
353 # Use search to get all matching records
354 if query:
355 records = self.search(query)
356 else:
357 # If no query, get all records
358 from .query import Query
359 records = self.search(Query())
361 # Yield records in batches for consistency
362 for i in range(0, len(records), config.batch_size):
363 batch = records[i:i + config.batch_size]
364 for record in batch:
365 yield record
367 def _default_stream_write(
368 self,
369 records: Iterator[Record],
370 config: Optional[StreamConfig] = None
371 ) -> StreamResult:
372 """
373 Default implementation of stream_write using create_batch method.
375 This provides batch writing functionality with graceful fallback
376 to individual record creation when batches fail.
377 """
378 config = config or StreamConfig()
379 result = StreamResult()
380 start_time = time.time()
381 quitting = False
382 batch_index = 0
384 batch = []
385 for record in records:
386 batch.append(record)
388 if len(batch) >= config.batch_size:
389 # Write batch with graceful fallback
390 continue_processing = process_batch_with_fallback(
391 batch,
392 self.create_batch,
393 self.create,
394 result,
395 config,
396 batch_index=batch_index
397 )
399 if not continue_processing:
400 quitting = True
401 break
403 batch = []
404 batch_index += 1
406 # Write remaining batch
407 if batch and not quitting:
408 process_batch_with_fallback(
409 batch,
410 self.create_batch,
411 self.create,
412 result,
413 config,
414 batch_index=batch_index
415 )
417 result.duration = time.time() - start_time
418 return result
421class AsyncStreamingMixin:
422 """Mixin class providing common streaming functionality for async databases."""
424 async def _default_stream_read(
425 self,
426 query: Optional["Query"] = None,
427 config: Optional[StreamConfig] = None
428 ) -> AsyncIterator[Record]:
429 """
430 Default implementation of async stream_read using search method.
432 This provides a simple streaming wrapper around the search method
433 that most backends can use without modification.
434 """
435 config = config or StreamConfig()
437 # Use search to get all matching records
438 if query:
439 records = await self.search(query)
440 else:
441 # If no query, get all records
442 from .query import Query
443 records = await self.search(Query())
445 # Yield records in batches for consistency
446 for i in range(0, len(records), config.batch_size):
447 batch = records[i:i + config.batch_size]
448 for record in batch:
449 yield record
451 async def _default_stream_write(
452 self,
453 records: AsyncIterator[Record],
454 config: Optional[StreamConfig] = None
455 ) -> StreamResult:
456 """
457 Default implementation of async stream_write using create_batch method.
459 This provides batch writing functionality with graceful fallback
460 to individual record creation when batches fail.
461 """
462 config = config or StreamConfig()
463 result = StreamResult()
464 start_time = time.time()
465 quitting = False
466 batch_index = 0
468 batch = []
469 async for record in records:
470 batch.append(record)
472 if len(batch) >= config.batch_size:
473 # Write batch with graceful fallback
474 continue_processing = await async_process_batch_with_fallback(
475 batch,
476 self.create_batch,
477 self.create,
478 result,
479 config,
480 batch_index=batch_index
481 )
483 if not continue_processing:
484 quitting = True
485 break
487 batch = []
488 batch_index += 1
490 # Write remaining batch
491 if batch and not quitting:
492 await async_process_batch_with_fallback(
493 batch,
494 self.create_batch,
495 self.create,
496 result,
497 config,
498 batch_index=batch_index
499 )
501 result.duration = time.time() - start_time
502 return result