Fix BatchWrite affinity and IPv6 routing URLs

alternator/core/handlers.py

@@ -22,6 +22,7 @@
     create_header_filter_handler,
     create_user_agent_header_handler,
 )
+from alternator.core.live_nodes import _format_host_port
 from alternator.core.query_plan import LazyQueryPlan
 from alternator.core.request import extract_operation_name, extract_request_params
 from alternator.exceptions import NoNodesAvailableError
@@ -83,20 +84,20 @@ def create_query_plan(
         """Create a URI iterator for a single request."""
         node_addresses = nodes.nodes
         if preferred_node is not None and preferred_node in node_addresses:
-            yield f"{scheme}://{preferred_node}:{port}"
+            yield f"{scheme}://{_format_host_port(preferred_node, port)}"
             remaining_nodes = tuple(
                 node for node in node_addresses if node != preferred_node
             )
             seed = _stable_seed(preferred_node)
             plan = LazyQueryPlan(nodes=remaining_nodes, seed=seed)
             for node in plan:
-                yield f"{scheme}://{node}:{port}"
+                yield f"{scheme}://{_format_host_port(node, port)}"
             return
 
         seed = random.getrandbits(64)
         plan = LazyQueryPlan(nodes=node_addresses, seed=seed)
         for node in plan:
-            yield f"{scheme}://{node}:{port}"
+            yield f"{scheme}://{_format_host_port(node, port)}"
 
     # Register event handler to update endpoint per-request
     def update_endpoint(

alternator/core/key_affinity.py

@@ -13,6 +13,7 @@
 
 from alternator._constants import PK_DISCOVERY_TIMEOUT_SECONDS
 from alternator.core.hashing import hash_attribute_value
+from alternator.core.query_plan import LazyQueryPlan
 
 if TYPE_CHECKING:
     from mypy_boto3_dynamodb import DynamoDBClient
@@ -205,7 +206,6 @@ def _select_batch_write_affinity_node(
     get_pk_name: Callable[[str], str | None],
 ) -> str | None:
     votes: Counter[str] = Counter()
-    selector = AffinitySelector()
 
     for candidate in _iter_batch_write_candidates(params):
         pk_name = get_pk_name(candidate.table_name)
@@ -226,7 +226,7 @@ def _select_batch_write_affinity_node(
         except (ValueError, TypeError, UnicodeEncodeError):
             continue
 
-        node = selector.select(nodes, hash_value)
+        node = _select_query_plan_first_node(nodes, hash_value)
         if node is not None:
             votes[node] += 1
 
@@ -240,6 +240,12 @@ def _select_batch_write_affinity_node(
     return winners[0]
 
 
+def _select_query_plan_first_node(nodes: NodeList, hash_value: int) -> str | None:
+    if not nodes:
+        return None
+    return next(LazyQueryPlan(nodes=tuple(sorted(nodes.nodes)), seed=hash_value))
+
+
 def _iter_batch_write_candidates(
     params: dict[str, Any],
 ) -> Iterable[_BatchWriteCandidate]:

alternator/core/live_nodes.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import contextlib
+import ipaddress
 import logging
 import threading
 import time
@@ -137,7 +138,8 @@ def build_localnodes_url(self, scope: RoutingScope, seed: str) -> str:
         Returns:
             Full URL with query string
         """
-        base = f"{self._config.scheme}://{seed}:{self._config.port}/localnodes"
+        authority = _format_host_port(seed, self._config.port)
+        base = f"{self._config.scheme}://{authority}/localnodes"
         query = scope.get_localnodes_query()
         return f"{base}?{query}" if query else base
 
@@ -236,6 +238,22 @@ def _no_nodes_for_scope_error(scope: RoutingScope) -> ConfigurationError:
     return ConfigurationError(f"No nodes found for routing scope: {scope.description}")
 
 
+def _format_host_port(host: str, port: int) -> str:
+    """Return a URL authority for a host-only address and shared port."""
+    if host.startswith("[") and host.endswith("]"):
+        return f"{host}:{port}"
+    if _is_ipv6_literal(host):
+        return f"[{host}]:{port}"
+    return f"{host}:{port}"
+
+
+def _is_ipv6_literal(host: str) -> bool:
+    try:
+        return ipaddress.ip_address(host).version == 6
+    except ValueError:
+        return False
+
+
 class SyncLiveNodesManager:
     """
     Synchronous LiveNodesManager with background refresh thread.
@@ -319,7 +337,8 @@ def next_node_uri(self) -> str:
                 scope_name=self._config.routing_scope.name,
                 attempted_hosts=list(self._config.seed_hosts),
             )
-        return f"{self._config.scheme}://{node}:{self._config.port}"
+        authority = _format_host_port(node, self._config.port)
+        return f"{self._config.scheme}://{authority}"
 
     def refresh_nodes(self) -> bool:
         """
@@ -509,7 +528,8 @@ def next_node_uri(self) -> str:
                 scope_name=self._config.routing_scope.name,
                 attempted_hosts=list(self._config.seed_hosts),
             )
-        return f"{self._config.scheme}://{node}:{self._config.port}"
+        authority = _format_host_port(node, self._config.port)
+        return f"{self._config.scheme}://{authority}"
 
     async def refresh_nodes(self) -> bool:
         """

tests/unit/test_async_manager.py

@@ -162,6 +162,32 @@ async def empty_fetch(url: str) -> list[str]:
         url = manager._core.build_localnodes_url(ClusterScope(), "192.168.1.1")
         assert url == "http://192.168.1.1:8000/localnodes"
 
+    @pytest.mark.asyncio
+    async def test_url_construction_brackets_ipv6_seed(self) -> None:
+        """Test raw IPv6 seed hosts are bracketed in discovery URLs."""
+
+        async def empty_fetch(url: str) -> list[str]:
+            return []
+
+        config = Config(seed_hosts=["2001:db8::1"], port=8000)
+        manager = AsyncLiveNodesManager(config, empty_fetch)
+
+        url = manager._core.build_localnodes_url(ClusterScope(), "2001:db8::1")
+        assert url == "http://[2001:db8::1]:8000/localnodes"
+
+    @pytest.mark.asyncio
+    async def test_next_node_uri_brackets_ipv6_node(self) -> None:
+        """Test raw IPv6 live nodes are bracketed in operation URLs."""
+
+        async def mock_fetch(url: str) -> list[str]:
+            return ["2001:db8::2"]
+
+        config = Config(seed_hosts=["seed"], port=8000)
+        manager = AsyncLiveNodesManager(config, mock_fetch)
+        await manager.refresh_nodes()
+
+        assert manager.next_node_uri() == "http://[2001:db8::2]:8000"
+
     @pytest.mark.asyncio
     async def test_url_construction_with_dc_scope(self, config: Config) -> None:
         """Test URL construction with datacenter scope."""

tests/unit/test_key_affinity.py

@@ -24,6 +24,7 @@
     should_use_affinity,
 )
 from alternator.core.live_nodes import NodeList
+from alternator.core.query_plan import LazyQueryPlan
 from alternator.core.request import extract_request_params
 
 Params = dict[str, Any]
@@ -51,6 +52,19 @@ def _pk_value_for_node(nodes: NodeList, target_node: str, prefix: str) -> str:
     raise AssertionError(f"could not find value for node {target_node}")
 
 
+def _query_plan_first_node(nodes: NodeList, hash_value: int) -> str:
+    return next(LazyQueryPlan(nodes=tuple(sorted(nodes.nodes)), seed=hash_value))
+
+
+def _pk_value_for_batch_node(nodes: NodeList, target_node: str, prefix: str) -> str:
+    for index in range(1000):
+        value = f"{prefix}-{index}"
+        hash_value = hash_attribute_value("S", value)
+        if _query_plan_first_node(nodes, hash_value) == target_node:
+            return value
+    raise AssertionError(f"could not find batch value for node {target_node}")
+
+
 class TestIsRmwOperation:
     """Tests for is_rmw_operation function."""
 
@@ -360,7 +374,7 @@ def capture_request(request: AWSPreparedRequest, **_: object) -> None:
     def test_batch_write_single_put_selects_node(self) -> None:
         """Test BatchWriteItem with a single PutRequest selects its node."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
-        value = _pk_value_for_node(nodes, "a", "batch-put")
+        value = _pk_value_for_batch_node(nodes, "a", "batch-put")
         params = {
             "RequestItems": {
                 "orders": [{"PutRequest": {"Item": {"pk": {"S": value}}}}],
@@ -378,10 +392,49 @@ def test_batch_write_single_put_selects_node(self) -> None:
             == "a"
         )
 
+    def test_batch_write_vote_uses_query_plan_first_pick(self) -> None:
+        """Test BatchWriteItem votes use canonical query-plan first pick."""
+        nodes = NodeList(
+            nodes=("node1", "node2", "node3", "node4", "node5", "node6"),
+            scope_name="test",
+        )
+        value = None
+        expected = None
+        modulo = None
+        for index in range(1000):
+            candidate = f"canonical-batch-{index}"
+            hash_value = hash_attribute_value("S", candidate)
+            query_plan_node = _query_plan_first_node(nodes, hash_value)
+            modulo_node = AffinitySelector().select(nodes, hash_value)
+            if query_plan_node != modulo_node:
+                value = candidate
+                expected = query_plan_node
+                modulo = modulo_node
+                break
+
+        assert value is not None
+        params = {
+            "RequestItems": {
+                "orders": [{"PutRequest": {"Item": {"pk": {"S": value}}}}],
+            }
+        }
+
+        assert (
+            select_affinity_node(
+                mode="ANY_WRITE",
+                operation_name="BatchWriteItem",
+                params=params,
+                nodes=nodes,
+                get_pk_name={"orders": "pk"}.get,
+            )
+            == expected
+        )
+        assert expected != modulo
+
     def test_batch_write_single_delete_selects_node(self) -> None:
         """Test BatchWriteItem with a single DeleteRequest selects its node."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
-        value = _pk_value_for_node(nodes, "c", "batch-delete")
+        value = _pk_value_for_batch_node(nodes, "c", "batch-delete")
         params = {
             "RequestItems": {
                 "orders": [{"DeleteRequest": {"Key": {"pk": {"S": value}}}}],
@@ -402,9 +455,9 @@ def test_batch_write_single_delete_selects_node(self) -> None:
     def test_batch_write_mixed_put_delete_unique_winner(self) -> None:
         """Test BatchWriteItem votes for the unique preferred node."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
-        b1 = _pk_value_for_node(nodes, "b", "b1")
-        b2 = _pk_value_for_node(nodes, "b", "b2")
-        c1 = _pk_value_for_node(nodes, "c", "c1")
+        b1 = _pk_value_for_batch_node(nodes, "b", "b1")
+        b2 = _pk_value_for_batch_node(nodes, "b", "b2")
+        c1 = _pk_value_for_batch_node(nodes, "c", "c1")
         params = {
             "RequestItems": {
                 "orders": [
@@ -429,9 +482,9 @@ def test_batch_write_mixed_put_delete_unique_winner(self) -> None:
     def test_batch_write_multi_table_reversed_order_same_winner(self) -> None:
         """Test batch voting is independent of table and request order."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
-        b1 = _pk_value_for_node(nodes, "b", "orders-b1")
-        b2 = _pk_value_for_node(nodes, "b", "sessions-b2")
-        a1 = _pk_value_for_node(nodes, "a", "orders-a1")
+        b1 = _pk_value_for_batch_node(nodes, "b", "orders-b1")
+        b2 = _pk_value_for_batch_node(nodes, "b", "sessions-b2")
+        a1 = _pk_value_for_batch_node(nodes, "a", "orders-a1")
         orders = [
             {"PutRequest": {"Item": {"pk": {"S": b1}}}},
             {"PutRequest": {"Item": {"pk": {"S": a1}}}},
@@ -543,8 +596,8 @@ def test_batch_write_no_nodes_falls_back(self) -> None:
     def test_batch_write_tied_votes_fall_back(self) -> None:
         """Test tied preferred-node votes produce no preferred node."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
-        a1 = _pk_value_for_node(nodes, "a", "tie-a")
-        b1 = _pk_value_for_node(nodes, "b", "tie-b")
+        a1 = _pk_value_for_batch_node(nodes, "a", "tie-a")
+        b1 = _pk_value_for_batch_node(nodes, "b", "tie-b")
         params = {
             "RequestItems": {
                 "orders": [
@@ -569,9 +622,8 @@ def test_batch_write_binary_pk_selects_stable_node(self) -> None:
         """Test binary partition-key values use stable hashing."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
         binary_value = b"\x00\x01stable"
-        expected = AffinitySelector().select(
-            nodes,
-            hash_attribute_value("B", binary_value),
+        expected = _query_plan_first_node(
+            nodes, hash_attribute_value("B", binary_value)
         )
         params = {
             "RequestItems": {
@@ -595,7 +647,7 @@ def test_batch_write_binary_pk_selects_stable_node(self) -> None:
     def test_batch_write_selection_does_not_mutate_params(self) -> None:
         """Test BatchWriteItem affinity selection leaves request params unchanged."""
         nodes = NodeList(nodes=("a", "b", "c"), scope_name="test")
-        value = _pk_value_for_node(nodes, "b", "no-mutate")
+        value = _pk_value_for_batch_node(nodes, "b", "no-mutate")
         params = {
             "RequestItems": {
                 "orders": [{"PutRequest": {"Item": {"pk": {"S": value}}}}],
@@ -661,6 +713,34 @@ def compute_affinity_node(
         assert first_url == "http://b:8000/"
         assert {second_url, third_url} == {"http://a:8000/", "http://c:8000/"}
 
+    def test_query_plan_brackets_ipv6_nodes(self) -> None:
+        """Test request handler formats raw IPv6 node addresses as URL authorities."""
+        config = Config(seed_hosts=["seed"], port=8000)
+        manager = MagicMock()
+        manager.nodes = NodeList(
+            nodes=("2001:db8::1", "2001:db8::2"), scope_name="cluster"
+        )
+        events = MagicMock()
+
+        _register_alternator_handlers(events, manager, config)
+        handlers = {
+            call[0][1].__name__: call[0][1] for call in events.register.call_args_list
+        }
+
+        request = MagicMock()
+        request.url = "http://seed:8000/"
+        request.headers = {"X-Amz-Target": "DynamoDB_20120810.ListTables"}
+        request.body = b"{}"
+        request._alternator_query_plan = None
+
+        update_endpoint = handlers["update_endpoint"]
+        update_endpoint(request)
+
+        assert request.url in {
+            "http://[2001:db8::1]:8000/",
+            "http://[2001:db8::2]:8000/",
+        }
+
 
 class TestExtractPartitionKey:
     """Tests for extract_partition_key function."""

tests/unit/test_sync_manager.py

@@ -159,6 +159,32 @@ def mock_fetch(url: str) -> Sequence[str]:
 
         assert captured_url[0] == "http://192.168.1.1:8000/localnodes"
 
+    def test_url_construction_brackets_ipv6_seed(self) -> None:
+        """Test raw IPv6 seed hosts are bracketed in discovery URLs."""
+        config = Config(seed_hosts=["2001:db8::1"], port=8000)
+        captured_url: list[str] = []
+
+        def mock_fetch(url: str) -> Sequence[str]:
+            captured_url.append(url)
+            return ["node1"]
+
+        manager = SyncLiveNodesManager(config, mock_fetch)
+        manager.refresh_nodes()
+
+        assert captured_url[0] == "http://[2001:db8::1]:8000/localnodes"
+
+    def test_next_node_uri_brackets_ipv6_node(self) -> None:
+        """Test raw IPv6 live nodes are bracketed in operation URLs."""
+        config = Config(seed_hosts=["seed"], port=8000)
+
+        def mock_fetch(url: str) -> Sequence[str]:
+            return ["2001:db8::2"]
+
+        manager = SyncLiveNodesManager(config, mock_fetch)
+        manager.refresh_nodes()
+
+        assert manager.next_node_uri() == "http://[2001:db8::2]:8000"
+
     def test_url_construction_with_dc_scope(self) -> None:
         """Test URL includes query string for datacenter scope."""
         config = Config(