Merge pull request #43 from hotdata-dev/feat/mode-and-key-on-managed-load

feat(managed): mode + key on managed-table loads

Author: shefeek-jinnah

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `load_managed_table(..., mode=...)` selects the load mode (`replace` (default), `append`, `delete`, `update`, `upsert`) instead of always replacing the table. `replace`/`append` apply the upload directly; `delete`/`update`/`upsert` match rows by the table's declared key. Backward compatible — omitting `mode` still replaces.
+- `create_managed_database(..., keys={table: [cols]})` and `add_managed_table(..., key=[cols])` declare a table's row-identity key, enabling the key-based load modes on it. Requires a `hotdata` client whose managed-table decl models carry `key` (see the dependency floor bump); tables declared without a key stay `replace`/`append`-only.
 
 ## [0.6.3] - 2026-07-08
 

hotdata_framework/client.py

@@ -241,14 +241,21 @@ def create_managed_database(
         *,
         schema: str = DEFAULT_SCHEMA,
         tables: list[str] | None = None,
+        keys: dict[str, list[str]] | None = None,
         expires_at: str | None = None,
     ) -> ManagedDatabase:
+        """Create a managed database. ``keys`` maps a table to its key columns
+        (enabling delete/update/upsert on it); omitted tables are keyless."""
+        keys = keys or {}
         schemas = None
         if tables:
             schemas = [
                 DatabaseDefaultSchemaDecl(
                     name=schema,
-                    tables=[DatabaseDefaultTableDecl(name=t) for t in tables],
+                    tables=[
+                        DatabaseDefaultTableDecl(name=t, key=list(keys.get(t, [])))
+                        for t in tables
+                    ],
                 )
             ]
         request = CreateDatabaseRequest(
@@ -314,6 +321,7 @@ def load_managed_table(
         schema: str = DEFAULT_SCHEMA,
         upload_id: str | None = None,
         file: str | None = None,
+        mode: str = "replace",
     ) -> LoadManagedTableResult:
         if (upload_id is None) == (file is None):
             raise ValueError("Exactly one of upload_id or file is required")
@@ -324,7 +332,7 @@ def load_managed_table(
             assert file is not None
             resolved_upload_id = self.upload_parquet(file)
         request = LoadManagedTableRequest(
-            mode="replace",
+            mode=mode,
             upload_id=resolved_upload_id,
         )
         try:
@@ -350,15 +358,17 @@ def add_managed_table(
         table: str,
         *,
         schema: str = DEFAULT_SCHEMA,
+        key: list[str] | None = None,
     ) -> ManagedTable:
         """Declare a new table on an existing managed database.
 
         The table is added empty (declared-but-unloaded); populate it with
         :meth:`load_managed_table`. Use this to evolve a managed database's
-        schema after creation without recreating it.
+        schema after creation without recreating it. ``key`` sets the
+        row-identity columns for delete/update/upsert; omit for keyless.
         """
         db = self.resolve_managed_database(database)
-        request = AddManagedTableRequest(name=table)
+        request = AddManagedTableRequest(name=table, key=list(key or []))
         try:
             self._databases_api().add_database_table(db.id, schema, request)
         except ApiException as e:

hotdata_framework/managed_client.py

@@ -203,13 +203,15 @@ def load_managed_table(
         *,
         schema: str,
         upload_id: str,
+        mode: str = "replace",
     ) -> LoadManagedTableResult:
         return self._request_with_retry(
             lambda: self._runtime.load_managed_table(
                 database,
                 table,
                 schema=schema,
                 upload_id=upload_id,
+                mode=mode,
             )
         )
 

pyproject.toml

@@ -26,7 +26,8 @@ classifiers = [
   "Typing :: Typed",
 ]
 dependencies = [
-  "hotdata>=0.6.0",
+  # 0.7.0 adds `key` to table decls (create_managed_database(keys=) / add_managed_table(key=))
+  "hotdata>=0.7.0",
   "pandas>=2.0",
   "pyarrow>=14.0",
 ]

tests/test_databases.py

@@ -5,6 +5,7 @@
 
 import pytest
 from hotdata.exceptions import ApiException
+from hotdata.models.database_default_table_decl import DatabaseDefaultTableDecl
 
 from hotdata_framework.client import HotdataClient
 from hotdata_framework.databases import (
@@ -13,6 +14,20 @@
 )
 
 
+def _decl_key_supported() -> bool:
+    # `key` ships with the regenerated client; the key tests activate once it does.
+    try:
+        return DatabaseDefaultTableDecl(name="t", key=["k"]).key == ["k"]
+    except Exception:
+        return False
+
+
+requires_key_field = pytest.mark.skipif(
+    not _decl_key_supported(),
+    reason="hotdata client without `key` on managed-table decls",
+)
+
+
 def _client() -> HotdataClient:
     return HotdataClient("k", "ws", host="https://api.hotdata.dev")
 
@@ -230,6 +245,69 @@ def test_load_managed_table_requires_exactly_one_source():
         )
 
 
+def _load_and_capture_request(client, **kwargs):
+    db = managed_database_from_detail(_detail())
+    loaded = SimpleNamespace(
+        connection_id="conn_1", schema_name="public", table_name="orders", row_count=1
+    )
+    with (
+        patch.object(client, "resolve_managed_database", return_value=db),
+        patch.object(client, "connections") as connections,
+    ):
+        connections.return_value.load_managed_table.return_value = loaded
+        client.load_managed_table("db_1", "orders", upload_id="upl_1", **kwargs)
+    return connections.return_value.load_managed_table.call_args.args[3]
+
+
+def test_load_managed_table_defaults_to_replace():
+    assert _load_and_capture_request(_client()).mode == "replace"
+
+
+@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
+
+
+@requires_key_field
+def test_create_managed_database_declares_keys():
+    client = _client()
+    with patch.object(client, "_databases_api") as dbs:
+        dbs.return_value.create_database.return_value = _detail(id="db_new")
+        client.create_managed_database(
+            "mydb", tables=["orders", "events"], keys={"orders": ["id"]}
+        )
+    req = dbs.return_value.create_database.call_args.args[0]
+    declared = {t.name: list(t.key) for t in req.schemas[0].tables}
+    assert declared == {"orders": ["id"], "events": []}
+
+
+@requires_key_field
+def test_add_managed_table_declares_key():
+    client = _client()
+    db = managed_database_from_detail(_detail())
+    with (
+        patch.object(client, "resolve_managed_database", return_value=db),
+        patch.object(client, "_databases_api") as dbs,
+    ):
+        client.add_managed_table("db_1", "line_items", key=["order_id", "sku"])
+    req = dbs.return_value.add_database_table.call_args.args[2]
+    assert req.name == "line_items"
+    assert list(req.key) == ["order_id", "sku"]
+
+
+@requires_key_field
+def test_add_managed_table_keyless_by_default():
+    client = _client()
+    db = managed_database_from_detail(_detail())
+    with (
+        patch.object(client, "resolve_managed_database", return_value=db),
+        patch.object(client, "_databases_api") as dbs,
+    ):
+        client.add_managed_table("db_1", "orders")
+    req = dbs.return_value.add_database_table.call_args.args[2]
+    assert list(req.key) == []
+
+
 def test_delete_managed_table_uses_default_connection_id():
     client = _client()
     db = managed_database_from_detail(_detail())