From 97ca9bb943700167d4222d4d213e183b19b1839c Mon Sep 17 00:00:00 2001
From: Wyatt Alt <wyatt.alt@gmail.com>
Date: Wed, 4 Mar 2026 11:11:36 -0800
Subject: [PATCH] feat: allow passing azure client/tenant ID through remote SDK
 (#3102)

Prior to this commit we supported passing the azure storage account name
to the lancedb remote SDK through headers. This adds support for client
ID and tenant ID as well.
---
 rust/lancedb/src/connection.rs    |  7 ++--
 rust/lancedb/src/remote/client.rs | 54 +++++++++++++++++++++++++++----
 rust/lancedb/src/remote/db.rs     |  7 +++-
 3 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs
index f442409d2..e745a921b 100644
--- a/rust/lancedb/src/connection.rs
+++ b/rust/lancedb/src/connection.rs
@@ -566,8 +566,11 @@ pub struct ConnectBuilder {
 }
 
 #[cfg(feature = "remote")]
-const ENV_VARS_TO_STORAGE_OPTS: [(&str, &str); 1] =
-    [("AZURE_STORAGE_ACCOUNT_NAME", "azure_storage_account_name")];
+const ENV_VARS_TO_STORAGE_OPTS: [(&str, &str); 3] = [
+    ("AZURE_STORAGE_ACCOUNT_NAME", "azure_storage_account_name"),
+    ("AZURE_CLIENT_ID", "azure_client_id"),
+    ("AZURE_TENANT_ID", "azure_tenant_id"),
+];
 
 impl ConnectBuilder {
     /// Create a new [`ConnectOptions`] with the given database URI.
diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs
index 7304c18c2..ac318c014 100644
--- a/rust/lancedb/src/remote/client.rs
+++ b/rust/lancedb/src/remote/client.rs
@@ -446,13 +446,23 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
                 })?,
             );
         }
-        if let Some(v) = options.0.get("azure_storage_account_name") {
-            headers.insert(
-                HeaderName::from_static("x-azure-storage-account-name"),
-                HeaderValue::from_str(v).map_err(|_| Error::InvalidInput {
-                    message: format!("non-ascii storage account name '{}' provided", db_name),
-                })?,
-            );
+        // Map azure storage options to x-azure-* headers.
+        // The option key uses underscores (e.g. "azure_client_id") while the
+        // header uses hyphens (e.g. "x-azure-client-id").
+        let azure_opts: [(&str, &str); 3] = [
+            ("azure_storage_account_name", "x-azure-storage-account-name"),
+            ("azure_client_id", "x-azure-client-id"),
+            ("azure_tenant_id", "x-azure-tenant-id"),
+        ];
+        for (opt_key, header_name) in azure_opts {
+            if let Some(v) = options.0.get(opt_key) {
+                headers.insert(
+                    HeaderName::from_static(header_name),
+                    HeaderValue::from_str(v).map_err(|_| Error::InvalidInput {
+                        message: format!("non-ascii value '{}' for option '{}'", v, opt_key),
+                    })?,
+                );
+            }
         }
 
         for (key, value) in &config.extra_headers {
@@ -1075,4 +1085,34 @@ mod tests {
             _ => panic!("Expected Runtime error"),
         }
     }
+
+    #[test]
+    fn test_default_headers_azure_opts() {
+        let mut opts = HashMap::new();
+        opts.insert(
+            "azure_storage_account_name".to_string(),
+            "myaccount".to_string(),
+        );
+        opts.insert("azure_client_id".to_string(), "my-client-id".to_string());
+        opts.insert("azure_tenant_id".to_string(), "my-tenant-id".to_string());
+        let remote_opts = RemoteOptions::new(opts);
+
+        let headers = RestfulLanceDbClient::<Sender>::default_headers(
+            "test-key",
+            "us-east-1",
+            "testdb",
+            false,
+            &remote_opts,
+            None,
+            &ClientConfig::default(),
+        )
+        .unwrap();
+
+        assert_eq!(
+            headers.get("x-azure-storage-account-name").unwrap(),
+            "myaccount"
+        );
+        assert_eq!(headers.get("x-azure-client-id").unwrap(), "my-client-id");
+        assert_eq!(headers.get("x-azure-tenant-id").unwrap(), "my-tenant-id");
+    }
 }
diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs
index b615fd3bc..b80c1cea1 100644
--- a/rust/lancedb/src/remote/db.rs
+++ b/rust/lancedb/src/remote/db.rs
@@ -777,7 +777,12 @@ impl RemoteOptions {
 
 impl From<StorageOptions> for RemoteOptions {
     fn from(options: StorageOptions) -> Self {
-        let supported_opts = vec!["account_name", "azure_storage_account_name"];
+        let supported_opts = vec![
+            "account_name",
+            "azure_storage_account_name",
+            "azure_client_id",
+            "azure_tenant_id",
+        ];
         let mut filtered = HashMap::new();
         for opt in supported_opts {
             if let Some(v) = options.0.get(opt) {