[cherry-pick] support connectType if use oss-bucket (#2988)

Signed-off-by: lentitude2tk <[email protected]>

Author: lentitude2tk

examples/bulk_import/example_stage_file_manager.py

@@ -1,10 +1,12 @@
+from pymilvus.bulk_writer.constants import ConnectType
 from pymilvus.bulk_writer.stage_file_manager import StageFileManager
 
 if __name__ == "__main__":
     stage_file_manager = StageFileManager(
         cloud_endpoint='https://api.cloud.zilliz.com',
         api_key='_api_key_for_cluster_org_',
         stage_name='_stage_name_for_project_',
+        connect_type=ConnectType.AUTO,
     )
     result = stage_file_manager.upload_file_to_stage("/Users/zilliz/data/", "data/")
     print(f"\nuploadFileToStage results: {result}")

pymilvus/bulk_writer/constants.py

@@ -10,7 +10,7 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 
-from enum import IntEnum
+from enum import Enum, IntEnum
 
 import numpy as np
 
@@ -88,3 +88,9 @@ class BulkFileType(IntEnum):
     JSON_RB = 2  # deprecated
     PARQUET = 3
     CSV = 4
+
+
+class ConnectType(Enum):
+    AUTO = "AUTO"
+    INTERNAL = "INTERNAL"
+    PUBLIC = "PUBLIC"

pymilvus/bulk_writer/endpoint_resolver.py

@@ -0,0 +1,64 @@
+import http
+import logging
+
+from pymilvus.bulk_writer.constants import ConnectType
+
+logger = logging.getLogger("EndpointResolver")
+logging.basicConfig(level=logging.INFO)
+
+
+class EndpointResolver:
+    @staticmethod
+    def resolve_endpoint(
+        default_endpoint: str, cloud: str, region: str, connect_type: ConnectType
+    ) -> str:
+        logger.info(
+            "Start resolving endpoint, cloud=%s, region=%s, connectType=%s",
+            cloud,
+            region,
+            connect_type,
+        )
+        if cloud == "ali":
+            default_endpoint = EndpointResolver._resolve_oss_endpoint(region, connect_type)
+        logger.info("Resolved endpoint: %s, reachable check passed", default_endpoint)
+        return default_endpoint
+
+    @staticmethod
+    def _resolve_oss_endpoint(region: str, connect_type: ConnectType) -> str:
+        internal_endpoint = f"oss-{region}-internal.aliyuncs.com"
+        public_endpoint = f"oss-{region}.aliyuncs.com"
+
+        if connect_type == ConnectType.INTERNAL:
+            logger.info("Forced INTERNAL endpoint selected: %s", internal_endpoint)
+            EndpointResolver._check_endpoint_reachable(internal_endpoint, True)
+            return internal_endpoint
+        if connect_type == ConnectType.PUBLIC:
+            logger.info("Forced PUBLIC endpoint selected: %s", public_endpoint)
+            EndpointResolver._check_endpoint_reachable(public_endpoint, True)
+            return public_endpoint
+        if EndpointResolver._check_endpoint_reachable(internal_endpoint, False):
+            logger.info("AUTO mode: internal endpoint reachable, using %s", internal_endpoint)
+            return internal_endpoint
+        logger.warning(
+            "AUTO mode: internal endpoint not reachable, fallback to public endpoint %s",
+            public_endpoint,
+        )
+        EndpointResolver._check_endpoint_reachable(public_endpoint, True)
+        return public_endpoint
+
+    @staticmethod
+    def _check_endpoint_reachable(endpoint: str, raise_error: bool) -> bool:
+        try:
+            conn = http.client.HTTPSConnection(endpoint, timeout=5)
+            conn.request("HEAD", "/")
+            resp = conn.getresponse()
+            code = resp.status
+            logger.debug("Checked endpoint %s, response code=%s", endpoint, code)
+        except Exception as e:
+            if raise_error:
+                logger.exception("Endpoint %s not reachable, throwing exception", endpoint)
+                raise RuntimeError(str(e)) from e
+            logger.warning("Endpoint %s not reachable, will fallback if needed", endpoint)
+            return False
+        else:
+            return 200 <= code < 400

pymilvus/bulk_writer/stage_bulk_writer.py

@@ -6,7 +6,7 @@
 from pymilvus.bulk_writer.stage_file_manager import StageFileManager
 from pymilvus.orm.schema import CollectionSchema
 
-from .constants import MB, BulkFileType
+from .constants import MB, BulkFileType, ConnectType
 from .local_bulk_writer import LocalBulkWriter
 
 logger = logging.getLogger(__name__)
@@ -35,7 +35,10 @@ def __init__(
         self._remote_files: List[List[str]] = []
         self._stage_name = stage_name
         self._stage_file_manager = StageFileManager(
-            cloud_endpoint=cloud_endpoint, api_key=api_key, stage_name=stage_name
+            cloud_endpoint=cloud_endpoint,
+            api_key=api_key,
+            stage_name=stage_name,
+            connect_type=ConnectType.AUTO,
         )
 
         logger.info(f"Remote buffer writer initialized, target path: {self._remote_path}")

pymilvus/bulk_writer/stage_file_manager.py

@@ -9,6 +9,8 @@
 from minio import Minio
 from minio.error import S3Error
 
+from pymilvus.bulk_writer.constants import ConnectType
+from pymilvus.bulk_writer.endpoint_resolver import EndpointResolver
 from pymilvus.bulk_writer.file_utils import FileUtils
 from pymilvus.bulk_writer.stage_restful import apply_stage
 
@@ -17,7 +19,13 @@
 
 
 class StageFileManager:
-    def __init__(self, cloud_endpoint: str, api_key: str, stage_name: str):
+    def __init__(
+        self,
+        cloud_endpoint: str,
+        api_key: str,
+        stage_name: str,
+        connect_type: ConnectType = ConnectType.AUTO,
+    ):
         """
         private preview feature. Please submit a request and contact us if you need it.
 
@@ -27,10 +35,16 @@ def __init__(self, cloud_endpoint: str, api_key: str, stage_name: str):
                 - For regions in China: https://api.cloud.zilliz.com.cn
             api_key (str): The API key associated with your organization
             stage_name (str): The name of the Stage.
+            connect_type: Current value is mainly for Aliyun OSS buckets, default is Auto.
+             - Default case, if the OSS bucket is reachable via the internal endpoint,
+               the internal endpoint will be used
+             - otherwise, the public endpoint will be used.
+             - You can also force the use of either the internal or public endpoint.
         """
         self.cloud_endpoint = cloud_endpoint
         self.api_key = api_key
         self.stage_name = stage_name
+        self.connect_type = connect_type
         self.local_file_paths = []
         self.total_bytes = 0
         self.stage_info = {}
@@ -51,8 +65,15 @@ def _refresh_stage_and_client(self, path: str):
 
         creds = self.stage_info["credentials"]
         http_client = urllib3.PoolManager(maxsize=100)
+
+        endpoint = EndpointResolver.resolve_endpoint(
+            self.stage_info["endpoint"],
+            self.stage_info["cloud"],
+            self.stage_info["region"],
+            self.connect_type,
+        )
         self._client = Minio(
-            endpoint=self.stage_info["endpoint"],
+            endpoint=endpoint,
             access_key=creds["tmpAK"],
             secret_key=creds["tmpSK"],
             session_token=creds["sessionToken"],

tests/test_bulk_writer_stage.py

@@ -5,7 +5,7 @@
 
 import pytest
 import requests
-from pymilvus.bulk_writer.constants import BulkFileType
+from pymilvus.bulk_writer.constants import BulkFileType, ConnectType
 from pymilvus.bulk_writer.stage_bulk_writer import StageBulkWriter
 from pymilvus.bulk_writer.stage_file_manager import StageFileManager
 from pymilvus.bulk_writer.stage_manager import StageManager
@@ -127,6 +127,7 @@ def test_apply_stage_success(
                 "endpoint": "s3.amazonaws.com",
                 "bucketName": "test-bucket",
                 "region": "us-west-2",
+                "cloud": "aws",
                 "condition": {"maxContentLength": 1073741824},
                 "credentials": {
                     "tmpAK": "test_access_key",
@@ -238,6 +239,7 @@ def stage_file_manager(self) -> StageFileManager:
             cloud_endpoint="https://api.cloud.zilliz.com",
             api_key="test_api_key",
             stage_name="test_stage",
+            connect_type=ConnectType.AUTO,
         )
 
     @pytest.fixture
@@ -249,6 +251,7 @@ def mock_stage_info(self) -> Dict[str, Any]:
             "endpoint": "s3.amazonaws.com",
             "bucketName": "test-bucket",
             "region": "us-west-2",
+            "cloud": "aws",
             "condition": {"maxContentLength": 1073741824},
             "credentials": {
                 "tmpAK": "test_access_key",
@@ -563,6 +566,7 @@ def mock_server_responses(self) -> Dict[str, Any]:
                     "endpoint": "s3.amazonaws.com",
                     "bucketName": "test-bucket",
                     "region": "us-west-2",
+                    "cloud": "aws",
                     "condition": {"maxContentLength": 1073741824},
                     "credentials": {
                         "tmpAK": "test_access_key",
@@ -614,6 +618,7 @@ def test_full_stage_workflow(
             cloud_endpoint="https://api.cloud.zilliz.com",
             api_key="test_api_key",
             stage_name="test_stage",
+            connect_type=ConnectType.AUTO,
         )
 
         # Verify stage info can be refreshed