rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support Hugging Face storage integrations with optional token Secrets (#4775)

Browse Source 
Adds `HUGGINGFACE` and an optional token Secret reference to the
existing storage integration methods. The same SQL transport handles
create, describe, list and drop. Only a Secret reference is stored in
integration metadata; mixed-provider credential options are rejected.

```python
from lancedb.storage_integration import SecretReference

info = db.create_storage_integration(
    "private_hub", "hf://datasets/acme/private/raw/",
    storage_provider="HUGGINGFACE",
    huggingface_token_secret=SecretReference(
        "hf_token", database="shared", namespace_path=("credentials",),
    ),
)
print(info.configuration_instructions)
info = db.describe_storage_integration("private_hub")
names = list(db.list_storage_integrations())
db.drop_storage_integration("private_hub")
```

Omit the token Secret for anonymous access at Hugging Face's default
anonymous
rate limits. A tokenless integration does not use a cluster's Hugging
Face token:

```python
info = db.create_storage_integration(
    "public_hub", "hf://datasets/acme/public/raw/",
    storage_provider="HUGGINGFACE",
)
```

```sql
CREATE STORAGE INTEGRATION public_hub
    STORAGE_PROVIDER = 'HUGGINGFACE'
    STORAGE_PREFIX = 'hf://datasets/acme/public/raw/';
```


```sql
CREATE STORAGE INTEGRATION private_hub
    STORAGE_PROVIDER = 'HUGGINGFACE'
    STORAGE_PREFIX = 'hf://datasets/acme/private/raw/'
    HUGGINGFACE_TOKEN_SECRET = shared.credentials.hf_token;
DESCRIBE STORAGE INTEGRATION private_hub;
SHOW STORAGE INTEGRATIONS;
SELECT * FROM 'hf://datasets/acme/private/raw/*.parquet' LIMIT 10;
DROP STORAGE INTEGRATION private_hub;
```

```mermaid
flowchart LR
    Python[Sync or async Python] --> SQL[Existing SQL transport]
    SQL --> Metadata[Integration and token Secret reference]
```

Adds Rust request construction, SQL rendering, typed metadata decoding,
sync/async Flight SQL regressions and public documentation. Local
builds/tests are excluded by request; CI provides execution validation.
This commit is contained in:
Jack Ye
2026-10-31 01:56:42 -07:00
committed by GitHub
parent 1fb2dd1551
commit 784d649380
7 changed files with 347 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/db.py
+31 -5
View File
@@ -394,12 +394,13 @@ 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", "GCP"] = "S3",
storage_provider: Literal["S3", "AZURE", "GCP", "HUGGINGFACE"] = "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,
gcp_credentials_json_secret: Optional[Union[str, SecretReference]] = None,
huggingface_token_secret: Optional[Union[str, SecretReference]] = None,
namespace_path: Optional[List[str]] = None,
) -> StorageIntegrationInfo:
raise NotImplementedError
@@ -2249,12 +2250,13 @@ 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", "GCP"] = "S3",
storage_provider: Literal["S3", "AZURE", "GCP", "HUGGINGFACE"] = "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,
gcp_credentials_json_secret: Optional[Union[str, SecretReference]] = None,
huggingface_token_secret: Optional[Union[str, SecretReference]] = None,
namespace_path: Optional[List[str]] = None,
) -> StorageIntegrationInfo:
return LOOP.run(
@@ -2271,6 +2273,7 @@ class LanceDBConnection(DBConnection):
azure_client_id_secret=azure_client_id_secret,
azure_client_secret_secret=azure_client_secret_secret,
gcp_credentials_json_secret=gcp_credentials_json_secret,
huggingface_token_secret=huggingface_token_secret,
namespace_path=namespace_path,
)
)
@@ -3542,12 +3545,13 @@ 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", "GCP"] = "S3",
storage_provider: Literal["S3", "AZURE", "GCP", "HUGGINGFACE"] = "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,
gcp_credentials_json_secret: Optional[Union[str, SecretReference]] = None,
huggingface_token_secret: Optional[Union[str, SecretReference]] = None,
namespace_path: Optional[List[str]] = None,
) -> StorageIntegrationInfo:
"""Register cloud storage access using database-owned credential Secrets.
@@ -3558,6 +3562,9 @@ class AsyncConnection(object):
and Secret references for tenant ID, client ID and client secret.
For GCP, set ``storage_provider="GCP"`` and ``gcp_credentials_json_secret``
to a Secret containing the application credentials JSON file content.
For Hugging Face, set ``storage_provider="HUGGINGFACE"`` and
optionally set ``huggingface_token_secret`` to a Secret containing the
token; omitting it selects anonymous access at the default rate limits.
Registration returns setup instructions without verifying storage access.
"""
azure_options = (
@@ -3565,9 +3572,27 @@ class AsyncConnection(object):
azure_client_id_secret,
azure_client_secret_secret,
)
if storage_provider != "HUGGINGFACE" and huggingface_token_secret is not None:
raise ValueError(
"Hugging Face options require storage_provider HUGGINGFACE"
)
if storage_provider != "GCP" and gcp_credentials_json_secret is not None:
raise ValueError("GCP credential options require storage_provider GCP")
if storage_provider == "GCP":
if storage_provider == "HUGGINGFACE":
if any(
value is not None
for value in azure_options
+ (
azure_storage_account,
aws_role_arn,
aws_access_key_id_secret,
aws_secret_access_key_secret,
aws_region,
gcp_credentials_json_secret,
)
):
raise ValueError("Hugging Face cannot contain other cloud options")
elif storage_provider == "GCP":
if gcp_credentials_json_secret is None:
raise ValueError("GCP requires gcp_credentials_json_secret")
if azure_storage_account is not None or any(
@@ -3622,7 +3647,7 @@ class AsyncConnection(object):
)
else:
raise ValueError("storage_provider must be S3, AZURE or GCP")
raise ValueError("storage_provider must be S3, AZURE, GCP or HUGGINGFACE")
def reference(value):
if value is None:
@@ -3645,6 +3670,7 @@ class AsyncConnection(object):
"azure_client_id_secret": reference(azure_client_id_secret),
"azure_client_secret_secret": reference(azure_client_secret_secret),
"gcp_credentials_json_secret": reference(gcp_credentials_json_secret),
"huggingface_token_secret": reference(huggingface_token_secret),
"aws_access_key_id_secret": reference(aws_access_key_id_secret),
"aws_secret_access_key_secret": reference(aws_secret_access_key_secret),
}
python/python/lancedb/remote/db.py
+3 -1
View File
@@ -1048,12 +1048,13 @@ 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", "GCP"] = "S3",
storage_provider: Literal["S3", "AZURE", "GCP", "HUGGINGFACE"] = "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,
gcp_credentials_json_secret: Optional[Union[str, SecretReference]] = None,
huggingface_token_secret: Optional[Union[str, SecretReference]] = None,
namespace_path: Optional[List[str]] = None,
) -> StorageIntegrationInfo:
return LOOP.run(
@@ -1070,6 +1071,7 @@ class RemoteDBConnection(DBConnection):
azure_client_id_secret=azure_client_id_secret,
azure_client_secret_secret=azure_client_secret_secret,
gcp_credentials_json_secret=gcp_credentials_json_secret,
huggingface_token_secret=huggingface_token_secret,
namespace_path=namespace_path,
)
)
python/python/lancedb/storage_integration.py
+3 -1
View File
@@ -30,7 +30,7 @@ class StorageIntegrationInfo:
name: str
database: str
namespace_path: list[str]
storage_provider: Literal["S3", "AZURE", "GCP"]
storage_provider: Literal["S3", "AZURE", "GCP", "HUGGINGFACE"]
storage_prefix: str
aws_region: Optional[str]
aws_account_id: str
@@ -49,6 +49,7 @@ class StorageIntegrationInfo:
azure_client_id_secret: Optional[SecretReference] = None
azure_client_secret_secret: Optional[SecretReference] = None
gcp_credentials_json_secret: Optional[SecretReference] = None
huggingface_token_secret: Optional[SecretReference] = None
def __post_init__(self):
for field in (
@@ -58,6 +59,7 @@ class StorageIntegrationInfo:
"azure_client_id_secret",
"azure_client_secret_secret",
"gcp_credentials_json_secret",
"huggingface_token_secret",
):
value = getattr(self, field)
if isinstance(value, dict):
python/python/tests/test_storage_integration.py
+135
View File
@@ -43,6 +43,29 @@ def remote_server(text_type=pa.string()):
elif b"DROP STORAGE INTEGRATION" in command:
result = pa.table({"dropped": [True]})
else:
if b"STORAGE_PROVIDER = 'HUGGINGFACE'" in command:
info.update(
storage_provider="HUGGINGFACE",
storage_prefix="hf://datasets/acme/private/raw/",
aws_role_arn=None,
aws_region=None,
aws_account_id="",
aws_iam_principal_arn="",
trust_policy=None,
permissions_policy=None,
configuration_instructions=(
"Grant the referenced identity access to the prefix."
),
huggingface_token_secret=(
dict(
database="shared",
namespace_path=["credentials"],
name="hf_token",
)
if b"HUGGINGFACE_TOKEN_SECRET" in command
else None
),
)
if b"GCP_CREDENTIALS_JSON_SECRET" in command:
info.update(
storage_provider="GCP",
@@ -103,6 +126,7 @@ def remote_server(text_type=pa.string()):
"azure_client_id_secret",
"azure_client_secret_secret",
"gcp_credentials_json_secret",
"huggingface_token_secret",
]:
if field in row and row[field] is not None:
row[field] = json.dumps(row[field])
@@ -392,3 +416,114 @@ async def test_async_gcp_requires_one_secret_and_rejects_other_cloud_options():
)
assert info.storage_provider == "GCP"
assert b'GCP_CREDENTIALS_JSON_SECRET = "public"."gcp"' in calls[0]
@pytest.mark.parametrize(
"text_type", [pa.string(), pa.large_string(), pa.string_view()]
)
def test_huggingface_lifecycle_uses_one_qualified_token_secret(text_type):
with remote_server(text_type) as (host, calls):
db = connect(host)
info = db.create_storage_integration(
"hub",
"hf://datasets/acme/private/raw/",
storage_provider="HUGGINGFACE",
huggingface_token_secret=SecretReference(
"hf_token", database="shared", namespace_path=("credentials",)
),
)
assert info.storage_provider == "HUGGINGFACE"
assert info.huggingface_token_secret.name == "hf_token"
assert info.huggingface_token_secret.database == "shared"
assert info.aws_role_arn is None
assert info.azure_storage_account is None
assert info.aws_iam_principal_arn == ""
assert info.trust_policy is None
assert not info.access_verified
assert db.describe_storage_integration("hub") == info
assert list(db.list_storage_integrations()) == ["raw", "zeta"]
db.drop_storage_integration("hub")
assert b"STORAGE_PROVIDER = 'HUGGINGFACE'" in calls[0]
assert (
b'HUGGINGFACE_TOKEN_SECRET = "shared"."credentials"."hf_token"' in calls[0]
)
assert b"STORAGE_AWS" not in calls[0]
assert b"STORAGE_AZURE" not in calls[0]
@pytest.mark.asyncio
async def test_async_huggingface_rejects_other_cloud_options():
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",
)
for invalid in (
{"huggingface_token_secret": "hf_token", "aws_role_arn": ROLE},
{
"huggingface_token_secret": "hf_token",
"azure_storage_account": "account",
},
):
with pytest.raises(ValueError):
await db.create_storage_integration(
"hub",
"hf://datasets/acme/private/raw/",
storage_provider="HUGGINGFACE",
**invalid,
)
with pytest.raises(ValueError):
await db.create_storage_integration(
"hub",
"hf://datasets/acme/private/raw/",
aws_role_arn=ROLE,
huggingface_token_secret="hf_token",
)
assert calls == []
info = await db.create_storage_integration(
"hub",
"hf://datasets/acme/private/raw/",
storage_provider="HUGGINGFACE",
huggingface_token_secret="hf_token",
)
assert info.storage_provider == "HUGGINGFACE"
assert b'HUGGINGFACE_TOKEN_SECRET = "public"."hf_token"' in calls[0]
@pytest.mark.parametrize(
"text_type", [pa.string(), pa.large_string(), pa.string_view()]
)
def test_anonymous_huggingface_lifecycle_omits_token_secret(text_type):
with remote_server(text_type) as (host, calls):
db = connect(host)
info = db.create_storage_integration(
"public_hub",
"hf://datasets/acme/public/raw/",
storage_provider="HUGGINGFACE",
)
assert info.storage_provider == "HUGGINGFACE"
assert info.huggingface_token_secret is None
assert db.describe_storage_integration("public_hub") == info
assert list(db.list_storage_integrations()) == ["raw", "zeta"]
db.drop_storage_integration("public_hub")
assert b"HUGGINGFACE_TOKEN_SECRET" not in calls[0]
@pytest.mark.asyncio
async def test_async_huggingface_allows_anonymous_access():
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",
)
info = await db.create_storage_integration(
"public_hub",
"hf://datasets/acme/public/raw/",
storage_provider="HUGGINGFACE",
)
assert info.huggingface_token_secret is None
assert b"HUGGINGFACE_TOKEN_SECRET" not in calls[0]