From 5e2176755f58dec7bdeccb290664062d9dfc28e7 Mon Sep 17 00:00:00 2001 From: Shefeek Jinnah Date: Mon, 20 Jul 2026 11:45:32 +0530 Subject: [PATCH 1/2] feat: add per-load key to load_managed_table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit load_managed_table gains an optional `key` — 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, which adds the `key` field to the load-request model. --- CHANGELOG.md | 7 ++++++ hotdata_framework/client.py | 2 ++ hotdata_framework/managed_client.py | 6 +++++ pyproject.toml | 2 +- tests/test_databases.py | 10 ++++++++ tests/test_managed_client.py | 39 ++++++++++++++++++++++++++++- 6 files changed, 64 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 422d02d..b3dbb3b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/hotdata_framework/client.py b/hotdata_framework/client.py index 62cad87..0a4dd6f 100644 --- a/hotdata_framework/client.py +++ b/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( diff --git a/hotdata_framework/managed_client.py b/hotdata_framework/managed_client.py index bac18cb..bf2ff69 100644 --- a/hotdata_framework/managed_client.py +++ b/hotdata_framework/managed_client.py @@ -205,10 +205,15 @@ 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, @@ -216,6 +221,7 @@ def load_managed_table( schema=schema, upload_id=upload_id, mode=mode, + key=key, ), retryable=(mode != "append"), ) diff --git a/pyproject.toml b/pyproject.toml index d5ecc92..8fb7fcb 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/tests/test_databases.py b/tests/test_databases.py index ab77958..c85a62c 100644 --- a/tests/test_databases.py +++ b/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 diff --git a/tests/test_managed_client.py b/tests/test_managed_client.py index 3966c0c..e831734 100644 --- a/tests/test_managed_client.py +++ b/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"]} From c266493706c1d61974cc220a37fab4d9e1533201 Mon Sep 17 00:00:00 2001 From: Shefeek Jinnah Date: Mon, 20 Jul 2026 12:02:02 +0530 Subject: [PATCH 2/2] chore: lock hotdata 0.8.0 --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index 45187be..f47e3a2 100644 --- a/uv.lock +++ b/uv.lock @@ -86,7 +86,7 @@ wheels = [ [[package]] name = "hotdata" -version = "0.7.0" +version = "0.8.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pydantic" }, @@ -94,9 +94,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/61/30/08681132e019f6c9ddb566dc38be6a302763c4b4848b2d97e3d7ace8d257/hotdata-0.7.0.tar.gz", hash = "sha256:ef4484a74c0f33ee543b0a0dbfb5a6c04e7b5812b277f42a6448000d41c2405d", size = 215130, upload-time = "2026-07-14T12:21:33.936Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8a/38/30ed3d1d99413e7684672fa424baef85a347db465b90c774443320cf1cea/hotdata-0.8.0.tar.gz", hash = "sha256:cdac515ffa193ed028491e4b7abcd1b6404d4e3f96986c9c0a2b2f407db42eb6", size = 216848, upload-time = "2026-07-20T06:30:37.659Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/ef/6c4236640629688074892759b1a7593c08f30e9d916b6951c2bbfec19d7f/hotdata-0.7.0-py3-none-any.whl", hash = "sha256:d9e3008e3084d22bcc27b7bf1d08d8c2c255fe3726c252c777bb614a3cae0576", size = 314103, upload-time = "2026-07-14T12:21:32.154Z" }, + { url = "https://files.pythonhosted.org/packages/bc/0f/f2ccaeb0f910c3f8cd84a911d934704277f0d6e5a20f8b49cf0eedb2a8aa/hotdata-0.8.0-py3-none-any.whl", hash = "sha256:5d64bfc185a0e7e2bf34ebdcea50787d3418fb250dacc53e00dc74d64e99f753", size = 316923, upload-time = "2026-07-20T06:30:35.778Z" }, ] [[package]] @@ -120,7 +120,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "hotdata", specifier = ">=0.7.0" }, + { name = "hotdata", specifier = ">=0.8.0" }, { name = "pandas", specifier = ">=2.0" }, { name = "pyarrow", specifier = ">=14.0" }, ]