perf(kv): push whole services to SQL vaults in one call

tests/core/test_kv_add_keys_batching.py

@@ -0,0 +1,66 @@
+"""`add_keys_with_progress` gives SQL vaults the whole service in one call and chunks network vaults."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Iterator, Optional, Union
+from uuid import UUID
+
+import pytest
+
+from unshackle.commands.kv import add_keys_with_progress
+from unshackle.core.vault import Vault
+from unshackle.vaults.SQLite import SQLite
+
+LOG = logging.getLogger("kv-test")
+
+
+def make_keys(n: int) -> dict[str, str]:
+    return {f"{i:032x}": f"{i + 1:032x}" for i in range(n)}
+
+
+class RecordingVault(Vault):
+    """Counts the batch sizes it receives, like a network vault would."""
+
+    def __init__(self) -> None:
+        super().__init__("net")
+        self.batches: list[int] = []
+
+    def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
+        return None
+
+    def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
+        return iter(())
+
+    def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool:
+        return False
+
+    def add_keys(self, service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
+        self.batches.append(len(kid_keys))
+        return len(kid_keys)
+
+    def get_services(self) -> Iterator[str]:
+        return iter(())
+
+
+def test_sql_vault_gets_one_call(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    vault = SQLite("local", tmp_path / "keys.db")
+    calls: list[int] = []
+    real = vault.add_keys
+
+    def spy(service: str, kid_keys: dict[Union[UUID, str], str]) -> int:
+        calls.append(len(kid_keys))
+        return real(service, kid_keys)
+
+    monkeypatch.setattr(vault, "add_keys", spy)
+    keys = make_keys(1201)
+    assert add_keys_with_progress(vault, "SVC", keys, LOG) == 1201
+    assert calls == [1201]
+    assert add_keys_with_progress(vault, "SVC", keys, LOG) == 0
+
+
+def test_network_vault_probes_then_batches() -> None:
+    vault = RecordingVault()
+    assert add_keys_with_progress(vault, "SVC", make_keys(1201), LOG) == 1201
+    assert vault.batches == [1, 500, 500, 200]

unshackle/commands/kv.py

@@ -2,7 +2,7 @@ import logging
 import re
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional, Union, cast
 from uuid import UUID
 
 import click
@@ -64,19 +64,20 @@ class _PaddedProgress(Progress):
 
 
 def add_keys_with_progress(vault: Vault, service: str, kid_keys: dict[str, str], log: logging.Logger) -> int:
-    """Add content keys to a vault in batches, with a progress bar of counts and time left."""
+    """Add content keys to a vault. Network vaults write in batches behind a progress bar."""
+    if type(vault).__name__ in ("MySQL", "SQLite"):
+        return vault.add_keys(service, cast(dict[Union[UUID, str], str], kid_keys))
 
     def chunk(i: int) -> int:
-        if type(vault).__name__ != "HTTP":
-            return 500
-        return 500 if i and getattr(vault, "batch_insert", False) else 1
+        # Probe with one key so a server that rejects batches costs one row, not 500.
+        return 500 if i and getattr(vault, "batch_insert", True) else 1
 
     kids = list(kid_keys)
     added = 0
     with _PaddedProgress(
         SpinnerColumn(finished_text=""),
         TextColumn("[bold]{task.description}"),
-        GradientPulseBarColumn(),
+        GradientPulseBarColumn(bar_width=None),
         MofNCompleteColumn(),
         "•",
         TextColumn("[green]{task.fields[added]} new"),
@@ -86,6 +87,7 @@ def add_keys_with_progress(vault: Vault, service: str, kid_keys: dict[str, str],
         TimeRemainingColumn(compact=True),
         console=console,
         transient=True,
+        expand=True,
     ) as progress:
         task = progress.add_task(f"{service} → {vault}", total=len(kids), added=0)
         i = 0
@@ -196,7 +198,7 @@ def copy(
             try:
                 services_to_copy = list(from_vault.get_services())
             except Exception as e:
-                log.warning(f"{from_vault.name}: cannot list services ({e}), skipped")
+                log.debug(f"{from_vault.name}: cannot list services ({e}), skipped")
                 continue
 
         if installed is not None:

unshackle/vaults/MySQL.py

@@ -138,9 +138,13 @@ class MySQL(Vault):
         cursor = conn.cursor()
 
         try:
-            placeholders = ",".join(["%s"] * len(kid_keys))
-            cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", list(kid_keys.keys()))
-            existing_kids = {row["kid"] for row in cursor.fetchall()}
+            existing_kids: set[str] = set()
+            kid_list = list(kid_keys)
+            for i in range(0, len(kid_list), 10000):
+                batch = kid_list[i : i + 10000]
+                placeholders = ",".join(["%s"] * len(batch))
+                cursor.execute(f"SELECT kid FROM `{service}` WHERE kid IN ({placeholders})", batch)
+                existing_kids.update(row["kid"] for row in cursor.fetchall())
 
             new_keys = {kid: key for kid, key in kid_keys.items() if kid not in existing_kids}