Pick up on SPARK_CONNECT_USER_AGENT in chat_databricks(). (#254)

`SPARK_CONNECT_USER_AGENT` is the standard environment variable used to set the user agent in various Databricks and Spark-related projects. This commit adds some detection for when this is set, much like we have in `odbc` and `pysparklyr`. Unit tests are included. Signed-off-by: Aaron Jacobs <[email protected]>
tidyverse · Jan 15, 2025 · f1117d8 · f1117d8
1 parent 2ced774
commit f1117d8
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -10,6 +10,9 @@
 * `chat_databricks()` now handles the `DATABRICKS_HOST` environment variable
   correctly whether it includes an HTTPS prefix or not (#252, @atheriel).
 
+* `chat_databricks()` now respects the `SPARK_CONNECT_USER_AGENT` environment
+  variable when making requests (#254, @atheriel).
+
 # ellmer 0.1.0
 
 * New `chat_vllm()` to chat with models served by vLLM (#140).

diff --git a/R/provider-databricks.R b/R/provider-databricks.R
@@ -89,6 +89,7 @@ method(chat_request, ProviderDatabricks) <- function(provider,
   req <- req_auth_bearer_token(req,
     databricks_token(provider@base_url, provider@token)
   )
+  req <- req_user_agent(req, databricks_user_agent())
   req <- req_retry(req, max_tries = 2)
   req <- req_error(req, body = function(resp) {
     if (resp_content_type(resp) == "application/json") {
@@ -174,6 +175,14 @@ databricks_workspace <- function() {
   host
 }
 
+databricks_user_agent <- function() {
+  user_agent <- paste0("r-ellmer/", utils::packageVersion("ellmer"))
+  if (nchar(Sys.getenv("SPARK_CONNECT_USER_AGENT")) != 0) {
+    user_agent <- paste(Sys.getenv("SPARK_CONNECT_USER_AGENT"), user_agent)
+  }
+  user_agent
+}
+
 # Try various ways to get Databricks credentials. This implements a subset of
 # the "Databricks client unified authentication" model.
 databricks_token <- function(workspace = databricks_workspace(), token = NULL) {

diff --git a/tests/testthat/test-provider-databricks.R b/tests/testthat/test-provider-databricks.R
@@ -100,3 +100,14 @@ test_that("workspace detection handles URLs with and without an https prefix", {
     )
   )
 })
+
+test_that("the user agent respects SPARK_CONNECT_USER_AGENT when set", {
+  withr::with_envvar(
+    c(SPARK_CONNECT_USER_AGENT = NA),
+    expect_match(databricks_user_agent(), "^r-ellmer")
+  )
+  withr::with_envvar(
+    c(SPARK_CONNECT_USER_AGENT = "testing"),
+    expect_match(databricks_user_agent(), "^testing r-ellmer")
+  )
+})