Add Azure-identity to pyproject.toml

.gitattributes

@@ -0,0 +1,13 @@
+# Thanks to: https://rehansaeed.com/gitattributes-best-practices/
+
+# Set default behavior to automatically normalize line endings.
+* text=auto
+
+# Force batch scripts to always use CRLF line endings so that if a repo is accessed
+# in Windows via a file share from Linux, the scripts will work.
+*.{cmd,[cC][mM][dD]} text eol=crlf
+*.{bat,[bB][aA][tT]} text eol=crlf
+
+# Force bash scripts to always use LF line endings so that if a repo is accessed
+# in Unix via a file share from Windows, the scripts will work.
+*.sh text eol=lf

.gitignore

@@ -9,4 +9,5 @@ __MACOSX
 build/
 venv/
 .idea/
-results/
+results/
+.venv/

README-MSSQLmd

@@ -0,0 +1,63 @@
+# Run VectorDBBench against MSSQL database
+
+VectorDBBench has been tested running on WSL2 + Ubuntu 22.04.4 LTS.
+
+## Install ODBC 
+
+Follow instructions here: https://learn.microsoft.com/en-us/sql/connect/odbc/linux-mac/installing-the-microsoft-odbc-driver-for-sql-server
+
+## Install Python 3.11
+
+Follow instructions here: https://ubuntuhandbook.org/index.php/2022/10/python-3-11-released-how-install-ubuntu/)
+
+## Install pip for Python3.11 :
+
+Use the following commands:
+
+```
+sudo apt install python3.11 python3.11-distutils python3.11-venv
+curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11
+```
+
+## Clone the repository
+
+```
+git clone https://github.com/MSSQL-VectorDBBench/VectorDBBench
+```
+
+Clone the repository into a local folder
+
+## Install VectorDBBench dependencies
+
+Change directoies into VectorDBBench and Install the VectorDBBench and its dependencies
+
+```
+cd VectorDBBench
+pip install pyodbc
+pip install .
+```
+
+
+## Run VectorDBBench on the Command Line Interface with help
+
+```
+vectordbbench mssql --help
+```
+
+## Run VectorDBBench on the Command Line Interface
+The database must exist and there must be enough room to build the index
+
+```
+vectordbbench mssql --database=vectordb --server=**IP_ADDRESS** --uid=sa --pwd=**PASSWORD_HERE**  --concurrency-duration=1800 --skip-search-concurrent --case-type=Performance1536D500K 
+```
+
+## Run VectorDBBench on the Command Line Interface with Existing Data
+
+```
+vectordbbench mssql --database=vectordb --server=**IP_ADDRESS** --uid=sa --pwd=**PASSWORD_HERE**  --concurrency-duration=1800 --skip-search-concurrent --case-type=Performance1536D500K 
+```
+
+## Start VectorDBBench in the GUI Mode
+```
+python -m vectordb_bench
+```

pyproject.toml

@@ -15,14 +15,14 @@ authors = [
   {name="XuanYang-cn", email="xuan.yang@zilliz.com"},
 ]
 description = "VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze."
-
 readme = "README.md"
 requires-python = ">=3.11"
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
 ]
+
 dependencies = [
     "click",
     "pytz",
@@ -35,7 +35,7 @@ dependencies = [
     "psutil",
     "polars",
     "plotly",
-    "environs",
+    "environs<14.0.1",
     "pydantic<v2",
     "scikit-learn",
     "pymilvus", # with pandas, numpy, ujson
@@ -68,6 +68,8 @@ all = [
     "memorydb",
     "alibabacloud_ha3engine_vector",
     "alibabacloud_searchengine20211025",
+    "pyodbc",
+    "azure-identity"
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -85,6 +87,7 @@ memorydb        = [ "memorydb" ]
 chromadb        = [ "chromadb" ]
 opensearch      = [ "opensearch-py" ]
 aliyun_opensearch = [ "alibabacloud_ha3engine_vector", "alibabacloud_searchengine20211025"]
+mssql = [ "pyodbc", "azure-identity" ]
 
 [project.urls]
 "repository" = "https://github.com/zilliztech/VectorDBBench"
@@ -207,3 +210,4 @@ builtins-ignorelist = [
     # "filter",
 ]
 
+

vectordb_bench/backend/clients/__init__.py

@@ -39,11 +39,16 @@ class DB(Enum):
     AWSOpenSearch = "OpenSearch"
     AliyunElasticsearch = "AliyunElasticsearch"
     Test = "test"
-    AliyunOpenSearch = "AliyunOpenSearch"
+    AliyunOpenSearch = "AliyunOpenSearch"    
+    MSSQL = "MSSQL"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912
         """Import while in use"""
+        if self == DB.MSSQL:
+            from .mssql.mssql import MSSQL
+            return MSSQL
+
         if self == DB.Milvus:
             from .milvus.milvus import Milvus
 
@@ -135,6 +140,10 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912
     @property
     def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912
         """Import while in use"""
+        if self == DB.MSSQL:
+            from .mssql.config import MSSQLConfig
+            return MSSQLConfig
+
         if self == DB.Milvus:
             from .milvus.config import MilvusConfig
 
@@ -227,6 +236,11 @@ def case_config_cls(  # noqa: PLR0911
         self,
         index_type: IndexType | None = None,
     ) -> type[DBCaseConfig]:
+
+        if self == DB.MSSQL:
+            from .mssql.config import MSSQLVectorIndexConfig
+            return MSSQLVectorIndexConfig
+
         if self == DB.Milvus:
             from .milvus.config import _milvus_case_config
 

vectordb_bench/backend/clients/mssql/cli.py

@@ -0,0 +1,56 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class MSSQLTypedDict(CommonTypedDict):
+    server: Annotated[
+        str, click.option("--server", type=str, help="server url", required=True)
+    ]
+    database: Annotated[
+        str,
+        click.option("--database", type=str, help="database name", required=True),
+    ]
+    uid: Annotated[
+        str,
+        click.option("--uid", type=str, help="User id", required=False),
+    ]
+    pwd: Annotated[
+        str,
+        click.option("--pwd", type=str, help="user password", required=False),
+    ]
+    entraid: Annotated[
+        str,
+        click.option("--entraid", type=str, help="Entra Id Authentication", required=False),
+    ]
+
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MSSQLTypedDict)
+def MSSQL(**parameters: Unpack[MSSQLTypedDict]):
+    from .config import MSSQLConfig, MSSQLVectorIndexConfig
+
+    run(
+        db=DB.MSSQL,
+        db_config=MSSQLConfig(
+            server=parameters["server"],
+            database=parameters["database"],
+            uid=parameters["uid"],
+            pwd=parameters["pwd"],
+            entraid=parameters["entraid"]
+        ),
+        db_case_config=MSSQLVectorIndexConfig(
+
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/mssql/config.py

@@ -0,0 +1,102 @@
+import pyodbc
+import struct
+import logging
+from azure.identity import ManagedIdentityCredential
+from pydantic import BaseModel, SecretStr
+from typing import Optional
+from ..api import DBConfig, DBCaseConfig, MetricType
+
+log = logging.getLogger(__name__)
+
+MSSQL_CONNECTION_STRING_PLACEHOLDER="DRIVER={ODBC Driver 18 for SQL Server};SERVER=%s;DATABASE=%s;UID=%s;PWD=%s;LongAsMax=yes;Connect Timeout=30;TrustServerCertificate=Yes"
+
+# --- Constants for Token Authentication ---
+SQL_COPT_SS_ACCESS_TOKEN = 1256 
+SQL_SERVER_TOKEN_SCOPE = "https://database.windows.net/.default"
+
+# --- Your Modified MSSQLConfig Class ---
+
+class MSSQLConfig(DBConfig):
+    server: str
+    database: str
+    uid: Optional[str] = None
+    pwd: Optional[SecretStr] = None
+    entraid: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        """
+        Prepares connection parameters. If entraid is provided, it fetches a token
+        manually and returns connection attributes for pyodbc.
+        """
+        # --- Case 1: Standard SQL Authentication ---
+        if self.entraid is None:
+            if not self.uid or not self.pwd:
+                log.error("UID and PWD must be provided for standard SQL auth.")
+
+            pwd_str = self.pwd.get_secret_value()
+            connection_string = (
+                f"DRIVER={{ODBC Driver 18 for SQL Server}};"
+                f"SERVER={self.server};"
+                f"DATABASE={self.database};"
+                f"UID={self.uid};"
+                f"PWD={pwd_str};"
+                "LongAsMax=yes;"
+                "Connect Timeout=30;"
+                "Encrypt=yes;"
+                "TrustServerCertificate=Yes"
+            )
+            return {"connection_string": connection_string}
+
+        # --- Case 2: Entra ID Managed Identity (Manual Token Auth) ---
+        log.info(f"Attempting to get token for User-Assigned Identity: {self.entraid}")
+
+        # 1. Get credentials and token using azure-identity
+        credential = ManagedIdentityCredential(client_id=self.entraid)
+        access_token = credential.get_token(SQL_SERVER_TOKEN_SCOPE)
+        token_bytes = access_token.token.encode("UTF-16-LE")
+
+        # 2. Pack the token for the driver
+        token_struct = struct.pack(f'<I{len(token_bytes)}s', len(token_bytes), token_bytes)
+
+        log.info("Token acquired successfully.")
+
+        # 3. Create the connection string WITHOUT auth keywords (UID, PWD, AUTHENTICATION)
+        connection_string = (
+            f"DRIVER={{ODBC Driver 18 for SQL Server}};"
+            f"SERVER={self.server};"
+            f"DATABASE={self.database};"
+            "LongAsMax=yes;"
+            "Connect Timeout=30;"
+            "Encrypt=yes;"
+            "TrustServerCertificate=Yes"
+        )
+
+        # 4. Return both the string and the token attributes
+        return {
+            "connection_string": connection_string,
+            "attrs_before": {SQL_COPT_SS_ACCESS_TOKEN: token_struct}
+        }
+
+
+class MSSQLVectorIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    efSearch: int | None = 48
+
+    def parse_metric(self) -> str: 
+        if self.metric_type == MetricType.L2:
+            return "euclidean"
+        elif self.metric_type == MetricType.IP:
+            return "dot"
+        return "cosine"
+
+    def index_param(self) -> dict:
+        return {
+            "lists" : self.lists,
+            "metric" : self.parse_metric()
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "efSearch" : self.efSearch,
+            "metric" : self.parse_metric()
+        }