Progress Bar Estimate

airbyte-commons-worker/src/main/java/io/airbyte/workers/internal/AirbyteMessageTracker.java

@@ -8,6 +8,7 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Charsets;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.hash.HashFunction;
@@ -19,6 +20,7 @@
 import io.airbyte.config.State;
 import io.airbyte.protocol.models.AirbyteControlConnectorConfigMessage;
 import io.airbyte.protocol.models.AirbyteControlMessage;
+import io.airbyte.protocol.models.AirbyteEstimateTraceMessage;
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import io.airbyte.protocol.models.AirbyteStateMessage;
@@ -53,13 +55,20 @@ public class AirbyteMessageTracker implements MessageTracker {
   private final BiMap<AirbyteStreamNameNamespacePair, Short> nameNamespacePairToIndex;
   private final Map<Short, Long> streamToTotalBytesEmitted;
   private final Map<Short, Long> streamToTotalRecordsEmitted;
+  private final Map<Short, Long> streamToTotalBytesEstimated;
+  private final Map<Short, Long> streamToTotalRecordsEstimated;
   private final StateDeltaTracker stateDeltaTracker;
   private final StateMetricsTracker stateMetricsTracker;
   private final List<AirbyteTraceMessage> destinationErrorTraceMessages;
   private final List<AirbyteTraceMessage> sourceErrorTraceMessages;
   private final StateAggregator stateAggregator;
   private final boolean logConnectorMessages = new EnvVariableFeatureFlags().logConnectorMessages();
 
+  // These variables support SYNC level estimates and are meant for sources where stream level
+  // estimates are not possible e.g. CDC sources.
+  private Long totalRecordsEstimatedSync;
+  private Long totalBytesEstimatedSync;
+
   private short nextStreamIndex;
 
   /**
@@ -95,6 +104,8 @@ protected AirbyteMessageTracker(final StateDeltaTracker stateDeltaTracker,
     this.hashFunction = Hashing.murmur3_32_fixed();
     this.streamToTotalBytesEmitted = new HashMap<>();
     this.streamToTotalRecordsEmitted = new HashMap<>();
+    this.streamToTotalBytesEstimated = new HashMap<>();
+    this.streamToTotalRecordsEstimated = new HashMap<>();
     this.stateDeltaTracker = stateDeltaTracker;
     this.stateMetricsTracker = stateMetricsTracker;
     this.nextStreamIndex = 0;
@@ -252,7 +263,7 @@ private void handleEmittedOrchestratorConnectorConfig(final AirbyteControlConnec
    */
   private void handleEmittedTrace(final AirbyteTraceMessage traceMessage, final ConnectorType connectorType) {
     switch (traceMessage.getType()) {
-      case ESTIMATE -> handleEmittedEstimateTrace(traceMessage, connectorType);
+      case ESTIMATE -> handleEmittedEstimateTrace(traceMessage.getEstimate());
       case ERROR -> handleEmittedErrorTrace(traceMessage, connectorType);
       default -> log.warn("Invalid message type for trace message: {}", traceMessage);
     }
@@ -266,8 +277,35 @@ private void handleEmittedErrorTrace(final AirbyteTraceMessage errorTraceMessage
     }
   }
 
-  @SuppressWarnings("PMD") // until method is implemented
-  private void handleEmittedEstimateTrace(final AirbyteTraceMessage estimateTraceMessage, final ConnectorType connectorType) {
+  /**
+   * There are several assumptions here:
+   * <p>
+   * - Assume the estimate is a whole number and not a sum i.e. each estimate replaces the previous
+   * estimate.
+   * <p>
+   * - Sources cannot emit both STREAM and SYNC estimates in a same sync. Error out if this happens.
+   */
+  @SuppressWarnings("PMD.AvoidDuplicateLiterals")
+  private void handleEmittedEstimateTrace(final AirbyteEstimateTraceMessage estimate) {
+    switch (estimate.getType()) {
+      case STREAM -> {
+        Preconditions.checkArgument(totalBytesEstimatedSync == null, "STREAM and SYNC estimates should not be emitted in the same sync.");
+        Preconditions.checkArgument(totalRecordsEstimatedSync == null, "STREAM and SYNC estimates should not be emitted in the same sync.");
+
+        log.info("Saving records estimates for namespace: {}, stream: {}", estimate.getNamespace(), estimate.getName());
+        final var index = getStreamIndex(new AirbyteStreamNameNamespacePair(estimate.getName(), estimate.getNamespace()));
+
+        streamToTotalRecordsEstimated.put(index, estimate.getRowEstimate());
+        streamToTotalBytesEstimated.put(index, estimate.getByteEstimate());
+      }
+      case SYNC -> {
+        Preconditions.checkArgument(streamToTotalBytesEstimated.isEmpty(), "STREAM and SYNC estimates should not be emitted in the same sync.");
+        Preconditions.checkArgument(streamToTotalRecordsEstimated.isEmpty(), "STREAM and SYNC estimates should not be emitted in the same sync.");
+
+        totalBytesEstimatedSync = estimate.getByteEstimate();
+        totalRecordsEstimatedSync = estimate.getRowEstimate();
+      }
+    }
 
   }
 
@@ -368,6 +406,15 @@ public Map<AirbyteStreamNameNamespacePair, Long> getStreamToEmittedRecords() {
         entry -> nameNamespacePairToIndex.inverse().get(entry.getKey()), Entry::getValue));
   }
 
+  /**
+   * Swap out stream indices for stream names and return total records estimated by stream.
+   */
+  @Override
+  public Map<AirbyteStreamNameNamespacePair, Long> getStreamToEstimatedRecords() {
+    return streamToTotalRecordsEstimated.entrySet().stream().collect(Collectors.toMap(
+        entry -> nameNamespacePairToIndex.inverse().get(entry.getKey()), Entry::getValue));
+  }
+
   /**
    * Swap out stream indices for stream names and return total bytes emitted by stream.
    */
@@ -377,6 +424,15 @@ public Map<AirbyteStreamNameNamespacePair, Long> getStreamToEmittedBytes() {
         entry -> nameNamespacePairToIndex.inverse().get(entry.getKey()), Entry::getValue));
   }
 
+  /**
+   * Swap out stream indices for stream names and return total bytes estimated by stream.
+   */
+  @Override
+  public Map<AirbyteStreamNameNamespacePair, Long> getStreamToEstimatedBytes() {
+    return streamToTotalBytesEstimated.entrySet().stream().collect(Collectors.toMap(
+        entry -> nameNamespacePairToIndex.inverse().get(entry.getKey()), Entry::getValue));
+  }
+
   /**
    * Compute sum of emitted record counts across all streams.
    */
@@ -385,6 +441,18 @@ public long getTotalRecordsEmitted() {
     return streamToTotalRecordsEmitted.values().stream().reduce(0L, Long::sum);
   }
 
+  /**
+   * Compute sum of estimated record counts across all streams.
+   */
+  @Override
+  public long getTotalRecordsEstimated() {
+    if (!streamToTotalRecordsEstimated.isEmpty()) {
+      return streamToTotalRecordsEstimated.values().stream().reduce(0L, Long::sum);
+    }
+
+    return totalRecordsEstimatedSync;
+  }
+
   /**
    * Compute sum of emitted bytes across all streams.
    */
@@ -393,6 +461,18 @@ public long getTotalBytesEmitted() {
     return streamToTotalBytesEmitted.values().stream().reduce(0L, Long::sum);
   }
 
+  /**
+   * Compute sum of estimated bytes across all streams.
+   */
+  @Override
+  public long getTotalBytesEstimated() {
+    if (!streamToTotalBytesEstimated.isEmpty()) {
+      return streamToTotalBytesEstimated.values().stream().reduce(0L, Long::sum);
+    }
+
+    return totalBytesEstimatedSync;
+  }
+
   /**
    * Compute sum of committed record counts across all streams. If the delta tracker has exceeded its
    * capacity, return empty because committed record counts cannot be reliably computed.

airbyte-commons-worker/src/main/java/io/airbyte/workers/internal/MessageTracker.java

@@ -66,6 +66,14 @@ public interface MessageTracker {
    */
   Map<AirbyteStreamNameNamespacePair, Long> getStreamToEmittedRecords();
 
+  /**
+   * Get the per-stream estimated record count provided by
+   * {@link io.airbyte.protocol.models.AirbyteEstimateTraceMessage}.
+   *
+   * @return returns a map of estimated record count by stream name.
+   */
+  Map<AirbyteStreamNameNamespacePair, Long> getStreamToEstimatedRecords();
+
   /**
    * Get the per-stream emitted byte count. This includes messages that were emitted by the source,
    * but never committed by the destination.
@@ -74,6 +82,14 @@ public interface MessageTracker {
    */
   Map<AirbyteStreamNameNamespacePair, Long> getStreamToEmittedBytes();
 
+  /**
+   * Get the per-stream estimated byte count provided by
+   * {@link io.airbyte.protocol.models.AirbyteEstimateTraceMessage}.
+   *
+   * @return returns a map of estimated bytes by stream name.
+   */
+  Map<AirbyteStreamNameNamespacePair, Long> getStreamToEstimatedBytes();
+
   /**
    * Get the overall emitted record count. This includes messages that were emitted by the source, but
    * never committed by the destination.
@@ -82,6 +98,13 @@ public interface MessageTracker {
    */
   long getTotalRecordsEmitted();
 
+  /**
+   * Get the overall estimated record count.
+   *
+   * @return returns the total count of estimated records across all streams.
+   */
+  long getTotalRecordsEstimated();
+
   /**
    * Get the overall emitted bytes. This includes messages that were emitted by the source, but never
    * committed by the destination.
@@ -90,6 +113,13 @@ public interface MessageTracker {
    */
   long getTotalBytesEmitted();
 
+  /**
+   * Get the overall estimated bytes.
+   *
+   * @return returns the total count of estimated bytes across all streams.
+   */
+  long getTotalBytesEstimated();
+
   /**
    * Get the overall committed record count.
    *

airbyte-commons-worker/src/main/java/io/airbyte/workers/test_utils/AirbyteMessageUtils.java

@@ -8,6 +8,7 @@
 import com.google.common.collect.ImmutableMap;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.protocol.models.AirbyteErrorTraceMessage;
+import io.airbyte.protocol.models.AirbyteEstimateTraceMessage;
 import io.airbyte.protocol.models.AirbyteGlobalState;
 import io.airbyte.protocol.models.AirbyteLogMessage;
 import io.airbyte.protocol.models.AirbyteMessage;
@@ -102,29 +103,60 @@ public static AirbyteStreamState createStreamState(final String streamName) {
     return new AirbyteStreamState().withStreamDescriptor(new StreamDescriptor().withName(streamName));
   }
 
+  public static AirbyteMessage createStreamEstimateMessage(final String name, final String namespace, final long byteEst, final long rowEst) {
+    return createEstimateMessage(AirbyteEstimateTraceMessage.Type.STREAM, name, namespace, byteEst, rowEst);
+  }
+
+  public static AirbyteMessage createSyncEstimateMessage(final long byteEst, final long rowEst) {
+    return createEstimateMessage(AirbyteEstimateTraceMessage.Type.SYNC, null, null, byteEst, rowEst);
+  }
+
+  public static AirbyteMessage createEstimateMessage(AirbyteEstimateTraceMessage.Type type,
+                                                     final String name,
+                                                     final String namespace,
+                                                     final long byteEst,
+                                                     final long rowEst) {
+    final var est = new AirbyteEstimateTraceMessage()
+        .withType(type)
+        .withByteEstimate(byteEst)
+        .withRowEstimate(rowEst);
+
+    if (name != null) {
+      est.withName(name);
+    }
+    if (namespace != null) {
+      est.withNamespace(namespace);
+    }
+
+    return new AirbyteMessage()
+        .withType(Type.TRACE)
+        .withTrace(new AirbyteTraceMessage().withType(AirbyteTraceMessage.Type.ESTIMATE)
+            .withEstimate(est));
+  }
+
+  public static AirbyteMessage createErrorMessage(final String message, final Double emittedAt) {
+    return new AirbyteMessage()
+        .withType(AirbyteMessage.Type.TRACE)
+        .withTrace(createErrorTraceMessage(message, emittedAt));
+  }
+
   public static AirbyteTraceMessage createErrorTraceMessage(final String message, final Double emittedAt) {
-    return new AirbyteTraceMessage()
-        .withType(io.airbyte.protocol.models.AirbyteTraceMessage.Type.ERROR)
-        .withEmittedAt(emittedAt)
-        .withError(new AirbyteErrorTraceMessage().withMessage(message));
+    return createErrorTraceMessage(message, emittedAt, null);
   }
 
   public static AirbyteTraceMessage createErrorTraceMessage(final String message,
                                                             final Double emittedAt,
                                                             final AirbyteErrorTraceMessage.FailureType failureType) {
-    return new AirbyteTraceMessage()
+    final var msg = new AirbyteTraceMessage()
         .withType(io.airbyte.protocol.models.AirbyteTraceMessage.Type.ERROR)
-        .withEmittedAt(emittedAt)
-        .withError(new AirbyteErrorTraceMessage().withMessage(message).withFailureType(failureType));
-  }
+        .withError(new AirbyteErrorTraceMessage().withMessage(message))
+        .withEmittedAt(emittedAt);
 
-  public static AirbyteMessage createTraceMessage(final String message, final Double emittedAt) {
-    return new AirbyteMessage()
-        .withType(AirbyteMessage.Type.TRACE)
-        .withTrace(new AirbyteTraceMessage()
-            .withType(AirbyteTraceMessage.Type.ERROR)
-            .withEmittedAt(emittedAt)
-            .withError(new AirbyteErrorTraceMessage().withMessage(message)));
+    if (failureType != null) {
+      msg.getError().withFailureType(failureType);
+    }
+
+    return msg;
   }
 
 }

airbyte-commons-worker/src/test/java/io/airbyte/workers/general/DefaultReplicationWorkerTest.java

@@ -307,7 +307,7 @@ void testReplicationRunnableWorkerFailure() throws Exception {
   @Test
   void testOnlyStateAndRecordMessagesDeliveredToDestination() throws Exception {
     final AirbyteMessage LOG_MESSAGE = AirbyteMessageUtils.createLogMessage(Level.INFO, "a log message");
-    final AirbyteMessage TRACE_MESSAGE = AirbyteMessageUtils.createTraceMessage("a trace message", 123456.0);
+    final AirbyteMessage TRACE_MESSAGE = AirbyteMessageUtils.createErrorMessage("a trace message", 123456.0);
     when(mapper.mapMessage(LOG_MESSAGE)).thenReturn(LOG_MESSAGE);
     when(mapper.mapMessage(TRACE_MESSAGE)).thenReturn(TRACE_MESSAGE);
     when(source.isFinished()).thenReturn(false, false, false, false, true);