diff --git a/tests/core/test_kv_add_keys_batching.py b/tests/core/test_kv_add_keys_batching.py new file mode 100644 index 00000000..e5c0a2c0 --- /dev/null +++ b/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] diff --git a/unshackle/commands/kv.py b/unshackle/commands/kv.py index 00986167..ea0bca07 100644 --- a/unshackle/commands/kv.py +++ b/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: diff --git a/unshackle/vaults/MySQL.py b/unshackle/vaults/MySQL.py index 4dccdfe1..d2c43eda 100644 --- a/unshackle/vaults/MySQL.py +++ b/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}