pinecone.openapi_support.rest_aiohttp

  1import ssl
  2import certifi
  3import json
  4from .rest_utils import RestClientInterface, RESTResponse, raise_exceptions_or_return
  5from .configuration import Configuration
  6
  7
  8class AiohttpRestClient(RestClientInterface):
  9    def __init__(self, configuration: Configuration) -> None:
 10        try:
 11            import aiohttp
 12        except ImportError:
 13            raise ImportError(
 14                "Additional dependencies are required to use Pinecone with asyncio. Include these extra dependencies in your project by installing `pinecone[asyncio]`."
 15            ) from None
 16
 17        if configuration.ssl_ca_cert is not None:
 18            ca_certs = configuration.ssl_ca_cert
 19        else:
 20            ca_certs = certifi.where()
 21
 22        ssl_context = ssl.create_default_context(cafile=ca_certs)
 23
 24        conn = aiohttp.TCPConnector(verify_ssl=configuration.verify_ssl, ssl=ssl_context)
 25
 26        if configuration.proxy:
 27            self._session = aiohttp.ClientSession(connector=conn, proxy=configuration.proxy)
 28        else:
 29            self._session = aiohttp.ClientSession(connector=conn)
 30
 31    async def close(self):
 32        await self._session.close()
 33
 34    async def request(
 35        self,
 36        method,
 37        url,
 38        query_params=None,
 39        headers=None,
 40        body=None,
 41        post_params=None,
 42        _preload_content=True,
 43        _request_timeout=None,
 44    ):
 45        if method in ["POST", "PUT", "PATCH", "OPTIONS"] and ("Content-Type" not in headers):
 46            headers["Content-Type"] = "application/json"
 47
 48        if "application/x-ndjson" in headers.get("Content-Type", "").lower():
 49            ndjson_data = "\n".join(json.dumps(record) for record in body)
 50
 51            async with self._session.request(
 52                method, url, params=query_params, headers=headers, data=ndjson_data
 53            ) as resp:
 54                content = await resp.read()
 55                return raise_exceptions_or_return(
 56                    RESTResponse(resp.status, content, resp.headers, resp.reason)
 57                )
 58
 59        else:
 60            async with self._session.request(
 61                method, url, params=query_params, headers=headers, json=body
 62            ) as resp:
 63                content = await resp.read()
 64                return raise_exceptions_or_return(
 65                    RESTResponse(resp.status, content, resp.headers, resp.reason)
 66                )
 67
 68    async def GET(
 69        self, url, headers=None, query_params=None, _preload_content=True, _request_timeout=None
 70    ):
 71        return await self.request(
 72            "GET",
 73            url,
 74            headers=headers,
 75            _preload_content=_preload_content,
 76            _request_timeout=_request_timeout,
 77            query_params=query_params,
 78        )
 79
 80    async def HEAD(
 81        self, url, headers=None, query_params=None, _preload_content=True, _request_timeout=None
 82    ):
 83        return await self.request(
 84            "HEAD",
 85            url,
 86            headers=headers,
 87            _preload_content=_preload_content,
 88            _request_timeout=_request_timeout,
 89            query_params=query_params,
 90        )
 91
 92    async def OPTIONS(
 93        self,
 94        url,
 95        headers=None,
 96        query_params=None,
 97        post_params=None,
 98        body=None,
 99        _preload_content=True,
100        _request_timeout=None,
101    ):
102        return await self.request(
103            "OPTIONS",
104            url,
105            headers=headers,
106            query_params=query_params,
107            post_params=post_params,
108            _preload_content=_preload_content,
109            _request_timeout=_request_timeout,
110            body=body,
111        )
112
113    async def DELETE(
114        self,
115        url,
116        headers=None,
117        query_params=None,
118        body=None,
119        _preload_content=True,
120        _request_timeout=None,
121    ):
122        return await self.request(
123            "DELETE",
124            url,
125            headers=headers,
126            query_params=query_params,
127            _preload_content=_preload_content,
128            _request_timeout=_request_timeout,
129            body=body,
130        )
131
132    async def POST(
133        self,
134        url,
135        headers=None,
136        query_params=None,
137        post_params=None,
138        body=None,
139        _preload_content=True,
140        _request_timeout=None,
141    ):
142        return await self.request(
143            "POST",
144            url,
145            headers=headers,
146            query_params=query_params,
147            post_params=post_params,
148            _preload_content=_preload_content,
149            _request_timeout=_request_timeout,
150            body=body,
151        )
152
153    async def PUT(
154        self,
155        url,
156        headers=None,
157        query_params=None,
158        post_params=None,
159        body=None,
160        _preload_content=True,
161        _request_timeout=None,
162    ):
163        return await self.request(
164            "PUT",
165            url,
166            headers=headers,
167            query_params=query_params,
168            post_params=post_params,
169            _preload_content=_preload_content,
170            _request_timeout=_request_timeout,
171            body=body,
172        )
173
174    async def PATCH(
175        self,
176        url,
177        headers=None,
178        query_params=None,
179        post_params=None,
180        body=None,
181        _preload_content=True,
182        _request_timeout=None,
183    ):
184        return await self.request(
185            "PATCH",
186            url,
187            headers=headers,
188            query_params=query_params,
189            post_params=post_params,
190            _preload_content=_preload_content,
191            _request_timeout=_request_timeout,
192            body=body,
193        )