Improve code quality - base class, exceptions, imports, tests

.gitignore

@@ -135,4 +135,10 @@ dmypy.json
 
 # Test reports
 TEST_REPORT_*.md
-docs/plans/*
+docs/plans/*
+# IDE
+.cursor/
+.kiro/
+
+# Test env (may contain tokens)
+milvus_cli/test/test.env

README.md

@@ -27,31 +27,40 @@ milvus_cli/
 ├── Core Modules
 │   ├── main.py          # Main entry point
 │   ├── Cli.py           # CLI command interface
-│   ├── Connection.py    # Milvus connection management
-│   ├── Collection.py    # Collection operations
-│   ├── Database.py      # Database management
-│   ├── Index.py         # Index management
-│   ├── Partition.py     # Partition management
-│   ├── Data.py          # Data import/export
-│   ├── Role.py          # Role management
-│   ├── User.py          # User management
-│   ├── Alias.py         # Alias management
-│   ├── Fs.py            # File system operations
-│   ├── Types.py         # Data type definitions
-│   ├── utils.py         # Utility functions
-│   └── Validation.py    # Input validation
+│   ├── BaseClient.py    # Base class for all client modules
+│   ├── ConnectionClient.py # Milvus connection management
+│   ├── CollectionClient.py # Collection operations
+│   ├── DatabaseClient.py   # Database management
+│   ├── IndexClient.py      # Index management
+│   ├── PartitionClient.py  # Partition management
+│   ├── DataClient.py       # Data import/export/search
+│   ├── RoleClient.py       # Role management
+│   ├── UserClient.py       # User management
+│   ├── AliasClient.py      # Alias management
+│   ├── ResourceGroup.py    # Resource group management
+│   ├── PrivilegeGroup.py   # Privilege group management
+│   ├── CliClient.py        # Main CLI client (aggregates all modules)
+│   ├── OutputFormatter.py  # Output formatting (table/json/csv)
+│   ├── Fs.py               # File system operations
+│   ├── Types.py            # Data type definitions
+│   ├── utils.py            # Utility functions
+│   └── Validation.py       # Input validation
 ├── scripts/             # CLI command implementations
 │   ├── milvus_cli.py    # Main CLI script
-│   ├── connection_cli.py # Connection-related commands
-│   ├── collection_cli.py # Collection-related commands
-│   ├── database_cli.py  # Database-related commands
-│   ├── index_cli.py     # Index-related commands
-│   ├── partition_cli.py # Partition-related commands
-│   ├── data_cli.py      # Data-related commands
-│   ├── role_cli.py      # Role-related commands
-│   ├── user_cli.py      # User-related commands
-│   ├── alias_cli.py     # Alias-related commands
-│   └── helper_cli.py    # Helper commands
+│   ├── init_client_cli.py # CLI initialization and global state
+│   ├── helper_client_cli.py # REPL loop, command groups, output settings
+│   ├── helper_cli.py    # Help, version, history commands
+│   ├── connection_client_cli.py # Connection commands
+│   ├── collection_client_cli.py # Collection commands
+│   ├── database_client_cli.py  # Database commands
+│   ├── index_client_cli.py     # Index commands
+│   ├── partition_client_cli.py # Partition commands
+│   ├── data_client_cli.py      # Data commands
+│   ├── role_client_cli.py      # Role commands
+│   ├── user_client_cli.py      # User commands
+│   ├── alias_client_cli.py     # Alias commands
+│   ├── resource_group_cli.py   # Resource group commands
+│   └── privilege_group_cli.py  # Privilege group commands
 ├── test/                # Unit tests (internal APIs)
 │   ├── test_config.py
 │   ├── test_connection_client.py

TESTING.md

@@ -7,8 +7,9 @@ This project has two types of tests:
 ### 1. **Unit Tests** (`milvus_cli/test/`)
 - **Framework**: `unittest`
 - **Purpose**: Test Python modules and classes directly
-- **Target**: Internal APIs (`ConnectionClient`, `CollectionClient`, etc.)
-- **Run with**: `python -m unittest discover milvus_cli/test`
+- **Target**: Internal APIs (`BaseClient`, `ConnectionClient`, `CollectionClient`, etc.)
+- **Config**: `milvus_cli/test/test.env` (set `MILVUS_TEST_URI`, see `test.env.example`)
+- **Run with**: `python -m pytest milvus_cli/test/ -v`
 
 ### 2. **Integration Tests** (`tests/`)
 - **Framework**: `pytest`

milvus_cli/AliasClient.py

@@ -1,40 +1,18 @@
+from __future__ import annotations
+
 from pymilvus import MilvusClient
+try:
+    from .BaseClient import BaseMilvusClient
+except ImportError:
+    from BaseClient import BaseMilvusClient
 
 
-class MilvusClientAlias(object):
-    """
-    Alias operations class based on MilvusClient API
-    Used to replace the original Alias operations based on utility functions
-    """
-
-    def __init__(self, connection_client=None):
-        """
-        Initialize Alias client
-
-        Args:
-            connection_client: MilvusClientConnection instance
-        """
-        self.connection_client = connection_client
-        self.alias = "default"  # Keep for compatibility
+class MilvusClientAlias(BaseMilvusClient):
+    """Alias operations based on MilvusClient API."""
 
-    def _get_client(self):
-        """
-        Get MilvusClient instance
-
-        Returns:
-            MilvusClient instance
-
-        Raises:
-            Exception: If not connected or connection is invalid
-        """
-        if not self.connection_client:
-            raise Exception("Connection client not set!")
-
-        client = self.connection_client.get_client()
-        if not client:
-            raise Exception("Not connected to Milvus! Please connect first.")
-
-        return client
+    def __init__(self, connection_client=None) -> None:
+        super().__init__(connection_client)
+        self.alias = "default"
 
     def create_alias(self, collectionName, aliasName):
         """
@@ -52,7 +30,7 @@ def create_alias(self, collectionName, aliasName):
 
             # Check if alias already exists
             if self.has_alias(aliasName):
-                raise Exception(f"Alias '{aliasName}' already exists")
+                raise ValueError(f"Alias '{aliasName}' already exists")
 
             # Create alias using MilvusClient API
             client.create_alias(
@@ -63,7 +41,7 @@ def create_alias(self, collectionName, aliasName):
             return f"Create alias {aliasName} successfully!"
 
         except Exception as e:
-            raise Exception(f"Create alias error!{str(e)}")
+            raise RuntimeError(f"Create alias error: {e}") from e
 
     def list_aliases(self, collectionName=None):
         """
@@ -95,7 +73,7 @@ def list_aliases(self, collectionName=None):
                 return self.list_all_aliases()
 
         except Exception as e:
-            raise Exception(f"List alias error!{str(e)}")
+            raise RuntimeError(f"List alias error: {e}") from e
 
     def drop_alias(self, aliasName):
         """
@@ -112,15 +90,15 @@ def drop_alias(self, aliasName):
 
             # Check if alias exists first
             if not self.has_alias(aliasName):
-                raise Exception(f"Alias '{aliasName}' does not exist")
+                raise ValueError(f"Alias '{aliasName}' does not exist")
 
             # Drop alias using MilvusClient API
             client.drop_alias(alias=aliasName)
 
             return f"Drop alias {aliasName} successfully!"
 
         except Exception as e:
-            raise Exception(f"Drop alias error!{str(e)}")
+            raise RuntimeError(f"Drop alias error: {e}") from e
 
     def alter_alias(self, aliasName, collectionName):
         """
@@ -145,7 +123,7 @@ def alter_alias(self, aliasName, collectionName):
             return f"Alter alias {aliasName} successfully!"
 
         except Exception as e:
-            raise Exception(f"Alter alias error!{str(e)}")
+            raise RuntimeError(f"Alter alias error: {e}") from e
 
     def describe_alias(self, aliasName):
         """
@@ -166,7 +144,7 @@ def describe_alias(self, aliasName):
             return alias_info
 
         except Exception as e:
-            raise Exception(f"Describe alias error!{str(e)}")
+            raise RuntimeError(f"Describe alias error: {e}") from e
 
     def has_alias(self, aliasName):
         """
@@ -201,7 +179,7 @@ def get_alias_collection(self, aliasName):
             return ''
 
         except Exception as e:
-            raise Exception(f"Get alias collection error!{str(e)}")
+            raise RuntimeError(f"Get alias collection error: {e}") from e
 
     def list_all_aliases(self):
         """
@@ -226,13 +204,13 @@ def list_all_aliases(self):
                                 'alias': alias,
                                 'collection': collection
                             })
-                except:
+                except Exception:
                     continue  # Skip collections without aliases or with errors
 
             return all_aliases
 
         except Exception as e:
-            raise Exception(f"List all aliases error!{str(e)}")
+            raise RuntimeError(f"List all aliases error: {e}") from e
 
     def validate_alias_name(self, aliasName):
         """

milvus_cli/BaseClient.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pymilvus import MilvusClient
+    from ConnectionClient import MilvusClientConnection
+
+
+class BaseMilvusClient:
+    """Base class for all Milvus client modules."""
+
+    def __init__(self, connection_client: MilvusClientConnection | None = None) -> None:
+        self.connection_client = connection_client
+
+    def _get_client(self) -> MilvusClient:
+        if not self.connection_client:
+            raise ConnectionError("Connection client not set!")
+        client = self.connection_client.get_client()
+        if not client:
+            raise ConnectionError("Not connected to Milvus! Please connect first.")
+        return client

milvus_cli/CliClient.py

@@ -1,91 +1,30 @@
-from ConnectionClient import MilvusClientConnection
-from DatabaseClient import MilvusClientDatabase
-from CollectionClient import MilvusClientCollection
-from IndexClient import MilvusClientIndex
-from DataClient import MilvusClientData
-from UserClient import MilvusClientUser
-from AliasClient import MilvusClientAlias
-from PartitionClient import MilvusClientPartition
-from RoleClient import MilvusClientRole
-from ResourceGroup import MilvusResourceGroup
-from PrivilegeGroup import MilvusPrivilegeGroup
+try:
+    from .ConnectionClient import MilvusClientConnection
+    from .DatabaseClient import MilvusClientDatabase
+    from .CollectionClient import MilvusClientCollection
+    from .IndexClient import MilvusClientIndex
+    from .DataClient import MilvusClientData
+    from .UserClient import MilvusClientUser
+    from .AliasClient import MilvusClientAlias
+    from .PartitionClient import MilvusClientPartition
+    from .RoleClient import MilvusClientRole
+    from .ResourceGroup import MilvusResourceGroup
+    from .PrivilegeGroup import MilvusPrivilegeGroup
+    from .OutputFormatter import OutputFormatter
+except ImportError:
+    from ConnectionClient import MilvusClientConnection
+    from DatabaseClient import MilvusClientDatabase
+    from CollectionClient import MilvusClientCollection
+    from IndexClient import MilvusClientIndex
+    from DataClient import MilvusClientData
+    from UserClient import MilvusClientUser
+    from AliasClient import MilvusClientAlias
+    from PartitionClient import MilvusClientPartition
+    from RoleClient import MilvusClientRole
+    from ResourceGroup import MilvusResourceGroup
+    from PrivilegeGroup import MilvusPrivilegeGroup
+    from OutputFormatter import OutputFormatter
 from pymilvus import __version__
-from tabulate import tabulate
-import json
-import csv
-import io
-
-
-class OutputFormatter:
-    """
-    Formatter class for CLI output in different formats (table, json, csv)
-    """
-
-    def __init__(self, format="table"):
-        self.format = format
-
-    def format_output(self, data, headers=None):
-        """Format data based on current format setting."""
-        if not data:
-            return "No data found."
-
-        if self.format == "json":
-            return json.dumps(data, indent=2, default=str)
-        elif self.format == "csv":
-            return self._to_csv(data)
-        else:  # table
-            return self._to_table(data, headers)
-
-    def format_list(self, items, header="Item"):
-        """Format a simple list of items."""
-        if not items:
-            return "No items found."
-
-        if self.format == "json":
-            return json.dumps(items, indent=2, default=str)
-        elif self.format == "csv":
-            output = io.StringIO()
-            writer = csv.writer(output)
-            writer.writerow([header])
-            for item in items:
-                writer.writerow([item])
-            return output.getvalue()
-        else:  # table
-            table_data = [[item] for item in items]
-            return tabulate(table_data, headers=[header], tablefmt="pretty")
-
-    def _to_table(self, data, headers=None):
-        """Convert data to ASCII table."""
-        if isinstance(data, list) and len(data) > 0:
-            if isinstance(data[0], dict):
-                headers = headers or list(data[0].keys())
-                table_data = [[row.get(h, "") for h in headers] for row in data]
-                return tabulate(table_data, headers=headers, tablefmt="pretty")
-            else:
-                return tabulate([[item] for item in data], tablefmt="pretty")
-        elif isinstance(data, dict):
-            return tabulate([[k, v] for k, v in data.items()], headers=["Key", "Value"], tablefmt="pretty")
-        else:
-            return str(data)
-
-    def _to_csv(self, data):
-        """Convert data to CSV format."""
-        output = io.StringIO()
-        if isinstance(data, list) and len(data) > 0:
-            if isinstance(data[0], dict):
-                writer = csv.DictWriter(output, fieldnames=data[0].keys())
-                writer.writeheader()
-                writer.writerows(data)
-            else:
-                writer = csv.writer(output)
-                for item in data:
-                    writer.writerow([item])
-        elif isinstance(data, dict):
-            writer = csv.writer(output)
-            writer.writerow(["Key", "Value"])
-            for k, v in data.items():
-                writer.writerow([k, v])
-        return output.getvalue()
 
 
 class MilvusClientCli(object):
@@ -116,7 +55,6 @@ def __init__(self):
         self.role = MilvusClientRole(self.connection)
         self.alias = MilvusClientAlias(self.connection)
         self.partition = MilvusClientPartition(self.connection)
-        self.formatter = OutputFormatter()
 
         # Resource and privilege group clients
         self.resource_group = MilvusResourceGroup(self.connection)