Fix: More accurate perf metrics for slow-starting streams (#385)

airbyte/progress.py

@@ -36,7 +36,6 @@
 from airbyte_protocol.models import (
  AirbyteMessage,
  AirbyteStreamStatus,
- Type,
 )
 
 from airbyte import logs
@@ -275,16 +274,17 @@ def tally_records_read(
  self.stream_read_counts[message.record.stream] += 1
 
  if message.record.stream not in self.stream_read_start_times:
- self._log_stream_read_start(stream_name=message.record.stream)
+ self.log_stream_start(stream_name=message.record.stream)
 
- elif (
- message.trace
- and message.trace.stream_status
- and message.trace.stream_status.status is AirbyteStreamStatus.COMPLETE
- ):
- self._log_stream_read_end(
- stream_name=message.trace.stream_status.stream_descriptor.name
- )
+ elif message.trace and message.trace.stream_status:
+ if message.trace.stream_status.status is AirbyteStreamStatus.STARTED:
+ self.log_stream_start(
+ stream_name=message.trace.stream_status.stream_descriptor.name
+ )
+ if message.trace.stream_status.status is AirbyteStreamStatus.COMPLETE:
+ self._log_stream_read_end(
+ stream_name=message.trace.stream_status.stream_descriptor.name
+ )
 
  # Bail if we're not due for a progress update.
  if count % update_period != 0:
@@ -345,12 +345,12 @@ def tally_confirmed_writes(
  """
  self._start_rich_view() # Start Rich's live view if not already running.
  for message in messages:
- if message.type is Type.STATE:
+ if message.state:
  # This is a state message from the destination. Tally the records written.
  if message.state.stream and message.state.destinationStats:
  stream_name = message.state.stream.stream_descriptor.name
- self.destination_stream_records_confirmed[stream_name] += (
- message.state.destinationStats.recordCount
+ self.destination_stream_records_confirmed[stream_name] += int(
+ message.state.destinationStats.recordCount or 0
  )
  self._update_display()
 
@@ -418,11 +418,14 @@ def _log_sync_start(self) -> None:
  event_type=EventType.SYNC,
  )
 
- def _log_stream_read_start(self, stream_name: str) -> None:
- self._print_info_message(
- f"Read started on stream `{stream_name}` at `{pendulum.now().format('HH:mm:ss')}`..."
- )
- self.stream_read_start_times[stream_name] = time.time()
+ def log_stream_start(self, stream_name: str) -> None:
+ """Log that a stream has started reading."""
+ if stream_name not in self.stream_read_start_times:
+ self._print_info_message(
+ f"Read started on stream `{stream_name}` at "
+ f"`{pendulum.now().format('HH:mm:ss')}`..."
+ )
+ self.stream_read_start_times[stream_name] = time.time()
 
  def _log_stream_read_end(self, stream_name: str) -> None:
  self._print_info_message(