Thanks to visit codestin.com
Credit goes to github.com

Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Oct 17, 2023
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
36e3cf5
feat: add `bigframes.options.bigquery.application_name` for partner a…
tswast Oct 17, 2023
95948ca
add tests for BigQueryOptions.application_name
tswast Oct 17, 2023
b7a082e
add tests for use_regional_endpoints
tswast Oct 17, 2023
63993b6
use correct path for _set_default_session_location_if_possible
tswast Oct 17, 2023
18249c2
Merge branch 'main' into b305950924-application_name
tswast Oct 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions bigframes/_config/bigquery_options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
4 changes: 3 additions & 1 deletion bigframes/pandas/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
170 changes: 8 additions & 162 deletions bigframes/session.py → bigframes/session/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,6 @@
ReadPickleBuffer,
StorageOptions,
)
import pydata_google_auth

import bigframes._config.bigquery_options as bigquery_options
import bigframes.constants as constants
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -319,7 +165,7 @@ def __init__(
ibis.bigquery.connect(
project_id=context.project,
client=self.bqclient,
storage_client=self.bqstorageclient,
storage_client=self.bqstoragereadclient,
),
)

Expand All @@ -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):
Expand Down
Loading