Add TiDB backend

README.md

@@ -55,6 +55,7 @@ All the database client supported
 | awsopensearch            | `pip install vectordb-bench[opensearch]` |
 | aliyun_opensearch        | `pip install vectordb-bench[aliyun_opensearch]` |
 | mongodb                  | `pip install vectordb-bench[mongodb]`       |
+| tidb                     | `pip install vectordb-bench[tidb]`          |
 
 ### Run
 

pyproject.toml

@@ -69,6 +69,7 @@ all = [
     "alibabacloud_ha3engine_vector",
     "alibabacloud_searchengine20211025",
     "mariadb",
+    "PyMySQL",
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -87,7 +88,8 @@ chromadb        = [ "chromadb" ]
 opensearch      = [ "opensearch-py" ]
 aliyun_opensearch = [ "alibabacloud_ha3engine_vector", "alibabacloud_searchengine20211025"]
 mongodb         = [ "pymongo" ]
-mariadb = [ "mariadb" ]
+mariadb         = [ "mariadb" ]
+tidb            = [ "PyMySQL" ]
 
 [project.urls]
 "repository" = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/__init__.py

@@ -42,6 +42,7 @@ class DB(Enum):
     Test = "test"
     AliyunOpenSearch = "AliyunOpenSearch"
     MongoDB = "MongoDB"
+    TiDB = "TiDB"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
@@ -141,6 +142,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
 
             return MariaDB
 
+        if self == DB.TiDB:
+            from .tidb.tidb import TiDB
+
+            return TiDB
+
         if self == DB.Test:
             from .test.test import Test
 
@@ -244,8 +250,14 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901
 
         if self == DB.MariaDB:
             from .mariadb.config import MariaDBConfig
+
             return MariaDBConfig
 
+        if self == DB.TiDB:
+            from .tidb.config import TiDBConfig
+
+            return TiDBConfig
+
         if self == DB.Test:
             from .test.config import TestConfig
 
@@ -333,6 +345,11 @@ def case_config_cls(  # noqa: PLR0911
 
             return _mariadb_case_config.get(index_type)
 
+        if self == DB.TiDB:
+            from .tidb.config import TiDBIndexConfig
+
+            return TiDBIndexConfig
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

vectordb_bench/backend/clients/tidb/cli.py

@@ -0,0 +1,98 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+
+from ....cli.cli import CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, run
+
+
+class TiDBTypedDict(CommonTypedDict):
+    user_name: Annotated[
+        str,
+        click.option(
+            "--username",
+            type=str,
+            help="Username",
+            default="root",
+            show_default=True,
+            required=True,
+        ),
+    ]
+    password: Annotated[
+        str,
+        click.option(
+            "--password",
+            type=str,
+            default="",
+            show_default=True,
+            help="Password",
+        ),
+    ]
+    host: Annotated[
+        str,
+        click.option(
+            "--host",
+            type=str,
+            default="127.0.0.1",
+            show_default=True,
+            required=True,
+            help="Db host",
+        ),
+    ]
+    port: Annotated[
+        int,
+        click.option(
+            "--port",
+            type=int,
+            default=4000,
+            show_default=True,
+            required=True,
+            help="Db Port",
+        ),
+    ]
+    db_name: Annotated[
+        str,
+        click.option(
+            "--db-name",
+            type=str,
+            default="test",
+            show_default=True,
+            required=True,
+            help="Db name",
+        ),
+    ]
+    ssl: Annotated[
+        bool,
+        click.option(
+            "--ssl/--no-ssl",
+            default=False,
+            show_default=True,
+            is_flag=True,
+            help="Enable or disable SSL, for TiDB Serverless SSL must be enabled",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(TiDBTypedDict)
+def TiDB(
+    **parameters: Unpack[TiDBTypedDict],
+):
+    from .config import TiDBConfig, TiDBIndexConfig
+
+    run(
+        db=DB.TiDB,
+        db_config=TiDBConfig(
+            db_label=parameters["db_label"],
+            user_name=parameters["username"],
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            port=parameters["port"],
+            db_name=parameters["db_name"],
+            ssl=parameters["ssl"],
+        ),
+        db_case_config=TiDBIndexConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/tidb/config.py

@@ -0,0 +1,49 @@
+from pydantic import SecretStr, BaseModel, validator
+from ..api import DBConfig, DBCaseConfig, MetricType
+
+
+class TiDBConfig(DBConfig):
+    user_name: str = "root"
+    password: SecretStr
+    host: str = "127.0.0.1"
+    port: int = 4000
+    db_name: str = "test"
+    ssl: bool = False
+
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        return v
+
+    def to_dict(self) -> dict:
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "user": self.user_name,
+            "password": pwd_str,
+            "database": self.db_name,
+            "ssl_verify_cert": self.ssl,
+            "ssl_verify_identity": self.ssl,
+        }
+
+
+class TiDBIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+
+    def get_metric_fn(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "vec_l2_distance"
+        elif self.metric_type == MetricType.COSINE:
+            return "vec_cosine_distance"
+        else:
+            raise ValueError(f"Unsupported metric type: {self.metric_type}")
+
+    def index_param(self) -> dict:
+        return {
+            "metric_fn": self.get_metric_fn(),
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_fn": self.get_metric_fn(),
+        }