diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 1639a7c83..7b4068715 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -1061,3 +1061,72 @@ options. The service interprets the stored JSON when accessing storage; registra does not validate its credentials format or verify access, and leaves `access_verified` false. Integration metadata contains the Secret reference, never its JSON content. GCP options cannot be combined with AWS or Azure options. + + +### Hugging Face storage integrations + +Store a Hugging Face token as one Secret, then register the repository prefix +with the existing integration methods. Only the Secret name is part of the +integration configuration: + +```python +import os + +from lancedb.storage_integration import SecretReference + +db.create_secret("hf_token", os.environ["HF_TOKEN"]) +info = db.create_storage_integration( + "private_hub", + "hf://datasets/acme/private/raw/", + storage_provider="HUGGINGFACE", + huggingface_token_secret=SecretReference("hf_token"), +) +print(info.configuration_instructions) +info = db.describe_storage_integration("private_hub") +names = list(db.list_storage_integrations()) +db.drop_storage_integration("private_hub") +``` + +The equivalent SQL references the existing Secret: + +```sql +CREATE STORAGE INTEGRATION private_hub + STORAGE_PROVIDER = 'HUGGINGFACE' + STORAGE_PREFIX = 'hf://datasets/acme/private/raw/' + HUGGINGFACE_TOKEN_SECRET = public.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; +``` + +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/'; +``` + +When provided, the Secret must contain the token itself, with repository read +permission. +Cross-database and namespace references use the same `SecretReference` rules as +other integrations. Sync and async APIs accept the same options. Registration +returns one sentence of instructions and leaves `access_verified` false; it does +not contact Hugging Face. Metadata contains the Secret reference, never its value. +Hugging Face options cannot be combined with AWS, Azure or GCP options. + +Repository prefixes use `hf://datasets//[/]` or +`hf://models//[/]`. An optional `@revision` follows the repository +name; omitting it allows matching reads of any revision. Prefixes must be concrete +paths without wildcards or traversal. This integration does not make Hugging Face +a writable database location for `ATTACH DATABASE`. Dataset SQL sources select a +matching integration automatically, with no integration argument in the query. diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 1576348c9..1c9d46558 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -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), } diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index c71f581ba..0c1f127c8 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -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, ) ) diff --git a/python/python/lancedb/storage_integration.py b/python/python/lancedb/storage_integration.py index 4f443cdc6..a9dda088b 100644 --- a/python/python/lancedb/storage_integration.py +++ b/python/python/lancedb/storage_integration.py @@ -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): diff --git a/python/python/tests/test_storage_integration.py b/python/python/tests/test_storage_integration.py index d295aa98e..ff7995ab3 100644 --- a/python/python/tests/test_storage_integration.py +++ b/python/python/tests/test_storage_integration.py @@ -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] diff --git a/rust/lancedb/src/remote/storage_integration.rs b/rust/lancedb/src/remote/storage_integration.rs index 19ad5e6e9..bdba4f1df 100644 --- a/rust/lancedb/src/remote/storage_integration.rs +++ b/rust/lancedb/src/remote/storage_integration.rs @@ -54,6 +54,41 @@ pub(super) fn create_sql( namespace: &[String], request: &CreateStorageIntegrationRequest, ) -> Result { + if request.storage_provider != StorageProvider::HuggingFace + && request.huggingface_token_secret.is_some() + { + return Err(invalid( + "Hugging Face credential options require storage_provider HUGGINGFACE", + )); + } + if request.storage_provider == StorageProvider::HuggingFace { + 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() + || 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() + || request.gcp_credentials_json_secret.is_some() + { + return Err(invalid( + "Hugging Face integrations cannot contain AWS, Azure or GCP options", + )); + } + let mut sql = format!( + "CREATE STORAGE INTEGRATION {} STORAGE_PROVIDER = 'HUGGINGFACE' STORAGE_PREFIX = {}", + object_name(name, namespace)?, + literal(&request.storage_prefix) + ); + if let Some(reference) = &request.huggingface_token_secret { + sql.push_str(&format!( + " HUGGINGFACE_TOKEN_SECRET = {}", + secret_name(reference)? + )); + } + return Ok(sql); + } if request.storage_provider != StorageProvider::Gcp && request.gcp_credentials_json_secret.is_some() { @@ -205,6 +240,7 @@ pub(super) fn decode_info(batches: &[RecordBatch]) -> Result, #[serde(default, skip_serializing_if = "Option::is_none")] pub gcp_credentials_json_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub huggingface_token_secret: Option, } /// Cloud storage providers supported by storage integrations. @@ -71,6 +73,8 @@ pub enum StorageProvider { Azure, #[serde(rename = "GCP")] Gcp, + #[serde(rename = "HUGGINGFACE")] + HuggingFace, } impl CreateStorageIntegrationRequest { @@ -88,6 +92,7 @@ impl CreateStorageIntegrationRequest { azure_client_id_secret: None, azure_client_secret_secret: None, gcp_credentials_json_secret: None, + huggingface_token_secret: None, } } @@ -109,6 +114,7 @@ impl CreateStorageIntegrationRequest { azure_client_id_secret: None, azure_client_secret_secret: None, gcp_credentials_json_secret: None, + huggingface_token_secret: None, } } @@ -137,6 +143,7 @@ impl CreateStorageIntegrationRequest { azure_client_id_secret: Some(client_id), azure_client_secret_secret: Some(client_secret), gcp_credentials_json_secret: None, + huggingface_token_secret: None, aws_role_arn: None, aws_access_key_id_secret: None, aws_secret_access_key_secret: None, @@ -157,6 +164,7 @@ impl CreateStorageIntegrationRequest { storage_provider: StorageProvider::Gcp, storage_prefix: storage_prefix.into(), gcp_credentials_json_secret: Some(credentials_json), + huggingface_token_secret: None, aws_role_arn: None, aws_access_key_id_secret: None, aws_secret_access_key_secret: None, @@ -168,6 +176,38 @@ impl CreateStorageIntegrationRequest { } } + /// Register anonymous Hugging Face access at the default anonymous rate limits. + /// Use `with_huggingface_token_secret` to reference a Secret containing a token. + /// + /// ``` + /// use lancedb::storage_integration::{CreateStorageIntegrationRequest, SecretReference}; + /// let config = CreateStorageIntegrationRequest::huggingface( + /// "hf://datasets/acme/private/data/", + /// ).with_huggingface_token_secret(SecretReference::new("hf_token")); + /// ``` + pub fn huggingface(storage_prefix: impl Into) -> Self { + Self { + storage_provider: StorageProvider::HuggingFace, + storage_prefix: storage_prefix.into(), + huggingface_token_secret: None, + gcp_credentials_json_secret: None, + aws_role_arn: None, + 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, + } + } + + /// Use a Hugging Face token stored in a Secret instead of anonymous access. + pub fn with_huggingface_token_secret(mut self, token: SecretReference) -> Self { + self.huggingface_token_secret = Some(token); + self + } + /// Set the AWS region of the storage prefix. pub fn with_aws_region(mut self, region: impl Into) -> Self { self.aws_region = Some(region.into()); @@ -205,6 +245,8 @@ pub struct StorageIntegrationInfo { pub azure_client_secret_secret: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub gcp_credentials_json_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub huggingface_token_secret: Option, pub aws_account_id: String, pub aws_iam_principal_arn: String, pub aws_external_id: String,