[r] use SOMATileDBContext (#256)

chanzuckerberg · Mar 16, 2023 · 303d555 · 303d555
1 parent 0bbe5ea
commit 303d555
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 15 deletions.
diff --git a/api/r/CellCensus/NAMESPACE b/api/r/CellCensus/NAMESPACE
@@ -9,5 +9,5 @@ importFrom(aws.s3,save_object)
 importFrom(httr,build_url)
 importFrom(httr,parse_url)
 importFrom(jsonlite,fromJSON)
-importFrom(tiledb,tiledb_ctx)
 importFrom(tiledbsoma,SOMACollection)
+importFrom(tiledbsoma,SOMATileDBContext)
diff --git a/api/r/CellCensus/R/open.R b/api/r/CellCensus/R/open.R
@@ -8,28 +8,40 @@ DEFAULT_TILEDB_CONFIGURATION <- c(
 #' @param census_version The version of the Census, e.g., "latest".
 #' @param uri The URI containing the Census SOMA objects. If specified, takes
 #'            precedence over `census_version`.
-#' @param s3_region Optional AWS region accompanying uri.
+#' @param tiledbsoma_ctx A custom `tiledbsoma::SOMATileDBContext`
 #'
 #' @return Top-level `tiledbsoma::SOMACollection` object
 #' @importFrom tiledbsoma SOMACollection
-#' @importFrom tiledb tiledb_ctx
+#' @importFrom tiledbsoma SOMATileDBContext
 #' @export
 #'
 #' @examples
-open_soma <- function(census_version = "latest", uri = "", s3_region = "") {
-  cfg <- DEFAULT_TILEDB_CONFIGURATION
+open_soma <- function(census_version = "latest", uri = NULL, tiledbsoma_ctx = NULL) {
+  s3_region <- NULL
 
-  if (uri == "") {
-    stopifnot("'s3_region' applicable only for specific uri" = s3_region == "")
+  if (is.null(uri)) {
     description <- get_census_version_description(census_version)
     uri <- description$soma.uri
     if ("soma.s3_region" %in% names(description) &&
       description$soma.s3_region != "") {
+      s3_region <- description$soma.s3_region
+    }
+  }
+
+  cfg <- DEFAULT_TILEDB_CONFIGURATION
+  if (is.null(tiledbsoma_ctx)) {
+    if (!is.null(s3_region)) {
       cfg <- c(cfg, c("vfs.s3.region" = description$soma.s3_region))
     }
-  } else if (s3_region != "") {
-    cfg <- c(cfg, c("vfs.s3.region" = s3_region))
+  } else {
+    # FIXME: we should use something like SOMATileDBContext$replace() (yet to
+    # exist) in case the user context has other important fields besides config
+    cfg <- as.vector(tiledb::config(tiledbsoma_ctx$get_tiledb_context()))
+    if (!is.null(s3_region)) {
+      cfg["vfs.s3.region"] <- s3_region
+    }
   }
+  tiledbsoma_ctx <- tiledbsoma::SOMATileDBContext$new(config = cfg)
 
-  return(tiledbsoma::SOMACollection$new(uri, ctx = tiledb::tiledb_ctx(cfg)))
+  return(tiledbsoma::SOMACollection$new(uri, tiledbsoma_ctx = tiledbsoma_ctx))
 }
diff --git a/api/r/CellCensus/R/source_h5ad.R b/api/r/CellCensus/R/source_h5ad.R
@@ -14,7 +14,9 @@ get_source_h5ad_uri <- function(dataset_id, census_version = "latest") {
   census <- open_soma(
     census_version,
     uri = description$soma.uri,
-    s3_region = description$soma.s3_region
+    tiledbsoma_ctx = tiledbsoma::SOMATileDBContext$new(
+      config = c("vfs.s3.region" = description$soma.s3_region)
+    )
   )
 
   # FIXME execution of value_filter:

diff --git a/api/r/CellCensus/man/open_soma.Rd b/api/r/CellCensus/man/open_soma.Rd
diff --git a/api/r/CellCensus/tests/testthat/test-open.R b/api/r/CellCensus/tests/testthat/test-open.R
@@ -10,3 +10,10 @@ test_that("open_soma latest/default", {
   coll_latest <- open_soma("latest")
   expect_equal(coll_default$uri, coll_latest$uri)
 })
+
+test_that("open_soma with custom context", {
+  ctx <- tiledbsoma::SOMATileDBContext$new(config = c("vfs.s3.region" = "bogus"))
+  # open_soma should override our bogus vfs.s3.region setting
+  coll <- open_soma(tiledbsoma_ctx = ctx)
+  expect_true(coll$exists())
+})
diff --git a/api/r/CellCensus/tests/testthat/test_source_h5ad.R b/api/r/CellCensus/tests/testthat/test_source_h5ad.R
@@ -1,5 +1,6 @@
 test_that("get_source_h5ad_uri", {
-  census <- open_soma()
+  census_region <- get_census_version_description("latest")$soma.s3_region
+  census <- open_soma("latest")
   datasets <- as.data.frame(census$get("census_info")$get("datasets")$read(
     column_names = c("dataset_id", "dataset_h5ad_path")
   ))
@@ -11,7 +12,7 @@ test_that("get_source_h5ad_uri", {
     loc <- get_source_h5ad_uri(dataset$dataset_id)
 
     expect_true(endsWith(loc$uri, paste("/", dataset$dataset_h5ad_path, sep = "")))
-    expect_equal(loc$s3_region, unname(tiledb::config(census$ctx)["vfs.s3.region"]))
+    expect_equal(loc$s3_region, census_region)
     # check URI joining with trailing slashes
     expect_false(endsWith(loc$uri, paste("//", dataset$dataset_h5ad_path, sep = "")))
   })