Fix: Improve progress updates, including list of streams finalized (#147

)

airbyte/_processors/sql/base.py

@@ -548,6 +548,8 @@ def write_stream_data(
  finally:
  self._drop_temp_table(temp_table_name, if_exists=True)
 
+ progress.log_stream_finalized(stream_name)
+
  # Return the batch handles as measure of work completed.
  return batches_to_finalize
 

airbyte/sources/base.py

@@ -611,7 +611,7 @@ def _tally_records(
 
  for message in messages:
  yield message
- progress.log_records_read(self._processed_records)
+ progress.log_records_read(new_total_count=self._processed_records)
 
  def _log_sync_start(
  self,

examples/run_faker_to_motherduck.py

@@ -35,7 +35,7 @@
  api_key=MOTHERDUCK_API_KEY,
 )
 
-result = source.read(cache=cache)
+result = source.read(cache=cache, force_full_refresh=True)
 
 for name, records in result.streams.items():
  print(f"Stream {name}: {len(records)} records")

examples/run_snowflake_faker.py

@@ -1,4 +1,10 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""
+Usage:
+ poetry install
+ poetry run python examples/run_snowflake_faker.py
+"""
+
 from __future__ import annotations
 
 import airbyte as ab

pyproject.toml

@@ -60,6 +60,7 @@ airbyte-source-faker = "^6.0.0"
 tomli = "^2.0"
 responses = "^0.25.0"
 airbyte-source-pokeapi = "^0.2.0"
+pytest-mock = "^3.14.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]

tests/conftest.py

@@ -325,9 +325,10 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
  can pass across all cache types.
  """
  all_cache_type_fixtures: dict[str, str] = {
- "BigQuery": "new_bigquery_cache",
+ # Ordered by priority (fastest first)
  "DuckDB": "new_duckdb_cache",
  "Postgres": "new_postgres_cache",
+ "BigQuery": "new_bigquery_cache",
  "Snowflake": "new_snowflake_cache",
  }
  if is_windows():

tests/integration_tests/test_all_cache_types.py

@@ -15,6 +15,7 @@
 
 import airbyte as ab
 from airbyte._executor import _get_bin_dir
+from airbyte.progress import progress, ReadProgress
 
 
 # Product count is always the same, regardless of faker scale.
@@ -117,6 +118,22 @@ def test_pokeapi_read(
  assert len(list(result.cache.streams["pokemon"])) == 1
 
 
+@pytest.fixture(scope="function")
+def progress_mock(
+ mocker: pytest.MockerFixture,
+) -> ReadProgress:
+ """Fixture to return a mocked version of progress.progress."""
+ # Mock the progress object.
+ mocker.spy(progress, 'reset')
+ mocker.spy(progress, 'log_records_read')
+ mocker.spy(progress, 'log_batch_written')
+ mocker.spy(progress, 'log_batches_finalizing')
+ mocker.spy(progress, 'log_batches_finalized')
+ mocker.spy(progress, 'log_stream_finalized')
+ mocker.spy(progress, 'log_success')
+ return progress
+
+
 # Uncomment this line if you want to see performance trace logs.
 # You can render perf traces using the viztracer CLI or the VS Code VizTracer Extension.
 #@viztracer.trace_and_save(output_dir=".pytest_cache/snowflake_trace/")
@@ -125,11 +142,33 @@ def test_pokeapi_read(
 def test_faker_read(
  source_faker_seed_a: ab.Source,
  new_generic_cache: ab.caches.CacheBase,
+ progress_mock: ReadProgress,
 ) -> None:
  """Test that the append strategy works as expected."""
  result = source_faker_seed_a.read(
  new_generic_cache, write_strategy="replace", force_full_refresh=True
  )
+ configured_count = source_faker_seed_a._config["count"]
+
+ # These numbers expect only 'users' stream selected:
+
+ assert progress_mock.total_records_read == configured_count
+ assert progress_mock.total_records_written == configured_count
+ assert progress_mock.log_records_read.call_count >= configured_count
+ assert progress_mock.reset.call_count == 1
+ assert progress_mock.log_batch_written.call_count == 1
+ assert progress_mock.total_batches_written == 1
+ assert progress_mock.log_batches_finalizing.call_count == 1
+ assert progress_mock.log_batches_finalized.call_count == 1
+ assert progress_mock.total_batches_finalized == 1
+ assert progress_mock.finalized_stream_names == {"users"}
+ assert progress_mock.log_stream_finalized.call_count == 1
+ assert progress_mock.log_success.call_count == 1
+
+ status_msg: str = progress_mock._get_status_message()
+ assert "Read **0** records" not in status_msg
+ assert f"Read **{configured_count}** records" in status_msg
+
  assert len(list(result.cache.streams["users"])) == FAKER_SCALE_A