Skip to content

Commit

Permalink
Add Milvus auth support through user_name and password fields (#416)
Browse files Browse the repository at this point in the history
* tweak(milvus): add auth support through user_name and password fields (like zilliz cloud)

* tweak(frontend): make username and password for milvus optional fields
  • Loading branch information
teynar authored Dec 4, 2024
1 parent 4bb2994 commit 6a63416
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 3 deletions.
26 changes: 25 additions & 1 deletion vectordb_bench/backend/clients/milvus/cli.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Annotated, TypedDict, Unpack
from typing import Annotated, TypedDict, Unpack, Optional

import click
from pydantic import SecretStr
Expand All @@ -21,6 +21,12 @@ class MilvusTypedDict(TypedDict):
uri: Annotated[
str, click.option("--uri", type=str, help="uri connection string", required=True)
]
user_name: Annotated[
Optional[str], click.option("--user-name", type=str, help="Db username", required=False)
]
password: Annotated[
Optional[str], click.option("--password", type=str, help="Db password", required=False)
]


class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict):
Expand All @@ -37,6 +43,8 @@ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=AutoIndexConfig(),
**parameters,
Expand All @@ -53,6 +61,8 @@ def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=FLATConfig(),
**parameters,
Expand All @@ -73,6 +83,8 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]) if parameters["password"] else None,
),
db_case_config=HNSWConfig(
M=parameters["m"],
Expand All @@ -97,6 +109,8 @@ def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=IVFFlatConfig(
nlist=parameters["nlist"],
Expand All @@ -116,6 +130,8 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=IVFSQ8Config(
nlist=parameters["nlist"],
Expand Down Expand Up @@ -143,6 +159,8 @@ def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=DISKANNConfig(
search_list=parameters["search_list"],
Expand Down Expand Up @@ -174,6 +192,8 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=GPUIVFFlatConfig(
nlist=parameters["nlist"],
Expand Down Expand Up @@ -208,6 +228,8 @@ def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=GPUIVFPQConfig(
nlist=parameters["nlist"],
Expand Down Expand Up @@ -274,6 +296,8 @@ def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
db_config=MilvusConfig(
db_label=parameters["db_label"],
uri=SecretStr(parameters["uri"]),
user=parameters["user_name"],
password=SecretStr(parameters["password"]),
),
db_case_config=GPUCAGRAConfig(
intermediate_graph_degree=parameters["intermediate_graph_degree"],
Expand Down
18 changes: 16 additions & 2 deletions vectordb_bench/backend/clients/milvus/config.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,26 @@
from pydantic import BaseModel, SecretStr
from pydantic import BaseModel, SecretStr, validator
from ..api import DBConfig, DBCaseConfig, MetricType, IndexType


class MilvusConfig(DBConfig):
uri: SecretStr = "http://localhost:19530"
user: str | None = None
password: SecretStr | None = None

def to_dict(self) -> dict:
return {"uri": self.uri.get_secret_value()}
return {
"uri": self.uri.get_secret_value(),
"user": self.user if self.user else None,
"password": self.password.get_secret_value() if self.password else None,
}

@validator("*")
def not_empty_field(cls, v, field):
if field.name in cls.common_short_configs() or field.name in cls.common_long_configs() or field.name in ["user", "password"]:
return v
if isinstance(v, (str, SecretStr)) and len(v) == 0:
raise ValueError("Empty string!")
return v


class MilvusIndexConfig(BaseModel):
Expand Down

0 comments on commit 6a63416

Please sign in to comment.