feat(java): support credential vending at write time

.github/workflows/java.yml

@@ -49,22 +49,9 @@ jobs:
       matrix:
         java-version: [8, 11, 17]
     name: Build and Test with Java ${{ matrix.java-version }}
-    services:
-      localstack:
-        image: localstack/localstack:4.0
-        ports:
-          - 4566:4566
-        env:
-          SERVICES: s3,dynamodb,kms
-          AWS_ACCESS_KEY_ID: ACCESS_KEY
-          AWS_SECRET_ACCESS_KEY: SECRET_KEY
-        options: >-
-          --health-cmd "curl -s http://localhost:4566/_localstack/health"
-          --health-interval 5s
-          --health-timeout 3s
-          --health-retries 3
-          --health-start-period 10s
     steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
       - name: Install dependencies
         run: |
           sudo apt update
@@ -76,8 +63,6 @@ jobs:
       - uses: rui314/setup-mold@v1
       - name: Install cargo-llvm-cov
         uses: taiki-e/install-action@cargo-llvm-cov
-      - name: Checkout repository
-        uses: actions/checkout@v4
       - uses: Swatinem/rust-cache@v2
         with:
           workspaces: java/lance-jni -> ../target/rust-maven-plugin/lance-jni
@@ -91,6 +76,9 @@ jobs:
         working-directory: java
         run: |
           mvn spotless:check
+      - name: Start localstack
+        run: |
+          docker compose -f docker-compose.yml up -d --wait
       - name: Running tests with Java ${{ matrix.java-version }}
         working-directory: java
         env:

java/lance-jni/src/blocking_dataset.rs

@@ -59,6 +59,11 @@ pub struct BlockingDataset {
 }
 
 impl BlockingDataset {
+    /// Get the storage options provider that was used when opening this dataset
+    pub fn get_storage_options_provider(&self) -> Option<Arc<dyn StorageOptionsProvider>> {
+        self.inner.storage_options_provider()
+    }
+
     pub fn drop(uri: &str, storage_options: HashMap<String, String>) -> Result<()> {
         RT.block_on(async move {
             let registry = Arc::new(ObjectStoreRegistry::default());
@@ -122,7 +127,7 @@ impl BlockingDataset {
             builder = builder.with_version(ver as u64);
         }
         builder = builder.with_storage_options(storage_options);
-        if let Some(provider) = storage_options_provider {
+        if let Some(provider) = storage_options_provider.clone() {
             builder = builder.with_storage_options_provider(provider)
         }
         if let Some(offset_seconds) = s3_credentials_refresh_offset_seconds {
@@ -284,14 +289,11 @@ impl BlockingDataset {
     pub fn commit_transaction(
         &mut self,
         transaction: Transaction,
-        write_params: HashMap<String, String>,
+        store_params: ObjectStoreParams,
     ) -> Result<Self> {
         let new_dataset = RT.block_on(
             CommitBuilder::new(Arc::new(self.clone().inner))
-                .with_store_params(ObjectStoreParams {
-                    storage_options: Some(write_params),
-                    ..Default::default()
-                })
+                .with_store_params(store_params)
                 .execute(transaction),
         )?;
         Ok(BlockingDataset { inner: new_dataset })
@@ -334,6 +336,7 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_createWithFfiSchema<'local
     enable_stable_row_ids: JObject, // Optional<Boolean>
     data_storage_version: JObject,  // Optional<String>
     storage_options_obj: JObject,   // Map<String, String>
+    s3_credentials_refresh_offset_seconds_obj: JObject, // Optional<Long>
 ) -> JObject<'local> {
     ok_or_throw!(
         env,
@@ -347,7 +350,8 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_createWithFfiSchema<'local
             mode,
             enable_stable_row_ids,
             data_storage_version,
-            storage_options_obj
+            storage_options_obj,
+            s3_credentials_refresh_offset_seconds_obj
         )
     )
 }
@@ -364,6 +368,7 @@ fn inner_create_with_ffi_schema<'local>(
     enable_stable_row_ids: JObject, // Optional<Boolean>
     data_storage_version: JObject,  // Optional<String>
     storage_options_obj: JObject,   // Map<String, String>
+    s3_credentials_refresh_offset_seconds_obj: JObject, // Optional<Long>
 ) -> Result<JObject<'local>> {
     let c_schema_ptr = arrow_schema_addr as *mut FFI_ArrowSchema;
     let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) };
@@ -380,6 +385,8 @@ fn inner_create_with_ffi_schema<'local>(
         enable_stable_row_ids,
         data_storage_version,
         storage_options_obj,
+        JObject::null(), // No provider for schema-only creation
+        s3_credentials_refresh_offset_seconds_obj,
         reader,
     )
 }
@@ -411,6 +418,7 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_createWithFfiStream<'local
     enable_stable_row_ids: JObject, // Optional<Boolean>
     data_storage_version: JObject,  // Optional<String>
     storage_options_obj: JObject,   // Map<String, String>
+    s3_credentials_refresh_offset_seconds_obj: JObject, // Optional<Long>
 ) -> JObject<'local> {
     ok_or_throw!(
         env,
@@ -424,7 +432,44 @@ pub extern "system" fn Java_com_lancedb_lance_Dataset_createWithFfiStream<'local
             mode,
             enable_stable_row_ids,
             data_storage_version,
-            storage_options_obj
+            storage_options_obj,
+            JObject::null(),
+            s3_credentials_refresh_offset_seconds_obj
+        )
+    )
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_Dataset_createWithFfiStreamAndProvider<'local>(
+    mut env: JNIEnv<'local>,
+    _obj: JObject,
+    arrow_array_stream_addr: jlong,
+    path: JString,
+    max_rows_per_file: JObject,            // Optional<Integer>
+    max_rows_per_group: JObject,           // Optional<Integer>
+    max_bytes_per_file: JObject,           // Optional<Long>
+    mode: JObject,                         // Optional<String>
+    enable_stable_row_ids: JObject,        // Optional<Boolean>
+    data_storage_version: JObject,         // Optional<String>
+    storage_options_obj: JObject,          // Map<String, String>
+    storage_options_provider_obj: JObject, // Optional<StorageOptionsProvider>
+    s3_credentials_refresh_offset_seconds_obj: JObject, // Optional<Long>
+) -> JObject<'local> {
+    ok_or_throw!(
+        env,
+        inner_create_with_ffi_stream(
+            &mut env,
+            arrow_array_stream_addr,
+            path,
+            max_rows_per_file,
+            max_rows_per_group,
+            max_bytes_per_file,
+            mode,
+            enable_stable_row_ids,
+            data_storage_version,
+            storage_options_obj,
+            storage_options_provider_obj,
+            s3_credentials_refresh_offset_seconds_obj
         )
     )
 }
@@ -434,13 +479,15 @@ fn inner_create_with_ffi_stream<'local>(
     env: &mut JNIEnv<'local>,
     arrow_array_stream_addr: jlong,
     path: JString,
-    max_rows_per_file: JObject,     // Optional<Integer>
-    max_rows_per_group: JObject,    // Optional<Integer>
-    max_bytes_per_file: JObject,    // Optional<Long>
-    mode: JObject,                  // Optional<String>
-    enable_stable_row_ids: JObject, // Optional<Boolean>
-    data_storage_version: JObject,  // Optional<String>
-    storage_options_obj: JObject,   // Map<String, String>
+    max_rows_per_file: JObject,            // Optional<Integer>
+    max_rows_per_group: JObject,           // Optional<Integer>
+    max_bytes_per_file: JObject,           // Optional<Long>
+    mode: JObject,                         // Optional<String>
+    enable_stable_row_ids: JObject,        // Optional<Boolean>
+    data_storage_version: JObject,         // Optional<String>
+    storage_options_obj: JObject,          // Map<String, String>
+    storage_options_provider_obj: JObject, // Optional<StorageOptionsProvider>
+    s3_credentials_refresh_offset_seconds_obj: JObject, // Optional<Long>
 ) -> Result<JObject<'local>> {
     let stream_ptr = arrow_array_stream_addr as *mut FFI_ArrowArrayStream;
     let reader = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }?;
@@ -454,6 +501,8 @@ fn inner_create_with_ffi_stream<'local>(
         enable_stable_row_ids,
         data_storage_version,
         storage_options_obj,
+        storage_options_provider_obj,
+        s3_credentials_refresh_offset_seconds_obj,
         reader,
     )
 }
@@ -469,6 +518,8 @@ fn create_dataset<'local>(
     enable_stable_row_ids: JObject,
     data_storage_version: JObject,
     storage_options_obj: JObject,
+    storage_options_provider_obj: JObject, // Optional<StorageOptionsProvider>
+    s3_credentials_refresh_offset_seconds_obj: JObject,
     reader: impl RecordBatchReader + Send + 'static,
 ) -> Result<JObject<'local>> {
     let path_str = path.extract(env)?;
@@ -482,6 +533,8 @@ fn create_dataset<'local>(
         &enable_stable_row_ids,
         &data_storage_version,
         &storage_options_obj,
+        &storage_options_provider_obj,
+        &s3_credentials_refresh_offset_seconds_obj,
     )?;
 
     let dataset = BlockingDataset::write(reader, &path_str, Some(write_params))?;