From 41d9c13fd3643f78d7282834ed3ac5453ee67882 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 01:34:53 +0000 Subject: [PATCH] fix(python): enable Azure default credentials --- python/python/lancedb/__init__.py | 14 ++++++++++++ python/python/lancedb/util.py | 4 ++++ python/python/tests/test_io.py | 36 ++++++++++++++++++++++++++++++- python/python/tests/test_util.py | 36 ++++++++++++++++++++++++++++++- 4 files changed, 88 insertions(+), 2 deletions(-) diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 235049f97..43c970ef7 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -162,6 +162,15 @@ def connect( ... }, ... ) + Azure managed identity authentication requires only the storage account name. + LanceDB acquires and refreshes managed identity tokens automatically. For a + user-assigned identity, also set ``azure_storage_client_id``: + + >>> db = lancedb.connect( # doctest: +SKIP + ... "az://my-container/lancedb", + ... storage_options={"account_name": "my-storage-account"}, + ... ) + For tests and temporary data, use an in-memory database: >>> db = lancedb.connect("memory://") @@ -455,6 +464,11 @@ async def connect_async( ... db = await lancedb.connect_async("s3://my-bucket/lancedb", ... storage_options={ ... "aws_access_key_id": "***"}) + ... # Azure managed identity tokens are acquired and refreshed automatically + ... db = await lancedb.connect_async( + ... "az://my-container/lancedb", + ... storage_options={"account_name": "my-storage-account"}, + ... ) ... # For tests and temporary data, use an in-memory database ... db = await lancedb.connect_async("memory://") ... # Connect to LanceDB cloud diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index f582be7b4..c0c0e89eb 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -102,6 +102,10 @@ def fs_from_uri(uri: str) -> Tuple[pa_fs.FileSystem, str]: az_blob_fs = adlfs.AzureBlobFileSystem( account_name=os.environ.get("AZURE_STORAGE_ACCOUNT_NAME"), account_key=os.environ.get("AZURE_STORAGE_ACCOUNT_KEY"), + # Without an explicit key, authenticate with DefaultAzureCredential + # instead of attempting anonymous access. In particular, this enables + # managed identity authentication on Azure hosts. + anon=False, ) fs = pa_fs.PyFileSystem(pa_fs.FSSpecHandler(az_blob_fs)) diff --git a/python/python/tests/test_io.py b/python/python/tests/test_io.py index dbf24d69d..aaf2078fe 100644 --- a/python/python/tests/test_io.py +++ b/python/python/tests/test_io.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import asyncio import os import lancedb @@ -12,11 +13,19 @@ import pytest # AWS_PROFILE=default TEST_S3_BASE_URL=s3://my_bucket/dataset pytest tests/test_io.py # # Azure: -# You need to setup Azure credentials an a base path to run this test. Example +# You need to set up Azure credentials and a base path to run this test. Examples: +# +# Account key: # export AZURE_STORAGE_ACCOUNT_NAME="" # export AZURE_STORAGE_ACCOUNT_KEY="" # export REMOTE_BASE_URL=az://my_blob/dataset # pytest tests/test_io.py +# +# Managed identity (system-assigned or user-assigned): +# export AZURE_STORAGE_ACCOUNT_NAME="" +# export AZURE_STORAGE_CLIENT_ID="" # user-assigned identity only +# export REMOTE_BASE_URL=az://my_blob/dataset +# pytest tests/test_io.py @pytest.fixture(autouse=True, scope="module") @@ -58,3 +67,28 @@ def test_remote_io(): assert len(db) == 1 assert db.open_table("test").name == db["test"].name + + +@pytest.mark.skipif( + (os.environ.get("REMOTE_BASE_URL") is None), + reason="please setup remote base url", +) +def test_remote_io_async(): + async def run(): + db = await lancedb.connect_async(os.environ["REMOTE_BASE_URL"]) + table = await db.create_table( + "test_async", + data=[ + {"vector": [3.1, 4.1], "item": "foo"}, + {"vector": [5.9, 26.5], "item": "bar"}, + ], + ) + + assert await table.count_rows() == 2 + assert (await db.open_table("test_async")).name == "test_async" + assert "test_async" in await db.table_names() + + await db.drop_table("test_async") + assert "test_async" not in await db.table_names() + + asyncio.run(run()) diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index a9b66b2dd..1b78555f9 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -26,7 +26,14 @@ import pandas as pd import polars as pl import pytest import lancedb -from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql +from lancedb import util +from lancedb.util import ( + flatten_columns, + fs_from_uri, + get_uri_scheme, + join_uri, + value_to_sql, +) from utils import exception_output @@ -78,6 +85,33 @@ def test_normalize_uri(): assert parsed_scheme == expected_scheme +def test_fs_from_uri_azure_uses_default_credential(monkeypatch): + azure_options = {} + filesystem = object() + + class MockAdlfs: + @staticmethod + def AzureBlobFileSystem(**kwargs): + azure_options.update(kwargs) + return object() + + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "account") + monkeypatch.delenv("AZURE_STORAGE_ACCOUNT_KEY", raising=False) + monkeypatch.setattr(util, "adlfs", MockAdlfs) + monkeypatch.setattr(util.pa_fs, "FSSpecHandler", lambda _: object()) + monkeypatch.setattr(util.pa_fs, "PyFileSystem", lambda _: filesystem) + + actual_filesystem, path = fs_from_uri("az://container/database") + + assert actual_filesystem is filesystem + assert path == "container/database" + assert azure_options == { + "account_name": "account", + "account_key": None, + "anon": False, + } + + def test_join_uri_remote(): schemes = ["s3", "az", "gs"] for scheme in schemes: