fix(vaults): look tables up by the name the vault stores

get_services() normalised table names through Services.get_tag() and then handed the result back as an identifier, so kv copy and sync skipped any table whose name differed in case or matched an alias. Yield stored names, and resolve tables case-insensitively.
This commit is contained in:
imSp4rky
2026-08-19 12:11:44 -06:00
parent c4c49ea672
commit 54adb82746
5 changed files with 283 additions and 66 deletions
+177
View File
@@ -0,0 +1,177 @@
"""SQLite vault table names resolve case-insensitively, and survive a `kv copy` round trip.
A vault table is named by the service's vault tag, whose case (and sometimes whose
spelling) differs from the installed service tag. Reads and writes have to find the real
table whatever its case, `get_services()` must yield names that can be fed straight back
in, and no lookup may create a second table differing only by case.
"""
from __future__ import annotations
import logging
import sqlite3
import threading
from pathlib import Path
from typing import Iterator, Optional, Union
from uuid import UUID
import pytest
from unshackle.core.vault import Vault
from unshackle.vaults.SQLite import SQLite
LOG = logging.getLogger("kv-test")
KID = "0" * 31 + "1"
KEY = "a" * 32
def tables(path: Path) -> list[str]:
conn = sqlite3.connect(path)
try:
return sorted(r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'"))
finally:
conn.close()
class UnlistableVault(Vault):
"""Stands in for the API vault, which supports neither listing nor bulk dumps."""
def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
return None
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
raise ValueError("Bulk dump is not supported.")
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:
return 0
def get_services(self) -> Iterator[str]:
raise ValueError("Listing services is not supported.")
@pytest.fixture
def aliased_service(monkeypatch: pytest.MonkeyPatch) -> str:
"""Register an SVCD service whose vault table is named after its alias, not its tag."""
from unshackle.core import services
monkeypatch.setattr(services, "SERVICES", [Path("SVCD") / "__init__.py"])
monkeypatch.setattr(services, "ALIASES", {"SVCD": ("svcdalias",)})
return "svcdalias"
@pytest.fixture
def vault(tmp_path: Path) -> SQLite:
"""Vault whose backing DB already holds a lowercase `svca` table with one key."""
v = SQLite("test", tmp_path / "kv.db")
v.add_keys("svca", {KID: KEY})
return v
@pytest.mark.parametrize("service", ["svca", "SVCA", "SvCa"])
def test_reads_find_table_in_any_case(vault: SQLite, service: str) -> None:
assert vault.resolve_table(service) == "svca"
assert vault.get_key(KID, service) == KEY
assert list(vault.get_keys(service)) == [(KID, KEY)]
def test_add_keys_writes_into_existing_table(vault: SQLite) -> None:
kid = "0" * 31 + "2"
assert vault.add_keys("SVCA", {kid: "b" * 32}) == 1
assert tables(vault.path) == ["sqlite_sequence", "svca"]
assert vault.get_key(kid, "svca") == "b" * 32
def test_add_key_writes_into_existing_table(vault: SQLite) -> None:
kid = "0" * 31 + "3"
assert vault.add_key("SVCA", kid, "c" * 32)
assert tables(vault.path) == ["sqlite_sequence", "svca"]
assert vault.get_key(kid, "svca") == "c" * 32
def test_unknown_service_is_not_an_error(vault: SQLite) -> None:
assert vault.resolve_table("SVCX") is None
assert vault.get_key(KID, "SVCX") is None
assert list(vault.get_keys("SVCX")) == []
def test_new_service_creates_table_with_the_given_case(tmp_path: Path) -> None:
v = SQLite("test", tmp_path / "kv.db")
v.add_keys("SVCB", {KID: KEY})
assert tables(v.path) == ["SVCB", "sqlite_sequence"]
def test_synced_lowercase_table_survives_the_service_being_added_later(tmp_path: Path) -> None:
"""A `kv sync` before the service exists writes `svcc`; once SVCC is installed every
lookup uses the canonical tag and must still land on that same table."""
v = SQLite("test", tmp_path / "kv.db")
v.add_keys("svcc", {KID: KEY})
kid = "0" * 31 + "2"
assert v.get_key(KID, "SVCC") == KEY
assert v.add_keys("SVCC", {kid: "b" * 32}) == 1
assert tables(v.path) == ["sqlite_sequence", "svcc"]
assert len(list(v.get_keys("SVCC"))) == 2
def test_null_keys_are_rejected_and_filtered(tmp_path: Path) -> None:
v = SQLite("test", tmp_path / "kv.db")
with pytest.raises(ValueError):
v.add_keys("SVCA", {KID: "0" * 32})
v.create_table("svca")
conn = sqlite3.connect(v.path)
conn.execute("INSERT INTO svca (kid, key_) VALUES (?, ?)", (KID, "0" * 32))
conn.commit()
conn.close()
assert v.get_key(KID, "SVCA") is None
assert list(v.get_keys("SVCA")) == []
def test_get_services_yields_stored_table_names(vault: SQLite, aliased_service: str) -> None:
"""Names are yielded verbatim. Normalising the `svcdalias` table to the installed SVCD
tag would hand back a name that addresses no table at all."""
vault.add_keys(aliased_service, {KID: KEY})
assert sorted(vault.get_services()) == ["svca", aliased_service]
def test_every_listed_service_can_be_copied_out_and_in(tmp_path: Path, aliased_service: str) -> None:
"""Every name `get_services()` lists must be readable by `get_keys()`."""
from unshackle.commands.kv import copy_service_data
src = SQLite("src", tmp_path / "src.db")
src.add_keys(aliased_service, {KID: KEY})
src.add_keys("svca", {KID: KEY})
dst = SQLite("dst", tmp_path / "dst.db")
added = [copy_service_data(dst, src, service, LOG) for service in src.get_services()]
assert added == [1, 1]
assert dst.get_key(KID, aliased_service) == KEY
assert dst.get_key(KID, "svca") == KEY
def test_copy_skips_a_service_it_cannot_read(tmp_path: Path) -> None:
"""A vault that refuses bulk dumps must not abort the copy."""
from unshackle.commands.kv import copy_service_data
dst = SQLite("dst", tmp_path / "dst.db")
assert copy_service_data(dst, UnlistableVault("unlistable"), "svca", LOG) == 0
def test_copy_skips_a_vault_that_cannot_list_services() -> None:
with pytest.raises(ValueError):
list(UnlistableVault("unlistable").get_services())
def test_lookups_work_from_another_thread(vault: SQLite) -> None:
"""`kv search` probes services in a ThreadPoolExecutor; connections are thread-local."""
found: list[Optional[str]] = []
thread = threading.Thread(target=lambda: found.append(vault.get_key(KID, "SVCA")))
thread.start()
thread.join()
assert found == [KEY]
+24 -11
View File
@@ -48,7 +48,12 @@ def process_service_keys(from_vault: Vault, service: str, log: logging.Logger) -
def copy_service_data(to_vault: Vault, from_vault: Vault, service: str, log: logging.Logger) -> int:
"""Copy data for a single service between vaults."""
content_keys = process_service_keys(from_vault, service, log)
try:
content_keys = process_service_keys(from_vault, service, log)
except Exception as e:
log.warning(f"{service}: Could not read from {from_vault} ({e}), skipped")
return 0
total_count = len(content_keys)
if total_count == 0:
@@ -124,7 +129,7 @@ def copy(
raise click.UsageError("--service and --local-only are mutually exclusive.")
if service:
service = Services.get_tag(service)
service = Services.get_vault_tag(service)
log.info(f"Filtering by service: {service}")
installed: Optional[set[str]] = None
@@ -134,11 +139,18 @@ def copy(
total_added = 0
for from_vault in from_vaults:
services_to_copy = [service] if service else list(from_vault.get_services())
if service:
services_to_copy = [service]
else:
try:
services_to_copy = list(from_vault.get_services())
except Exception as e:
log.warning(f"{from_vault.name}: cannot list services ({e}), skipped")
continue
if installed is not None:
before = len(services_to_copy)
services_to_copy = [s for s in services_to_copy if s and s.upper() in installed]
services_to_copy = [s for s in services_to_copy if s and Services.get_tag(s).upper() in installed]
skipped = before - len(services_to_copy)
if skipped:
log.info(f"{from_vault.name}: skipping {skipped} service(s) not installed locally")
@@ -218,7 +230,7 @@ def add(file: Path, service: str, vaults: list[str]) -> None:
raise click.ClickException("You must provide at least one Vault.")
log = logging.getLogger("kv")
service = Services.get_tag(service)
service = Services.get_vault_tag(service)
vaults_ = load_vaults(list(vaults))
@@ -305,7 +317,7 @@ def search(kid: str, service: Optional[str], vault_name: Optional[str]) -> None:
vaults_ = load_vaults(vault_names)
service_tag = Services.get_tag(service) if service else None
service_tag = Services.get_vault_tag(service) if service else None
hits: list[tuple[str, str, str]] = []
for vault in vaults_:
@@ -349,13 +361,14 @@ def prepare(vaults: list[str]) -> None:
vaults_ = load_vaults(vaults)
for vault in vaults_:
if hasattr(vault, "has_table") and hasattr(vault, "create_table"):
for service_tag in Services.get_tags():
if vault.has_table(service_tag):
log.info(f"{vault} already has a {service_tag} Table")
if hasattr(vault, "resolve_table") and hasattr(vault, "create_table"):
for service_tag in dict.fromkeys(Services.get_vault_tag(tag) for tag in Services.get_tags()):
existing = vault.resolve_table(service_tag)
if existing:
log.info(f"{vault} already has a {existing} Table")
else:
try:
vault.create_table(service_tag, commit=True)
vault.create_table(service_tag)
log.info(f"{vault}: Created {service_tag} Table")
except PermissionError:
log.error(f"{vault} user has no create table permission, skipping...")
+6 -1
View File
@@ -43,7 +43,12 @@ class Vault(metaclass=ABCMeta):
@abstractmethod
def get_services(self) -> Iterator[str]:
"""Get a list of Service Tags from Vault."""
"""
Get a list of Service Tags from Vault.
Tags come back exactly as the Vault stores them, so passing one to
get_key/get_keys/add_key/add_keys reaches the namespace it came from.
"""
__all__ = ("Vault",)
+51 -32
View File
@@ -5,7 +5,6 @@ from uuid import UUID
import pymysql
from pymysql.cursors import DictCursor
from unshackle.core.services import Services
from unshackle.core.vault import Vault
@@ -31,35 +30,8 @@ class MySQL(Vault):
if isinstance(kid, UUID):
kid = kid.hex
service_variants = [service]
if service != service.lower():
service_variants.append(service.lower())
if service != service.upper():
service_variants.append(service.upper())
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
for service_name in service_variants:
if not self.has_table(service_name):
continue
cursor.execute(
# TODO: SQL injection risk
f"SELECT `id`, `key_` FROM `{service_name}` WHERE `kid`=%s AND `key_`!=%s",
(kid, "0" * 32),
)
cek = cursor.fetchone()
if cek:
return cek["key_"]
return None
finally:
cursor.close()
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
if not self.has_table(service):
table = self.resolve_table(service)
if not table:
return None
conn = self.conn_factory.get()
@@ -68,7 +40,26 @@ class MySQL(Vault):
try:
cursor.execute(
# TODO: SQL injection risk
f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=%s",
f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=%s AND `key_`!=%s",
(kid, "0" * 32),
)
cek = cursor.fetchone()
return cek["key_"] if cek else None
finally:
cursor.close()
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
table = self.resolve_table(service)
if not table:
return None
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
# TODO: SQL injection risk
f"SELECT `kid`, `key_` FROM `{table}` WHERE `key_`!=%s",
("0" * 32,),
)
for row in cursor.fetchall():
@@ -80,6 +71,8 @@ class MySQL(Vault):
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
service = self.resolve_table(service) or service
if not self.has_permission("INSERT", table=service):
raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.")
@@ -119,6 +112,8 @@ class MySQL(Vault):
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
service = self.resolve_table(service) or service
if not self.has_permission("INSERT", table=service):
raise PermissionError(f"MySQL vault {self.slug} has no INSERT permission.")
@@ -169,10 +164,34 @@ class MySQL(Vault):
cursor.execute("SHOW TABLES")
for table in cursor.fetchall():
# each entry has a key named `Tables_in_<db name>`
yield Services.get_tag(list(table.values())[0])
yield list(table.values())[0]
finally:
cursor.close()
def resolve_table(self, name: str) -> Optional[str]:
"""
Get the actual Table name matching `name` case-insensitively, if it exists.
The UPPER() comparison is explicit so the result does not depend on the server's
`lower_case_table_names` setting. Where both spellings exist, exact case wins.
"""
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(
"SELECT TABLE_NAME FROM information_schema.TABLES"
" WHERE TABLE_SCHEMA=%s AND UPPER(TABLE_NAME)=UPPER(%s)",
(conn.db, name),
)
names = [list(row.values())[0] for row in cursor.fetchall()]
finally:
cursor.close()
if name in names:
return name
return names[0] if names else None
def has_table(self, name: str) -> bool:
"""Check if the Vault has a Table with the specified name."""
conn = self.conn_factory.get()
+25 -22
View File
@@ -5,7 +5,6 @@ from sqlite3 import Connection
from typing import Iterator, Optional, Union
from uuid import UUID
from unshackle.core.services import Services
from unshackle.core.vault import Vault
@@ -22,40 +21,30 @@ class SQLite(Vault):
if isinstance(kid, UUID):
kid = kid.hex
table = self.resolve_table(service)
if not table:
return None
conn = self.conn_factory.get()
cursor = conn.cursor()
service_variants = [service]
if service != service.lower():
service_variants.append(service.lower())
if service != service.upper():
service_variants.append(service.upper())
try:
for service_name in service_variants:
if not self.has_table(service_name):
continue
cursor.execute(
f"SELECT `id`, `key_` FROM `{service_name}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32)
)
cek = cursor.fetchone()
if cek:
return cek[1]
return None
cursor.execute(f"SELECT `id`, `key_` FROM `{table}` WHERE `kid`=? AND `key_`!=?", (kid, "0" * 32))
cek = cursor.fetchone()
return cek[1] if cek else None
finally:
cursor.close()
def get_keys(self, service: str) -> Iterator[tuple[str, str]]:
if not self.has_table(service):
table = self.resolve_table(service)
if not table:
return None
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute(f"SELECT `kid`, `key_` FROM `{service}` WHERE `key_`!=?", ("0" * 32,))
cursor.execute(f"SELECT `kid`, `key_` FROM `{table}` WHERE `key_`!=?", ("0" * 32,))
for kid, key_ in cursor.fetchall():
yield kid, key_
finally:
@@ -65,6 +54,7 @@ class SQLite(Vault):
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
service = self.resolve_table(service) or service
if not self.has_table(service):
self.create_table(service)
@@ -98,6 +88,7 @@ class SQLite(Vault):
if not key or key.count("0") == len(key):
raise ValueError("You cannot add a NULL Content Key to a Vault.")
service = self.resolve_table(service) or service
if not self.has_table(service):
self.create_table(service)
@@ -158,7 +149,19 @@ class SQLite(Vault):
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
for (name,) in cursor.fetchall():
if name != "sqlite_sequence":
yield Services.get_tag(name)
yield name
finally:
cursor.close()
def resolve_table(self, name: str) -> Optional[str]:
"""Get the actual Table name matching `name` case-insensitively, if it exists."""
conn = self.conn_factory.get()
cursor = conn.cursor()
try:
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=? COLLATE NOCASE", (name,))
row = cursor.fetchone()
return row[0] if row else None
finally:
cursor.close()