Fix sanity tests

Author: jhamon

pinecone/__init__.py

@@ -41,7 +41,7 @@
     "QueryResponse": ("pinecone.db_data.models", "QueryResponse"),
     "UpsertResponse": ("pinecone.db_data.models", "UpsertResponse"),
     "UpdateRequest": ("pinecone.db_data.models", "UpdateRequest"),
-    "ImportErrorMode": ("pinecone.core.openapi.db_data.models", "ImportErrorMode"),
+    "ImportErrorMode": ("pinecone.db_data.resources.sync.bulk_import", "ImportErrorMode"),
     "VectorDictionaryMissingKeysError": (
         "pinecone.db_data.errors",
         "VectorDictionaryMissingKeysError",

pinecone/__init__.pyi

@@ -29,7 +29,7 @@ from pinecone.db_data.models import (
     UpsertResponse,
     UpdateRequest,
 )
-from pinecone.core.openapi.db_data.models import ImportErrorMode
+from pinecone.db_data.resources.sync.bulk_import import ImportErrorMode
 from pinecone.db_data.errors import (
     VectorDictionaryMissingKeysError,
     VectorDictionaryExcessKeysError,

pinecone/data/__init__.py

@@ -1,7 +1,5 @@
 import warnings
 
-from pinecone.db_data import *
-
 warnings.warn(
     "The module at `pinecone.data` has moved to `pinecone.db_data`. "
     "Please update your imports. "

pinecone/data/features/__init__.py

@@ -1,10 +1,11 @@
 import warnings
 
-from pinecone.db_data.features import *
+from .bulk_imports import *
+from .inference import *
+
 
 warnings.warn(
-    "The module at `pinecone.data.features` has moved to `pinecone.db_data.features`. "
-    "Please update your imports. "
+    "The module at `pinecone.data.features` has been removed. Code has been refactored and integrated into other parts of the client. "
     "This warning will become an error in a future version of the Pinecone Python SDK.",
     DeprecationWarning,
 )

pinecone/data/features/bulk_imports/__init__.py

@@ -1,6 +1,9 @@
 import warnings
 
-from pinecone.db_data.features.bulk_import import *
+from pinecone.db_data.resources.asyncio.bulk_import_asyncio import *
+from pinecone.db_data.resources.sync.bulk_import import *
+from pinecone.db_data.resources.sync.bulk_import_request_factory import *
+
 
 warnings.warn(
     "The module at `pinecone.data.features.bulk_import` has moved to `pinecone.db_data.features.bulk_import`. "

pinecone/db_data/__init__.py

@@ -25,7 +25,7 @@
     MetadataDictionaryExpectedError,
 )
 
-from .features.bulk_import import ImportErrorMode
+from .resources.sync.bulk_import import ImportErrorMode
 
 
 import warnings

pinecone/db_data/features/bulk_import/__init__.py

@@ -2,7 +2,7 @@
 import warnings
 import logging
 import json
-from typing import Union, List, Optional, Dict, Any, Literal, TYPE_CHECKING
+from typing import Union, List, Optional, Dict, Any, Literal, Iterator, TYPE_CHECKING
 
 from pinecone.config import ConfigBuilder
 
@@ -19,7 +19,6 @@
 from .dataclasses import Vector, SparseValues, FetchResponse, SearchQuery, SearchRerank
 from .interfaces import IndexInterface
 from .request_factory import IndexRequestFactory
-from .features.bulk_import import ImportFeatureMixin
 from .types import (
     SparseVectorTypedDict,
     VectorTypedDict,
@@ -47,6 +46,15 @@
 
 if TYPE_CHECKING:
     from pinecone.config import Config, OpenApiConfiguration
+    from .resources.sync.bulk_import import BulkImportResource
+
+    from pinecone.core.openapi.db_data.models import (
+        StartImportResponse,
+        ListImportsResponse,
+        ImportModel,
+    )
+
+    from .resources.sync.bulk_import import ImportErrorMode
 
 logger = logging.getLogger(__name__)
 """ @private """
@@ -58,12 +66,15 @@ def parse_query_response(response: QueryResponse):
     return response
 
 
-class Index(PluginAware, IndexInterface, ImportFeatureMixin):
+class Index(PluginAware, IndexInterface):
     """
     A client for interacting with a Pinecone index via REST API.
     For improved performance, use the Pinecone GRPC index client.
     """
 
+    _bulk_import_resource: Optional["BulkImportResource"]
+    """ @private """
+
     def __init__(
         self,
         api_key: str,
@@ -101,6 +112,9 @@ def __init__(
 
         self._api_client = self._vector_api.api_client
 
+        self._bulk_import_resource = None
+        """ @private """
+
         # Pass the same api_client to the ImportFeatureMixin
         super().__init__(api_client=self._api_client)
 
@@ -129,6 +143,15 @@ def pool_threads(self) -> int:
         )
         return self._pool_threads
 
+    @property
+    def bulk_import(self) -> "BulkImportResource":
+        """@private"""
+        if self._bulk_import_resource is None:
+            from .resources.sync.bulk_import import BulkImportResource
+
+            self._bulk_import_resource = BulkImportResource(api_client=self._api_client)
+        return self._bulk_import_resource
+
     def _openapi_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
         return filter_dict(kwargs, OPENAPI_ENDPOINT_PARAMS)
 
@@ -457,3 +480,124 @@ def list(self, **kwargs):
                 kwargs.update({"pagination_token": results.pagination.next})
             else:
                 done = True
+
+    @validate_and_convert_errors
+    def start_import(
+        self,
+        uri: str,
+        integration_id: Optional[str] = None,
+        error_mode: Optional[
+            Union["ImportErrorMode", Literal["CONTINUE", "ABORT"], str]
+        ] = "CONTINUE",
+    ) -> "StartImportResponse":
+        """
+        Args:
+            uri (str): The URI of the data to import. The URI must start with the scheme of a supported storage provider.
+            integration_id (Optional[str], optional): If your bucket requires authentication to access, you need to pass the id of your storage integration using this property. Defaults to None.
+            error_mode: Defaults to "CONTINUE". If set to "CONTINUE", the import operation will continue even if some
+                records fail to import. Pass "ABORT" to stop the import operation if any records fail to import.
+
+        Returns:
+            `StartImportResponse`: Contains the id of the import operation.
+
+        Import data from a storage provider into an index. The uri must start with the scheme of a supported
+        storage provider. For buckets that are not publicly readable, you will also need to separately configure
+        a storage integration and pass the integration id.
+
+        Examples:
+            >>> from pinecone import Pinecone
+            >>> index = Pinecone().Index('my-index')
+            >>> index.start_import(uri="s3://bucket-name/path/to/data.parquet")
+            { id: "1" }
+        """
+        return self.bulk_import.start(uri=uri, integration_id=integration_id, error_mode=error_mode)
+
+    @validate_and_convert_errors
+    def list_imports(self, **kwargs) -> Iterator["ImportModel"]:
+        """
+        Args:
+            limit (Optional[int]): The maximum number of operations to fetch in each network call. If unspecified, the server will use a default value. [optional]
+            pagination_token (Optional[str]): When there are multiple pages of results, a pagination token is returned in the response. The token can be used
+                to fetch the next page of results. [optional]
+
+        Returns:
+            Returns a generator that yields each import operation. It automatically handles pagination tokens on your behalf so you can
+            easily iterate over all results. The `list_imports` method accepts all of the same arguments as list_imports_paginated
+
+        ```python
+        for op in index.list_imports():
+            print(op)
+        ```
+
+        You can convert the generator into a list by wrapping the generator in a call to the built-in `list` function:
+
+        ```python
+        operations = list(index.list_imports())
+        ```
+
+        You should be cautious with this approach because it will fetch all operations at once, which could be a large number
+        of network calls and a lot of memory to hold the results.
+        """
+        for i in self.bulk_import.list(**kwargs):
+            yield i
+
+    @validate_and_convert_errors
+    def list_imports_paginated(
+        self, limit: Optional[int] = None, pagination_token: Optional[str] = None, **kwargs
+    ) -> "ListImportsResponse":
+        """
+        Args:
+            limit (Optional[int]): The maximum number of ids to return. If unspecified, the server will use a default value. [optional]
+            pagination_token (Optional[str]): A token needed to fetch the next page of results. This token is returned
+                in the response if additional results are available. [optional]
+
+        Returns: ListImportsResponse object which contains the list of operations as ImportModel objects, pagination information,
+            and usage showing the number of read_units consumed.
+
+        The list_imports_paginated operation returns information about import operations.
+        It returns operations in a paginated form, with a pagination token to fetch the next page of results.
+
+        Consider using the `list_imports` method to avoid having to handle pagination tokens manually.
+
+        Examples:
+            >>> results = index.list_imports_paginated(limit=5)
+            >>> results.pagination.next
+            eyJza2lwX3Bhc3QiOiI5OTMiLCJwcmVmaXgiOiI5OSJ9
+            >>> results.data[0]
+            {
+                "id": "6",
+                "uri": "s3://dev-bulk-import-datasets-pub/10-records-dim-10/",
+                "status": "Completed",
+                "percent_complete": 100.0,
+                "records_imported": 10,
+                "created_at": "2024-09-06T14:52:02.567776+00:00",
+                "finished_at": "2024-09-06T14:52:28.130717+00:00"
+            }
+            >>> next_results = index.list_imports_paginated(limit=5, pagination_token=results.pagination.next)
+        """
+        return self.bulk_import.list_paginated(
+            limit=limit, pagination_token=pagination_token, **kwargs
+        )
+
+    @validate_and_convert_errors
+    def describe_import(self, id: str) -> "ImportModel":
+        """
+        Args:
+            id (str): The id of the import operation. This value is returned when
+                starting an import, and can be looked up using list_imports.
+
+        Returns:
+            `ImportModel`: An object containing operation id, status, and other details.
+
+        describe_import is used to get detailed information about a specific import operation.
+        """
+        return self.bulk_import.describe(id=id)
+
+    @validate_and_convert_errors
+    def cancel_import(self, id: str):
+        """Cancel an import operation.
+
+        Args:
+            id (str): The id of the import operation to cancel.
+        """
+        return self.bulk_import.cancel(id=id)