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):