feat: add bigframes.options.bigquery.application_name for partner attribution (#117)

Because `session.py` was getting long, this also refactors `session.py` to separate client construction in a separate module.

Fixes internal issue 305950924
🦕

Author: tswast

bigframes/_config/bigquery_options.py

@@ -37,14 +37,33 @@ def __init__(
         location: Optional[str] = None,
         bq_connection: Optional[str] = None,
         use_regional_endpoints: bool = False,
+        application_name: Optional[str] = None,
     ):
         self._credentials = credentials
         self._project = project
         self._location = location
         self._bq_connection = bq_connection
         self._use_regional_endpoints = use_regional_endpoints
+        self._application_name = application_name
         self._session_started = False
 
+    @property
+    def application_name(self) -> Optional[str]:
+        """The application name to amend to the user-agent sent to Google APIs.
+
+        Recommended format is ``"appplication-name/major.minor.patch_version"``
+        or ``"(gpn:PartnerName;)"`` for official Google partners.
+        """
+        return self._application_name
+
+    @application_name.setter
+    def application_name(self, value: Optional[str]):
+        if self._session_started and self._application_name != value:
+            raise ValueError(
+                SESSION_STARTED_MESSAGE.format(attribute="application_name")
+            )
+        self._application_name = value
+
     @property
     def credentials(self) -> Optional[google.auth.credentials.Credentials]:
         """The OAuth2 Credentials to use for this client."""

bigframes/pandas/__init__.py

@@ -51,6 +51,7 @@
 import bigframes.dataframe
 import bigframes.series
 import bigframes.session
+import bigframes.session.clients
 import third_party.bigframes_vendored.pandas.core.reshape.concat as vendored_pandas_concat
 import third_party.bigframes_vendored.pandas.core.reshape.merge as vendored_pandas_merge
 import third_party.bigframes_vendored.pandas.core.reshape.tile as vendored_pandas_tile
@@ -180,11 +181,12 @@ def _set_default_session_location_if_possible(query):
     ):
         return
 
-    clients_provider = bigframes.session.ClientsProvider(
+    clients_provider = bigframes.session.clients.ClientsProvider(
         project=options.bigquery.project,
         location=options.bigquery.location,
         use_regional_endpoints=options.bigquery.use_regional_endpoints,
         credentials=options.bigquery.credentials,
+        application_name=options.bigquery.application_name,
     )
 
     bqclient = clients_provider.bqclient

bigframes/session.py renamed to bigframes/session/__init__.py

@@ -61,7 +61,6 @@
     ReadPickleBuffer,
     StorageOptions,
 )
-import pydata_google_auth
 
 import bigframes._config.bigquery_options as bigquery_options
 import bigframes.constants as constants
@@ -75,6 +74,7 @@
 import bigframes.formatting_helpers as formatting_helpers
 from bigframes.remote_function import read_gbq_function as bigframes_rgf
 from bigframes.remote_function import remote_function as bigframes_rf
+import bigframes.session.clients
 import bigframes.version
 
 # Even though the ibis.backends.bigquery.registry import is unused, it's needed
@@ -85,18 +85,6 @@
 import third_party.bigframes_vendored.pandas.io.parsers.readers as third_party_pandas_readers
 import third_party.bigframes_vendored.pandas.io.pickle as third_party_pandas_pickle
 
-_ENV_DEFAULT_PROJECT = "GOOGLE_CLOUD_PROJECT"
-_APPLICATION_NAME = f"bigframes/{bigframes.version.__version__}"
-_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
-
-# BigQuery is a REST API, which requires the protocol as part of the URL.
-_BIGQUERY_REGIONAL_ENDPOINT = "https://{location}-bigquery.googleapis.com"
-
-# BigQuery Connection and Storage are gRPC APIs, which don't support the
-# https:// protocol in the API endpoint URL.
-_BIGQUERYCONNECTION_REGIONAL_ENDPOINT = "{location}-bigqueryconnection.googleapis.com"
-_BIGQUERYSTORAGE_REGIONAL_ENDPOINT = "{location}-bigquerystorage.googleapis.com"
-
 _BIGFRAMES_DEFAULT_CONNECTION_ID = "bigframes-default-connection"
 
 _MAX_CLUSTER_COLUMNS = 4
@@ -122,149 +110,6 @@ def _is_query(query_or_table: str) -> bool:
     return re.search(r"\s", query_or_table.strip(), re.MULTILINE) is not None
 
 
-def _get_default_credentials_with_project():
-    return pydata_google_auth.default(scopes=_SCOPES, use_local_webserver=False)
-
-
-class ClientsProvider:
-    """Provides client instances necessary to perform cloud operations."""
-
-    def __init__(
-        self,
-        project: Optional[str],
-        location: Optional[str],
-        use_regional_endpoints: Optional[bool],
-        credentials: Optional[google.auth.credentials.Credentials],
-    ):
-        credentials_project = None
-        if credentials is None:
-            credentials, credentials_project = _get_default_credentials_with_project()
-
-        # Prefer the project in this order:
-        # 1. Project explicitly specified by the user
-        # 2. Project set in the environment
-        # 3. Project associated with the default credentials
-        project = (
-            project
-            or os.getenv(_ENV_DEFAULT_PROJECT)
-            or typing.cast(Optional[str], credentials_project)
-        )
-
-        if not project:
-            raise ValueError(
-                "Project must be set to initialize BigQuery client. "
-                "Try setting `bigframes.options.bigquery.project` first."
-            )
-
-        self._project = project
-        self._location = location
-        self._use_regional_endpoints = use_regional_endpoints
-        self._credentials = credentials
-
-        # cloud clients initialized for lazy load
-        self._bqclient = None
-        self._bqconnectionclient = None
-        self._bqstorageclient = None
-        self._cloudfunctionsclient = None
-        self._resourcemanagerclient = None
-
-    @property
-    def bqclient(self):
-        if not self._bqclient:
-            bq_options = None
-            if self._use_regional_endpoints:
-                bq_options = google.api_core.client_options.ClientOptions(
-                    api_endpoint=_BIGQUERY_REGIONAL_ENDPOINT.format(
-                        location=self._location
-                    ),
-                )
-            bq_info = google.api_core.client_info.ClientInfo(
-                user_agent=_APPLICATION_NAME
-            )
-            self._bqclient = bigquery.Client(
-                client_info=bq_info,
-                client_options=bq_options,
-                credentials=self._credentials,
-                project=self._project,
-                location=self._location,
-            )
-
-        return self._bqclient
-
-    @property
-    def bqconnectionclient(self):
-        if not self._bqconnectionclient:
-            bqconnection_options = None
-            if self._use_regional_endpoints:
-                bqconnection_options = google.api_core.client_options.ClientOptions(
-                    api_endpoint=_BIGQUERYCONNECTION_REGIONAL_ENDPOINT.format(
-                        location=self._location
-                    )
-                )
-            bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo(
-                user_agent=_APPLICATION_NAME
-            )
-            self._bqconnectionclient = (
-                google.cloud.bigquery_connection_v1.ConnectionServiceClient(
-                    client_info=bqconnection_info,
-                    client_options=bqconnection_options,
-                    credentials=self._credentials,
-                )
-            )
-
-        return self._bqconnectionclient
-
-    @property
-    def bqstorageclient(self):
-        if not self._bqstorageclient:
-            bqstorage_options = None
-            if self._use_regional_endpoints:
-                bqstorage_options = google.api_core.client_options.ClientOptions(
-                    api_endpoint=_BIGQUERYSTORAGE_REGIONAL_ENDPOINT.format(
-                        location=self._location
-                    )
-                )
-            bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo(
-                user_agent=_APPLICATION_NAME
-            )
-            self._bqstorageclient = google.cloud.bigquery_storage_v1.BigQueryReadClient(
-                client_info=bqstorage_info,
-                client_options=bqstorage_options,
-                credentials=self._credentials,
-            )
-
-        return self._bqstorageclient
-
-    @property
-    def cloudfunctionsclient(self):
-        if not self._cloudfunctionsclient:
-            functions_info = google.api_core.gapic_v1.client_info.ClientInfo(
-                user_agent=_APPLICATION_NAME
-            )
-            self._cloudfunctionsclient = (
-                google.cloud.functions_v2.FunctionServiceClient(
-                    client_info=functions_info,
-                    credentials=self._credentials,
-                )
-            )
-
-        return self._cloudfunctionsclient
-
-    @property
-    def resourcemanagerclient(self):
-        if not self._resourcemanagerclient:
-            resourcemanager_info = google.api_core.gapic_v1.client_info.ClientInfo(
-                user_agent=_APPLICATION_NAME
-            )
-            self._resourcemanagerclient = (
-                google.cloud.resourcemanager_v3.ProjectsClient(
-                    credentials=self._credentials, client_info=resourcemanager_info
-                )
-            )
-
-        return self._resourcemanagerclient
-
-
 class Session(
     third_party_pandas_gbq.GBQIOMixin,
     third_party_pandas_parquet.ParquetIOMixin,
@@ -279,14 +124,14 @@ class Session(
             Configuration adjusting how to connect to BigQuery and related
             APIs. Note that some options are ignored if ``clients_provider`` is
             set.
-        clients_provider (bigframes.session.ClientsProvider):
+        clients_provider (bigframes.session.bigframes.session.clients.ClientsProvider):
             An object providing client library objects.
     """
 
     def __init__(
         self,
         context: Optional[bigquery_options.BigQueryOptions] = None,
-        clients_provider: Optional[ClientsProvider] = None,
+        clients_provider: Optional[bigframes.session.clients.ClientsProvider] = None,
     ):
         if context is None:
             context = bigquery_options.BigQueryOptions()
@@ -306,11 +151,12 @@ def __init__(
         if clients_provider:
             self._clients_provider = clients_provider
         else:
-            self._clients_provider = ClientsProvider(
+            self._clients_provider = bigframes.session.clients.ClientsProvider(
                 project=context.project,
                 location=self._location,
                 use_regional_endpoints=context.use_regional_endpoints,
                 credentials=context.credentials,
+                application_name=context.application_name,
             )
 
         self._create_and_bind_bq_session()
@@ -319,7 +165,7 @@ def __init__(
             ibis.bigquery.connect(
                 project_id=context.project,
                 client=self.bqclient,
-                storage_client=self.bqstorageclient,
+                storage_client=self.bqstoragereadclient,
             ),
         )
 
@@ -338,8 +184,8 @@ def bqconnectionclient(self):
         return self._clients_provider.bqconnectionclient
 
     @property
-    def bqstorageclient(self):
-        return self._clients_provider.bqstorageclient
+    def bqstoragereadclient(self):
+        return self._clients_provider.bqstoragereadclient
 
     @property
     def cloudfunctionsclient(self):