diff --git a/dlt/common/schema/typing.py b/dlt/common/schema/typing.py index 8707202f59..855f26706a 100644 --- a/dlt/common/schema/typing.py +++ b/dlt/common/schema/typing.py @@ -237,6 +237,8 @@ class TMergeDispositionDict(TWriteDispositionDict, total=False): active_record_timestamp: Optional[TAnyDateTime] boundary_timestamp: Optional[TAnyDateTime] row_version_column_name: Optional[str] + retire_if_absent: Optional[bool] + natural_key: Optional[str] TWriteDispositionConfig = Union[TWriteDisposition, TWriteDispositionDict, TMergeDispositionDict] diff --git a/dlt/extract/hints.py b/dlt/extract/hints.py index 7a30812e74..22a7574796 100644 --- a/dlt/extract/hints.py +++ b/dlt/extract/hints.py @@ -457,9 +457,9 @@ def _merge_merge_disposition_dict(dict_: Dict[str, Any]) -> None: if "boundary_timestamp" in md_dict: dict_["x-boundary-timestamp"] = md_dict["boundary_timestamp"] if "retire_if_absent" in md_dict: - dict_["x-retire-if-absent"] = md_dict["retire_if_absent"] # type: ignore[typeddict-item] + dict_["x-retire-if-absent"] = md_dict["retire_if_absent"] if "natural_key" in md_dict: - nk = md_dict["natural_key"] # type: ignore[typeddict-item] + nk = md_dict["natural_key"] if nk in dict_["columns"]: dict_["columns"][nk]["x-natural-key"] = True else: @@ -546,5 +546,9 @@ def validate_write_disposition_hint(wd: TTableHintTemplate[TWriteDispositionConf f'could not parse `{ts}` value "{wd[ts]}"' # type: ignore[literal-required] ) - if "retire_if_absent" in wd and not wd["retire_if_absent"] and "natural_key" not in wd: # type: ignore[typeddict-item] + if ( + "retire_if_absent" in wd + and not wd["retire_if_absent"] + and "natural_key" not in wd + ): raise ValueError("`natural_key` is required when `retire_if_absent=False`") diff --git a/tests/load/pipeline/test_scd2.py b/tests/load/pipeline/test_scd2.py index 3ac30e3664..8bbf30a47f 100644 --- a/tests/load/pipeline/test_scd2.py +++ b/tests/load/pipeline/test_scd2.py @@ -725,7 +725,7 @@ def test_retire_if_absent( ) -> None: p = destination_config.setup_pipeline("abstract", dev_mode=True) - @dlt.resource( # type: ignore[call-overload] + @dlt.resource( table_name="dim_test", write_disposition={ "disposition": "merge", @@ -792,11 +792,11 @@ def r(data): "retire_if_absent": True, } ) - assert r.compute_table_schema()["x-retire-if-absent"] + assert r.compute_table_schema()["x-retire-if-absent"] # type: ignore[typeddict-item] # user-provided hints for `natural_key` column should be respected r.apply_hints( - columns={"nk": {"x-foo": "foo"}}, + columns={"nk": {"x-foo": "foo"}}, # type: ignore[typeddict-unknown-key] write_disposition={ "disposition": "merge", "strategy": "scd2",