diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index 386ad325..4127a985 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -491,6 +491,8 @@ class Config: metadata_service_url = ConfigAttribute(env='DATABRICKS_METADATA_SERVICE_URL', auth='metadata-service', sensitive=True) + max_connection_pools: int = ConfigAttribute() + max_connections_per_pool: int = ConfigAttribute() def __init__(self, *, @@ -893,7 +895,31 @@ def __init__(self, cfg: Config = None): self._session = requests.Session() self._session.auth = self._authenticate - self._session.mount("https://", HTTPAdapter(max_retries=retry_strategy)) + + # Number of urllib3 connection pools to cache before discarding the least + # recently used pool. Python requests default value is 10. + pool_connections = cfg.max_connection_pools + if pool_connections is None: + pool_connections = 20 + + # The maximum number of connections to save in the pool. Improves performance + # in multithreaded situations. For now, we're setting it to the same value + # as connection_pool_size. + pool_maxsize = cfg.max_connections_per_pool + if cfg.max_connections_per_pool is None: + pool_maxsize = pool_connections + + # If pool_block is False, then more connections will are created, + # but not saved after the first use. Blocks when no free connections are available. + # urllib3 ensures that no more than pool_maxsize connections are used at a time. + # Prevents platform from flooding. By default, requests library doesn't block. + pool_block = True + + http_adapter = HTTPAdapter(max_retries=retry_strategy, + pool_connections=pool_connections, + pool_maxsize=pool_maxsize, + pool_block=pool_block) + self._session.mount("https://", http_adapter) @property def account_id(self) -> str: