pinecone.config.openapi

  1import sys
  2from typing import List, Optional
  3
  4import certifi
  5import socket
  6import copy
  7
  8from urllib3.connection import HTTPConnection
  9
 10from pinecone.core.openapi.shared.configuration import (
 11    Configuration as OpenApiConfiguration,
 12)
 13
 14TCP_KEEPINTVL = 60  # Sec
 15TCP_KEEPIDLE = 300  # Sec
 16TCP_KEEPCNT = 4
 17
 18
 19class OpenApiConfigFactory:
 20    @classmethod
 21    def build(cls, api_key: str, host: Optional[str] = None, **kwargs):
 22        openapi_config = OpenApiConfiguration()
 23        openapi_config.api_key = {"ApiKeyAuth": api_key}
 24        openapi_config.host = host
 25        openapi_config.ssl_ca_cert = certifi.where()
 26        openapi_config.socket_options = cls._get_socket_options()
 27        openapi_config.discard_unknown_keys = True
 28
 29        return openapi_config
 30
 31    @classmethod
 32    def copy(cls, openapi_config: OpenApiConfiguration, api_key: str, host: str) -> OpenApiConfiguration:
 33        """
 34        Copy a user-supplied openapi configuration and update it with the user's api key and host.
 35        If they have not specified other socket configuration, we will use the default values.
 36        We expect these objects are being passed mainly a vehicle for proxy configuration, so
 37        we don't modify those settings.
 38        """
 39        copied = copy.deepcopy(openapi_config)
 40
 41        copied.api_key = {"ApiKeyAuth": api_key}
 42        copied.host = host
 43
 44        # Set sensible defaults if the user hasn't set them
 45        if not copied.socket_options:
 46            copied.socket_options = cls._get_socket_options()
 47
 48        # We specifically do not modify the user's ssl_ca_cert or proxy settings, as
 49        # they may have set them intentionally. This is the main reason somebody would
 50        # pass an openapi_config in the first place.
 51
 52        return copied
 53
 54    @classmethod
 55    def _get_socket_options(
 56        self,
 57        do_keep_alive: bool = True,
 58        keep_alive_idle_sec: int = TCP_KEEPIDLE,
 59        keep_alive_interval_sec: int = TCP_KEEPINTVL,
 60        keep_alive_tries: int = TCP_KEEPCNT,
 61    ) -> List[tuple]:
 62        """
 63        Returns the socket options to pass to OpenAPI's Rest client
 64        Args:
 65            do_keep_alive: Whether to enable TCP keep alive mechanism
 66            keep_alive_idle_sec: Time in seconds of connection idleness before starting to send keep alive probes
 67            keep_alive_interval_sec: Interval time in seconds between keep alive probe messages
 68            keep_alive_tries: Number of failed keep alive tries (unanswered KA messages) before terminating the connection
 69
 70        Returns:
 71            A list of socket options for the Rest client's connection pool
 72        """
 73        # Source: https://www.finbourne.com/blog/the-mysterious-hanging-client-tcp-keep-alives
 74
 75        socket_params = HTTPConnection.default_socket_options
 76        if not do_keep_alive:
 77            return socket_params
 78
 79        socket_params += [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
 80
 81        # TCP Keep Alive Probes for different platforms
 82        platform = sys.platform
 83        # TCP Keep Alive Probes for Linux
 84        if (
 85            platform == "linux"
 86            and hasattr(socket, "TCP_KEEPIDLE")
 87            and hasattr(socket, "TCP_KEEPINTVL")
 88            and hasattr(socket, "TCP_KEEPCNT")
 89        ):
 90            socket_params += [(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, keep_alive_idle_sec)]
 91            socket_params += [
 92                (
 93                    socket.IPPROTO_TCP,
 94                    socket.TCP_KEEPINTVL,
 95                    keep_alive_interval_sec,
 96                )
 97            ]
 98            socket_params += [(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, keep_alive_tries)]
 99
100        # TCP Keep Alive Probes for Windows OS
101        # NOTE: Changing TCP KA params on windows is done via a different mechanism which OpenAPI's Rest client doesn't expose.
102        # Since the default values work well, it seems setting `(socket.SO_KEEPALIVE, 1)` is sufficient.
103        # Leaving this code here for future reference.
104        # elif platform == 'win32' and hasattr(socket, "SIO_KEEPALIVE_VALS"):
105        #     socket.ioctl((socket.SIO_KEEPALIVE_VALS, (1, keep_alive_idle_sec * 1000, keep_alive_interval_sec * 1000)))
106
107        # TCP Keep Alive Probes for Mac OS
108        elif platform == "darwin":
109            TCP_KEEPALIVE = 0x10
110            socket_params += [(socket.IPPROTO_TCP, TCP_KEEPALIVE, keep_alive_interval_sec)]
111
112        return socket_params