From d842e9ac3f52bb33956b5d476fabf829bbeed4cc Mon Sep 17 00:00:00 2001
From: "junjie.jiang" <junjie.jiang@zilliz.com>
Date: Mon, 5 Aug 2024 10:57:59 +0800
Subject: [PATCH] Limit milvus-lite files to end with .db

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
---
 pymilvus/orm/connections.py | 12 +++++-------
 tests/test_milvus_lite.py   |  8 ++++++++
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/pymilvus/orm/connections.py b/pymilvus/orm/connections.py
index 48fde7ef2..1aa119a2f 100644
--- a/pymilvus/orm/connections.py
+++ b/pymilvus/orm/connections.py
@@ -241,11 +241,7 @@ def add_connection(self, **kwargs):
             self._alias[alias] = alias_config
 
     def __get_full_address(
-        self,
-        address: str = "",
-        uri: str = "",
-        host: str = "",
-        port: str = "",
+        self, address: str = "", uri: str = "", host: str = "", port: str = "",
     ) -> (str, parse.ParseResult):
         if address != "":
             if not is_legal_address(address):
@@ -365,10 +361,11 @@ def connect(
             "tcp",
         ]:
             # start and connect milvuslite
-            if kwargs["uri"].endswith("/"):
+            if not kwargs["uri"].endswith(".db"):
                 raise ConnectionConfigException(
-                    message=f"Open local milvus failed, {kwargs['uri']} is not a local file path"
+                    message=f"Open local milvus failed, {kwargs['uri']} is not endswith .db"
                 )
+
             parent_path = pathlib.Path(kwargs["uri"]).parent
             if not parent_path.is_dir():
                 raise ConnectionConfigException(
@@ -377,6 +374,7 @@ def connect(
 
             from milvus_lite.server_manager import server_manager_instance
 
+            logger.info(f"Pass in the local path {kwargs['uri']}, and run it using milvus-lite")
             local_uri = server_manager_instance.start_and_get_uri(kwargs["uri"])
             if local_uri is None:
                 raise ConnectionConfigException(message="Open local milvus failed")
diff --git a/tests/test_milvus_lite.py b/tests/test_milvus_lite.py
index 0149fd5bc..7f000aa89 100644
--- a/tests/test_milvus_lite.py
+++ b/tests/test_milvus_lite.py
@@ -5,6 +5,7 @@
 import pytest
 
 from pymilvus.milvus_client import MilvusClient
+from pymilvus.exceptions import ConnectionConfigException
 
 
 @pytest.mark.skipif(sys.platform.startswith('win'), reason="Milvus Lite is not supported on Windows")
@@ -56,3 +57,10 @@ def test_milvus_lite(self):
                     filter="subject == 'history'",
                 )
             assert len(res) == 3
+
+    def test_illegal_name(self):
+        try:
+            MilvusClient("localhost")
+            assert False
+        except ConnectionConfigException as e:
+            assert e.message == "Open local milvus failed, localhost is not endswith .db"