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.
This commit is contained in:
Jack Ye
2026-10-31 14:39:46 +08:00
committed by GitHub
parent e56150568c
commit b15fa9d6e4
7 changed files with 369 additions and 21 deletions
+37
View File
@@ -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/`.
+72 -14
View File
@@ -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),
}
+11
View File
@@ -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,
)
)
+14 -4
View File
@@ -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))
@@ -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]
+76 -1
View File
@@ -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<String> {
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<StorageIntegrationI
"namespace_path"
| "aws_access_key_id_secret"
| "aws_secret_access_key_secret"
| "azure_tenant_id_secret"
| "azure_client_id_secret"
| "azure_client_secret_secret"
| "trust_policy"
| "permissions_policy"
) && let Some(json) = value.as_str()
@@ -183,6 +233,31 @@ mod tests {
assert!(create_sql("a/b", &[], &role).is_err());
}
#[test]
fn azure_sql_uses_only_explicit_account_and_three_secret_objects() {
let reference = SecretReference::new("credential")
.with_database("shared")
.with_namespace(vec!["credentials".into()]);
let mut config = CreateStorageIntegrationRequest::azure(
"az://container/raw/",
"customeraccount",
reference.clone(),
reference.clone(),
reference,
);
let sql = create_sql("blobs", &[], &config).unwrap();
assert!(sql.contains("STORAGE_PROVIDER = 'AZURE'"));
assert!(sql.contains("STORAGE_AZURE_STORAGE_ACCOUNT = 'customeraccount'"));
assert!(sql.contains(
"STORAGE_AZURE_CLIENT_SECRET_SECRET = \"shared\".\"credentials\".\"credential\""
));
assert!(!sql.contains("STORAGE_AWS"));
config.azure_tenant_id_secret = None;
assert!(create_sql("blobs", &[], &config).is_err());
config.aws_region = Some("us-east-1".into());
assert!(create_sql("blobs", &[], &config).is_err());
}
#[test]
fn sql_metadata_rejects_missing_rows_and_invalid_json() {
assert!(decode_info(&[]).is_err());
+59 -2
View File
@@ -36,7 +36,7 @@ impl SecretReference {
}
}
/// Configuration for registering an AWS S3 storage integration.
/// Configuration for registering an AWS S3 or Azure Blob storage integration.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct CreateStorageIntegrationRequest {
@@ -50,6 +50,14 @@ pub struct CreateStorageIntegrationRequest {
pub aws_secret_access_key_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub aws_region: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_storage_account: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_tenant_id_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_client_id_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_client_secret_secret: Option<SecretReference>,
}
/// 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<String>,
storage_account: impl Into<String>,
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<SecretReference>,
pub aws_region: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_storage_account: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_tenant_id_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_client_id_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub azure_client_secret_secret: Option<SecretReference>,
pub aws_account_id: String,
pub aws_iam_principal_arn: String,
pub aws_external_id: String,