Coverage for src/dataknobs_data/migration/migrator.py: 10%
164 statements
« prev ^ index » next coverage.py v7.10.3, created at 2025-08-17 16:10 -0500
1"""
2Enhanced data migrator with streaming support.
3"""
5import asyncio
6import concurrent.futures
7from typing import Callable, Iterator, Optional, Union, List
9from dataknobs_data.database import AsyncDatabase, SyncDatabase
10from dataknobs_data.query import Query
11from dataknobs_data.records import Record
12from dataknobs_data.streaming import StreamConfig, StreamResult
14from .transformer import Transformer
15from .migration import Migration
16from .progress import MigrationProgress
19class Migrator:
20 """
21 Data migration orchestrator with streaming support.
23 Provides memory-efficient migration between databases using streaming,
24 with support for transformations, progress tracking, and parallel processing.
25 """
27 def migrate(
28 self,
29 source: SyncDatabase,
30 target: SyncDatabase,
31 transform: Optional[Union[Transformer, Migration]] = None,
32 query: Optional[Query] = None,
33 batch_size: int = 1000,
34 on_progress: Optional[Callable[[MigrationProgress], None]] = None,
35 on_error: Optional[Callable[[Exception, Record], bool]] = None
36 ) -> MigrationProgress:
37 """
38 Migrate data between databases with optional transformation.
40 Args:
41 source: Source database
42 target: Target database
43 transform: Optional transformer or migration to apply
44 query: Optional query to filter source records
45 batch_size: Number of records to process per batch
46 on_progress: Optional callback for progress updates
47 on_error: Optional error handler (return True to continue)
49 Returns:
50 MigrationProgress with final statistics
51 """
52 progress = MigrationProgress().start()
54 # Get total count for progress tracking
55 all_records = source.search(query or Query())
56 progress.total = len(all_records)
58 batch = []
59 for record in all_records:
60 try:
61 # Apply transformation if provided
62 if transform is not None:
63 if isinstance(transform, Transformer):
64 original_id = record.id # Preserve ID before transformation
65 transformed = transform.transform(record)
66 if transformed is None:
67 # Record filtered out
68 progress.record_skip("Filtered by transformer", original_id)
69 continue
70 record = transformed
71 elif isinstance(transform, Migration):
72 record = transform.apply(record)
74 batch.append(record)
76 # Process batch when full
77 if len(batch) >= batch_size:
78 self._write_batch(target, batch, progress, on_error)
79 batch = []
81 if on_progress:
82 on_progress(progress)
84 except Exception as e:
85 progress.record_failure(str(e), record.id if hasattr(record, 'id') else None, e)
86 if on_error:
87 if not on_error(e, record):
88 # Handler says stop - re-raise to stop processing immediately
89 raise
90 # Handler says continue - keep going
91 else:
92 # No handler - stop processing immediately
93 raise
95 # Process final batch
96 if batch:
97 self._write_batch(target, batch, progress, on_error)
99 progress.finish()
101 if on_progress:
102 on_progress(progress)
104 return progress
106 def migrate_stream(
107 self,
108 source: SyncDatabase,
109 target: SyncDatabase,
110 transform: Optional[Union[Transformer, Migration]] = None,
111 query: Optional[Query] = None,
112 config: Optional[StreamConfig] = None,
113 on_progress: Optional[Callable[[MigrationProgress], None]] = None
114 ) -> MigrationProgress:
115 """
116 Stream-based migration for memory efficiency.
118 Never loads full dataset into memory.
120 Args:
121 source: Source database with streaming support
122 target: Target database with streaming support
123 transform: Optional transformer or migration to apply
124 query: Optional query to filter source records
125 config: Streaming configuration
126 on_progress: Optional callback for progress updates
128 Returns:
129 MigrationProgress with final statistics
130 """
131 config = config or StreamConfig()
132 progress = MigrationProgress().start()
134 # Estimate total (if possible)
135 try:
136 progress.total = source.count(query)
137 except Exception:
138 # Count not available, will track as we go
139 pass
141 # Create streaming pipeline
142 def transform_stream(records: Iterator[Record]) -> Iterator[Record]:
143 """Apply transformation to streaming records."""
144 for record in records:
145 progress.processed += 1 # Track that we've processed this record
146 try:
147 if transform is not None:
148 if isinstance(transform, Transformer):
149 original_id = record.id # Preserve ID before transformation
150 transformed = transform.transform(record)
151 if transformed:
152 yield transformed
153 else:
154 progress.record_skip("Filtered by transformer", original_id)
155 elif isinstance(transform, Migration):
156 yield transform.apply(record)
157 else:
158 yield record
159 except Exception as e:
160 if config.on_error and config.on_error(e, record):
161 progress.record_failure(str(e), record.id if hasattr(record, 'id') else None, e)
162 continue
163 else:
164 progress.record_failure(str(e), record.id if hasattr(record, 'id') else None, e)
165 raise
167 # Stream from source through transformation to target
168 source_stream = source.stream_read(query, config)
169 transformed_stream = transform_stream(source_stream)
171 # Write stream to target
172 result = target.stream_write(transformed_stream, config)
174 # Update progress from result
175 # Note: processed was already tracked in transform_stream
176 # Result contains only write successes/failures
177 progress.succeeded += result.successful
178 progress.failed += result.failed
179 progress.errors.extend(result.errors)
181 progress.finish()
183 if on_progress:
184 on_progress(progress)
186 return progress
188 def migrate_parallel(
189 self,
190 source: SyncDatabase,
191 target: SyncDatabase,
192 transform: Optional[Union[Transformer, Migration]] = None,
193 partitions: int = 4,
194 partition_field: str = "partition_id",
195 on_progress: Optional[Callable[[MigrationProgress], None]] = None
196 ) -> MigrationProgress:
197 """
198 Parallel streaming migration.
200 Partition data and migrate in parallel streams.
202 Args:
203 source: Source database
204 target: Target database
205 transform: Optional transformer or migration
206 partitions: Number of parallel partitions
207 partition_field: Field to use for partitioning
208 on_progress: Optional callback for progress updates
210 Returns:
211 Combined MigrationProgress
212 """
213 def migrate_partition(partition_id: int) -> MigrationProgress:
214 """Migrate a single partition."""
215 query = Query().filter(partition_field, "=", partition_id)
216 return self.migrate_stream(source, target, transform, query)
218 total_progress = MigrationProgress().start()
220 with concurrent.futures.ThreadPoolExecutor(max_workers=partitions) as executor:
221 futures = [
222 executor.submit(migrate_partition, i)
223 for i in range(partitions)
224 ]
226 for future in concurrent.futures.as_completed(futures):
227 partition_progress = future.result()
228 total_progress.merge(partition_progress)
230 if on_progress:
231 on_progress(total_progress)
233 total_progress.finish()
234 return total_progress
236 async def migrate_async(
237 self,
238 source: AsyncDatabase,
239 target: AsyncDatabase,
240 transform: Optional[Union[Transformer, Migration]] = None,
241 query: Optional[Query] = None,
242 config: Optional[StreamConfig] = None,
243 on_progress: Optional[Callable[[MigrationProgress], None]] = None
244 ) -> MigrationProgress:
245 """
246 Async stream-based migration.
248 Args:
249 source: Async source database
250 target: Async target database
251 transform: Optional transformer or migration
252 query: Optional query to filter source records
253 config: Streaming configuration
254 on_progress: Optional callback for progress updates
256 Returns:
257 MigrationProgress with final statistics
258 """
259 config = config or StreamConfig()
260 progress = MigrationProgress().start()
262 # Estimate total (if possible)
263 try:
264 progress.total = await source.count(query)
265 except Exception:
266 pass
268 # Create async streaming pipeline
269 async def transform_stream(records):
270 """Apply transformation to async streaming records."""
271 async for record in records:
272 progress.processed += 1 # Track that we've processed this record
273 try:
274 if transform is not None:
275 if isinstance(transform, Transformer):
276 original_id = record.id # Preserve ID before transformation
277 transformed = transform.transform(record)
278 if transformed:
279 yield transformed
280 else:
281 progress.record_skip("Filtered by transformer", original_id)
282 elif isinstance(transform, Migration):
283 yield transform.apply(record)
284 else:
285 yield record
286 except Exception as e:
287 if config.on_error and config.on_error(e, record):
288 progress.record_failure(str(e), record.id if hasattr(record, 'id') else None, e)
289 continue
290 else:
291 progress.record_failure(str(e), record.id if hasattr(record, 'id') else None, e)
292 raise
294 # Stream from source through transformation to target
295 source_stream = source.stream_read(query, config)
296 transformed_stream = transform_stream(source_stream)
298 # Write stream to target
299 result = await target.stream_write(transformed_stream, config)
301 # Update progress from result
302 # Note: processed was already tracked in transform_stream
303 # Result contains only write successes/failures
304 progress.succeeded += result.successful
305 progress.failed += result.failed
306 progress.errors.extend(result.errors)
308 progress.finish()
310 if on_progress:
311 on_progress(progress)
313 return progress
315 def _write_batch(
316 self,
317 target: SyncDatabase,
318 batch: List[Record],
319 progress: MigrationProgress,
320 on_error: Optional[Callable[[Exception, Record], bool]] = None
321 ) -> None:
322 """
323 Write a batch of records to target database.
325 Args:
326 target: Target database
327 batch: Batch of records to write
328 progress: Progress tracker to update
329 on_error: Optional error handler
330 """
331 for record in batch:
332 try:
333 # Ensure record has an ID
334 if not record.id:
335 record.generate_id()
337 target.create(record)
338 progress.record_success(record.id)
339 except Exception as e:
340 progress.record_failure(str(e), record.id, e)
341 if on_error:
342 if not on_error(e, record):
343 # Handler says stop - re-raise to stop processing immediately
344 raise
345 # Handler says continue - keep going
346 else:
347 # No handler - stop processing immediately
348 raise
350 def validate_migration(
351 self,
352 source: SyncDatabase,
353 target: SyncDatabase,
354 query: Optional[Query] = None,
355 sample_size: Optional[int] = None
356 ) -> tuple[bool, List[str]]:
357 """
358 Validate that migration was successful.
360 Args:
361 source: Source database
362 target: Target database
363 query: Optional query used for migration
364 sample_size: Optional number of records to sample for validation
366 Returns:
367 Tuple of (is_valid, list_of_issues)
368 """
369 issues = []
371 # Get counts
372 source_records = source.search(query or Query())
373 target_records = target.search(Query())
375 source_count = len(source_records)
376 target_count = len(target_records)
378 if source_count != target_count:
379 issues.append(
380 f"Record count mismatch: source={source_count}, target={target_count}"
381 )
383 # Sample validation
384 if sample_size:
385 sample = source_records[:sample_size]
386 else:
387 sample = source_records
389 for source_record in sample:
390 if source_record.id:
391 target_record = target.read(source_record.id)
392 if not target_record:
393 issues.append(f"Record {source_record.id} not found in target")
395 return len(issues) == 0, issues