Merge pull request #49 from hotdata-dev/feat/per-load-key

feat: add per-load key to load_managed_table

Author: shefeek-jinnah

CHANGELOG.md

@@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- `load_managed_table` accepts a `key` argument — the merge key columns for
+  `delete`/`update`/`upsert` loads, matched per-load instead of requiring a key
+  declared at table creation. Omit it to use the table's declared key; ignored
+  for `replace`/`append`. Requires `hotdata>=0.8.0`.
+
 
 ## [0.7.3] - 2026-07-16
 

hotdata_framework/client.py

@@ -340,6 +340,7 @@ def load_managed_table(
         upload_id: str | None = None,
         file: str | None = None,
         mode: ManagedLoadMode = "replace",
+        key: list[str] | None = None,
     ) -> LoadManagedTableResult:
         if (upload_id is None) == (file is None):
             raise ValueError("Exactly one of upload_id or file is required")
@@ -352,6 +353,7 @@ def load_managed_table(
         request = LoadManagedTableRequest(
             mode=mode,
             upload_id=resolved_upload_id,
+            key=key,
         )
         try:
             loaded = self.connections().load_managed_table(

hotdata_framework/managed_client.py

@@ -205,17 +205,23 @@ def load_managed_table(
         schema: str,
         upload_id: str,
         mode: ManagedLoadMode = "replace",
+        key: list[str] | None = None,
     ) -> LoadManagedTableResult:
         # append is the only non-idempotent mode: if the server commits the load
         # but the response is lost, a retry re-appends the same rows. Run it
         # at-most-once; every other mode is safe to retry.
+        #
+        # `key` is the merge key for delete/update/upsert loads: when set it is
+        # matched per-load instead of a key declared at table creation. Omit it
+        # to use the table's declared key. Ignored for replace/append.
         return self._request_with_retry(
             lambda: self._runtime.load_managed_table(
                 database,
                 table,
                 schema=schema,
                 upload_id=upload_id,
                 mode=mode,
+                key=key,
             ),
             retryable=(mode != "append"),
         )

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 ]
 dependencies = [
   # 0.7.0 adds `key` to table decls (create_managed_database(keys=) / add_managed_table(key=))
-  "hotdata>=0.7.0",
+  "hotdata>=0.8.0",
   "pandas>=2.0",
   "pyarrow>=14.0",
 ]

tests/test_databases.py

@@ -308,6 +308,16 @@ def test_load_managed_table_defaults_to_replace():
     assert _load_and_capture_request(_client()).mode == "replace"
 
 
+def test_load_managed_table_passes_key():
+    request = _load_and_capture_request(_client(), mode="delete", key=["id"])
+    assert request.mode == "delete"
+    assert request.key == ["id"]
+
+
+def test_load_managed_table_key_defaults_to_none():
+    assert _load_and_capture_request(_client()).key is None
+
+
 @pytest.mark.parametrize("mode", ["append", "delete", "update", "upsert"])
 def test_load_managed_table_forwards_mode(mode: str):
     assert _load_and_capture_request(_client(), mode=mode).mode == mode

tests/test_managed_client.py

@@ -163,7 +163,13 @@ def _load_recording_runtime(calls: list[str]) -> SimpleNamespace:
     with a transient error, so retry behaviour is observable via ``calls``."""
 
     def load_managed_table(
-        database: str, table: str, *, schema: str, upload_id: str, mode: str
+        database: str,
+        table: str,
+        *,
+        schema: str,
+        upload_id: str,
+        mode: str,
+        key: list[str] | None = None,
     ) -> SimpleNamespace:
         calls.append(mode)
         raise TimeoutError("commit succeeded but response was lost")
@@ -209,3 +215,34 @@ def test_idempotent_load_retries_on_transient(monkeypatch: pytest.MonkeyPatch) -
         client.load_managed_table("db", "orders", schema="public", upload_id="u1", mode="replace")
 
     assert calls == ["replace", "replace", "replace"]  # retried up to max_retries
+
+
+def test_load_managed_table_forwards_key(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A per-load ``key`` is passed straight through to the runtime client."""
+    monkeypatch.setattr(mc.time, "sleep", lambda _seconds: None)
+    captured: dict[str, Any] = {}
+
+    def load_managed_table(
+        database: str,
+        table: str,
+        *,
+        schema: str,
+        upload_id: str,
+        mode: str,
+        key: list[str] | None = None,
+    ) -> SimpleNamespace:
+        captured["mode"] = mode
+        captured["key"] = key
+        return SimpleNamespace(
+            connection_id="c", schema_name=schema, table_name=table, row_count=0
+        )
+
+    client = _managed_client(max_retries=1)
+    runtime = _fake_runtime()
+    runtime.load_managed_table = load_managed_table
+    client._runtime = runtime
+
+    client.load_managed_table(
+        "db", "orders", schema="public", upload_id="u1", mode="delete", key=["id"]
+    )
+    assert captured == {"mode": "delete", "key": ["id"]}