include new keys in typing

dlt-hub · Sep 17, 2024 · 70d036d · 70d036d
1 parent bcbd0c7
commit 70d036d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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",