SQL Connector Bindings + SQL Execution API (#44)

Changes * Adding tests for sql connector and sql execution api * Adjusting db_host(), no longer required to specify scheme (https/http) as brickster now will override to https. * Incrementing to 0.2.2 * Skip sql connector tests when env isn't available * Adjusting docs for R CMD Check --------- Co-authored-by: Zac Davies <[email protected]>
databrickslabs · May 6, 2024 · f172e45 · f172e45
1 parent a677cae
commit f172e45
Show file tree

Hide file tree

Showing 27 changed files with 1,779 additions and 15 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -1,6 +1,6 @@
 Package: brickster
 Title: R Toolkit for Databricks
-Version: 0.2.1
+Version: 0.2.2
 Authors@R: 
   c(
     person(given = "Zac",
@@ -18,6 +18,7 @@ License: Apache License (>= 2)
 Encoding: UTF-8
 LazyData: true
 Imports:
+    arrow,
     base64enc,
     cli,
     curl,
@@ -29,7 +30,11 @@ Imports:
     jsonlite,
     magrittr,
     purrr,
-    rlang
+    reticulate,
+    R6 (>= 2.4.0),
+    rlang,
+    tibble,
+    utils
 Suggests: 
     testthat (>= 3.0.0),
     htmltools,

diff --git a/NAMESPACE b/NAMESPACE
@@ -1,5 +1,6 @@
 # Generated by roxygen2: do not edit by hand
 
+export(DatabricksSqlClient)
 export(access_control_req_group)
 export(access_control_req_user)
 export(access_control_request)
@@ -90,6 +91,11 @@ export(db_secrets_scope_acl_put)
 export(db_secrets_scope_create)
 export(db_secrets_scope_delete)
 export(db_secrets_scope_list_all)
+export(db_sql_client)
+export(db_sql_exec_cancel)
+export(db_sql_exec_query)
+export(db_sql_exec_result)
+export(db_sql_exec_status)
 export(db_sql_global_warehouse_get)
 export(db_sql_query_history)
 export(db_sql_warehouse_create)
@@ -116,6 +122,7 @@ export(db_workspace_list)
 export(db_workspace_mkdirs)
 export(db_wsid)
 export(dbfs_storage_info)
+export(determine_brickster_venv)
 export(docker_image)
 export(email_notifications)
 export(file_storage_info)
@@ -126,6 +133,7 @@ export(get_latest_dbr)
 export(git_source)
 export(in_databricks_nb)
 export(init_script_info)
+export(install_db_sql_connector)
 export(is.access_control_req_group)
 export(is.access_control_req_user)
 export(is.access_control_request)
@@ -174,6 +182,7 @@ export(notebook_task)
 export(notebook_use_posit_repo)
 export(open_workspace)
 export(pipeline_task)
+export(py_db_sql_connector)
 export(python_wheel_task)
 export(remove_lib_path)
 export(s3_storage_info)
@@ -183,8 +192,11 @@ export(spark_jar_task)
 export(spark_python_task)
 export(spark_submit_task)
 export(wait_for_lib_installs)
+import(R6)
+import(arrow)
 import(cli)
 import(httr2)
+import(tibble)
 importFrom(glue,glue)
 importFrom(magrittr,`%>%`)
 importFrom(rlang,.data)

diff --git a/R/data-structures.R b/R/data-structures.R
@@ -718,7 +718,7 @@ is.lib_pypi <- function(x) {
 #' `org.jsoup:jsoup:1.7.2`.
 #' @param repo Maven repo to install the Maven package from. If omitted, both
 #' Maven Central Repository and Spark Packages are searched.
-#' @param exclusions List of dependences to exclude. For example:
+#' @param exclusions List of dependencies to exclude. For example:
 #' `list("slf4j:slf4j", "*:hadoop-client")`.
 #' [Maven dependency exclusions](https://maven.apache.org/guides/introduction/introduction-to-optional-and-excludes-dependencies.html).
 #'

diff --git a/R/databricks-helpers.R b/R/databricks-helpers.R
@@ -0,0 +1,31 @@
+on_databricks <- function() {
+  dbr <- Sys.getenv("DATABRICKS_RUNTIME_VERSION")
+  dbr != ""
+}
+
+in_databricks_nb <- function() {
+  ("/databricks/spark/R/lib"  %in% .libPaths()) &&
+    exists("DATABRICKS_GUID", envir = .GlobalEnv)
+}
+
+use_posit_repo <- function() {
+  if (in_databricks_nb()) {
+    codename <- system("lsb_release -c --short", intern = T)
+    mirror <- paste0("https://packagemanager.posit.co/cran/__linux__/", codename, "/latest")
+    options(repos = c(POSIT = mirror))
+  }
+}
+
+#' Determine brickster virtualenv
+#'
+#' @details Returns `NULL` when running within Databricks,
+#' otherwise `"r-brickster"`
+#'
+#' @export
+determine_brickster_venv <- function() {
+  if (on_databricks()) {
+    NULL
+  } else {
+    "r-brickster"
+  }
+}
diff --git a/R/package-auth.R b/R/package-auth.R
@@ -39,9 +39,25 @@ db_host <- function(id = NULL, prefix = NULL, profile = getOption("db_profile",
 
   if (is.null(id) && is.null(prefix)) {
     host <- read_env_var(key = "host", profile = profile)
+    parsed_url <- httr2::url_parse(host)
+
+    # inject scheme if not present then re-build with https
+    if (is.null(parsed_url$schema)) {
+      parsed_url$scheme <- "https"
+    }
+
+    # if hostname is missing change path to host
+    if (is.null(parsed_url$host)) {
+      parsed_url$hostname <- parsed_url$path
+      parsed_url$path <- NULL
+    }
+
+    host <- httr2::url_build(parsed_url)
+    host <- httr2::url_parse(host)$hostname
+
   } else {
     # otherwise construct host string
-    host <- paste0("https://", prefix, id, ".cloud.databricks.com")
+    host <- paste0(prefix, id, ".cloud.databricks.com")
   }
 
   host

diff --git a/R/request-helpers.R b/R/request-helpers.R
@@ -17,10 +17,19 @@
 #' @importFrom magrittr `%>%`
 db_request <- function(endpoint, method, version = NULL, body = NULL, host, token, ...) {
 
-  req <- httr2::request(base_url = paste0(host, "api", "/", version, "/")) %>%
+  url <- list(
+    scheme = "https",
+    hostname = host,
+    path = paste0("/api/", version)
+  )
+
+  url <- httr2::url_build(url)
+  user_agent_str <- paste0("brickster/", utils::packageVersion("brickster"))
+
+  req <- httr2::request(base_url = url) %>%
     httr2::req_auth_bearer_token(token) %>%
-    httr2::req_headers("User-Agent" = "brickster/1.0") %>%
-    httr2::req_user_agent(string = "brickster/1.0") %>%
+    httr2::req_headers("User-Agent" = user_agent_str) %>%
+    httr2::req_user_agent(string = user_agent_str) %>%
     httr2::req_url_path_append(endpoint) %>%
     httr2::req_method(method) %>%
     httr2::req_retry(max_tries = 3, backoff = ~ 2)