diff --git a/src/freshdata/enterprise/privacy.py b/src/freshdata/enterprise/privacy.py index c3f60eee..3ebcd523 100644 --- a/src/freshdata/enterprise/privacy.py +++ b/src/freshdata/enterprise/privacy.py @@ -34,12 +34,13 @@ import json import os import re +import threading import warnings from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Literal +from typing import IO, Any, Literal import pandas as pd @@ -586,29 +587,137 @@ def __len__(self) -> int: return len(self._map) +def _os_file_lock_module() -> tuple[str, Any] | None: + """Return the available OS file-lock module: ``fcntl``, then ``msvcrt``, else None.""" + try: + import fcntl + except ImportError: + pass + else: + return "fcntl", fcntl + try: + import msvcrt + except ImportError: + return None + return "msvcrt", msvcrt + + +@contextlib.contextmanager +def _locked_file(handle: IO[str], *, exclusive: bool) -> Iterator[None]: + """Hold an OS lock on the open ``handle`` for the duration of the block. + + Uses ``fcntl.flock`` (shared or exclusive) where available. On Windows it locks + byte 0 of the file with ``msvcrt.locking``, which is always exclusive. When + neither module can be imported the block runs without a cross-process lock. + """ + found = _os_file_lock_module() + if found is None: + yield + return + kind, module = found + fd = handle.fileno() + if kind == "fcntl": + module.flock(fd, module.LOCK_EX if exclusive else module.LOCK_SH) + try: + yield + finally: + module.flock(fd, module.LOCK_UN) + return + handle.seek(0) + module.locking(fd, module.LK_LOCK, 1) + try: + yield + finally: + handle.seek(0) + module.locking(fd, module.LK_UNLCK, 1) + + class JsonTokenVault(TokenVault): """A token vault persisted to an explicit JSON file. The file holds the sensitive token→value mapping, so protect it like any secret. Nothing is written until :meth:`put` (or :meth:`save`) is called. + + Several instances (in one process or in several) may share the same path: + writes are merged and serialised. Each write takes an exclusive lock on the + file (``fcntl.flock`` on POSIX, ``msvcrt.locking`` on Windows), re-reads the + file, merges its entries with this instance's map and rewrites the file in + place before releasing the lock. :meth:`get` reloads the file when a token is + missing and the file changed since it was last read. On platforms with neither + lock module only the per-instance thread lock applies, so separate instances + and processes are not serialised there. + + The file is rewritten in place, so a crash mid-write can leave it incomplete + (empty or truncated JSON). An empty file loads as an empty vault; a truncated + one raises on load. """ def __init__(self, path: str | Path) -> None: self.path = Path(path) self._map: dict[str, str] = {} + self._stamp: tuple[int, int] | None = None + self._lock = threading.RLock() if self.path.exists(): - self._map = json.loads(self.path.read_text(encoding="utf-8")) + self._reload() + + @staticmethod + def _parse(text: str) -> dict[str, str]: + return json.loads(text) if text.strip() else {} + + @staticmethod + def _stat_stamp(fd: int) -> tuple[int, int]: + st = os.fstat(fd) + return (st.st_mtime_ns, st.st_size) + + def _current_stamp(self) -> tuple[int, int] | None: + try: + st = self.path.stat() + except OSError: + return None + return (st.st_mtime_ns, st.st_size) + + def _reload(self) -> None: + """Merge the file contents into the in-memory map, under a shared lock.""" + try: + with open(self.path, encoding="utf-8") as handle: # noqa: SIM117 + with _locked_file(handle, exclusive=False): + handle.seek(0) + disk = self._parse(handle.read()) + self._map = {**self._map, **disk} + self._stamp = self._stat_stamp(handle.fileno()) + except FileNotFoundError: + return def get(self, token: str) -> str | None: - return self._map.get(token) + with self._lock: + value = self._map.get(token) + if value is None and self._current_stamp() not in (None, self._stamp): + self._reload() + value = self._map.get(token) + return value def put(self, token: str, value: str) -> None: - self._map[token] = value - self.save() + self._write({token: value}, always=False) def save(self) -> None: - self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text(json.dumps(self._map, indent=2), encoding="utf-8") + self._write({}, always=True) + + def _write(self, updates: dict[str, str], *, always: bool) -> None: + with self._lock: + self.path.parent.mkdir(parents=True, exist_ok=True) + with open(self.path, "a+", encoding="utf-8") as handle: # noqa: SIM117 + with _locked_file(handle, exclusive=True): + handle.seek(0) + disk = self._parse(handle.read()) + merged = {**self._map, **disk, **updates} + if always or merged != disk: + handle.seek(0) + handle.truncate() + handle.write(json.dumps(merged, indent=2)) + handle.flush() + os.fsync(handle.fileno()) + self._map = merged + self._stamp = self._stat_stamp(handle.fileno()) class SqliteTokenVault(TokenVault): @@ -617,36 +726,45 @@ class SqliteTokenVault(TokenVault): Suited to larger reversible-tokenization runs where holding the whole map in memory (or rewriting a JSON file on every ``put``) is undesirable. The table stores only the ``token -> value`` mapping; protect the file like any secret. + + The vault is thread-safe: its connection may be used from any thread (for + example a vault created in the main thread and passed to a worker pool), and + a lock serialises every call on it. Other processes sharing the file wait up + to 30 seconds for SQLite's own database lock. """ def __init__(self, path: str | Path) -> None: import sqlite3 + self._lock = threading.RLock() self.path = Path(path) if str(self.path) != ":memory:": self.path.parent.mkdir(parents=True, exist_ok=True) - self._conn = sqlite3.connect(str(self.path)) + self._conn = sqlite3.connect(str(self.path), check_same_thread=False, timeout=30.0) self._conn.execute( "CREATE TABLE IF NOT EXISTS tokens (token TEXT PRIMARY KEY, value TEXT NOT NULL)" ) self._conn.commit() def get(self, token: str) -> str | None: - cur = self._conn.execute("SELECT value FROM tokens WHERE token = ?", (token,)) - row = cur.fetchone() + with self._lock: + cur = self._conn.execute("SELECT value FROM tokens WHERE token = ?", (token,)) + row = cur.fetchone() return None if row is None else str(row[0]) def put(self, token: str, value: str) -> None: - self._conn.execute( - "INSERT OR REPLACE INTO tokens (token, value) VALUES (?, ?)", (token, value) - ) - self._conn.commit() + with self._lock: + self._conn.execute( + "INSERT OR REPLACE INTO tokens (token, value) VALUES (?, ?)", (token, value) + ) + self._conn.commit() def __len__(self) -> int: - return int(self._conn.execute("SELECT COUNT(*) FROM tokens").fetchone()[0]) + with self._lock: + return int(self._conn.execute("SELECT COUNT(*) FROM tokens").fetchone()[0]) def close(self) -> None: - with contextlib.suppress(Exception): + with contextlib.suppress(Exception), self._lock: self._conn.close() def __del__(self) -> None: # best-effort: avoid leaking the DB handle diff --git a/tests/test_token_vault_concurrency.py b/tests/test_token_vault_concurrency.py new file mode 100644 index 00000000..d9c491b0 --- /dev/null +++ b/tests/test_token_vault_concurrency.py @@ -0,0 +1,287 @@ +"""Token vaults shared across instances and threads (#279). + +JsonTokenVault merges and serialises writes from several instances on one path; +SqliteTokenVault's connection is usable from any thread. +""" + +from __future__ import annotations + +import json +import os +import sys +import threading +import types +from concurrent.futures import ThreadPoolExecutor + +import pandas as pd +import pytest + +from freshdata.enterprise import ( + JsonTokenVault, + PrivacyPolicy, + PrivacyRule, + SqliteTokenVault, + apply_privacy_policy, + detokenize_series, +) + +KEY = "unit-test-key" +N_THREADS = 8 +PUTS_PER_THREAD = 10 +JOIN_TIMEOUT = 30.0 + + +def _run_threads(target, n: int = N_THREADS) -> None: + """Start ``n`` threads together (via a barrier) and re-raise any worker error.""" + barrier = threading.Barrier(n) + errors: list[BaseException] = [] + + def worker(i: int) -> None: + try: + barrier.wait(timeout=JOIN_TIMEOUT) + target(i) + except BaseException as exc: # pragma: no cover - surfaced below + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(n)] + for t in threads: + t.start() + for t in threads: + t.join(JOIN_TIMEOUT) + assert not any(t.is_alive() for t in threads), "worker thread did not finish" + if errors: + raise errors[0] + + +# -------------------------------------------------------------------------- +# JsonTokenVault +# -------------------------------------------------------------------------- + + +def test_json_two_instances_keep_both_mappings(tmp_path): + # Issue #279 repro: the second instance used to overwrite the first's entry. + path = tmp_path / "vault.json" + a, b = JsonTokenVault(path), JsonTokenVault(path) + a.put("tok_a", "111-11-1111") + b.put("tok_b", "222-22-2222") + assert JsonTokenVault(path).get("tok_a") == "111-11-1111" + assert JsonTokenVault(path).get("tok_b") == "222-22-2222" + + +def test_json_interleaved_puts_from_two_instances(tmp_path): + path = tmp_path / "vault.json" + a, b = JsonTokenVault(path), JsonTokenVault(path) + expected = {} + for i in range(6): + vault = a if i % 2 == 0 else b + vault.put(f"tok_{i}", f"value-{i}") + expected[f"tok_{i}"] = f"value-{i}" + assert json.loads(path.read_text(encoding="utf-8")) == expected + for token, value in expected.items(): + assert a.get(token) == value + assert b.get(token) == value + + +def test_json_get_sees_entry_written_by_other_instance(tmp_path): + path = tmp_path / "vault.json" + b = JsonTokenVault(path) # created before the file exists + a = JsonTokenVault(path) + assert b.get("tok_x") is None + a.put("tok_x", "x-value") + assert b.get("tok_x") == "x-value" + a.put("tok_y", "y-value") + assert b.get("tok_y") == "y-value" + + +def test_json_threads_with_own_instances_keep_all_entries(tmp_path): + path = tmp_path / "vault.json" + + def work(i: int) -> None: + vault = JsonTokenVault(path) + for j in range(PUTS_PER_THREAD): + vault.put(f"tok_{i}_{j}", f"value-{i}-{j}") + + _run_threads(work) + on_disk = json.loads(path.read_text(encoding="utf-8")) + assert len(on_disk) == N_THREADS * PUTS_PER_THREAD + fresh = JsonTokenVault(path) + for i in range(N_THREADS): + for j in range(PUTS_PER_THREAD): + assert fresh.get(f"tok_{i}_{j}") == f"value-{i}-{j}" + + +def test_json_threads_sharing_one_instance(tmp_path): + path = tmp_path / "vault.json" + vault = JsonTokenVault(path) + + def work(i: int) -> None: + for j in range(PUTS_PER_THREAD): + vault.put(f"tok_{i}_{j}", f"value-{i}-{j}") + + _run_threads(work) + assert len(json.loads(path.read_text(encoding="utf-8"))) == N_THREADS * PUTS_PER_THREAD + + +def test_json_empty_file_loads_as_empty_vault(tmp_path): + path = tmp_path / "vault.json" + path.write_text("", encoding="utf-8") + vault = JsonTokenVault(path) + assert vault.get("tok_missing") is None + vault.put("tok_a", "a") + assert json.loads(path.read_text(encoding="utf-8")) == {"tok_a": "a"} + + +def test_json_nothing_written_until_put(tmp_path): + path = tmp_path / "sub" / "vault.json" + vault = JsonTokenVault(path) + assert vault.get("tok_a") is None + assert not path.exists() + vault.put("tok_a", "a") + assert path.read_text(encoding="utf-8") == json.dumps({"tok_a": "a"}, indent=2) + + +def test_json_repeated_put_of_same_value_skips_rewrite(tmp_path, monkeypatch): + path = tmp_path / "vault.json" + vault = JsonTokenVault(path) + vault.put("tok_a", "a") + calls = [] + real_fsync = os.fsync + monkeypatch.setattr(os, "fsync", lambda fd: (calls.append(fd), real_fsync(fd))) + vault.put("tok_a", "a") + JsonTokenVault(path).put("tok_a", "a") + assert calls == [] + vault.put("tok_b", "b") + assert len(calls) == 1 + + +def test_json_save_keeps_entries_from_other_instances(tmp_path): + path = tmp_path / "vault.json" + a, b = JsonTokenVault(path), JsonTokenVault(path) + a.put("tok_a", "a") + b.save() + assert json.loads(path.read_text(encoding="utf-8")) == {"tok_a": "a"} + + +def test_json_tokenize_round_trip_through_two_vault_instances(tmp_path): + path = tmp_path / "vault.json" + rule = PrivacyRule(id="ssn", action="tokenize", reversible=True, columns=("ssn",)) + policy = PrivacyPolicy(name="p", rules=(rule,), key=KEY) + out_a, _ = apply_privacy_policy( + pd.DataFrame({"ssn": ["123-45-6789"]}), policy, vault=JsonTokenVault(path) + ) + out_b, _ = apply_privacy_policy( + pd.DataFrame({"ssn": ["987-65-4321"]}), policy, vault=JsonTokenVault(path) + ) + fresh = JsonTokenVault(path) + assert list(detokenize_series(out_a["ssn"], fresh, KEY)) == ["123-45-6789"] + assert list(detokenize_series(out_b["ssn"], fresh, KEY)) == ["987-65-4321"] + + +def test_json_uses_exclusive_flock_for_writes_and_shared_for_reads(tmp_path, monkeypatch): + fcntl = pytest.importorskip("fcntl") + ops = [] + + def recording_flock(fd, op): + ops.append(op) + return fcntl.flock(fd, op) + + stub = types.SimpleNamespace( + LOCK_EX=fcntl.LOCK_EX, LOCK_SH=fcntl.LOCK_SH, LOCK_UN=fcntl.LOCK_UN, flock=recording_flock + ) + monkeypatch.setitem(sys.modules, "fcntl", stub) + path = tmp_path / "vault.json" + JsonTokenVault(path).put("tok_a", "a") + assert ops == [fcntl.LOCK_EX, fcntl.LOCK_UN] + ops.clear() + JsonTokenVault(path) + assert ops == [fcntl.LOCK_SH, fcntl.LOCK_UN] + + +def test_json_windows_lock_covers_byte_zero(tmp_path, monkeypatch): + calls = [] + + def locking(fd, mode, nbytes): + calls.append((mode, nbytes, os.lseek(fd, 0, os.SEEK_CUR))) + + stub = types.SimpleNamespace(LK_LOCK=1, LK_UNLCK=0, locking=locking) + monkeypatch.setitem(sys.modules, "fcntl", None) + monkeypatch.setitem(sys.modules, "msvcrt", stub) + path = tmp_path / "vault.json" + a, b = JsonTokenVault(path), JsonTokenVault(path) + a.put("tok_a", "a") + b.put("tok_b", "b") + assert calls == [(1, 1, 0), (0, 1, 0)] * 2 + assert json.loads(path.read_text(encoding="utf-8")) == {"tok_a": "a", "tok_b": "b"} + + +def test_json_without_lock_modules_falls_back_to_thread_lock(tmp_path, monkeypatch): + monkeypatch.setitem(sys.modules, "fcntl", None) + monkeypatch.setitem(sys.modules, "msvcrt", None) + path = tmp_path / "vault.json" + vault = JsonTokenVault(path) + + def work(i: int) -> None: + for j in range(PUTS_PER_THREAD): + vault.put(f"tok_{i}_{j}", f"value-{i}-{j}") + + _run_threads(work) + other = JsonTokenVault(path) + other.put("tok_other", "other") + assert len(json.loads(path.read_text(encoding="utf-8"))) == N_THREADS * PUTS_PER_THREAD + 1 + + +# -------------------------------------------------------------------------- +# SqliteTokenVault +# -------------------------------------------------------------------------- + + +def test_sqlite_vault_works_from_worker_thread(tmp_path): + # Issue #279 repro: a vault created in the main thread used in a thread pool. + vault = SqliteTokenVault(tmp_path / "v.db") + policy = PrivacyPolicy( + rules=(PrivacyRule(id="t", action="tokenize", columns=("ssn",), key="k"),) + ) + frame = pd.DataFrame({"ssn": ["123-45-6789"]}) + with ThreadPoolExecutor(1) as ex: + out, report = ex.submit(apply_privacy_policy, frame, policy, vault=vault).result() + assert out["ssn"].iloc[0].startswith("tok_") + assert report.vault_info["entries"] == 1 + assert len(vault) == 1 + vault.close() + + +def test_sqlite_get_put_len_from_other_thread(tmp_path): + vault = SqliteTokenVault(tmp_path / "v.db") + vault.put("tok_main", "main") + + def use() -> tuple[str | None, int]: + vault.put("tok_worker", "worker") + return vault.get("tok_main"), len(vault) + + with ThreadPoolExecutor(1) as ex: + assert ex.submit(use).result() == ("main", 2) + assert vault.get("tok_worker") == "worker" + vault.close() + + +def test_sqlite_eight_threads_share_one_vault(tmp_path): + vault = SqliteTokenVault(tmp_path / "v.db") + + def work(i: int) -> None: + for j in range(PUTS_PER_THREAD): + vault.put(f"tok_{i}_{j}", f"value-{i}-{j}") + assert vault.get(f"tok_{i}_{j}") == f"value-{i}-{j}" + assert len(vault) >= 1 + + _run_threads(work) + assert len(vault) == N_THREADS * PUTS_PER_THREAD + vault.close() + reopened = SqliteTokenVault(tmp_path / "v.db") + assert reopened.get("tok_7_9") == "value-7-9" + reopened.close() + + +def test_sqlite_close_twice_is_harmless(tmp_path): + vault = SqliteTokenVault(tmp_path / "v.db") + vault.close() + vault.close()