pinecone.data.features.bulk_import

  1from enum import Enum
  2from typing import Optional, Literal, Iterator, List, Type, cast
  3
  4from pinecone.config.config import ConfigBuilder
  5from pinecone.core_ea.openapi.db_data import ApiClient
  6from pinecone.core_ea.openapi.db_data.api.bulk_operations_api import BulkOperationsApi
  7from pinecone.core_ea.openapi.shared import API_VERSION
  8
  9from pinecone.utils import parse_non_empty_args, install_json_repr_override, setup_openapi_client
 10
 11from pinecone.core_ea.openapi.db_data.models import (
 12    StartImportRequest,
 13    StartImportResponse,
 14    ImportListResponse,
 15    ImportModel,
 16    ImportErrorMode as ImportErrorModeClass,
 17)
 18
 19for m in [StartImportResponse, ImportListResponse, ImportModel]:
 20    install_json_repr_override(m)
 21
 22ImportErrorMode: Type[Enum] = cast(
 23    Type[Enum], Enum("ImportErrorMode", ImportErrorModeClass.allowed_values[("on_error",)])
 24)
 25
 26
 27class ImportFeatureMixin:
 28    def __init__(self, **kwargs):
 29        config = ConfigBuilder.build(
 30            **kwargs,
 31        )
 32        openapi_config = ConfigBuilder.build_openapi_config(config, kwargs.get("openapi_config", None))
 33
 34        if kwargs.get("__import_operations_api", None):
 35            self.__import_operations_api = kwargs.get("__import_operations_api")
 36        else:
 37            self.__import_operations_api = setup_openapi_client(
 38                api_client_klass=ApiClient,
 39                api_klass=BulkOperationsApi,
 40                config=config,
 41                openapi_config=openapi_config,
 42                pool_threads=kwargs.get("pool_threads", 1),
 43                api_version=API_VERSION,
 44            )
 45
 46    def start_import(
 47        self,
 48        uri: str,
 49        integration_id: Optional[str] = None,
 50        error_mode: Optional[Literal["CONTINUE", "ABORT"]] = "CONTINUE",
 51    ) -> StartImportResponse:
 52        """Import data from a storage provider into an index. The uri must start with the scheme of a supported
 53        storage provider. For buckets that are not publicly readable, you will also need to separately configure
 54        a storage integration and pass the integration id.
 55
 56        Examples:
 57            >>> from pinecone import Pinecone
 58            >>> index = Pinecone().Index('my-index')
 59            >>> index.start_import(uri="s3://bucket-name/path/to/data.parquet")
 60            { id: "1" }
 61
 62        Args:
 63            uri (str): The URI of the data to import. The URI must start with the scheme of a supported storage provider.
 64            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.
 65            error_mode: Defaults to "CONTINUE". If set to "CONTINUE", the import operation will continue even if some
 66                records fail to import. Pass "ABORT" to stop the import operation if any records fail to import.
 67
 68        Returns:
 69            StartImportResponse: Contains the id of the import operation.
 70        """
 71        if isinstance(error_mode, ImportErrorMode):
 72            error_mode = error_mode.value
 73        elif isinstance(error_mode, str):
 74            try:
 75                error_mode = ImportErrorMode(error_mode.lower()).value
 76            except ValueError:
 77                raise ValueError(f"Invalid error_mode value: {error_mode}")
 78
 79        args_dict = parse_non_empty_args(
 80            [
 81                ("uri", uri),
 82                ("integration_id", integration_id),
 83                ("error_mode", ImportErrorModeClass(on_error=error_mode)),
 84            ]
 85        )
 86
 87        return self.__import_operations_api.start_import(StartImportRequest(**args_dict))
 88
 89    def list_imports(self, **kwargs) -> Iterator[List[ImportModel]]:
 90        """
 91        Returns a generator that yields each import operation. It automatically handles pagination tokens on your behalf so you can
 92        easily iterate over all results. The `list_imports` method accepts all of the same arguments as list_imports_paginated
 93
 94        ```python
 95        for op in index.list_imports():
 96            print(op)
 97        ```
 98
 99        You can convert the generator into a list by wrapping the generator in a call to the built-in `list` function:
100
101        ```python
102        operations = list(index.list_imports())
103        ```
104
105        You should be cautious with this approach because it will fetch all operations at once, which could be a large number
106        of network calls and a lot of memory to hold the results.
107
108        Args:
109            limit (Optional[int]): The maximum number of operations to fetch in each network call. If unspecified, the server will use a default value. [optional]
110            pagination_token (Optional[str]): When there are multiple pages of results, a pagination token is returned in the response. The token can be used
111                to fetch the next page of results. [optional]
112        """
113        done = False
114        while not done:
115            results = self.list_imports_paginated(**kwargs)
116            if len(results.data) > 0:
117                for op in results.data:
118                    yield op
119
120            if results.pagination:
121                kwargs.update({"pagination_token": results.pagination.next})
122            else:
123                done = True
124
125    def list_imports_paginated(
126        self,
127        limit: Optional[int] = None,
128        pagination_token: Optional[str] = None,
129        **kwargs,
130    ) -> ImportListResponse:
131        """
132        The list_imports_paginated operation returns information about import operations.
133        It returns operations in a paginated form, with a pagination token to fetch the next page of results.
134
135        Consider using the `list_imports` method to avoid having to handle pagination tokens manually.
136
137        Examples:
138            >>> results = index.list_imports_paginated(limit=5)
139            >>> results.pagination.next
140            eyJza2lwX3Bhc3QiOiI5OTMiLCJwcmVmaXgiOiI5OSJ9
141            >>> results.data[0]
142            {
143                "id": "6",
144                "uri": "s3://dev-bulk-import-datasets-pub/10-records-dim-10/",
145                "status": "Completed",
146                "percent_complete": 100.0,
147                "records_imported": 10,
148                "created_at": "2024-09-06T14:52:02.567776+00:00",
149                "finished_at": "2024-09-06T14:52:28.130717+00:00"
150            }
151            >>> next_results = index.list_imports_paginated(limit=5, pagination_token=results.pagination.next)
152
153        Args:
154            limit (Optional[int]): The maximum number of ids to return. If unspecified, the server will use a default value. [optional]
155            pagination_token (Optional[str]): A token needed to fetch the next page of results. This token is returned
156                in the response if additional results are available. [optional]
157
158        Returns: ImportListResponse object which contains the list of operations as ImportModel objects, pagination information,
159            and usage showing the number of read_units consumed.
160        """
161        args_dict = parse_non_empty_args(
162            [
163                ("limit", limit),
164                ("pagination_token", pagination_token),
165            ]
166        )
167        return self.__import_operations_api.list_imports(**args_dict)
168
169    def describe_import(self, id: str) -> ImportModel:
170        """
171        describe_import is used to get detailed information about a specific import operation.
172
173        Args:
174            id (str): The id of the import operation. This value is returned when
175            starting an import, and can be looked up using list_imports.
176
177        Returns:
178            ImportModel: An object containing operation id, status, and other details.
179        """
180        if isinstance(id, int):
181            id = str(id)
182        return self.__import_operations_api.describe_import(id=id)
183
184    def cancel_import(self, id: str):
185        """Cancel an import operation.
186
187        Args:
188            id (str): The id of the import operation to cancel.
189        """
190        if isinstance(id, int):
191            id = str(id)
192        return self.__import_operations_api.cancel_import(id=id)