Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 30 additions & 13 deletions dlt/extract/resource.py
Original file line number Diff line number Diff line change
Expand Up @@ -470,11 +470,13 @@ def add_step(
self._pipe.insert_step(item_transform, insert_at)
return self

def _remove_incremental_step(self) -> None:
self._pipe.remove_by_type(Incremental, IncrementalResourceWrapper)
def _remove_incremental_step(self) -> int:
return self._pipe.remove_by_type(Incremental, IncrementalResourceWrapper)

def set_incremental(
self, new_incremental: Union[Incremental[Any], IncrementalResourceWrapper]
self,
new_incremental: Union[Incremental[Any], IncrementalResourceWrapper],
insert_at: Optional[int] = None,
) -> Optional[Union[Incremental[Any], IncrementalResourceWrapper]]:
"""Set/replace the incremental transform for the resource.

Expand Down Expand Up @@ -508,7 +510,7 @@ def set_incremental(
if new_incremental:
if not isinstance(new_incremental, IncrementalResourceWrapper):
new_incremental = Incremental.ensure_instance(new_incremental)
self.add_step(new_incremental)
self.add_step(new_incremental, insert_at=insert_at)
return new_incremental

def _set_hints(
Expand All @@ -519,7 +521,11 @@ def _set_hints(
if not create_table_variant:
# try to late assign incremental
if table_schema_template.get("incremental") is not None:
self.set_incremental(table_schema_template["incremental"])
incremental_step_index = self._pipe.find(Incremental, IncrementalResourceWrapper)
self.set_incremental(
table_schema_template["incremental"],
insert_at=incremental_step_index if incremental_step_index >= 0 else None,
)
elif incremental := self.incremental:
# set primary key as dedup key on new incremental
resource_primary_key = self._hints.get("primary_key")
Expand Down Expand Up @@ -685,7 +691,7 @@ def _set_explicit_args(
except Exception:
pass

def _eject_config(self) -> bool:
def _eject_config(self) -> tuple[bool, int]:
"""Unwraps the pipe generator step from config injection and incremental wrappers by restoring the original step.

Removes the step with incremental wrapper. Should be used before a subsequent _inject_config is called on the
Expand All @@ -696,12 +702,16 @@ def _eject_config(self) -> bool:
if not self._pipe.is_empty and not self._args_bound:
orig_gen = getattr(self._pipe.gen, "__GEN__", None)
if orig_gen:
self._remove_incremental_step()
incremental_step_index = self._remove_incremental_step()
self._pipe.replace_gen(orig_gen)
return True
return False
return True, incremental_step_index
return False, -1

def _inject_config(self, incremental_from_hints_override: Optional[bool] = None) -> Self:
def _inject_config(
self,
incremental_from_hints_override: Optional[bool] = None,
incremental_step_index: int = -1,
) -> Self:
"""Wraps the pipe generation step in incremental and config injection wrappers and adds pipe step with
Incremental transform.
"""
Expand All @@ -723,7 +733,10 @@ def _inject_config(self, incremental_from_hints_override: Optional[bool] = None)
),
)
incr_f = incremental.wrap(sig, gen)
self.set_incremental(incremental)
self.set_incremental(
incremental,
insert_at=incremental_step_index if incremental_step_index >= 0 else None,
)
else:
incr_f = gen
resource_sections = (known_sections.SOURCES, self.section, self.name)
Expand Down Expand Up @@ -792,8 +805,12 @@ def _clone(
# try to eject and then inject configuration and incremental wrapper when resource is cloned
# this makes sure that a take config values from a right section and wrapper has a separated
# instance in the pipeline
if cloned_r._eject_config():
cloned_r._inject_config(incremental_from_hints_override=incremental_from_hints)
ejected, incremental_step_index = cloned_r._eject_config()
if ejected:
cloned_r._inject_config(
incremental_from_hints_override=incremental_from_hints,
incremental_step_index=incremental_step_index,
)
return cloned_r

def _update_wrapper(self) -> None:
Expand Down
30 changes: 29 additions & 1 deletion tests/extract/test_incremental.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@
IncrementalPrimaryKeyMissing,
)
from dlt.extract.incremental.lag import apply_lag
from dlt.extract.items_transform import ValidateItem
from dlt.extract.items_transform import MetricsItem, ValidateItem
from dlt.extract.resource import DltResource
from dlt.pipeline.exceptions import PipelineStepFailed
from dlt.sources.helpers.transform import take_first
Expand Down Expand Up @@ -2969,6 +2969,34 @@ def updated_is_int(updated_at=dlt.sources.incremental("updated_at", initial_valu
assert pip_ex.value.__cause__.cursor_path == "updated_at"


def test_clone_preserves_inserted_step_after_hints_incremental() -> None:
@dlt.source
def my_source():
@dlt.resource
def events():
yield [
{"id": 1, "updated_at": 1},
{"id": 2, "updated_at": 2},
{"id": 3, "updated_at": 3},
]

return events

def count_rows(items: TDataItems, meta: Any, metrics: Dict[str, Any]) -> None:
if isinstance(items, list):
metrics["row_count"] = metrics.get("row_count", 0) + len(items)

resource = my_source().events
resource.apply_hints(incremental=dlt.sources.incremental("updated_at", initial_value=2))
resource.add_metrics(count_rows, insert_at=len(resource._pipe))

assert resource._pipe.find(Incremental) < resource._pipe.find(MetricsItem)

cloned_resource = resource._clone()

assert cloned_resource._pipe.find(Incremental) < cloned_resource._pipe.find(MetricsItem)


def test_incremental_merge_native_representation():
incremental = Incremental(cursor_path="some_path", lag=10) # type: ignore

Expand Down