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(**kwargs)
 30        openapi_config = ConfigBuilder.build_openapi_config(
 31            config, kwargs.get("openapi_config", None)
 32        )
 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, limit: Optional[int] = None, pagination_token: Optional[str] = None, **kwargs
127    ) -> ImportListResponse:
128        """
129        The list_imports_paginated operation returns information about import operations.
130        It returns operations in a paginated form, with a pagination token to fetch the next page of results.
131
132        Consider using the `list_imports` method to avoid having to handle pagination tokens manually.
133
134        Examples:
135            >>> results = index.list_imports_paginated(limit=5)
136            >>> results.pagination.next
137            eyJza2lwX3Bhc3QiOiI5OTMiLCJwcmVmaXgiOiI5OSJ9
138            >>> results.data[0]
139            {
140                "id": "6",
141                "uri": "s3://dev-bulk-import-datasets-pub/10-records-dim-10/",
142                "status": "Completed",
143                "percent_complete": 100.0,
144                "records_imported": 10,
145                "created_at": "2024-09-06T14:52:02.567776+00:00",
146                "finished_at": "2024-09-06T14:52:28.130717+00:00"
147            }
148            >>> next_results = index.list_imports_paginated(limit=5, pagination_token=results.pagination.next)
149
150        Args:
151            limit (Optional[int]): The maximum number of ids to return. If unspecified, the server will use a default value. [optional]
152            pagination_token (Optional[str]): A token needed to fetch the next page of results. This token is returned
153                in the response if additional results are available. [optional]
154
155        Returns: ImportListResponse object which contains the list of operations as ImportModel objects, pagination information,
156            and usage showing the number of read_units consumed.
157        """
158        args_dict = parse_non_empty_args([("limit", limit), ("pagination_token", pagination_token)])
159        return self.__import_operations_api.list_imports(**args_dict)
160
161    def describe_import(self, id: str) -> ImportModel:
162        """
163        describe_import is used to get detailed information about a specific import operation.
164
165        Args:
166            id (str): The id of the import operation. This value is returned when
167            starting an import, and can be looked up using list_imports.
168
169        Returns:
170            ImportModel: An object containing operation id, status, and other details.
171        """
172        if isinstance(id, int):
173            id = str(id)
174        return self.__import_operations_api.describe_import(id=id)
175
176    def cancel_import(self, id: str):
177        """Cancel an import operation.
178
179        Args:
180            id (str): The id of the import operation to cancel.
181        """
182        if isinstance(id, int):
183            id = str(id)
184        return self.__import_operations_api.cancel_import(id=id)