|
| 1 | +package com.scalar.db.dataloader.cli.command.dataimport; |
| 2 | + |
| 3 | +import com.scalar.db.dataloader.core.dataimport.ImportEventListener; |
| 4 | +import com.scalar.db.dataloader.core.dataimport.datachunk.ImportDataChunkStatus; |
| 5 | +import com.scalar.db.dataloader.core.dataimport.task.result.ImportTaskResult; |
| 6 | +import com.scalar.db.dataloader.core.dataimport.transactionbatch.ImportTransactionBatchResult; |
| 7 | +import com.scalar.db.dataloader.core.dataimport.transactionbatch.ImportTransactionBatchStatus; |
| 8 | +import java.time.Duration; |
| 9 | +import java.util.Map; |
| 10 | +import java.util.concurrent.*; |
| 11 | +import java.util.concurrent.atomic.AtomicLong; |
| 12 | + |
| 13 | +public class ConsoleImportProgressListener implements ImportEventListener { |
| 14 | + |
| 15 | + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); |
| 16 | + private final Duration updateInterval; |
| 17 | + private final long startTime; |
| 18 | + private final Map<Integer, String> chunkLogs = new ConcurrentHashMap<>(); |
| 19 | + private final Map<Integer, String> chunkFailureLogs = new ConcurrentHashMap<>(); |
| 20 | + private final AtomicLong totalRecords = new AtomicLong(); |
| 21 | + private volatile boolean completed = false; |
| 22 | + |
| 23 | + public ConsoleImportProgressListener(Duration updateInterval) { |
| 24 | + this.updateInterval = updateInterval; |
| 25 | + this.startTime = System.currentTimeMillis(); |
| 26 | + scheduler.scheduleAtFixedRate( |
| 27 | + this::render, 0, updateInterval.toMillis(), TimeUnit.MILLISECONDS); |
| 28 | + } |
| 29 | + |
| 30 | + @Override |
| 31 | + public void onDataChunkStarted(ImportDataChunkStatus status) { |
| 32 | + chunkLogs.put( |
| 33 | + status.getDataChunkId(), |
| 34 | + String.format( |
| 35 | + "🔄 Chunk %d: Processing... %d records so far", |
| 36 | + status.getDataChunkId(), status.getTotalRecords())); |
| 37 | + } |
| 38 | + |
| 39 | + @Override |
| 40 | + public void onDataChunkCompleted(ImportDataChunkStatus status) { |
| 41 | + long elapsed = System.currentTimeMillis() - status.getStartTime().toEpochMilli(); |
| 42 | + totalRecords.addAndGet(status.getTotalRecords()); |
| 43 | + if (status.getSuccessCount() > 0) { |
| 44 | + chunkLogs.put( |
| 45 | + status.getDataChunkId(), |
| 46 | + String.format( |
| 47 | + "✓ Chunk %d: %,d records imported (%.1fs), %d records imported successfully, import of %d records failed", |
| 48 | + status.getDataChunkId(), |
| 49 | + status.getTotalRecords(), |
| 50 | + elapsed / 1000.0, |
| 51 | + status.getSuccessCount(), |
| 52 | + status.getFailureCount())); |
| 53 | + } |
| 54 | + // if (status.getFailureCount() > 0) { |
| 55 | + // chunkFailureLogs.put( |
| 56 | + // status.getDataChunkId(), |
| 57 | + // String.format( |
| 58 | + // "❌ Chunk %d: Failed - %d records failed to be imported) ", |
| 59 | + // status.getDataChunkId(), status.getFailureCount())); |
| 60 | + // } |
| 61 | + } |
| 62 | + |
| 63 | + @Override |
| 64 | + public void onAllDataChunksCompleted() { |
| 65 | + completed = true; |
| 66 | + scheduler.shutdown(); |
| 67 | + render(); // Final render |
| 68 | + } |
| 69 | + |
| 70 | + @Override |
| 71 | + public void onTransactionBatchStarted(ImportTransactionBatchStatus batchStatus) { |
| 72 | + // Optional: Implement if you want to show more granular batch progress |
| 73 | + } |
| 74 | + |
| 75 | + @Override |
| 76 | + public void onTransactionBatchCompleted(ImportTransactionBatchResult batchResult) { |
| 77 | + if (!batchResult.isSuccess()) { |
| 78 | + chunkFailureLogs.put( |
| 79 | + batchResult.getDataChunkId(), |
| 80 | + String.format( |
| 81 | + "❌ Chunk %d: Transaction batch %d Failed - %d records failed to be imported) ", |
| 82 | + batchResult.getDataChunkId(), |
| 83 | + batchResult.getTransactionBatchId(), |
| 84 | + batchResult.getRecords().size())); |
| 85 | + } |
| 86 | + // Optional: Implement error reporting or success/failure count |
| 87 | + } |
| 88 | + |
| 89 | + @Override |
| 90 | + public void onTaskComplete(ImportTaskResult taskResult) { |
| 91 | + // Optional: Summary or stats after final chunk |
| 92 | + } |
| 93 | + |
| 94 | + private void render() { |
| 95 | + StringBuilder builder = new StringBuilder(); |
| 96 | + long now = System.currentTimeMillis(); |
| 97 | + long elapsed = now - startTime; |
| 98 | + double recPerSec = (totalRecords.get() * 1000.0) / (elapsed == 0 ? 1 : elapsed); |
| 99 | + |
| 100 | + builder.append( |
| 101 | + String.format( |
| 102 | + "\rImporting... %,d records | %.0f rec/s | %s\n", |
| 103 | + totalRecords.get(), recPerSec, formatElapsed(elapsed))); |
| 104 | + |
| 105 | + chunkLogs.values().stream() |
| 106 | + .sorted() // Optional: stable ordering |
| 107 | + .forEach(line -> builder.append(line).append("\n")); |
| 108 | + chunkFailureLogs.values().stream() |
| 109 | + .sorted() // Optional: stable ordering |
| 110 | + .forEach(line -> builder.append(line).append("\n")); |
| 111 | + |
| 112 | + clearConsole(); |
| 113 | + System.out.print(builder); |
| 114 | + System.out.flush(); |
| 115 | + } |
| 116 | + |
| 117 | + private String formatElapsed(long elapsedMillis) { |
| 118 | + long seconds = (elapsedMillis / 1000) % 60; |
| 119 | + long minutes = (elapsedMillis / 1000) / 60; |
| 120 | + return String.format("%dm %ds elapsed", minutes, seconds); |
| 121 | + } |
| 122 | + |
| 123 | + private void clearConsole() { |
| 124 | + // Clear screen for updated multiline rendering |
| 125 | + System.out.print("\033[H\033[2J"); // ANSI escape for clearing screen |
| 126 | + System.out.flush(); |
| 127 | + } |
| 128 | +} |
0 commit comments