lusid-sdk 2.0.50b0__py3-none-any.whl → 2.1.684__py3-none-any.whl

lusid/extensions/retry.py

@@ -1,19 +1,33 @@
 import time
-from typing import overload
+from typing import Optional
 
-from lusid import ApiException
 import asyncio
 
+from lusid.configuration import Configuration
+from lusid.exceptions import ApiException
+
 
 class RetryingRestWrapper:
     """Wrapper for HTTP requests
     Which retries on failure
     And waits the amount of time specified in the Retry After header.
     """
-    def __init__(self, rest_object, retries: int = 3):
+    def __init__(
+        self, 
+        rest_object, 
+        retries: int = Configuration.DEFAULT_RETRIES, 
+        rate_limit_retries: int = Configuration.DEFAULT_RATE_LIMIT_RETRIES
+    ):
         if not isinstance(retries, int):
-            raise ValueError(f"retries should be an int, found {type(self.retries)}")
+            raise TypeError(f"retries should be an int, found {type(retries)}")
+        if retries < 0:
+            raise ValueError(f"retries should be greater than or equal to zero but was '{retries}'")
+        if not isinstance(rate_limit_retries, int):
+            raise TypeError(f"rate_limit_retries should be an int, found {type(rate_limit_retries)}")
+        if rate_limit_retries < 0:
+            raise ValueError(f"rate_limit_retries should be greater than or equal to zero but was '{rate_limit_retries}'")
         self.retries: int = retries
+        self.rate_limit_retries: int = rate_limit_retries
         self.rest_object = rest_object
 
     def request(
@@ -26,9 +40,12 @@ class RetryingRestWrapper:
         post_params=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
-        tries = 0
-        while tries < self.retries + 1:
+        retries_count = 0
+        rate_limit_retries_count = 0
+
+        while True:
             try:
                 return self.rest_object.request(
                     method,
@@ -39,17 +56,26 @@ class RetryingRestWrapper:
                     post_params,
                     _preload_content,
                     _request_timeout,
+                    opts,
                 )
             except ApiException as ex:
-                
-                retry_after = ex.headers.get("Retry-After")
 
-                # have done max number of retries
-                if tries >= self.retries:
-                    raise
+                if ex.status == 429 and ((opts != None and opts.rate_limit_retries != None) or self.rate_limit_retries != None):
+                    # check for limit of rate limit retries
+                    limit = opts.rate_limit_retries if (opts and opts.rate_limit_retries != None) else self.rate_limit_retries
+                    if rate_limit_retries_count >= limit:
+                        raise
+                    rate_limit_retries_count += 1
+                else:
+                    # check for limit of all other retries
+                    if retries_count >= self.retries:
+                        raise
+                    retries_count += 1
+
+                retry_after = ex.headers.get("Retry-After")
 
                 # try after delay
-                elif retry_after is not None:
+                if retry_after is not None:
                     if not isinstance(retry_after, float):
                         try:
                             retry_after = float(retry_after)
@@ -61,7 +87,6 @@ class RetryingRestWrapper:
                 # no retry header
                 else:
                     raise
-                tries += 1
 
     def get_request(
         self,
@@ -70,12 +95,14 @@ class RetryingRestWrapper:
         query_params=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("GET", url,
                             headers=headers,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            query_params=query_params)
+                            query_params=query_params,
+                            opts=opts)
 
     def head_request(
         self,
@@ -84,12 +111,14 @@ class RetryingRestWrapper:
         query_params=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("HEAD", url,
                             headers=headers,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            query_params=query_params)
+                            query_params=query_params,
+                            opts=opts)
 
     def options_request(
         self,
@@ -100,6 +129,7 @@ class RetryingRestWrapper:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("OPTIONS", url,
                             headers=headers,
@@ -107,7 +137,8 @@ class RetryingRestWrapper:
                             post_params=post_params,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            body=body)
+                            body=body,
+                            opts=opts)
 
     def delete_request(
         self,
@@ -117,13 +148,15 @@ class RetryingRestWrapper:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("DELETE", url,
                             headers=headers,
                             query_params=query_params,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            body=body)
+                            body=body,
+                            opts=opts)
 
     def post_request(
         self,
@@ -134,6 +167,7 @@ class RetryingRestWrapper:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("POST", url,
                             headers=headers,
@@ -141,7 +175,8 @@ class RetryingRestWrapper:
                             post_params=post_params,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            body=body)
+                            body=body,
+                            opts=opts)
 
     def put_request(
         self,
@@ -152,6 +187,7 @@ class RetryingRestWrapper:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("PUT", url,
                             headers=headers,
@@ -159,7 +195,8 @@ class RetryingRestWrapper:
                             post_params=post_params,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            body=body)
+                            body=body,
+                            opts=opts)
 
     def patch_request(
         self,
@@ -170,6 +207,7 @@ class RetryingRestWrapper:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return self.request("PATCH", url,
                             headers=headers,
@@ -177,7 +215,8 @@ class RetryingRestWrapper:
                             post_params=post_params,
                             _preload_content=_preload_content,
                             _request_timeout=_request_timeout,
-                            body=body)
+                            body=body,
+                            opts=opts)
 
 
 class RetryingRestWrapperAsync:
@@ -185,10 +224,17 @@ class RetryingRestWrapperAsync:
     Which retries on failure
     And waits the amount of time specified in the Retry After header.
     """
-    def __init__(self, rest_object, retries: int = 3):
+    def __init__(self, rest_object, retries: int = 3, rate_limit_retries: Optional[int] = Configuration.DEFAULT_RATE_LIMIT_RETRIES):
         if not isinstance(retries, int):
-            raise ValueError(f"retries should be an int, found {type(self.retries)}")
+            raise TypeError(f"retries should be an int, found {type(retries)}")
+        if retries < 0:
+            raise ValueError(f"retries should be greater than or equal to zero but was '{retries}'")
+        if not isinstance(rate_limit_retries, int):
+            raise TypeError(f"rate_limit_retries should be an int, found {type(rate_limit_retries)}")
+        if rate_limit_retries < 0:
+            raise ValueError(f"rate_limit_retries should be greater than or equal to zero but was '{rate_limit_retries}'")
         self.retries: int = retries
+        self.rate_limit_retries: Optional[int] = rate_limit_retries
         self.rest_object = rest_object
 
     async def close(self):
@@ -204,9 +250,12 @@ class RetryingRestWrapperAsync:
         post_params=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
-        tries = 0
-        while tries < self.retries + 1:
+        retries_count = 0
+        rate_limit_retries_count = 0
+        
+        while True:
             try:
                 return await self.rest_object.request(
                     method,
@@ -217,16 +266,26 @@ class RetryingRestWrapperAsync:
                     post_params,
                     _preload_content,
                     _request_timeout,
+                    opts
                 )
             except ApiException as ex:
-                retry_after = ex.headers.get("Retry-After")
 
-                # have done max number of retries
-                if tries >= self.retries:
-                    raise
+                if ex.status == 429 and ((opts and opts.rate_limit_retries != None) or self.rate_limit_retries != None):
+                    # check for limit of rate limit retries
+                    limit = opts.rate_limit_retries if (opts and opts.rate_limit_retries != None) else self.rate_limit_retries
+                    if rate_limit_retries_count >= limit:
+                        raise
+                    rate_limit_retries_count += 1
+                else:
+                    # check for limit of all other retries
+                    if retries_count >= self.retries:
+                        raise
+                    retries_count += 1
+
+                retry_after = ex.headers.get("Retry-After")
 
                 # try after delay
-                elif retry_after is not None:
+                if retry_after is not None:
                     if not isinstance(retry_after, float):
                         try:
                             retry_after = float(retry_after)
@@ -238,7 +297,6 @@ class RetryingRestWrapperAsync:
                 # no retry header
                 else:
                     raise
-                tries += 1
 
     async def get_request(
         self,
@@ -247,12 +305,14 @@ class RetryingRestWrapperAsync:
         query_params=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("GET", url,
                                    headers=headers,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   query_params=query_params))
+                                   query_params=query_params,
+                                   opts=opts))
 
     async def head_request(
         self,
@@ -261,12 +321,14 @@ class RetryingRestWrapperAsync:
         query_params=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("HEAD", url,
                                    headers=headers,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   query_params=query_params))
+                                   query_params=query_params,
+                                   opts=opts))
 
     async def options_request(
         self,
@@ -277,6 +339,7 @@ class RetryingRestWrapperAsync:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("OPTIONS", url,
                                    headers=headers,
@@ -284,7 +347,8 @@ class RetryingRestWrapperAsync:
                                    post_params=post_params,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   body=body))
+                                   body=body,
+                                   opts=opts))
 
     async def delete_request(
         self,
@@ -294,13 +358,15 @@ class RetryingRestWrapperAsync:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("DELETE", url,
                                    headers=headers,
                                    query_params=query_params,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   body=body))
+                                   body=body,
+                                   opts=opts))
 
     async def post_request(
         self,
@@ -311,6 +377,7 @@ class RetryingRestWrapperAsync:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("POST", url,
                                    headers=headers,
@@ -318,7 +385,8 @@ class RetryingRestWrapperAsync:
                                    post_params=post_params,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   body=body))
+                                   body=body,
+                                   opts=opts))
 
     async def put_request(
         self,
@@ -329,6 +397,7 @@ class RetryingRestWrapperAsync:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("PUT", url,
                                    headers=headers,
@@ -336,7 +405,8 @@ class RetryingRestWrapperAsync:
                                    post_params=post_params,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   body=body))
+                                   body=body,
+                                   opts=opts))
 
     async def patch_request(
         self,
@@ -347,6 +417,7 @@ class RetryingRestWrapperAsync:
         body=None,
         _preload_content=True,
         _request_timeout=None,
+        opts=None,
     ):
         return (await self.request("PATCH", url,
                                    headers=headers,
@@ -354,4 +425,5 @@ class RetryingRestWrapperAsync:
                                    post_params=post_params,
                                    _preload_content=_preload_content,
                                    _request_timeout=_request_timeout,
-                                   body=body))
+                                   body=body,
+                                   opts=opts))

lusid/extensions/socket_keep_alive.py

@@ -5,6 +5,15 @@ from urllib3.connection import HTTPConnection
 
 __logger__ = logging.getLogger(__name__)
 
+# The content to send on Mac OS in the TCP Keep Alive probe
+TCP_KEEPALIVE = 0x10
+# The maximum time to keep the connection idle before sending probes
+TCP_KEEP_IDLE = 60
+# The interval between probes
+TCP_KEEPALIVE_INTERVAL = 60
+# The maximum number of failed probes before terminating the connection
+TCP_KEEP_CNT = 3
+
 
 def keep_alive_socket_options() -> Sequence:
     """Returns default socket options for all platforms for
@@ -15,15 +24,6 @@ def keep_alive_socket_options() -> Sequence:
     Sequence
         Set of socket options.
     """
-    # The content to send on Mac OS in the TCP Keep Alive probe
-    TCP_KEEPALIVE = 0x10
-    # The maximum time to keep the connection idle before sending probes
-    TCP_KEEP_IDLE = 60
-    # The interval between probes
-    TCP_KEEPALIVE_INTERVAL = 60
-    # The maximum number of failed probes before terminating the connection
-    TCP_KEEP_CNT = 3
-
     try:
         # linux and some windows runtimes
         return HTTPConnection.default_socket_options + [
@@ -35,18 +35,17 @@ def keep_alive_socket_options() -> Sequence:
     except AttributeError:
         pass
     try:
-        # other windows runtimes
+        # darwin
         return HTTPConnection.default_socket_options + [
-            socket.SIO_KEEPALIVE_VALS,
-            (1, TCP_KEEP_IDLE * 1000, TCP_KEEPALIVE_INTERVAL * 1000),
+            (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+            (socket.IPPROTO_TCP, TCP_KEEPALIVE, TCP_KEEPALIVE_INTERVAL),
         ]
     except AttributeError:
         pass
     try:
-        # darwin
+        # windows
         return HTTPConnection.default_socket_options + [
-            (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
-            (socket.IPPROTO_TCP, TCP_KEEPALIVE, TCP_KEEPALIVE_INTERVAL),
+            (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
         ]
     except AttributeError:
         __logger__.exception("Unable to set TCP Keep-alive socket options")

lusid/extensions/tcp_keep_alive_connector.py

@@ -1,7 +1,29 @@
-from ssl import SSLContext
 import aiohttp
-import asyncio
-from typing import List, Tuple, Any, Union, Optional
+from typing import List, Tuple, Any, Union
+import socket
+from urllib3 import HTTPSConnectionPool, HTTPConnectionPool
+from lusid.extensions.socket_keep_alive import TCP_KEEPALIVE_INTERVAL, TCP_KEEP_IDLE
+
+
+def adjust_connection_socket(conn):
+    """
+    Adjusts the socket settings so that the client sends a TCP keep alive probe over the connection. This is only
+    applied where possible, if the ability to set the socket options is not available, for example using Anaconda,
+    then the settings will be left as is.
+
+    :param conn: The connection to update the socket settings for
+    :param str protocol: The protocol of the connection
+
+    :return: None
+    """
+    try:
+        # set TCP keep alive interval on windows connections
+        conn.ioctl(
+            socket.SIO_KEEPALIVE_VALS,
+            (1, TCP_KEEP_IDLE * 1000, TCP_KEEPALIVE_INTERVAL * 1000),
+        )
+    except AttributeError:
+        pass
 
 
 class TcpKeepAliveConnector(aiohttp.TCPConnector):
@@ -11,56 +33,81 @@ class TcpKeepAliveConnector(aiohttp.TCPConnector):
     """
     def __init__(
         self,
-        *,
-        verify_ssl: bool = True,
-        fingerprint: Optional[bytes] = None,
-        use_dns_cache: bool = True,
-        ttl_dns_cache: Optional[int] = 10,
-        family: int = 0,
-        ssl_context: Optional[SSLContext] = None,
-        ssl: Union[None, bool, aiohttp.Fingerprint, SSLContext] = None,
-        local_addr: Optional[Tuple[str, int]] = None,
-        resolver: Optional[aiohttp.abc.AbstractResolver] = None,
-        keepalive_timeout: Union[None, float, object] = aiohttp.helpers.sentinel,
-        force_close: bool = False,
-        limit: int = 100,
-        limit_per_host: int = 0,
-        enable_cleanup_closed: bool = False,
-        loop: Optional[asyncio.AbstractEventLoop] = None,
+        connector: aiohttp.TCPConnector,
         socket_options: Union[Tuple[Any, Any, Any], Tuple[Any, Any, None, int]]
     ) -> None:
-        super().__init__(
-            verify_ssl=verify_ssl,
-            fingerprint=fingerprint,
-            use_dns_cache=use_dns_cache,
-            ttl_dns_cache=ttl_dns_cache,
-            family=family,
-            ssl_context=ssl_context,
-            ssl=ssl,
-            local_addr=local_addr,
-            resolver=resolver,
-            keepalive_timeout=keepalive_timeout,
-            force_close=force_close,
-            limit=limit,
-            limit_per_host=limit_per_host,
-            enable_cleanup_closed=enable_cleanup_closed,
-            loop=loop,
-        )
+        self.__connector = connector
         self.socket_options = socket_options or []
 
-    async def _create_connection(
+    @property
+    def _timeout_ceil_threshold(self):
+        return self.__connector._timeout_ceil_threshold
+
+    @property
+    def _loop(self):
+        return self.__connector._loop
+
+    @property
+    def closed(self):
+        return self.__connector.closed
+
+    async def close(self) -> None:
+        await self.__connector.close()
+
+    async def connect(
         self,
         req: aiohttp.ClientRequest,
         traces: List[aiohttp.tracing.Trace],
         timeout: aiohttp.ClientTimeout,
-    ) -> aiohttp.client_proto.ResponseHandler:
-        if req.proxy:
-            transport, proto = await self._create_proxy_connection(req, traces, timeout)
-        else:
-            transport, proto = await self._create_direct_connection(
-                req, traces, timeout
-            )
-        sock = transport.get_extra_info("socket")
+    ) -> aiohttp.connector.Connection:
+        """Wraps TCP connector, each new connection will have socket options
+        and windows ioctl for keep alives applied
+
+        Parameters
+        ----------
+        req : aiohttp.ClientRequest
+        traces : List[aiohttp.tracing.Trace]
+        timeout : aiohttp.ClientTimeout
+        """
+        connection = await self.__connector.connect(req, traces, timeout)
+        sock = connection.protocol.transport.get_extra_info("socket")
         for option in self.socket_options:
             sock.setsockopt(*option)
-        return proto
+        adjust_connection_socket(sock)
+        return connection
+
+
+class TCPKeepAliveHTTPSConnectionPool(HTTPSConnectionPool):
+    """
+    This class overrides the _validate_conn method in the HTTPSConnectionPool class. This is the entry point to use
+    for modifying the socket as it is called after the socket is created and before the request is made.
+    """
+
+    def _validate_conn(self, conn):
+        """
+        Called right before a request is made, after the socket is created.
+        """
+        # Call the method on the base class
+        super()._validate_conn(conn)
+
+        # Set up TCP Keep Alive probes, this is the only line added to this function
+        adjust_connection_socket(conn)
+
+
+class TCPKeepAliveHTTPConnectionPool(HTTPConnectionPool):
+    """
+    This class overrides the _validate_conn method in the HTTPSConnectionPool class. This is the entry point to use
+    for modifying the socket as it is called after the socket is created and before the request is made.
+
+    In the base class this method is passed completely.
+    """
+
+    def _validate_conn(self, conn):
+        """
+        Called right before a request is made, after the socket is created.
+        """
+        # Call the method on the base class
+        super()._validate_conn(conn)
+
+        # Set up TCP Keep Alive probes, this is the only line added to this function
+        adjust_connection_socket(conn)