Skip to content

Commit

Permalink
adds the capability to include customer user agent string
Browse files Browse the repository at this point in the history
  • Loading branch information
chalmerlowe committed Oct 9, 2024
1 parent b3faa0b commit 0c78acc
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 2 deletions.
71 changes: 69 additions & 2 deletions pandas_gbq/gbq.py
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,8 @@ def __init__(
auth_redirect_uri=None,
client_id=None,
client_secret=None,
user_agent=None,
rfc9110_delimiter=False,
):
global context
from google.api_core.exceptions import ClientError, GoogleAPIError
Expand All @@ -284,6 +286,8 @@ def __init__(
self.auth_redirect_uri = auth_redirect_uri
self.client_id = client_id
self.client_secret = client_secret
self.user_agent = user_agent
self.rfc9110_delimiter = rfc9110_delimiter

default_project = None

Expand Down Expand Up @@ -337,11 +341,17 @@ def log_elapsed_seconds(self, prefix="Elapsed", postfix="s.", overlong=6):

def get_client(self):
import google.api_core.client_info
import pandas

# import pandas # noqa: F401 # TODO is this line needed here?

bigquery = FEATURES.bigquery_try_import()

user_agent = create_user_agent(
user_agent=self.user_agent, rfc9110_delimiter=self.rfc9110_delimiter
)

client_info = google.api_core.client_info.ClientInfo(
user_agent="pandas-{}".format(pandas.__version__)
user_agent=user_agent,
)
return bigquery.Client(
project=self.project_id,
Expand Down Expand Up @@ -961,6 +971,8 @@ def to_gbq(
auth_redirect_uri=None,
client_id=None,
client_secret=None,
user_agent=None,
rfc9110_delimiter=False,
):
"""Write a DataFrame to a Google BigQuery table.
Expand Down Expand Up @@ -1072,6 +1084,13 @@ def to_gbq(
client_secret : str
The Client Secret associated with the Client ID for the Google Cloud Project
the user is attempting to connect to.
user_agent : str
Custom user agent string used as a prefix to the pandas version.
rfc9110_delimiter : bool
Sets user agent delimiter to a hyphen or a slash.
Default is False, meaning a hyphen will be used.
.. versionadded:: 0.23.3
"""

_test_google_api_imports()
Expand Down Expand Up @@ -1130,6 +1149,8 @@ def to_gbq(
auth_redirect_uri=auth_redirect_uri,
client_id=client_id,
client_secret=client_secret,
user_agent=user_agent,
rfc9110_delimiter=rfc9110_delimiter,
)
bqclient = connector.client

Expand Down Expand Up @@ -1409,3 +1430,49 @@ def create(self, dataset_id):
self.client.create_dataset(dataset)
except self.http_error as ex:
self.process_http_error(ex)


def create_user_agent(
user_agent: Optional[str] = None, rfc9110_delimiter: bool = False
) -> str:
"""Creates a user agent string.
The legacy format of our the user agent string was: `product-x.y.z` (where x,
y, and z are the major, minor, and micro version numbers).
Users are able to prepend this string with their own user agent identifier
to render something similar to `<my_user_agent> pandas-x.y.z`.
The legacy format used a hyphen to separate the product from the product
version which differs slightly from the format recommended by RFC9110, which is:
`product/x.y.z`. To produce a user agent more in line with the RFC, set
rfc9110_delimiter to True. This setting does not depend on whether a
user_agent is also supplied.
Reference:
https://www.rfc-editor.org/info/rfc9110
Args:
user_agent (Optional[str]): User agent string.
rfc9110_delimiter (Optional[bool]): Sets delimiter to a hyphen or a slash.
Default is False, meaning a hyphen will be used.
Returns (str):
Customized user agent string.
"""
import pandas as pd

if rfc9110_delimiter:
delimiter = "/"
else:
delimiter = "-"

identity = f"pandas{delimiter}{pd.__version__}"

if user_agent is None:
user_agent = identity
else:
user_agent = f"{user_agent} {identity}"

return user_agent
25 changes: 25 additions & 0 deletions tests/unit/test_to_gbq.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import google.api_core.exceptions
import google.cloud.bigquery
import pandas as pd
from pandas import DataFrame
import pytest

Expand Down Expand Up @@ -158,3 +159,27 @@ def test_to_gbq_with_if_exists_unknown():
project_id="myproj",
if_exists="unknown",
)


@pytest.mark.parametrize(
"user_agent,rfc9110_delimiter,expected",
[
(
"test_user_agent/2.0.42",
False,
f"test_user_agent/2.0.42 pandas-{pd.__version__}",
),
(None, False, f"pandas-{pd.__version__}"),
(
"test_user_agent/2.0.42",
True,
f"test_user_agent/2.0.42 pandas/{pd.__version__}",
),
(None, True, f"pandas/{pd.__version__}"),
],
)
def test_create_user_agent(user_agent, rfc9110_delimiter, expected):
from pandas_gbq.gbq import create_user_agent

result = create_user_agent(user_agent, rfc9110_delimiter)
assert result == expected

0 comments on commit 0c78acc

Please sign in to comment.