Skip to content

Commit

Permalink
[SPARK-44312][CONNECT][PYTHON] Allow to set a user agent with an envi…
Browse files Browse the repository at this point in the history
…ronment variable

### What changes were proposed in this pull request?
Use the `SPARK_CONNECT_USER_AGENT` environment variable as a fallback for the prepended user agent string if one did not set the user agent in `ChannelBuilder`.

### Why are the changes needed?
Currently one has to specify a custom user agent string in `ChannelBuilder`. It would be useful to be able to set this string with an environment variable.

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
Manual testing + existing tests.

Closes apache#41866 from dillitz/SPARK-44312-user-agent-environment.

Authored-by: Robert Dillitz <robert.dillitz@databricks.com>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
  • Loading branch information
dillitz authored and HyukjinKwon committed Jul 7, 2023
1 parent e3c2bf5 commit 485e8bc
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion python/pyspark/sql/connect/client/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,10 @@ def userAgent(self) -> str:
or "_SPARK_CONNECT_PYTHON" when not specified.
The returned value will be percent encoded.
"""
user_agent = self.params.get(ChannelBuilder.PARAM_USER_AGENT, "_SPARK_CONNECT_PYTHON")
user_agent = self.params.get(
ChannelBuilder.PARAM_USER_AGENT,
os.getenv("SPARK_CONNECT_USER_AGENT", "_SPARK_CONNECT_PYTHON"),
)
ua_len = len(urllib.parse.quote(user_agent))
if ua_len > 2048:
raise SparkConnectException(
Expand Down

0 comments on commit 485e8bc

Please sign in to comment.