fix: Updating milvus connect function to work with remote instance (#5401)

Fiona-Waters · web-flow · commit b89faddf4308 · 2025-05-30T09:50:51.000-04:00
* Updating milvus connect function to work with remote instance

Signed-off-by: Fiona Waters &lt;fiwaters6@gmail.com&gt;

* Update test configuration to use path for db

Signed-off-by: Fiona Waters &lt;fiwaters6@gmail.com&gt;

---------

Signed-off-by: Fiona Waters &lt;fiwaters6@gmail.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py b/sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py
@@ -88,8 +88,8 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
     """
 
     type: Literal["milvus"] = "milvus"
-    path: Optional[StrictStr] = "online_store.db"
-    host: Optional[StrictStr] = "localhost"
+    path: Optional[StrictStr] = ""
+    host: Optional[StrictStr] = "https://localhost"
     port: Optional[int] = 19530
     index_type: Optional[str] = "FLAT"
     metric_type: Optional[str] = "COSINE"
@@ -126,13 +126,16 @@ def _get_db_path(self, config: RepoConfig) -> str:
 
     def _connect(self, config: RepoConfig) -> MilvusClient:
         if not self.client:
-            if config.provider == "local":
+            if config.provider == "local" and config.online_store.path:
                 db_path = self._get_db_path(config)
                 print(f"Connecting to Milvus in local mode using {db_path}")
                 self.client = MilvusClient(db_path)
             else:
+                print(
+                    f"Connecting to Milvus remotely at {config.online_store.host}:{config.online_store.port}"
+                )
                 self.client = MilvusClient(
-                    url=f"{config.online_store.host}:{config.online_store.port}",
+                    uri=f"{config.online_store.host}:{config.online_store.port}",
                     token=f"{config.online_store.username}:{config.online_store.password}"
                     if config.online_store.username and config.online_store.password
                     else "",
diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py
@@ -89,7 +89,7 @@
 )
 
 DYNAMO_CONFIG = {"type": "dynamodb", "region": "us-west-2"}
-MILVUS_CONFIG = {"type": "milvus", "embedding_dim": "2"}
+MILVUS_CONFIG = {"type": "milvus", "embedding_dim": 2, "path": "online_store.db"}
 REDIS_CONFIG = {"type": "redis", "connection_string": "localhost:6379,db=0"}
 REDIS_CLUSTER_CONFIG = {
     "type": "redis",
diff --git a/sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py b/sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py
@@ -1,9 +1,5 @@
 from typing import Any, Dict
 
-import docker
-from testcontainers.core.container import DockerContainer
-from testcontainers.core.waiting_utils import wait_for_logs
-
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
@@ -12,33 +8,14 @@
 class MilvusOnlineStoreCreator(OnlineStoreCreator):
     def __init__(self, project_name: str, **kwargs):
         super().__init__(project_name)
-        self.fixed_port = 19530
-        self.container = DockerContainer("milvusdb/milvus:v2.4.9").with_exposed_ports(
-            self.fixed_port
-        )
-        self.client = docker.from_env()
 
     def create_online_store(self) -> Dict[str, Any]:
-        self.container.start()
-        # Wait for Milvus server to be ready
-        # log_string_to_wait_for = "Ready to accept connections"
-        log_string_to_wait_for = ""
-        wait_for_logs(
-            container=self.container, predicate=log_string_to_wait_for, timeout=30
-        )
-        host = "localhost"
-        port = self.container.get_exposed_port(self.fixed_port)
         return {
             "type": "milvus",
-            "host": host,
-            "port": int(port),
             "path": "online_store.db",
             "index_type": "IVF_FLAT",
             "metric_type": "L2",
             "embedding_dim": 2,
             "vector_enabled": True,
             "nlist": 1,
         }
-
-    def teardown(self):
-        self.container.stop()

Original file line number	Diff line number	Diff line change
`@@ -89,7 +89,7 @@`
`89`	`89`	`)`
`90`	`90`
`91`	`91`	`DYNAMO_CONFIG = {"type": "dynamodb", "region": "us-west-2"}`
`92`		`-MILVUS_CONFIG = {"type": "milvus", "embedding_dim": "2"}`
	`92`	`+MILVUS_CONFIG = {"type": "milvus", "embedding_dim": 2, "path": "online_store.db"}`
`93`	`93`	`REDIS_CONFIG = {"type": "redis", "connection_string": "localhost:6379,db=0"}`
`94`	`94`	`REDIS_CLUSTER_CONFIG = {`
`95`	`95`	`"type": "redis",`