[SPT-1668] Add proxy, custom CA, and connection pool options

README.md

@@ -135,6 +135,99 @@ print(len(options_chain))
 
 Also, please refer to the API documentation to get a full understanding of how the API works and the supported arguments. All required arguments are annotated with red asterisks " * " and argument examples are set.
 
+## Network Configuration
+
+Options for corporate networks, proxies, and multithreaded use. They apply to both `RESTClient` and `WebSocketClient` unless noted.
+
+### Custom CA certificates (`certificate verify failed`)
+
+If you're behind a corporate firewall that inspects TLS traffic, requests can fail with `SSLError: [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed`. Point the client at your company's CA bundle.
+
+**With environment variables (no code changes):**
+
+```bash
+export SSL_CERT_FILE=/path/to/corporate-ca-bundle.pem
+# or, if you already set this for the requests library:
+export REQUESTS_CA_BUNDLE=/path/to/corporate-ca-bundle.pem
+```
+
+**In code:**
+
+```python
+from massive import RESTClient, WebSocketClient
+
+client = RESTClient(ca_certs="/path/to/corporate-ca-bundle.pem")
+ws_client = WebSocketClient(ca_certs="/path/to/corporate-ca-bundle.pem")
+```
+
+The client uses the first one it finds: `ca_certs`, then `SSL_CERT_FILE`, then `REQUESTS_CA_BUNDLE`, then the bundled [certifi](https://pypi.org/project/certifi/) certificates. If an environment variable points to a file that doesn't exist, it's skipped and a warning is logged.
+
+The bundle replaces certifi, so it must also include the public root certificates. If you only have your company's root certificate, add it to certifi's bundle:
+
+```bash
+cat "$(python -m certifi)" corporate-root.pem > corporate-ca-bundle.pem
+```
+
+### Proxies
+
+**REST:** pass the proxy URL. `RESTClient` doesn't read `HTTPS_PROXY` on its own, so existing setups keep connecting the same way.
+
+```python
+import os
+from massive import RESTClient
+
+client = RESTClient(proxy="http://proxy.example.com:8080")
+
+# with a username and password
+client = RESTClient(proxy="http://user:password@proxy.example.com:8080")
+
+# use the proxy from your environment, like the requests library does
+client = RESTClient(proxy=os.environ.get("HTTPS_PROXY"))
+```
+
+**WebSocket:** `WebSocketClient` already uses the proxy from `HTTPS_PROXY` / `HTTP_PROXY` (and respects `NO_PROXY`). To set one explicitly, or to turn it off, pass `proxy` to `run()` or `connect()`:
+
+```python
+ws_client = WebSocketClient(subscriptions=["T.AAPL"])
+ws_client.run(handle_msg, proxy="http://user:password@proxy.example.com:8080")
+# ws_client.run(handle_msg, proxy=None)  # connect directly, ignoring HTTPS_PROXY
+```
+
+**Behind a proxy that inspects TLS**, you'll usually need both:
+
+```python
+client = RESTClient(
+    proxy="http://proxy.example.com:8080",
+    ca_certs="/path/to/corporate-ca-bundle.pem",
+)
+```
+
+### Using the REST client from multiple threads
+
+`RESTClient` can be shared across threads. It keeps up to `maxsize` connections (default 10) open for reuse, so set it to at least the number of threads making requests at the same time. Otherwise, extra threads open new connections and redo the TLS handshake on most requests.
+
+```python
+from concurrent.futures import ThreadPoolExecutor
+from massive import RESTClient
+
+client = RESTClient(maxsize=20)
+tickers = ["AAPL", "MSFT", "NVDA", "AMZN"]
+
+with ThreadPoolExecutor(max_workers=20) as pool:
+    results = list(pool.map(lambda t: client.get_previous_close_agg(t), tickers))
+```
+
+Note that `num_pools` doesn't help here: it's the number of different *hosts* the client keeps connections for, and the client only talks to one.
+
+### Other connection options
+
+Any other [urllib3 pool option](https://urllib3.readthedocs.io/en/stable/reference/urllib3.connectionpool.html#urllib3.HTTPConnectionPool) can be passed with `connection_pool_kw`. These take precedence over the client's defaults, so avoid overriding `headers`, which carries your API key.
+
+```python
+# wait for a free connection instead of opening extra ones beyond maxsize
+client = RESTClient(maxsize=4, connection_pool_kw={"block": True})
+```
+
 ## Debugging with RESTClient
 
 Sometimes you may find it useful to see the actual request and response details while working with the API. The `RESTClient` allows for this through its `trace=True` option.

massive/certs.py

@@ -0,0 +1,27 @@
+import logging
+import os
+from typing import Optional
+
+import certifi
+
+
+def ca_bundle_path(logger: Optional[logging.Logger] = None) -> str:
+    """
+    Path of the CA bundle used to verify TLS connections.
+
+    Honors the standard ``SSL_CERT_FILE`` and ``REQUESTS_CA_BUNDLE`` environment
+    variables (used for example behind corporate TLS inspection), falling back to
+    the certifi bundle. A variable that doesn't point to an existing file is
+    skipped with a warning, so a stale value doesn't break connections.
+    """
+    for env_var in ("SSL_CERT_FILE", "REQUESTS_CA_BUNDLE"):
+        path = os.environ.get(env_var)
+        if not path:
+            continue
+        if os.path.isfile(path):
+            return path
+        if logger:
+            logger.warning(
+                "%s is set to %r, which isn't a file; ignoring it", env_var, path
+            )
+    return certifi.where()

massive/rest/__init__.py

@@ -20,7 +20,9 @@
     ContractsClient,
 )
 from .vX import VXClient
-from typing import Optional, Any
+from .base import logger
+from ..certs import ca_bundle_path
+from typing import Optional, Any, Dict
 import os
 
 BASE = "https://api.massive.com"
@@ -60,7 +62,22 @@ def __init__(
         verbose: bool = False,
         trace: bool = False,
         custom_json: Optional[Any] = None,
+        maxsize: int = 10,
+        ca_certs: Optional[str] = None,
+        proxy: Optional[str] = None,
+        connection_pool_kw: Optional[Dict[str, Any]] = None,
     ):
+        """
+        :param maxsize: Connections kept open per host, for reuse across threads.
+        :param ca_certs: Path to a CA bundle to verify TLS with. Defaults to the
+            SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable, then certifi.
+        :param proxy: Proxy URL to send requests through, e.g. "http://user:pass@proxy:8080".
+            Environment proxy variables (HTTPS_PROXY etc.) are not used.
+        :param connection_pool_kw: Other urllib3 PoolManager options. They override
+            the client's defaults.
+        """
+        # resolve once, so the main and vx clients share it and warn only once
+        ca_certs = ca_certs or ca_bundle_path(logger)
         super().__init__(
             api_key=api_key,
             connect_timeout=connect_timeout,
@@ -72,6 +89,10 @@ def __init__(
             verbose=verbose,
             trace=trace,
             custom_json=custom_json,
+            maxsize=maxsize,
+            ca_certs=ca_certs,
+            proxy=proxy,
+            connection_pool_kw=connection_pool_kw,
         )
         self.vx = VXClient(
             api_key=api_key,
@@ -84,4 +105,8 @@ def __init__(
             verbose=verbose,
             trace=trace,
             custom_json=custom_json,
+            maxsize=maxsize,
+            ca_certs=ca_certs,
+            proxy=proxy,
+            connection_pool_kw=connection_pool_kw,
         )

massive/rest/base.py

@@ -1,4 +1,3 @@
-import certifi
 import json
 import urllib3
 import inspect
@@ -8,6 +7,7 @@
 from datetime import datetime
 from importlib.metadata import version, PackageNotFoundError
 from .models.request import RequestOptionBuilder
+from ..certs import ca_bundle_path
 from ..logging import get_logger
 import logging
 from urllib.parse import urlencode, urlparse
@@ -34,6 +34,10 @@ def __init__(
         verbose: bool,
         trace: bool,
         custom_json: Optional[Any] = None,
+        maxsize: int = 10,
+        ca_certs: Optional[str] = None,
+        proxy: Optional[str] = None,
+        connection_pool_kw: Optional[Dict[str, Any]] = None,
     ):
         if api_key is None:
             raise AuthError(
@@ -73,14 +77,30 @@ def __init__(
 
         # https://urllib3.readthedocs.io/en/stable/reference/urllib3.poolmanager.html
         # https://urllib3.readthedocs.io/en/stable/reference/urllib3.connectionpool.html#urllib3.HTTPConnectionPool
-        self.client = urllib3.PoolManager(
-            num_pools=num_pools,
-            headers=self.headers,  # default headers sent with each request.
-            ca_certs=certifi.where(),
-            cert_reqs="CERT_REQUIRED",
-            retries=retry_strategy,  # use the customized Retry instance
-            timeout=self.timeout,  # set timeout for each request
-        )
+        pool_kw: Dict[str, Any] = {
+            "num_pools": num_pools,
+            # connections kept per host; urllib3 defaults to 1, which makes threaded
+            # clients reconnect (and redo the TLS handshake) for most requests
+            "maxsize": maxsize,
+            "headers": self.headers,  # default headers sent with each request.
+            "ca_certs": ca_certs or ca_bundle_path(logger),
+            "cert_reqs": "CERT_REQUIRED",
+            "retries": retry_strategy,  # use the customized Retry instance
+            "timeout": self.timeout,  # set timeout for each request
+        }
+        # any other urllib3 pool options, which take precedence over the defaults
+        pool_kw.update(connection_pool_kw or {})
+        self.client: urllib3.PoolManager
+        if proxy:
+            # https://urllib3.readthedocs.io/en/stable/advanced-usage.html#proxies
+            proxy_auth = urllib3.util.parse_url(proxy).auth
+            if proxy_auth:
+                pool_kw.setdefault(
+                    "proxy_headers", urllib3.make_headers(proxy_basic_auth=proxy_auth)
+                )
+            self.client = urllib3.ProxyManager(proxy, **pool_kw)
+        else:
+            self.client = urllib3.PoolManager(**pool_kw)
 
         if verbose:
             logger.setLevel(logging.DEBUG)

massive/websocket/__init__.py

@@ -7,10 +7,10 @@
 import asyncio
 import ssl
 import time
-import certifi
 from .models import *
 from websockets.asyncio.client import connect, ClientConnection
 from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError
+from ..certs import ca_bundle_path
 from ..logging import get_logger
 import logging
 from ..exceptions import AuthError
@@ -35,6 +35,7 @@ def __init__(
         max_reconnects: Optional[int] = 5,
         secure: bool = True,
         custom_json: Optional[Any] = None,
+        ca_certs: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -47,6 +48,7 @@ def __init__(
         :param subscriptions: List of subscription parameters.
         :param max_reconnects: How many times to reconnect on network outage before ending .connect event loop.
         :param custom_json: Optional module exposing loads/dumps functions (similar to Python's json module) to be used for JSON conversions.
+        :param ca_certs: Path to a CA bundle to verify TLS with. Defaults to the SSL_CERT_FILE or REQUESTS_CA_BUNDLE environment variable, then certifi.
         :return: A client.
         """
         if api_key is None:
@@ -78,6 +80,7 @@ def __init__(
         self.schedule_resub = True
         self._server_errors: Counter = Counter()
         self._server_errors_since = 0.0
+        self.ca_certs = ca_certs
         if custom_json:
             self.json = custom_json
         else:
@@ -102,11 +105,7 @@ async def connect(
         """
         reconnects = 0
         logger.debug("connect: %s", self.url)
-        # darwin needs some extra <3
-        ssl_context = None
-        if self.url.startswith("wss://"):
-            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-            ssl_context.load_verify_locations(certifi.where())
+        ssl_context = self._ssl_context()
 
         last_exc = None
         async for s in connect(
@@ -190,6 +189,14 @@ async def connect(
         ):
             raise last_exc
 
+    def _ssl_context(self) -> Optional[ssl.SSLContext]:
+        # darwin needs some extra <3
+        if not self.url.startswith("wss://"):
+            return None
+        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        context.load_verify_locations(self.ca_certs or ca_bundle_path(logger))
+        return context
+
     def run(
         self,
         handle_msg: Union[

test_rest/test_connection.py

@@ -0,0 +1,90 @@
+import os
+import tempfile
+import unittest
+from unittest import mock
+
+import certifi
+import urllib3
+
+from massive import RESTClient
+
+
+class ConnectionOptionsTest(unittest.TestCase):
+    def setUp(self):
+        tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(tmp.cleanup)
+        self.ssl_ca = os.path.join(tmp.name, "ssl-ca.pem")
+        self.requests_ca = os.path.join(tmp.name, "requests-ca.pem")
+        for path in (self.ssl_ca, self.requests_ca):
+            open(path, "w").close()
+
+    def client(self, env=None, **kwargs):
+        with mock.patch.dict(os.environ, env or {}, clear=True):
+            return RESTClient(api_key="key", **kwargs)
+
+    def test_ca_bundle_defaults_to_certifi(self):
+        c = self.client()
+        self.assertEqual(c.client.connection_pool_kw["ca_certs"], certifi.where())
+
+    def test_ca_bundle_from_env(self):
+        c = self.client({"REQUESTS_CA_BUNDLE": self.requests_ca})
+        self.assertEqual(c.client.connection_pool_kw["ca_certs"], self.requests_ca)
+        c = self.client(
+            {"SSL_CERT_FILE": self.ssl_ca, "REQUESTS_CA_BUNDLE": self.requests_ca}
+        )
+        self.assertEqual(c.client.connection_pool_kw["ca_certs"], self.ssl_ca)
+
+    def test_missing_env_bundle_is_skipped_with_a_warning(self):
+        with self.assertLogs("RESTClient", level="WARNING") as logs:
+            c = self.client({"SSL_CERT_FILE": "/does-not-exist.pem"})
+        self.assertEqual(c.client.connection_pool_kw["ca_certs"], certifi.where())
+        self.assertEqual(c.vx.client.connection_pool_kw["ca_certs"], certifi.where())
+        self.assertEqual(len(logs.output), 1, logs.output)
+        self.assertIn("SSL_CERT_FILE", logs.output[0])
+        with self.assertLogs("RESTClient", level="WARNING"):
+            c = self.client(
+                {
+                    "SSL_CERT_FILE": "/does-not-exist.pem",
+                    "REQUESTS_CA_BUNDLE": self.requests_ca,
+                }
+            )
+        self.assertEqual(c.client.connection_pool_kw["ca_certs"], self.requests_ca)
+
+    def test_ca_certs_argument_wins_over_env(self):
+        c = self.client({"SSL_CERT_FILE": self.ssl_ca}, ca_certs="/explicit-ca.pem")
+        self.assertEqual(c.client.connection_pool_kw["ca_certs"], "/explicit-ca.pem")
+
+    def test_ca_certs_applies_to_vx(self):
+        c = self.client(ca_certs="/explicit-ca.pem")
+        self.assertEqual(c.vx.client.connection_pool_kw["ca_certs"], "/explicit-ca.pem")
+
+    def test_maxsize(self):
+        self.assertEqual(self.client().client.connection_pool_kw["maxsize"], 10)
+        c = self.client(maxsize=32)
+        self.assertEqual(c.client.connection_pool_kw["maxsize"], 32)
+
+    def test_connection_pool_kw_overrides_defaults(self):
+        timeout = urllib3.Timeout(connect=1, read=2)
+        c = self.client(connection_pool_kw={"block": True, "timeout": timeout})
+        self.assertTrue(c.client.connection_pool_kw["block"])
+        self.assertIs(c.client.connection_pool_kw["timeout"], timeout)
+        # defaults that weren't overridden are kept
+        self.assertEqual(c.client.connection_pool_kw["cert_reqs"], "CERT_REQUIRED")
+
+    def test_no_proxy_by_default(self):
+        c = self.client({"HTTPS_PROXY": "http://proxy.example:8080"})
+        self.assertNotIsInstance(c.client, urllib3.ProxyManager)
+
+    def test_proxy(self):
+        c = self.client(proxy="http://user:secret@proxy.example:8080")
+        self.assertIsInstance(c.client, urllib3.ProxyManager)
+        self.assertEqual(c.client.proxy.host, "proxy.example")
+        self.assertEqual(
+            c.client.proxy_headers["proxy-authorization"],
+            urllib3.make_headers(proxy_basic_auth="user:secret")["proxy-authorization"],
+        )
+        self.assertEqual(c.client.connection_pool_kw["maxsize"], 10)
+
+
+if __name__ == "__main__":
+    unittest.main()