From 688fa0680bd1a854578785cd8a20c42b8d16dc93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E5=9F=BA=E9=AD=81?= <1412414664@qq.com> Date: Thu, 25 Jun 2026 01:10:18 +0800 Subject: [PATCH] fix: start extract resource timing before first item --- dlt/extract/extract.py | 10 ++++++++++ dlt/extract/pipe_iterator.py | 31 +++++++++++++++++++++++++++++-- tests/extract/test_extract.py | 20 ++++++++++++++++++++ 3 files changed, 59 insertions(+), 2 deletions(-) diff --git a/dlt/extract/extract.py b/dlt/extract/extract.py index 87e78ae9fb..bbd3c38023 100644 --- a/dlt/extract/extract.py +++ b/dlt/extract/extract.py @@ -57,6 +57,7 @@ from dlt.extract.incremental import IncrementalResourceWrapper from dlt.extract.items import TableNameMeta from dlt.extract.items_transform import ItemTransform +from dlt.extract.pipe import Pipe from dlt.extract.pipe_iterator import PipeIterator from dlt.extract.source import DltSource from dlt.extract.reference import SourceReference, SourceFactory @@ -503,6 +504,14 @@ def _extract_single_source( load_id, self.extract_storage.item_storages["model"], schema, collector=collector ), } + object_extractor = extractors["object"] + + def start_resource_counter(pipe: Pipe) -> None: + resource = source.resources.with_pipe(pipe) + if resource.has_dynamic_table_name: + return + collector.update(object_extractor._get_static_table_name(resource, None), inc=0) + # make sure we close storage on exception with collector(f"Extract {source.name}"): with self.manage_writers(load_id, source): @@ -511,6 +520,7 @@ def _extract_single_source( source.resources.selected_pipes, max_parallel_items=max_parallel_items, workers=workers, + on_source_started=start_resource_counter, ) as pipes: left_gens = total_gens = len(pipes._sources) collector.update("Resources", 0, total_gens) diff --git a/dlt/extract/pipe_iterator.py b/dlt/extract/pipe_iterator.py index af7b0092cf..29f2f7dffb 100644 --- a/dlt/extract/pipe_iterator.py +++ b/dlt/extract/pipe_iterator.py @@ -2,8 +2,10 @@ import types from typing import ( AsyncIterator, + Callable, ClassVar, Dict, + Optional, Sequence, Union, Iterator, @@ -61,11 +63,14 @@ def __init__( futures_poll_interval: float, sources: List[SourcePipeItem], next_item_mode: TPipeNextItemMode, + on_source_started: Optional[Callable[[Pipe], None]] = None, ) -> None: self._sources = sources self._next_item_mode: TPipeNextItemMode = next_item_mode self._initial_sources_count = len(sources) self._current_source_index: int = 0 + self._started_pipes: set[str] = set() + self._on_source_started = on_source_started self._futures_pool = FuturesPool( workers=workers, poll_interval=futures_poll_interval, @@ -82,6 +87,7 @@ def from_pipe( workers: int = 5, futures_poll_interval: float = 0.01, next_item_mode: TPipeNextItemMode = "round_robin", + on_source_started: Optional[Callable[[Pipe], None]] = None, ) -> "PipeIterator": # join all dependent pipes if pipe.parent: @@ -95,7 +101,14 @@ def from_pipe( # create extractor sources = [SourcePipeItem(pipe.gen, 0, pipe, None)] - return cls(max_parallel_items, workers, futures_poll_interval, sources, next_item_mode) + return cls( + max_parallel_items, + workers, + futures_poll_interval, + sources, + next_item_mode, + on_source_started, + ) @classmethod @with_config(spec=PipeIteratorConfiguration) @@ -109,6 +122,7 @@ def from_pipes( futures_poll_interval: float = 0.01, copy_on_fork: bool = False, next_item_mode: TPipeNextItemMode = "round_robin", + on_source_started: Optional[Callable[[Pipe], None]] = None, ) -> "PipeIterator": # print(f"max_parallel_items: {max_parallel_items} workers: {workers}") sources: List[SourcePipeItem] = [] @@ -141,7 +155,14 @@ def _fork_pipeline(pipe: Pipe) -> None: _fork_pipeline(pipe) # create extractor - return cls(max_parallel_items, workers, futures_poll_interval, sources, next_item_mode) + return cls( + max_parallel_items, + workers, + futures_poll_interval, + sources, + next_item_mode, + on_source_started, + ) def __next__(self) -> PipeItem: pipe_item: Union[ResolvablePipeItem, SourcePipeItem] = None @@ -270,6 +291,12 @@ def _get_source_item(self) -> ResolvablePipeItem: return None # get next item from the current source gen, step, pipe, meta = self._sources[self._current_source_index] + if ( + self._on_source_started is not None + and pipe.instance_id not in self._started_pipes + ): + self._started_pipes.add(pipe.instance_id) + self._on_source_started(pipe) with pipe_context(pipe): pipe_item = next(gen) if pipe_item is not None: diff --git a/tests/extract/test_extract.py b/tests/extract/test_extract.py index 1d69c94072..50bf69f266 100644 --- a/tests/extract/test_extract.py +++ b/tests/extract/test_extract.py @@ -1,4 +1,5 @@ from typing import Dict, List, NamedTuple, Optional, Any +import io import pytest import os @@ -11,6 +12,7 @@ SchemaStorageConfiguration, NormalizeStorageConfiguration, ) +from dlt.common.runtime.collector import LogCollector from dlt.common.storages.schema_storage import SchemaStorage from dlt.common.schema.typing import TColumnSchema, TWriteDisposition from dlt.common.typing import TTableNames, TDataItems @@ -91,6 +93,24 @@ def test_storage_reuse_package() -> None: } +def test_resource_log_time_includes_wait_before_first_item(extract_step: Extract) -> None: + clock = [0.0] + buffer = io.StringIO() + collector = LogCollector(log_period=1000.0, logger=buffer, dump_system_stats=False) + collector._clock = lambda: clock[0] # type: ignore[assignment] + extract_step.collector = collector + + @dlt.resource + def slow_resource(): + clock[0] = 5.0 + yield [{"id": 1}] + + source = DltSource(dlt.Schema("timed"), "module", [slow_resource()]) + extract_step.extract(source, 20, 1) + + assert "slow_resource: 1 | Time: 5.00s" in buffer.getvalue() + + def test_extract_select_tables_mark(extract_step: Extract) -> None: n_f = lambda i: ("odd" if i % 2 == 1 else "even") + "_table"