From b15fa9d6e44758e5470275d17edcedca2176f903 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 30 Oct 2026 23:39:46 -0700 Subject: [PATCH] feat: accept Azure inputs in storage integration clients (#4759) Register Azure Blob integrations through Flight SQL using an explicit storage account and three Secret references for tenant ID, client ID and client secret. Extends the database-owned integration APIs with provider-specific inputs while keeping the same create, describe, list and drop operations. ```python info = db.create_storage_integration( "customer_blobs", "az://customer-data/raw/", storage_provider="AZURE", azure_storage_account="customeraccount", azure_tenant_id_secret="tenant", azure_client_id_secret="client", azure_client_secret_secret="credential", ) print(info.configuration_instructions) info = db.describe_storage_integration("customer_blobs") names = list(db.list_storage_integrations()) db.drop_storage_integration("customer_blobs") ``` Sync and async connections accept the same options. Use `SecretReference` for cross-database or namespace qualification. Values remain in Secrets; returned metadata includes only references and setup guidance. ```sql CREATE STORAGE INTEGRATION customer_blobs STORAGE_PROVIDER = 'AZURE' STORAGE_PREFIX = 'az://customer-data/raw/' STORAGE_AZURE_STORAGE_ACCOUNT = 'customeraccount' STORAGE_AZURE_TENANT_ID_SECRET = public.tenant STORAGE_AZURE_CLIENT_ID_SECRET = public.client STORAGE_AZURE_CLIENT_SECRET_SECRET = public.credential; DESCRIBE STORAGE INTEGRATION customer_blobs; SHOW STORAGE INTEGRATIONS; DROP STORAGE INTEGRATION customer_blobs; ``` ```mermaid flowchart LR Python[Sync or async Python] --> SQL[Flight SQL] SQL --> Metadata[Configuration and Secret references] ``` Azure supports service principal client credentials only. Adds SQL rendering, provider validation and Python Flight SQL regressions across Utf8, LargeUtf8 and Utf8View results. Local builds and tests are excluded as requested; validation runs in CI. --- docs/src/python/python.md | 37 +++++++ python/python/lancedb/db.py | 86 ++++++++++++--- python/python/lancedb/remote/db.py | 11 ++ python/python/lancedb/storage_integration.py | 18 +++- .../python/tests/test_storage_integration.py | 100 ++++++++++++++++++ .../lancedb/src/remote/storage_integration.rs | 77 +++++++++++++- rust/lancedb/src/storage_integration.rs | 61 ++++++++++- 7 files changed, 369 insertions(+), 21 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index b3d4efa95..9219b67e6 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -986,3 +986,40 @@ are never included in integration metadata. change. Use the structured identity and policy fields for programmatic access. Registration does not verify AWS access. `AsyncConnection` provides the same methods; await create, describe, and drop, and use `async for` for listing. + + +### Azure storage integrations + +Create three Secrets containing the Azure tenant ID, client ID and client secret in an ordinary database. Reference their names when registering Azure Blob access: + +```python +integration = db.create_storage_integration( + "customer_blobs", + "az://customer-data/raw/", + storage_provider="AZURE", + azure_storage_account="customeraccount", + azure_tenant_id_secret="azure_tenant", + azure_client_id_secret="azure_client", + azure_client_secret_secret="azure_credential", +) +print(integration.configuration_instructions) +``` + +Use `SecretReference` to qualify another database or namespace. Sync and async connections use the same options. Secret values are never part of integration metadata. Azure supports service principal client credentials; managed identity, account key and SAS authentication are not integration modes. + +The equivalent SQL is: + +```sql +CREATE STORAGE INTEGRATION customer_blobs + STORAGE_PROVIDER = 'AZURE' + STORAGE_PREFIX = 'az://customer-data/raw/' + STORAGE_AZURE_STORAGE_ACCOUNT = 'customeraccount' + STORAGE_AZURE_TENANT_ID_SECRET = public.azure_tenant + STORAGE_AZURE_CLIENT_ID_SECRET = public.azure_client + STORAGE_AZURE_CLIENT_SECRET_SECRET = public.azure_credential; +DESCRIBE STORAGE INTEGRATION customer_blobs; +SHOW STORAGE INTEGRATIONS; +DROP STORAGE INTEGRATION customer_blobs; +``` + +Registration returns setup instructions and leaves `access_verified` false. It stores an account-qualified prefix such as `az://customer-data@customeraccount.blob.core.windows.net/raw/`. diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index b9aae1a29..032fcb71e 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -394,6 +394,11 @@ class DBConnection(EnforceOverrides): aws_access_key_id_secret: Optional[Union[str, SecretReference]] = None, aws_secret_access_key_secret: Optional[Union[str, SecretReference]] = None, aws_region: Optional[str] = None, + storage_provider: Literal["S3", "AZURE"] = "S3", + azure_storage_account: Optional[str] = None, + azure_tenant_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_secret_secret: Optional[Union[str, SecretReference]] = None, namespace_path: Optional[List[str]] = None, ) -> StorageIntegrationInfo: raise NotImplementedError @@ -2243,6 +2248,11 @@ class LanceDBConnection(DBConnection): aws_access_key_id_secret: Optional[Union[str, SecretReference]] = None, aws_secret_access_key_secret: Optional[Union[str, SecretReference]] = None, aws_region: Optional[str] = None, + storage_provider: Literal["S3", "AZURE"] = "S3", + azure_storage_account: Optional[str] = None, + azure_tenant_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_secret_secret: Optional[Union[str, SecretReference]] = None, namespace_path: Optional[List[str]] = None, ) -> StorageIntegrationInfo: return LOOP.run( @@ -2253,6 +2263,11 @@ class LanceDBConnection(DBConnection): aws_access_key_id_secret=aws_access_key_id_secret, aws_secret_access_key_secret=aws_secret_access_key_secret, aws_region=aws_region, + storage_provider=storage_provider, + azure_storage_account=azure_storage_account, + azure_tenant_id_secret=azure_tenant_id_secret, + azure_client_id_secret=azure_client_id_secret, + azure_client_secret_secret=azure_client_secret_secret, namespace_path=namespace_path, ) ) @@ -3524,29 +3539,68 @@ class AsyncConnection(object): aws_access_key_id_secret: Optional[Union[str, SecretReference]] = None, aws_secret_access_key_secret: Optional[Union[str, SecretReference]] = None, aws_region: Optional[str] = None, + storage_provider: Literal["S3", "AZURE"] = "S3", + azure_storage_account: Optional[str] = None, + azure_tenant_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_secret_secret: Optional[Union[str, SecretReference]] = None, namespace_path: Optional[List[str]] = None, ) -> StorageIntegrationInfo: - """Register database-owned S3 access using a role or two Secret references. + """Register S3 or Azure Blob access using database-owned credential objects. ``namespace_path`` follows Secret naming rules. A SecretReference can name another database; a string names a root Secret in this database. + For Azure, set ``storage_provider="AZURE"``, ``azure_storage_account``, + and Secret references for tenant ID, client ID and client secret. Registration returns setup instructions without verifying storage access. """ - role_mode = aws_role_arn is not None - key_mode = ( - aws_access_key_id_secret is not None - and aws_secret_access_key_secret is not None + azure_options = ( + azure_tenant_id_secret, + azure_client_id_secret, + azure_client_secret_secret, ) - if role_mode == key_mode or ( - role_mode - and ( + if storage_provider == "AZURE": + if not azure_storage_account or any( + value is None for value in azure_options + ): + raise ValueError( + "Azure requires a storage account and " + "all three credential Secret references" + ) + if any( + value is not None + for value in ( + aws_role_arn, + aws_access_key_id_secret, + aws_secret_access_key_secret, + aws_region, + ) + ): + raise ValueError("Azure integrations cannot contain AWS options") + elif storage_provider == "S3": + if azure_storage_account is not None or any( + value is not None for value in azure_options + ): + raise ValueError("S3 integrations cannot contain Azure options") + role_mode = aws_role_arn is not None + key_mode = ( aws_access_key_id_secret is not None - or aws_secret_access_key_secret is not None - ) - ): - raise ValueError( - "Specify either aws_role_arn or both AWS credential Secret references" + and aws_secret_access_key_secret is not None ) + if role_mode == key_mode or ( + role_mode + and ( + aws_access_key_id_secret is not None + or aws_secret_access_key_secret is not None + ) + ): + raise ValueError( + "Specify either aws_role_arn or " + "both AWS credential Secret references" + ) + + else: + raise ValueError("storage_provider must be S3 or AZURE") def reference(value): if value is None: @@ -3560,10 +3614,14 @@ class AsyncConnection(object): return asdict(value) request = { - "storage_provider": "S3", + "storage_provider": storage_provider, "storage_prefix": storage_prefix, "aws_role_arn": aws_role_arn, "aws_region": aws_region, + "azure_storage_account": azure_storage_account, + "azure_tenant_id_secret": reference(azure_tenant_id_secret), + "azure_client_id_secret": reference(azure_client_id_secret), + "azure_client_secret_secret": reference(azure_client_secret_secret), "aws_access_key_id_secret": reference(aws_access_key_id_secret), "aws_secret_access_key_secret": reference(aws_secret_access_key_secret), } diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 99bd9bca9..25aef11b9 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -15,6 +15,7 @@ from typing import ( Dict, Iterable, List, + Literal, Optional, Sequence, Tuple, @@ -1047,6 +1048,11 @@ class RemoteDBConnection(DBConnection): aws_access_key_id_secret: Optional[Union[str, SecretReference]] = None, aws_secret_access_key_secret: Optional[Union[str, SecretReference]] = None, aws_region: Optional[str] = None, + storage_provider: Literal["S3", "AZURE"] = "S3", + azure_storage_account: Optional[str] = None, + azure_tenant_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_id_secret: Optional[Union[str, SecretReference]] = None, + azure_client_secret_secret: Optional[Union[str, SecretReference]] = None, namespace_path: Optional[List[str]] = None, ) -> StorageIntegrationInfo: return LOOP.run( @@ -1057,6 +1063,11 @@ class RemoteDBConnection(DBConnection): aws_access_key_id_secret=aws_access_key_id_secret, aws_secret_access_key_secret=aws_secret_access_key_secret, aws_region=aws_region, + storage_provider=storage_provider, + azure_storage_account=azure_storage_account, + azure_tenant_id_secret=azure_tenant_id_secret, + azure_client_id_secret=azure_client_id_secret, + azure_client_secret_secret=azure_client_secret_secret, namespace_path=namespace_path, ) ) diff --git a/python/python/lancedb/storage_integration.py b/python/python/lancedb/storage_integration.py index f3ede6631..025fbba12 100644 --- a/python/python/lancedb/storage_integration.py +++ b/python/python/lancedb/storage_integration.py @@ -18,7 +18,7 @@ class SecretReference: @dataclass(frozen=True) class StorageIntegrationInfo: - """Registered S3 configuration and policies to give to your AWS administrator. + """Registered cloud configuration and setup instructions. ``configuration_instructions`` is a single sentence of setup guidance for display; its wording may change. Use the structured identity and policy fields @@ -30,7 +30,7 @@ class StorageIntegrationInfo: name: str database: str namespace_path: list[str] - storage_provider: Literal["S3"] + storage_provider: Literal["S3", "AZURE"] storage_prefix: str aws_region: Optional[str] aws_account_id: str @@ -39,14 +39,24 @@ class StorageIntegrationInfo: created_at_millis: int access_verified: bool trust_policy: Optional[Dict[str, Any]] - permissions_policy: Dict[str, Any] + permissions_policy: Optional[Dict[str, Any]] configuration_instructions: str aws_role_arn: Optional[str] = None aws_access_key_id_secret: Optional[SecretReference] = None aws_secret_access_key_secret: Optional[SecretReference] = None + azure_storage_account: Optional[str] = None + azure_tenant_id_secret: Optional[SecretReference] = None + azure_client_id_secret: Optional[SecretReference] = None + azure_client_secret_secret: Optional[SecretReference] = None def __post_init__(self): - for field in ("aws_access_key_id_secret", "aws_secret_access_key_secret"): + for field in ( + "aws_access_key_id_secret", + "aws_secret_access_key_secret", + "azure_tenant_id_secret", + "azure_client_id_secret", + "azure_client_secret_secret", + ): value = getattr(self, field) if isinstance(value, dict): object.__setattr__(self, field, SecretReference(**value)) diff --git a/python/python/tests/test_storage_integration.py b/python/python/tests/test_storage_integration.py index 064ab6fe1..345343848 100644 --- a/python/python/tests/test_storage_integration.py +++ b/python/python/tests/test_storage_integration.py @@ -43,6 +43,23 @@ def remote_server(text_type=pa.string()): elif b"DROP STORAGE INTEGRATION" in command: result = pa.table({"dropped": [True]}) else: + if b"STORAGE_AZURE_STORAGE_ACCOUNT" in command: + info.update( + storage_provider="AZURE", + storage_prefix="az://customer-data@customeraccount.blob.core.windows.net/raw/", + azure_storage_account="customeraccount", + aws_role_arn=None, + trust_policy=None, + permissions_policy=None, + ) + for field, name in ( + ("azure_tenant_id_secret", "tenant"), + ("azure_client_id_secret", "client"), + ("azure_client_secret_secret", "credential"), + ): + info[field] = dict( + database="shared", namespace_path=["credentials"], name=name + ) if b"STORAGE_AWS_ACCESS_KEY_ID_SECRET" in command: info["aws_role_arn"] = None info["trust_policy"] = None @@ -63,6 +80,9 @@ def remote_server(text_type=pa.string()): "permissions_policy", "aws_access_key_id_secret", "aws_secret_access_key_secret", + "azure_tenant_id_secret", + "azure_client_id_secret", + "azure_client_secret_secret", ]: if field in row and row[field] is not None: row[field] = json.dumps(row[field]) @@ -205,3 +225,83 @@ def test_namespace_connections_keep_existing_apis_available(tmp_path): db = lancedb.connect_namespace("dir", {"root": str(tmp_path)}) with pytest.raises(NotImplementedError): db.create_storage_integration("raw", "s3://customer/raw/", aws_role_arn=ROLE) + + +@pytest.mark.parametrize( + "text_type", [pa.string(), pa.large_string(), pa.string_view()] +) +def test_azure_lifecycle_uses_three_qualified_secret_references(text_type): + with remote_server(text_type) as (host, calls): + db = connect(host) + refs = [ + SecretReference(name, database="shared", namespace_path=("credentials",)) + for name in ("tenant", "client", "credential") + ] + info = db.create_storage_integration( + "blobs", + "az://customer-data/raw/", + storage_provider="AZURE", + azure_storage_account="customeraccount", + azure_tenant_id_secret=refs[0], + azure_client_id_secret=refs[1], + azure_client_secret_secret=refs[2], + ) + assert info.storage_provider == "AZURE" + assert info.azure_storage_account == "customeraccount" + assert info.azure_tenant_id_secret.name == "tenant" + assert info.azure_client_id_secret.database == "shared" + assert info.azure_client_secret_secret.name == "credential" + assert info.trust_policy is None + assert info.permissions_policy is None + assert not info.access_verified + assert db.describe_storage_integration("blobs") == info + assert list(db.list_storage_integrations()) == ["raw", "zeta"] + db.drop_storage_integration("blobs") + sql = calls[0] + assert b"STORAGE_PROVIDER = 'AZURE'" in sql + assert b"STORAGE_AZURE_STORAGE_ACCOUNT = 'customeraccount'" in sql + assert ( + b'STORAGE_AZURE_CLIENT_SECRET_SECRET = "shared"."credentials"."credential"' + in sql + ) + assert b"STORAGE_AWS" not in sql + + +@pytest.mark.asyncio +async def test_async_azure_requires_account_and_all_three_secret_references(): + with remote_server() as (host, calls): + db = await lancedb.connect_async( + "db://default", + api_key="fake", + sql_host_override=host, + host_override="http://127.0.0.1:1", + ) + config = dict( + storage_provider="AZURE", + azure_storage_account="customeraccount", + azure_tenant_id_secret="tenant", + azure_client_id_secret="client", + azure_client_secret_secret="credential", + ) + for missing in ( + "azure_storage_account", + "azure_tenant_id_secret", + "azure_client_id_secret", + "azure_client_secret_secret", + ): + partial = dict(config) + partial.pop(missing) + with pytest.raises(ValueError): + await db.create_storage_integration( + "blobs", "az://customer-data/raw/", **partial + ) + with pytest.raises(ValueError): + await db.create_storage_integration( + "blobs", "az://customer-data/raw/", aws_role_arn=ROLE, **config + ) + assert calls == [] + info = await db.create_storage_integration( + "blobs", "az://customer-data/raw/", **config + ) + assert info.storage_provider == "AZURE" + assert b'STORAGE_AZURE_TENANT_ID_SECRET = "public"."tenant"' in calls[0] diff --git a/rust/lancedb/src/remote/storage_integration.rs b/rust/lancedb/src/remote/storage_integration.rs index 48ac51f27..51e0ccf8b 100644 --- a/rust/lancedb/src/remote/storage_integration.rs +++ b/rust/lancedb/src/remote/storage_integration.rs @@ -7,7 +7,7 @@ use arrow_array::{ use serde_json::Value; use crate::storage_integration::{ - CreateStorageIntegrationRequest, SecretReference, StorageIntegrationInfo, + CreateStorageIntegrationRequest, SecretReference, StorageIntegrationInfo, StorageProvider, }; use crate::utils::{validate_database_name, validate_namespace, validate_secret_name}; use crate::{Error, Result}; @@ -54,6 +54,53 @@ pub(super) fn create_sql( namespace: &[String], request: &CreateStorageIntegrationRequest, ) -> Result { + if request.storage_provider == StorageProvider::Azure { + if request.aws_role_arn.is_some() + || request.aws_access_key_id_secret.is_some() + || request.aws_secret_access_key_secret.is_some() + || request.aws_region.is_some() + { + return Err(invalid("Azure integrations cannot contain AWS options")); + } + let account = request + .azure_storage_account + .as_ref() + .filter(|account| !account.is_empty()) + .ok_or_else(|| invalid("Azure integrations require azure_storage_account"))?; + let mut sql = format!( + "CREATE STORAGE INTEGRATION {} STORAGE_PROVIDER = 'AZURE' STORAGE_PREFIX = {} STORAGE_AZURE_STORAGE_ACCOUNT = {}", + object_name(name, namespace)?, + literal(&request.storage_prefix), + literal(account) + ); + for (key, reference) in [ + ( + "STORAGE_AZURE_TENANT_ID_SECRET", + &request.azure_tenant_id_secret, + ), + ( + "STORAGE_AZURE_CLIENT_ID_SECRET", + &request.azure_client_id_secret, + ), + ( + "STORAGE_AZURE_CLIENT_SECRET_SECRET", + &request.azure_client_secret_secret, + ), + ] { + let reference = reference.as_ref().ok_or_else(|| { + invalid("Azure integrations require all three credential Secret references") + })?; + sql.push_str(&format!(" {key} = {}", secret_name(reference)?)); + } + return Ok(sql); + } + if request.azure_storage_account.is_some() + || request.azure_tenant_id_secret.is_some() + || request.azure_client_id_secret.is_some() + || request.azure_client_secret_secret.is_some() + { + return Err(invalid("S3 integrations cannot contain Azure options")); + } let mut sql = format!( "CREATE STORAGE INTEGRATION {} STORAGE_PROVIDER = 'S3' STORAGE_PREFIX = {}", object_name(name, namespace)?, @@ -122,6 +169,9 @@ pub(super) fn decode_info(batches: &[RecordBatch]) -> Result, #[serde(default, skip_serializing_if = "Option::is_none")] pub aws_region: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_storage_account: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_tenant_id_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_client_id_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_client_secret_secret: Option, } /// Cloud storage providers supported by storage integrations. @@ -57,6 +65,8 @@ pub struct CreateStorageIntegrationRequest { pub enum StorageProvider { #[serde(rename = "S3")] S3, + #[serde(rename = "AZURE")] + Azure, } impl CreateStorageIntegrationRequest { @@ -69,6 +79,10 @@ impl CreateStorageIntegrationRequest { aws_access_key_id_secret: None, aws_secret_access_key_secret: None, aws_region: None, + azure_storage_account: None, + azure_tenant_id_secret: None, + azure_client_id_secret: None, + azure_client_secret_secret: None, } } @@ -85,6 +99,41 @@ impl CreateStorageIntegrationRequest { aws_access_key_id_secret: Some(access_key), aws_secret_access_key_secret: Some(secret_key), aws_region: None, + azure_storage_account: None, + azure_tenant_id_secret: None, + azure_client_id_secret: None, + azure_client_secret_secret: None, + } + } + + /// Register Azure Blob access using three Secret references and an explicit account. + /// + /// ``` + /// use lancedb::storage_integration::{CreateStorageIntegrationRequest, SecretReference}; + /// let config = CreateStorageIntegrationRequest::azure( + /// "az://customer-data/raw/", "customeraccount", + /// SecretReference::new("tenant"), SecretReference::new("client"), + /// SecretReference::new("credential"), + /// ); + /// ``` + pub fn azure( + storage_prefix: impl Into, + storage_account: impl Into, + tenant_id: SecretReference, + client_id: SecretReference, + client_secret: SecretReference, + ) -> Self { + Self { + storage_provider: StorageProvider::Azure, + storage_prefix: storage_prefix.into(), + azure_storage_account: Some(storage_account.into()), + azure_tenant_id_secret: Some(tenant_id), + azure_client_id_secret: Some(client_id), + azure_client_secret_secret: Some(client_secret), + aws_role_arn: None, + aws_access_key_id_secret: None, + aws_secret_access_key_secret: None, + aws_region: None, } } @@ -95,7 +144,7 @@ impl CreateStorageIntegrationRequest { } } -/// Persisted configuration and complete AWS setup instructions. +/// Persisted configuration and cloud setup instructions. /// /// Registration does not verify AWS access or bind the integration to a storage /// operation. The policies and external ID are configuration metadata, never @@ -115,6 +164,14 @@ pub struct StorageIntegrationInfo { #[serde(default, skip_serializing_if = "Option::is_none")] pub aws_secret_access_key_secret: Option, pub aws_region: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_storage_account: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_tenant_id_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_client_id_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub azure_client_secret_secret: Option, pub aws_account_id: String, pub aws_iam_principal_arn: String, pub aws_external_id: String,