mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-14 21:15:06 +00:00
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:
@@ -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
@@ -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...")
|
||||
|
||||
@@ -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
@@ -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
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user