Skip to content

Commit c5efecd

Browse files
simonhearneclaude
authored and
ntkathole
committed
feat: Support token, uri and db_name in the Milvus online store
The Milvus client was always built from host:port with a username:password token, so Zilliz Cloud API keys could not be used and every collection landed in the default database. Adds three optional config fields: - token: an API key or username:password; takes precedence over username/password. - uri: a full endpoint; takes precedence over host/port (and over a Milvus Lite path). - db_name: the Milvus database to use. Existing configurations build exactly the same client arguments. Co-Authored-By: Claude Opus 5.5 <[email protected]> Signed-off-by: Simon Hearne <[email protected]>
1 parent b8989cb commit c5efecd

4 files changed

Lines changed: 150 additions & 24 deletions

File tree

‎docs/reference/online-stores/milvus.md‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,14 +58,36 @@ online_store:
5858
```
5959
{% endcode %}
6060
61+
Connecting to [Zilliz Cloud](https://zilliz.com/cloud) (managed Milvus) with an API key.
62+
Read the token from an environment variable rather than committing it:
63+
64+
{% code title="feature_store.yaml" %}
65+
```yaml
66+
project: my_feature_repo
67+
registry: data/registry.db
68+
provider: local
69+
online_store:
70+
type: milvus
71+
uri: "https://<your-cluster-endpoint>" # Public Endpoint from the Zilliz Cloud console
72+
token: ${ZILLIZ_TOKEN} # pragma: allowlist secret
73+
db_name: "default"
74+
embedding_dim: 768
75+
index_type: "AUTOINDEX"
76+
metric_type: "COSINE"
77+
```
78+
{% endcode %}
79+
6180
## Configuration options
6281
6382
| Option | Default | Description |
6483
|:-------|:--------|:------------|
6584
| `path` | `""` | Path to a Milvus Lite database file. Used when `provider: local` and `path` is set. |
6685
| `host` | `http://localhost` | Milvus server host, including the scheme. |
6786
| `port` | `19530` | Milvus server port. |
87+
| `uri` | unset | Full endpoint, e.g. `https://<cluster>.zillizcloud.com:19530`. Takes precedence over `host`/`port`, and over `path`. |
6888
| `username` / `password` | `""` | Credentials, sent as the token `username:password`. |
89+
| `token` | unset | API key or `username:password`. Takes precedence over `username`/`password`. |
90+
| `db_name` | unset | Milvus database to use. The database must already exist. Defaults to the server's `default` database. |
6991
| `embedding_dim` | `128` | Dimension of vector fields. |
7092
| `index_type` | `FLAT` | Index type for vector fields with `vector_index=True`. |
7193
| `metric_type` | `COSINE` | Default metric when a field does not set `vector_search_metric`. |

‎sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,12 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
212212
path: Optional[StrictStr] = ""
213213
host: Optional[StrictStr] = "http://localhost"
214214
port: Optional[int] = 19530
215+
# Full endpoint, e.g. a Zilliz Cloud https URI. Takes precedence over host/port.
216+
uri: Optional[StrictStr] = None
217+
# API key or "username:password". Takes precedence over username/password.
218+
token: Optional[StrictStr] = None
219+
# Milvus database to use. The database must already exist.
220+
db_name: Optional[StrictStr] = None
215221
index_type: Optional[str] = "FLAT"
216222
metric_type: Optional[str] = "COSINE"
217223
embedding_dim: Optional[int] = 128
@@ -258,22 +264,25 @@ def _get_db_path(self, config: RepoConfig) -> str:
258264

259265
def _connect(self, config: RepoConfig) -> MilvusClient:
260266
if not self.client:
261-
if config.provider == "local" and config.online_store.path:
267+
online_config = config.online_store
268+
if (
269+
config.provider == "local"
270+
and online_config.path
271+
and not online_config.uri
272+
):
262273
db_path = self._get_db_path(config)
263274
logger.info("Connecting to Milvus in local mode using %s", db_path)
264275
self.client = MilvusClient(db_path)
265276
else:
266-
logger.info(
267-
"Connecting to Milvus remotely at %s:%s",
268-
config.online_store.host,
269-
config.online_store.port,
270-
)
271-
self.client = MilvusClient(
272-
uri=f"{config.online_store.host}:{config.online_store.port}",
273-
token=f"{config.online_store.username}:{config.online_store.password}"
274-
if config.online_store.username and config.online_store.password
275-
else "",
276-
)
277+
uri = online_config.uri or f"{online_config.host}:{online_config.port}"
278+
logger.info("Connecting to Milvus remotely at %s", uri)
279+
client_kwargs: Dict[str, Any] = {
280+
"uri": uri,
281+
"token": _milvus_token(online_config),
282+
}
283+
if online_config.db_name:
284+
client_kwargs["db_name"] = online_config.db_name
285+
self.client = MilvusClient(**client_kwargs)
277286
return self.client
278287

279288
def _get_or_create_collection(
@@ -1019,6 +1028,15 @@ def _table_id(project: str, table: FeatureView, enable_versioning: bool = False)
10191028
return compute_table_id(project, table, enable_versioning)
10201029

10211030

1031+
def _milvus_token(online_config: MilvusOnlineStoreConfig) -> str:
1032+
"""Return the token to authenticate with: ``token``, else ``username:password``."""
1033+
if online_config.token:
1034+
return online_config.token
1035+
if online_config.username and online_config.password:
1036+
return f"{online_config.username}:{online_config.password}"
1037+
return ""
1038+
1039+
10221040
def _get_composite_key_name(table: FeatureView) -> str:
10231041
return "_".join([field.name for field in table.entity_columns]) + "_pk"
10241042

‎sdk/python/tests/integration/online_store/test_milvus_remote.py‎

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@
55
ZILLIZ_URI=https://<cluster>.zillizcloud.com ZILLIZ_TOKEN=<user:password> \
66
pytest --integration sdk/python/tests/integration/online_store/test_milvus_remote.py
77
8-
They also run against a self-hosted Milvus server, e.g.
8+
``ZILLIZ_TOKEN`` can be an API key or ``username:password``. The tests also run
9+
against a self-hosted Milvus server, e.g.
910
``ZILLIZ_URI=http://localhost:19530 ZILLIZ_TOKEN=root:Milvus``.
1011
"""
1112

@@ -16,7 +17,6 @@
1617
from pathlib import Path
1718
from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar
1819
from unittest.mock import patch
19-
from urllib.parse import urlparse
2020

2121
import pytest
2222

@@ -44,16 +44,7 @@
4444

4545

4646
def _connection_config() -> Dict[str, Any]:
47-
assert ZILLIZ_URI and ZILLIZ_TOKEN
48-
parsed = urlparse(ZILLIZ_URI)
49-
default_port = 443 if parsed.scheme == "https" else 19530
50-
username, _, password = ZILLIZ_TOKEN.partition(":")
51-
return {
52-
"host": f"{parsed.scheme}://{parsed.hostname}",
53-
"port": parsed.port or default_port,
54-
"username": username,
55-
"password": password,
56-
}
47+
return {"uri": ZILLIZ_URI, "token": ZILLIZ_TOKEN}
5748

5849

5950
def _repo_config(tmp_path: Path, project: str, **online_store: Any) -> RepoConfig:
@@ -277,3 +268,24 @@ def test_released_collection_is_loaded_once(
277268

278269
assert rows[0] is not None and rows[0]["city"].string_val == "Paris"
279270
assert load_spy.call_count == 1
271+
272+
273+
def test_db_name(tmp_path: Path, project: str, store: MilvusOnlineStore) -> None:
274+
from pymilvus import MilvusClient
275+
276+
db_name = os.environ.get("ZILLIZ_DB_NAME", "feast_it_db")
277+
admin = MilvusClient(uri=ZILLIZ_URI, token=ZILLIZ_TOKEN)
278+
try:
279+
if db_name not in admin.list_databases():
280+
admin.create_database(db_name)
281+
except Exception as e:
282+
pytest.skip(f"Cannot create database {db_name!r} on this cluster: {e}")
283+
284+
config = _repo_config(tmp_path, project, db_name=db_name)
285+
fv = _scalar_feature_view()
286+
store.update(config, [], [fv], [], [], partial=False)
287+
288+
collection_name = f"{project}_{fv.name}"
289+
db_client = MilvusClient(uri=ZILLIZ_URI, token=ZILLIZ_TOKEN, db_name=db_name)
290+
assert collection_name in db_client.list_collections()
291+
assert collection_name not in admin.list_collections()

‎sdk/python/tests/unit/infra/online_store/test_milvus_online_store.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,3 +301,77 @@ def test_load_collection_not_called_per_query(tmp_path: Path) -> None:
301301
)
302302

303303
assert load_spy.call_count == 0
304+
305+
306+
@pytest.mark.parametrize(
307+
"online_store, expected_kwargs",
308+
[
309+
# Existing configs build exactly the same client arguments as before.
310+
({}, {"uri": "http://localhost:19530", "token": ""}),
311+
(
312+
{
313+
"host": "https://milvus.internal",
314+
"port": 443,
315+
"username": "u",
316+
"password": "p",
317+
},
318+
{"uri": "https://milvus.internal:443", "token": "u:p"},
319+
),
320+
# token takes precedence over username/password.
321+
(
322+
{"username": "u", "password": "p", "token": "api-key"},
323+
{"uri": "http://localhost:19530", "token": "api-key"},
324+
),
325+
# uri takes precedence over host/port.
326+
(
327+
{
328+
"uri": "https://in01-abc.zillizcloud.com",
329+
"host": "http://ignored",
330+
"token": "k",
331+
},
332+
{"uri": "https://in01-abc.zillizcloud.com", "token": "k"},
333+
),
334+
(
335+
{
336+
"uri": "https://in01-abc.zillizcloud.com",
337+
"token": "k",
338+
"db_name": "catalog",
339+
},
340+
{
341+
"uri": "https://in01-abc.zillizcloud.com",
342+
"token": "k",
343+
"db_name": "catalog",
344+
},
345+
),
346+
],
347+
)
348+
@patch(f"{MILVUS_MODULE}.MilvusClient")
349+
def test_remote_client_arguments(
350+
mock_client_cls: MagicMock,
351+
online_store: Dict[str, Any],
352+
expected_kwargs: Dict[str, Any],
353+
) -> None:
354+
config = _mock_config(**online_store)
355+
config.provider = "gcp"
356+
357+
MilvusOnlineStore()._connect(config)
358+
359+
mock_client_cls.assert_called_once_with(**expected_kwargs)
360+
361+
362+
@patch(f"{MILVUS_MODULE}.MilvusClient")
363+
def test_uri_takes_precedence_over_lite_path(mock_client_cls: MagicMock) -> None:
364+
config = _mock_config(path="online_store.db", uri="http://milvus:19530", token="k")
365+
366+
MilvusOnlineStore()._connect(config)
367+
368+
mock_client_cls.assert_called_once_with(uri="http://milvus:19530", token="k")
369+
370+
371+
@patch(f"{MILVUS_MODULE}.MilvusClient")
372+
def test_lite_path_used_without_uri(mock_client_cls: MagicMock) -> None:
373+
config = _mock_config(path="/tmp/online_store.db", token="ignored")
374+
375+
MilvusOnlineStore()._connect(config)
376+
377+
mock_client_cls.assert_called_once_with("/tmp/online_store.db")

0 commit comments

Comments
 (0)