feat: add bpd.options.bigquery.requests_transport_adapters option

This allows for overriding requests-based settings such as the maximum
connection pool size.

Author: tswast

bigframes/_config/bigquery_options.py

@@ -16,10 +16,11 @@
 
 from __future__ import annotations
 
-from typing import Literal, Optional
+from typing import Literal, Optional, Sequence, Tuple
 import warnings
 
 import google.auth.credentials
+import requests.adapters
 
 import bigframes.enums
 import bigframes.exceptions as bfe
@@ -90,6 +91,9 @@ def __init__(
         allow_large_results: bool = False,
         ordering_mode: Literal["strict", "partial"] = "strict",
         client_endpoints_override: Optional[dict] = None,
+        requests_transport_adapters: Sequence[
+            Tuple[str, requests.adapters.BaseAdapter]
+        ] = (),
     ):
         self._credentials = credentials
         self._project = project
@@ -100,6 +104,7 @@ def __init__(
         self._kms_key_name = kms_key_name
         self._skip_bq_connection_check = skip_bq_connection_check
         self._allow_large_results = allow_large_results
+        self._requests_transport_adapters = requests_transport_adapters
         self._session_started = False
         # Determines the ordering strictness for the session.
         self._ordering_mode = _validate_ordering_mode(ordering_mode)
@@ -379,3 +384,32 @@ def client_endpoints_override(self, value: dict):
             )
 
         self._client_endpoints_override = value
+
+    @property
+    def requests_transport_adapters(
+        self,
+    ) -> Sequence[Tuple[str, requests.adapters.BaseAdapter]]:
+        """Transport adapters for requests-based REST clients such as the
+        google-cloud-bigquery package.
+
+        For more details, see the explanation in `requests guide to transport
+        adapters
+        <https://requests.readthedocs.io/en/latest/user/advanced/#transport-adapters>`_.
+
+        Returns:
+            Sequence[Tuple[str, requests.adapters.BaseAdapter]]:
+                Prefixes and corresponding transport adapters to `mount
+                <https://requests.readthedocs.io/en/latest/api/#requests.Session.mount>`_
+                in requests-based REST clients.
+        """
+        return self._requests_transport_adapters
+
+    @requests_transport_adapters.setter
+    def requests_transport_adapters(
+        self, value: Sequence[Tuple[str, requests.adapters.BaseAdapter]]
+    ) -> None:
+        if self._session_started and self._requests_transport_adapters != value:
+            raise ValueError(
+                SESSION_STARTED_MESSAGE.format(attribute="requests_transport_adapters")
+            )
+        self._requests_transport_adapters = value

bigframes/pandas/io/api.py

@@ -496,6 +496,7 @@ def _set_default_session_location_if_possible(query):
             application_name=config.options.bigquery.application_name,
             bq_kms_key_name=config.options.bigquery.kms_key_name,
             client_endpoints_override=config.options.bigquery.client_endpoints_override,
+            requests_transport_adapters=config.options.bigquery.requests_transport_adapters,
         )
 
         bqclient = clients_provider.bqclient

bigframes/session/__init__.py

@@ -171,6 +171,7 @@ def __init__(
                 application_name=context.application_name,
                 bq_kms_key_name=self._bq_kms_key_name,
                 client_endpoints_override=context.client_endpoints_override,
+                requests_transport_adapters=context.requests_transport_adapters,
             )
 
         # TODO(shobs): Remove this logic after https://github.com/ibis-project/ibis/issues/8494

bigframes/session/clients.py

@@ -17,18 +17,20 @@
 import os
 import threading
 import typing
-from typing import Optional
+from typing import Optional, Sequence, Tuple
 
 import google.api_core.client_info
 import google.api_core.client_options
 import google.api_core.gapic_v1.client_info
 import google.auth.credentials
+import google.auth.transport.requests
 import google.cloud.bigquery as bigquery
 import google.cloud.bigquery_connection_v1
 import google.cloud.bigquery_storage_v1
 import google.cloud.functions_v2
 import google.cloud.resourcemanager_v3
 import pydata_google_auth
+import requests
 
 import bigframes.constants
 import bigframes.version
@@ -79,6 +81,10 @@ def __init__(
         application_name: Optional[str] = None,
         bq_kms_key_name: Optional[str] = None,
         client_endpoints_override: dict = {},
+        *,
+        requests_transport_adapters: Sequence[
+            Tuple[str, requests.adapters.BaseAdapter]
+        ] = (),
     ):
         credentials_project = None
         if credentials is None:
@@ -124,6 +130,7 @@ def __init__(
                 )
         self._location = location
         self._use_regional_endpoints = use_regional_endpoints
+        self._requests_transport_adapters = requests_transport_adapters
 
         self._credentials = credentials
         self._bq_kms_key_name = bq_kms_key_name
@@ -173,12 +180,21 @@ def _create_bigquery_client(self):
             user_agent=self._application_name
         )
 
+        requests_session = google.auth.transport.requests.AuthorizedSession(
+            self._credentials
+        )
+        for prefix, adapter in self._requests_transport_adapters:
+            requests_session.mount(prefix, adapter)
+
         bq_client = bigquery.Client(
             client_info=bq_info,
             client_options=bq_options,
-            credentials=self._credentials,
             project=self._project,
             location=self._location,
+            # Instead of credentials, use _http so that users can override
+            # requests options with transport adapters. See internal issue
+            # b/419106112.
+            _http=requests_session,
         )
 
         if self._bq_kms_key_name: