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 Configuration as OpenApiConfiguration
 11
 12TCP_KEEPINTVL = 60  # Sec
 13TCP_KEEPIDLE = 300  # Sec
 14TCP_KEEPCNT = 4
 15
 16
 17class OpenApiConfigFactory:
 18    @classmethod
 19    def build(cls, api_key: str, host: Optional[str] = None, **kwargs):
 20        openapi_config = OpenApiConfiguration()
 21        openapi_config.api_key = {"ApiKeyAuth": api_key}
 22        openapi_config.host = host
 23        openapi_config.ssl_ca_cert = certifi.where()
 24        openapi_config.socket_options = cls._get_socket_options()
 25        openapi_config.discard_unknown_keys = True
 26
 27        return openapi_config
 28
 29    @classmethod
 30    def copy(
 31        cls, openapi_config: OpenApiConfiguration, api_key: str, host: str
 32    ) -> 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 += [(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, keep_alive_interval_sec)]
 92            socket_params += [(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, keep_alive_tries)]
 93
 94        # TCP Keep Alive Probes for Windows OS
 95        # NOTE: Changing TCP KA params on windows is done via a different mechanism which OpenAPI's Rest client doesn't expose.
 96        # Since the default values work well, it seems setting `(socket.SO_KEEPALIVE, 1)` is sufficient.
 97        # Leaving this code here for future reference.
 98        # elif platform == 'win32' and hasattr(socket, "SIO_KEEPALIVE_VALS"):
 99        #     socket.ioctl((socket.SIO_KEEPALIVE_VALS, (1, keep_alive_idle_sec * 1000, keep_alive_interval_sec * 1000)))
100
101        # TCP Keep Alive Probes for Mac OS
102        elif platform == "darwin":
103            TCP_KEEPALIVE = 0x10
104            socket_params += [(socket.IPPROTO_TCP, TCP_KEEPALIVE, keep_alive_interval_sec)]
105
106        return socket_params