diff --git a/tests/test_utils.py b/tests/test_utils.py index 69a2f7953..c2e469c99 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -109,3 +109,21 @@ def test_use_shuffled(self, use_shuffled): for t in trains: assert "shuffle" not in t assert "train" in t + + +class TestRedactSensitive: + @pytest.mark.parametrize("key", ["password", "Password", "api_key", "token"]) + def test_sensitive_keys_are_masked(self, key): + assert utils.redact_sensitive({"host": "h", key: "s3cret"}) == {"host": "h", key: utils.MASK} + + def test_http_auth_keeps_user_and_masks_secret(self): + assert utils.redact_sensitive({"http_auth": ("admin", "s3cret")}) == {"http_auth": ("admin", utils.MASK)} + # serverless OpenSearch passes an opaque auth object (e.g. AWS4Auth) instead of a pair + assert utils.redact_sensitive({"http_auth": object()}) == {"http_auth": utils.MASK} + + def test_nested_values_and_non_sensitive_keys(self): + config = {"hosts": [{"host": "h", "port": 443, "password": "s3cret"}], "use_ssl": True} + assert utils.redact_sensitive(config) == { + "hosts": [{"host": "h", "port": 443, "password": utils.MASK}], + "use_ssl": True, + } diff --git a/vectordb_bench/backend/clients/adbpg/adbpg.py b/vectordb_bench/backend/clients/adbpg/adbpg.py index c2e9bcc04..5453f159c 100644 --- a/vectordb_bench/backend/clients/adbpg/adbpg.py +++ b/vectordb_bench/backend/clients/adbpg/adbpg.py @@ -12,6 +12,7 @@ from psycopg import Connection, Cursor, sql from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import VectorDB from .config import AdbpgConfigDict, AdbpgIndexConfig @@ -68,7 +69,7 @@ def __init__( # construct basic units self.conn, self.cursor = self._create_connection(**self.connect_config) - log.info(f"{self.name} config values: {self.connect_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.connect_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -77,7 +78,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.connect_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.connect_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/clients/aliyun_opensearch/aliyun_opensearch.py b/vectordb_bench/backend/clients/aliyun_opensearch/aliyun_opensearch.py index 138e10bba..5ea61b2c8 100644 --- a/vectordb_bench/backend/clients/aliyun_opensearch/aliyun_opensearch.py +++ b/vectordb_bench/backend/clients/aliyun_opensearch/aliyun_opensearch.py @@ -7,6 +7,7 @@ from alibabacloud_ha3engine_vector.client import Client from alibabacloud_ha3engine_vector.models import QueryRequest +from ...utils import redact_sensitive from ..api import MetricType, VectorDB from .config import AliyunOpenSearchIndexConfig @@ -52,7 +53,7 @@ def __init__( ), ) - log.info(f"Aliyun_OpenSearch client config: {self.db_config}") + log.info(f"Aliyun_OpenSearch client config: {redact_sensitive(self.db_config)}") if drop_old: log.info(f"aliyun_OpenSearch client drop old index: {self.collection_name}") diff --git a/vectordb_bench/backend/clients/alloydb/alloydb.py b/vectordb_bench/backend/clients/alloydb/alloydb.py index b9808ce54..773ca000b 100644 --- a/vectordb_bench/backend/clients/alloydb/alloydb.py +++ b/vectordb_bench/backend/clients/alloydb/alloydb.py @@ -10,6 +10,7 @@ from pgvector.psycopg import register_vector from psycopg import Connection, Cursor, sql +from ...utils import redact_sensitive from ..api import VectorDB from .config import AlloyDBConfigDict, AlloyDBIndexConfig @@ -51,7 +52,7 @@ def __init__( self.cursor.execute("CREATE EXTENSION IF NOT EXISTS alloydb_scann CASCADE") self.conn.commit() - log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, diff --git a/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py b/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py index f6774b1b3..00f1ac298 100644 --- a/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py +++ b/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py @@ -7,6 +7,7 @@ from vectordb_bench import config from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import VectorDB from .config import AWSOpenSearchIndexConfig, AWSOS_Engine @@ -48,7 +49,7 @@ def __init__( self.with_scalar_labels = with_scalar_labels self._insert_batch_size = kwargs.get("insert_batch_size", config.DEFAULT_INSERT_BATCH_SIZE) - log.info(f"AWS_OpenSearch client config: {self.db_config}") + log.info(f"AWS_OpenSearch client config: {redact_sensitive(self.db_config)}") log.info(f"AWS_OpenSearch db case config : {self.case_config}") self._is_serverless = ".aoss." in self.db_config.get("hosts", [{}])[0].get("host", "") client = OpenSearch(**self.db_config) diff --git a/vectordb_bench/backend/clients/cockroachdb/cockroachdb.py b/vectordb_bench/backend/clients/cockroachdb/cockroachdb.py index 89fdf2826..514297933 100644 --- a/vectordb_bench/backend/clients/cockroachdb/cockroachdb.py +++ b/vectordb_bench/backend/clients/cockroachdb/cockroachdb.py @@ -12,6 +12,7 @@ from psycopg_pool import ConnectionPool from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import VectorDB from .config import CockroachDBIndexConfig @@ -91,7 +92,7 @@ def __init__( # noqa: PLR0915 self.conn: Connection | None = None self.cursor: Cursor | None = None - log.info(f"{self.name} config: {self.connect_config}, pool_size={self.pool_size}") + log.info(f"{self.name} config: {redact_sensitive(self.connect_config)}, pool_size={self.pool_size}") # Allow manual index creation (both flags can be False) # This is useful when CREATE INDEX times out in subprocess on multi-node clusters diff --git a/vectordb_bench/backend/clients/hologres/hologres.py b/vectordb_bench/backend/clients/hologres/hologres.py index c396e23b0..e57dde41a 100644 --- a/vectordb_bench/backend/clients/hologres/hologres.py +++ b/vectordb_bench/backend/clients/hologres/hologres.py @@ -14,6 +14,7 @@ from psycopg.adapt import Dumper from psycopg.pq import Format +from ...utils import redact_sensitive from ..api import VectorDB from .config import HologresConfig, HologresIndexConfig @@ -89,7 +90,7 @@ def __init__( cursor.execute("CREATE EXTENSION proxima;") conn.commit() - log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -98,7 +99,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.db_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/clients/lindorm/lindorm_search.py b/vectordb_bench/backend/clients/lindorm/lindorm_search.py index 7901cb35f..a00fbb3fd 100644 --- a/vectordb_bench/backend/clients/lindorm/lindorm_search.py +++ b/vectordb_bench/backend/clients/lindorm/lindorm_search.py @@ -8,6 +8,7 @@ from opensearchpy import OpenSearch from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import IndexType, VectorDB from .config import LindormConfig, LindormIndexConfig @@ -50,7 +51,7 @@ def __init__( self.vector_col_name = vector_col_name self.with_scalar_labels = with_scalar_labels self._pending_post_optimize_sleep = False # need sleep after optimize - log.debug(f"Lindorm client config: {self.db_config}") + log.debug(f"Lindorm client config: {redact_sensitive(self.db_config)}") log.debug(f"index name: {self.index_name}") log.debug(f"Lindorm index config: {self.case_config}") client = OpenSearch(**self.db_config) diff --git a/vectordb_bench/backend/clients/memorydb/memorydb.py b/vectordb_bench/backend/clients/memorydb/memorydb.py index d816d0fd4..907909b39 100644 --- a/vectordb_bench/backend/clients/memorydb/memorydb.py +++ b/vectordb_bench/backend/clients/memorydb/memorydb.py @@ -12,6 +12,7 @@ from redis.commands.search.indexDefinition import IndexDefinition, IndexType from redis.commands.search.query import Query +from ...utils import redact_sensitive from ..api import VectorDB from .config import MemoryDBIndexConfig @@ -36,7 +37,7 @@ def __init__( self.dbsize = kwargs.get("num_rows") # Create a MemoryDB connection, if db has password configured, add it to the connection here and in init(): - log.info(f"Establishing connection to: {self.db_config}") + log.info(f"Establishing connection to: {redact_sensitive(self.db_config)}") conn = self.get_client(primary=True) log.info(f"Connection established: {conn}") log.info(conn.execute_command("INFO server")) diff --git a/vectordb_bench/backend/clients/oceanbase/oceanbase.py b/vectordb_bench/backend/clients/oceanbase/oceanbase.py index b109d558d..7831a444b 100644 --- a/vectordb_bench/backend/clients/oceanbase/oceanbase.py +++ b/vectordb_bench/backend/clients/oceanbase/oceanbase.py @@ -9,6 +9,7 @@ import mysql.connector as mysql from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import IndexType, VectorDB from .config import OceanBaseConfigDict, OceanBaseHNSWConfig @@ -47,7 +48,9 @@ def __init__( self._vector_field = "embedding" log.info( - f"{self.name} initialized with config:\nDatabase: {self.db_config}\nCase Config: {self.db_case_config}" + f"{self.name} initialized with config:\n" + f"Database: {redact_sensitive(self.db_config)}\n" + f"Case Config: {self.db_case_config}" ) self._conn = None diff --git a/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py b/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py index 12617a640..45eb8999f 100644 --- a/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py +++ b/vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py @@ -10,6 +10,7 @@ from vectordb_bench.backend.filter import Filter, FilterOp from vectordb_bench.backend.payload import PayloadProfile +from vectordb_bench.backend.utils import redact_sensitive from ..api import VectorDB from .config import OSSOpenSearchFtsConfig, OSSOpenSearchIndexConfig, OSSOS_Engine @@ -226,7 +227,7 @@ def __init__( self.filter: dict[str, Any] | None = None self.routing_key: str | None = None - log.info(f"OSS_OpenSearch client config: {self.db_config}") + log.info(f"OSS_OpenSearch client config: {redact_sensitive(self.db_config)}") log.info(f"OSS_OpenSearch db case config: {self.case_config}") client = OpenSearch(**self.db_config) self._handle_index_initialization(client, drop_old) diff --git a/vectordb_bench/backend/clients/pgdiskann/pgdiskann.py b/vectordb_bench/backend/clients/pgdiskann/pgdiskann.py index 46e8fabd4..7b151b175 100644 --- a/vectordb_bench/backend/clients/pgdiskann/pgdiskann.py +++ b/vectordb_bench/backend/clients/pgdiskann/pgdiskann.py @@ -11,6 +11,7 @@ from psycopg import Connection, Cursor, sql from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import VectorDB from .config import PgDiskANNConfigDict, PgDiskANNIndexConfig @@ -57,7 +58,7 @@ def __init__( self.conn, self.cursor = self._create_connection(**self.db_config) - log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -66,7 +67,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.db_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py b/vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py index 3006b861a..776937c51 100644 --- a/vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py +++ b/vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py @@ -10,6 +10,7 @@ from pgvecto_rs.psycopg import register_vector from psycopg import Connection, Cursor, sql +from ...utils import redact_sensitive from ..api import VectorDB from .config import PgVectoRSConfig, PgVectoRSIndexConfig @@ -46,7 +47,7 @@ def __init__( # construct basic units self.conn, self.cursor = self._create_connection(**self.db_config) - log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -55,7 +56,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.db_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/clients/pgvector/pgvector.py b/vectordb_bench/backend/clients/pgvector/pgvector.py index 13f471afe..ac9153ffd 100644 --- a/vectordb_bench/backend/clients/pgvector/pgvector.py +++ b/vectordb_bench/backend/clients/pgvector/pgvector.py @@ -12,6 +12,7 @@ from psycopg import Connection, Cursor, sql from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import VectorDB from .config import PgVectorConfigDict, PgVectorIndexConfig @@ -62,7 +63,7 @@ def __init__( self.cursor.execute("CREATE EXTENSION IF NOT EXISTS vector") self.conn.commit() - log.info(f"{self.name} config values: {self.connect_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.connect_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -71,7 +72,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.connect_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.connect_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/clients/pgvectorscale/pgvectorscale.py b/vectordb_bench/backend/clients/pgvectorscale/pgvectorscale.py index 3985c0716..be0d3c723 100644 --- a/vectordb_bench/backend/clients/pgvectorscale/pgvectorscale.py +++ b/vectordb_bench/backend/clients/pgvectorscale/pgvectorscale.py @@ -10,6 +10,7 @@ from pgvector.psycopg import register_vector from psycopg import Connection, Cursor, sql +from ...utils import redact_sensitive from ..api import VectorDB from .config import PgVectorScaleConfigDict, PgVectorScaleIndexConfig @@ -46,7 +47,7 @@ def __init__( self.conn, self.cursor = self._create_connection(**self.db_config) - log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -55,7 +56,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.db_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/clients/seekdb/seekdb.py b/vectordb_bench/backend/clients/seekdb/seekdb.py index 69a27bd4b..74e77776b 100644 --- a/vectordb_bench/backend/clients/seekdb/seekdb.py +++ b/vectordb_bench/backend/clients/seekdb/seekdb.py @@ -9,6 +9,7 @@ import mysql.connector as mysql from vectordb_bench.backend.filter import Filter, FilterOp +from vectordb_bench.backend.utils import redact_sensitive from ..api import IndexType, VectorDB from .config import SeekDBConfigDict, SeekDBHNSWConfig @@ -71,7 +72,9 @@ def __init__( self.expr = "" log.info( - f"{self.name} initialized with config:\nDatabase: {self.db_config}\nCase Config: {self.db_case_config}" + f"{self.name} initialized with config:\n" + f"Database: {redact_sensitive(self.db_config)}\n" + f"Case Config: {self.db_case_config}" ) self._conn = None diff --git a/vectordb_bench/backend/clients/vectorchord/vectorchord.py b/vectordb_bench/backend/clients/vectorchord/vectorchord.py index 35dfa76d8..e4b2d8d8b 100644 --- a/vectordb_bench/backend/clients/vectorchord/vectorchord.py +++ b/vectordb_bench/backend/clients/vectorchord/vectorchord.py @@ -12,6 +12,7 @@ from psycopg import Connection, Cursor, sql from ...filter import Filter, FilterOp +from ...utils import redact_sensitive from ..api import VectorDB from .config import VectorChordConfigDict, VectorChordIndexConfig @@ -62,7 +63,7 @@ def __init__( self.cursor.execute("CREATE EXTENSION IF NOT EXISTS vchord CASCADE") self.conn.commit() - log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}") + log.info(f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}") if not any( ( self.case_config.create_index_before_load, @@ -71,7 +72,7 @@ def __init__( ): msg = ( f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.db_config}\n{self.case_config}" + f"{self.name} config values: {redact_sensitive(self.db_config)}\n{self.case_config}" ) log.error(msg) raise RuntimeError(msg) diff --git a/vectordb_bench/backend/utils.py b/vectordb_bench/backend/utils.py index 8c9c99a1f..b361a676d 100644 --- a/vectordb_bench/backend/utils.py +++ b/vectordb_bench/backend/utils.py @@ -6,6 +6,7 @@ from collections.abc import Callable from contextlib import contextmanager from functools import wraps +from typing import Any import psutil @@ -147,3 +148,44 @@ def compose_gt_file(filters: float | str | None = None) -> str: msg = f"Filters not supported: {filters}" raise ValueError(msg) + + +MASK = "**********" + +SENSITIVE_KEYS = frozenset({"api_key", "password", "token"}) + + +def redact_sensitive(value: Any) -> Any: + """Recursively replace credentials in a config structure with a fixed mask. + + ``DBConfig.to_dict()`` hands the driver its plaintext credentials, so its result + must be redacted before it reaches a log record or a result file. + + Examples: + >>> redact_sensitive({"host": "h", "password": "s3cret"}) + {'host': 'h', 'password': '**********'} + >>> redact_sensitive({"http_auth": ("admin", "s3cret")}) + {'http_auth': ('admin', '**********')} + """ + if isinstance(value, dict): + return {key: _redact_entry(key, item) for key, item in value.items()} + if isinstance(value, list): + return [redact_sensitive(item) for item in value] + return value + + +def _redact_entry(key: Any, value: Any) -> Any: + name = key.lower() if isinstance(key, str) else key + if name in SENSITIVE_KEYS and value: + return MASK + if name == "http_auth": + return _redact_http_auth(value) + return redact_sensitive(value) + + +def _redact_http_auth(value: Any) -> Any: + """http_auth is either a (user, secret) pair or an opaque auth object (e.g. AWS4Auth).""" + if isinstance(value, tuple | list) and len(value) == 2: + redacted = (value[0], MASK) + return list(redacted) if isinstance(value, list) else redacted + return MASK if value else value diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 4e471363a..27b633c10 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -3,12 +3,13 @@ from dataclasses import asdict from datetime import date, datetime from enum import Enum, StrEnum -from typing import Any, ClassVar, Literal, Self +from typing import Literal, Self import ujson from pydantic import ConfigDict, PositiveInt, field_validator, model_validator from vectordb_bench.backend.dataset import DatasetWithSizeMap +from vectordb_bench.backend.utils import redact_sensitive from . import config from .backend.cases import Case, CaseType, PerformanceCase, type2case @@ -388,22 +389,6 @@ class TestResult(BaseModel): file_fmt: str = "result_{}_{}_{}.json" # result_20230718_statndard_milvus.json timestamp: float = 0.0 - sensitive_output_fields: ClassVar[set[str]] = {"api_key", "password", "token"} - - @classmethod - def _redact_sensitive_fields(cls, value: Any) -> Any: - if isinstance(value, dict): - return { - key: ( - "**********" - if isinstance(key, str) and key.lower() in cls.sensitive_output_fields and item - else cls._redact_sensitive_fields(item) - ) - for key, item in value.items() - } - if isinstance(value, list): - return [cls._redact_sensitive_fields(item) for item in value] - return value @staticmethod def _output_metrics_for_case(case_result: CaseResult) -> dict: @@ -449,7 +434,7 @@ def model_dump_for_output(self) -> dict: for idx, case_result in enumerate(self.results): output["results"][idx]["metrics"] = self._output_metrics_for_case(case_result) output["results"][idx]["task_config"]["case_config"] = self._output_case_config_for_case(case_result) - return self._redact_sensitive_fields(output) + return redact_sensitive(output) def flush(self): db2case = self.get_db_results()