Align BatchWriteItem affinity query plan

alternator/async_client.py

@@ -25,6 +25,7 @@
 from alternator.core.auth import apply_auth
 from alternator.core.handlers import _register_alternator_handlers
 from alternator.core.key_affinity import (
+    AffinityTarget,
     select_affinity_node,
 )
 from alternator.core.live_nodes import AsyncLiveNodesManager, NodeList
@@ -203,7 +204,7 @@ def preload(self, table_pk_map: dict[str, str]) -> None:
 def _create_async_affinity_node_computer(
     config: Config,
     pk_cache: AsyncPartitionKeyCache | None,
-) -> Callable[[str, dict[str, Any], NodeList], str | None] | None:
+) -> Callable[[str, dict[str, Any], NodeList], AffinityTarget | None] | None:
     """
     Create a function that selects the preferred key-affinity node.
 
@@ -240,7 +241,7 @@ def compute_affinity_node(
         operation_name: str,
         params: dict[str, Any],
         nodes: NodeList,
-    ) -> str | None:
+    ) -> AffinityTarget | None:
         """Select the preferred key-affinity node for this request."""
         return select_affinity_node(
             mode=affinity_mode.name,

alternator/client.py

@@ -20,6 +20,7 @@
 from alternator.core.auth import apply_auth
 from alternator.core.handlers import _register_alternator_handlers
 from alternator.core.key_affinity import (
+    AffinityTarget,
     PartitionKeyCache,
     select_affinity_node,
 )
@@ -167,7 +168,7 @@ def _create_boto_config(config: Config, *, auth_enabled: bool) -> BotoConfig:
 def _create_affinity_node_computer(
     config: Config,
     client: DynamoDBClient,
-) -> Callable[[str, dict[str, Any], NodeList], str | None] | None:
+) -> Callable[[str, dict[str, Any], NodeList], AffinityTarget | None] | None:
     """
     Create a function that selects the preferred key-affinity node.
 
@@ -198,7 +199,7 @@ def compute_affinity_node(
         operation_name: str,
         params: dict[str, Any],
         nodes: NodeList,
-    ) -> str | None:
+    ) -> AffinityTarget | None:
         """Select the preferred key-affinity node for this request."""
         return select_affinity_node(
             mode=affinity_mode.name,

alternator/core/handlers.py

@@ -22,6 +22,7 @@
     create_header_filter_handler,
     create_user_agent_header_handler,
 )
+from alternator.core.key_affinity import AffinityTarget
 from alternator.core.query_plan import LazyQueryPlan
 from alternator.core.request import extract_operation_name, extract_request_params
 from alternator.exceptions import NoNodesAvailableError
@@ -49,7 +50,9 @@ def _register_alternator_handlers(
     events: BaseEventHooks,
     manager: _HasNodes,
     config: Config,
-    compute_affinity_node: Callable[[str, DynamoDBParams, NodeList], str | None]
+    compute_affinity_node: Callable[
+        [str, DynamoDBParams, NodeList], AffinityTarget | None
+    ]
     | None = None,
     *,
     auth_enabled: bool = False,
@@ -78,10 +81,21 @@ def _register_alternator_handlers(
 
     def create_query_plan(
         nodes: NodeList,
-        preferred_node: str | None,
+        preferred_node: AffinityTarget | None,
     ) -> Iterator[str]:
         """Create a URI iterator for a single request."""
         node_addresses = nodes.nodes
+        if isinstance(preferred_node, tuple):
+            emitted: set[str] = set()
+            for node in preferred_node:
+                if node in node_addresses and node not in emitted:
+                    emitted.add(node)
+                    yield f"{scheme}://{node}:{port}"
+            for node in node_addresses:
+                if node not in emitted:
+                    yield f"{scheme}://{node}:{port}"
+            return
+
         if preferred_node is not None and preferred_node in node_addresses:
             yield f"{scheme}://{preferred_node}:{port}"
             remaining_nodes = tuple(
@@ -146,7 +160,7 @@ def update_endpoint(
 
     def _create_request_query_plan(
         request: AWSRequest | AWSPreparedRequest,
-        preferred_node: str | None | object = _PREFERRED_NODE_UNSET,
+        preferred_node: AffinityTarget | None | object = _PREFERRED_NODE_UNSET,
     ) -> Iterator[str]:
         nodes = manager.nodes
         if not nodes:
@@ -155,7 +169,7 @@ def _create_request_query_plan(
                 scope_name=scope_name,
             )
 
-        selected_preferred_node: str | None
+        selected_preferred_node: AffinityTarget | None
         if preferred_node is _PREFERRED_NODE_UNSET:
             selected_preferred_node = None
             if compute_affinity_node is not None:
@@ -170,7 +184,7 @@ def _create_request_query_plan(
                         nodes,
                     )
         else:
-            selected_preferred_node = cast("str | None", preferred_node)
+            selected_preferred_node = cast("AffinityTarget | None", preferred_node)
         return create_query_plan(nodes, selected_preferred_node)
 
     def _store_query_plan(

alternator/core/key_affinity.py

@@ -20,6 +20,7 @@
     from alternator.core.live_nodes import NodeList
 
 logger = logging.getLogger("alternator")
+AffinityTarget = str | tuple[str, ...]
 
 
 class _BatchWriteRoutingTarget(NamedTuple):
@@ -31,6 +32,7 @@ class _BatchWriteRoutingTarget(NamedTuple):
 class _BatchWriteCandidate(NamedTuple):
     table_name: str
     attributes: dict[str, Any]
+    operation: str
 
 
 class AffinitySelector:
@@ -122,7 +124,7 @@ def select_affinity_node(
     params: dict[str, Any],
     nodes: NodeList,
     get_pk_name: Callable[[str], str | None],
-) -> str | None:
+) -> AffinityTarget | None:
     """Select the preferred affinity node for a request, or None for fallback."""
     if not should_use_affinity(mode, operation_name, params):
         return None
@@ -203,7 +205,7 @@ def _select_batch_write_affinity_node(
     params: dict[str, Any],
     nodes: NodeList,
     get_pk_name: Callable[[str], str | None],
-) -> str | None:
+) -> tuple[str, ...] | None:
     votes: Counter[str] = Counter()
     selector = AffinitySelector()
 
@@ -233,11 +235,12 @@ def _select_batch_write_affinity_node(
     if not votes:
         return None
 
-    top_count = max(votes.values())
-    winners = [node for node, count in votes.items() if count == top_count]
-    if len(winners) != 1:
-        return None
-    return winners[0]
+    return tuple(
+        sorted(
+            votes,
+            key=lambda node: (-votes[node], node),
+        )
+    )
 
 
 def _iter_batch_write_candidates(
@@ -254,22 +257,37 @@ def _iter_batch_write_candidates(
         if not isinstance(table_name, str) or not isinstance(writes, list):
             continue
         for write in writes:
-            if not isinstance(write, dict):
-                continue
+            candidate = _batch_write_candidate(table_name, write)
+            if candidate is not None:
+                candidates.append(candidate)
 
-            put_request = write.get("PutRequest")
-            if isinstance(put_request, dict):
-                item = put_request.get("Item")
-                if isinstance(item, dict):
-                    candidates.append(_BatchWriteCandidate(table_name, item))
+    return tuple(candidates)
 
-            delete_request = write.get("DeleteRequest")
-            if isinstance(delete_request, dict):
-                key = delete_request.get("Key")
-                if isinstance(key, dict):
-                    candidates.append(_BatchWriteCandidate(table_name, key))
 
-    return tuple(candidates)
+def _batch_write_candidate(
+    table_name: str,
+    write: object,
+) -> _BatchWriteCandidate | None:
+    if not isinstance(write, dict):
+        return None
+
+    operations = [
+        operation for operation in ("PutRequest", "DeleteRequest") if operation in write
+    ]
+    if len(operations) != 1:
+        return None
+
+    operation = operations[0]
+    request = write[operation]
+    if not isinstance(request, dict):
+        return None
+
+    attribute_field = "Item" if operation == "PutRequest" else "Key"
+    attributes = request.get(attribute_field)
+    if not isinstance(attributes, dict):
+        return None
+
+    return _BatchWriteCandidate(table_name, attributes, operation)
 
 
 def _non_empty_string(value: object) -> bool:
@@ -324,34 +342,22 @@ def _find_batch_write_routing_target(
         if not isinstance(table_name, str) or not isinstance(writes, list):
             continue
         for write in writes:
-            if not isinstance(write, dict):
+            candidate = _batch_write_candidate(table_name, write)
+            if candidate is None:
                 continue
 
-            put_request = write.get("PutRequest")
-            if isinstance(put_request, dict):
-                item = put_request.get("Item")
-                if isinstance(item, dict):
-                    target = _min_batch_write_target(
-                        target,
-                        _BatchWriteRoutingTarget(
-                            table_name,
-                            item,
-                            _batch_write_sort_key(table_name, "PutRequest", item),
-                        ),
-                    )
-
-            delete_request = write.get("DeleteRequest")
-            if isinstance(delete_request, dict):
-                key = delete_request.get("Key")
-                if isinstance(key, dict):
-                    target = _min_batch_write_target(
-                        target,
-                        _BatchWriteRoutingTarget(
-                            table_name,
-                            key,
-                            _batch_write_sort_key(table_name, "DeleteRequest", key),
-                        ),
-                    )
+            target = _min_batch_write_target(
+                target,
+                _BatchWriteRoutingTarget(
+                    table_name,
+                    candidate.attributes,
+                    _batch_write_sort_key(
+                        table_name,
+                        candidate.operation,
+                        candidate.attributes,
+                    ),
+                ),
+            )
 
     return target