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,