From 6d126a3601acc0d400cc4fbe5cf45371e08f7e2e Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 30 Oct 2026 00:26:42 -0700 Subject: [PATCH] feat: add database storage integration APIs (#4618) Adds database-owned S3 storage integrations with Secret-style namespace paths in a separate keyspace. Use a regular database under the cluster root for bootstrap integrations and Secrets. `default` is an example name, with no enforced or automatically provisioned database; `system` remains read-only. Authentication uses either an IAM role or two existing Secret references, including cross-database references; configuration metadata never contains credential values. ### Python experience ```python catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret") db = catalog.connect_database("default") info = db.create_storage_integration( "raw", "s3://customer-bucket/raw/", aws_role_arn="arn:aws:iam::123456789012:role/lancedb-access", aws_region="us-east-1", ) print(info.configuration_instructions) info = db.describe_storage_integration("raw") print(list(db.list_storage_integrations())) db.drop_storage_integration("raw") ``` For static credentials, replace `aws_role_arn` with `aws_access_key_id_secret=SecretReference("s3_key_id", database="default")` and `aws_secret_access_key_secret=SecretReference("s3_secret_key", database="default")`, imported from `lancedb.storage_integration`. Pass `namespace_path` to address nested namespaces. Async database connections expose the same APIs; listing is lazy. All four client operations execute SQL through the existing SQL transport, with typed Arrow result decoding and name-based listing cursors. ### SQL experience ```sql CREATE STORAGE INTEGRATION default.public.raw STORAGE_PROVIDER = 'S3' STORAGE_PREFIX = 's3://customer-bucket/raw/' STORAGE_AWS_ROLE_ARN = 'arn:aws:iam::123456789012:role/lancedb-access' STORAGE_AWS_REGION = 'us-east-1'; DESCRIBE STORAGE INTEGRATION default.public.raw; SHOW STORAGE INTEGRATIONS IN default.public; DROP STORAGE INTEGRATION IF EXISTS default.public.raw; CREATE STORAGE INTEGRATION default.public.raw STORAGE_PROVIDER = 'S3' STORAGE_PREFIX = 's3://customer-bucket/raw/' STORAGE_AWS_ACCESS_KEY_ID_SECRET = default.public.s3_key_id STORAGE_AWS_SECRET_ACCESS_KEY_SECRET = default.public.s3_secret_key; ``` CREATE and DESCRIBE return configuration and a single sentence of setup guidance. IAM mode also returns the deployment role principal, stable external ID, customer trust policy, and prefix-scoped permissions policy. Registration does not verify access. Creation requires namespace CREATE_STORAGE_INTEGRATION plus parent USAGE; Secret references require source database, namespace, and Secret USAGE. Describe requires object USAGE; drop requires OWNERSHIP. SQL supports GRANT/REVOKE for these objects and their creation privilege. Includes client documentation and Flight SQL regression coverage; no dedicated storage-integration REST endpoints are used. Local builds/tests are excluded by request. --- docs/src/python/python.md | 81 +++++++ python/python/lancedb/__init__.py | 1 + python/python/lancedb/_lancedb.pyi | 20 ++ python/python/lancedb/authz.py | 12 + python/python/lancedb/catalog.py | 2 +- python/python/lancedb/db.py | 187 ++++++++++++++++ python/python/lancedb/remote/db.py | 57 +++++ python/python/lancedb/storage_integration.py | 52 +++++ .../python/tests/test_storage_integration.py | 207 ++++++++++++++++++ python/src/authz.rs | 13 ++ python/src/connection.rs | 66 ++++++ rust/lancedb/src/authz/mod.rs | 2 +- rust/lancedb/src/authz/object.rs | 67 ++++++ rust/lancedb/src/authz/privilege.rs | 34 ++- rust/lancedb/src/catalog.rs | 5 +- rust/lancedb/src/connection.rs | 63 ++++++ rust/lancedb/src/database.rs | 34 +++ rust/lancedb/src/lib.rs | 1 + rust/lancedb/src/remote.rs | 1 + rust/lancedb/src/remote/db.rs | 82 ++++++- rust/lancedb/src/remote/sql_test.rs | 114 ++++++++++ .../lancedb/src/remote/storage_integration.rs | 201 +++++++++++++++++ rust/lancedb/src/storage_integration.rs | 126 +++++++++++ .../tests/fixtures/storage_integration.json | 36 +++ 24 files changed, 1453 insertions(+), 11 deletions(-) create mode 100644 python/python/lancedb/storage_integration.py create mode 100644 python/python/tests/test_storage_integration.py create mode 100644 rust/lancedb/src/remote/storage_integration.rs create mode 100644 rust/lancedb/src/storage_integration.rs create mode 100644 rust/lancedb/tests/fixtures/storage_integration.json diff --git a/docs/src/python/python.md b/docs/src/python/python.md index b3c0c4672..e737223bf 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -42,6 +42,12 @@ table names. `db.list_tables().tables` is a list of the first page; `limit` controls its page size and `page_token` resumes at a saved token. The same patterns apply to `list_namespaces()`. Other resource listings use `page_limit`. +Resource listings (`list_storage_integrations`, `list_secrets`, `list_views`, `list_jobs`, `list_functions`, and +`list_materialized_views`) return lazy iterators. Use `list(db.list_views())` to +collect synchronous results, or `[name async for name in db.list_views()]` with +an asynchronous connection. `page_limit` controls each request, and `page_token` +resumes a listing. Drain the iterator's cached results before saving its token. + ::: lancedb.listing.Listing ::: lancedb.listing.TableListing @@ -162,6 +168,11 @@ registered functions, for example, namespace. These objects identify resources for ACL operations; a secret object contains the credential's name, not its value. +Storage integration objects identify database-owned storage configuration, for +example `Object.storage_integration(database="default", name="raw")`. Pass +`namespace_path` for nested namespaces. Their keyspace is separate from Secrets, +Functions, and relations; the wire name is `storage_integration:default:public$raw`. + More object types may be added in the future. ### Access control lists @@ -205,6 +216,27 @@ for privilege in [Privilege.CREATE_SECRET, Privilege.CREATE_FUNCTION]: Both privileges also accept their string names in synchronous and asynchronous ACL methods. +Storage integration creation requires `CREATE_STORAGE_INTEGRATION` on its containing +namespace, together with database and namespace `USAGE`: + +```python +catalog.authz.add_acl( + object=Object.namespace(database="default"), + subject=Subject.role("developer"), + privilege=Privilege.CREATE_STORAGE_INTEGRATION, +) +catalog.authz.add_acl( + object=Object.storage_integration(database="default", name="raw"), + subject=Subject.role("reader"), + privilege=Privilege.USAGE, +) +``` + +Creation grants ownership to the creator. Describing an integration requires +`USAGE` on that integration; dropping it requires `OWNERSHIP`. Listing integrations +requires `USAGE` on the selected database and namespace. Describing and dropping +also require parent `USAGE` grants. Manage these grants through `catalog.authz`. + By default, the `list_acls` function lists all the access control entries in the system. By adding a function parameter, the output is filtered to just ACLs matching that parameter. For example, this code lists all the ACLs for the @@ -851,3 +883,52 @@ rows nearest to a query vector and can be created with the ::: lancedb.query.AsyncTakeQuery options: inherited_members: true + +## Storage integrations + +::: lancedb.storage_integration.StorageIntegrationInfo + +::: lancedb.storage_integration.SecretReference + +Integrations belong to a database and follow the same namespace-path naming rules +as Secrets, with their own keyspace. Use a regular database under the cluster root for storage configuration needed +before opening an external database. `default` is a convention; no database with +that name is required or created automatically. The system database is read-only. +Client methods execute SQL through the existing SQL transport. + +```python +catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret") +db = catalog.connect_database("default") +integration = db.create_storage_integration( + "raw", "s3://customer-bucket/raw/", + aws_role_arn="arn:aws:iam::123456789012:role/lancedb-access", +) +print(integration.configuration_instructions) +integration = db.describe_storage_integration("raw") +for name in db.list_storage_integrations(): + print(name) +db.drop_storage_integration("raw") +``` + +Static credentials use two existing, single-value Secrets instead of a role: + +```python +from lancedb.storage_integration import SecretReference + +integration = db.create_storage_integration( + "raw", "s3://customer-bucket/raw/", + aws_access_key_id_secret=SecretReference("s3-key-id", database="default"), + aws_secret_access_key_secret=SecretReference("s3-secret-key", database="default"), +) +``` + +Pass `namespace_path` on the object APIs and Secret references to select a namespace. +A string Secret reference selects the root namespace of the integration's database. +Cross-database references require `USAGE` on the source database, namespace, and +Secret. Role and static-credential modes are mutually exclusive. Credential values +are never included in integration metadata. + +`configuration_instructions` is one sentence of display guidance whose wording may +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. diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 74f37fc2e..4c38f374c 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -41,6 +41,7 @@ from .functions import ( ) from .secrets import EnvVarSecret as EnvVarSecret from .secrets import SecretInfo as SecretInfo +from .storage_integration import StorageIntegrationInfo as StorageIntegrationInfo from .materialized_view import ( AsyncMaterializedView, MaterializedView, diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 2b8d1655a..4c9ad38de 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -171,6 +171,22 @@ class Connection(object): async def drop_function_async( self, name: str, version: str, namespace_path: Optional[List[str]] = None ) -> Tuple[bool, Job]: ... + async def create_storage_integration( + self, name: str, request: str, namespace_path: Optional[List[str]] = None + ) -> str: ... + async def describe_storage_integration( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> str: ... + def list_storage_integrations( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> NameListing: ... + async def drop_storage_integration( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> None: ... async def create_secret( self, name: str, value: str, namespace_path: Optional[List[str]] = None ) -> None: ... @@ -1080,6 +1096,10 @@ class AuthzObject: database: str, namespace_path: Sequence[str], name: str ) -> AuthzObject: ... @staticmethod + def storage_integration( + database: str, namespace: Sequence[str], name: str + ) -> AuthzObject: ... + @staticmethod def secret( database: str, namespace_path: Sequence[str], name: str ) -> AuthzObject: ... diff --git a/python/python/lancedb/authz.py b/python/python/lancedb/authz.py index e62a276cf..cb7d04ed2 100644 --- a/python/python/lancedb/authz.py +++ b/python/python/lancedb/authz.py @@ -172,6 +172,17 @@ class Object: """A secret, defaulting to the database's public namespace.""" return cls(str(_lancedb.AuthzObject.secret(database, namespace_path, name))) + @classmethod + def storage_integration( + cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",) + ) -> "Object": + """A database storage integration, defaulting to its public namespace.""" + return cls( + str( + _lancedb.AuthzObject.storage_integration(database, namespace_path, name) + ) + ) + @classmethod def function( cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",) @@ -217,6 +228,7 @@ class Privilege(str, Enum): CREATE_SECRET = "CREATE_SECRET" CREATE_FUNCTION = "CREATE_FUNCTION" OPERATE = "OPERATE" + CREATE_STORAGE_INTEGRATION = "CREATE_STORAGE_INTEGRATION" @dataclass(frozen=True) diff --git a/python/python/lancedb/catalog.py b/python/python/lancedb/catalog.py index 9a0dc3412..338be98b8 100644 --- a/python/python/lancedb/catalog.py +++ b/python/python/lancedb/catalog.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""Remote catalogs manage databases through a server's root namespace.""" +"""Remote catalogs manage databases.""" from functools import cached_property from datetime import timedelta diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index e13e3dc9c..d50927936 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -8,6 +8,8 @@ from abc import abstractmethod import asyncio from datetime import timedelta from pathlib import Path +from dataclasses import asdict +import json import sys from typing import ( TYPE_CHECKING, @@ -58,6 +60,7 @@ from . import __version__ from ._lancedb import connect as lancedb_connect # type: ignore from .functions import FunctionVersion, UdfDefinition from .job import AsyncJob, Job, _typed_job +from .storage_integration import SecretReference, StorageIntegrationInfo from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery from .sql import QueryDescription, QueryParameters, to_parameter_batch @@ -375,6 +378,38 @@ class DBConnection(EnforceOverrides): """ pass + def create_storage_integration( + self, + name: str, + storage_prefix: str, + *, + aws_role_arn: Optional[str] = None, + 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, + namespace_path: Optional[List[str]] = None, + ) -> StorageIntegrationInfo: + raise NotImplementedError + + def describe_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> StorageIntegrationInfo: + raise NotImplementedError + + def list_storage_integrations( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + raise NotImplementedError + + def drop_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + raise NotImplementedError + @abstractmethod def create_table( self, @@ -2099,6 +2134,62 @@ class LanceDBConnection(DBConnection): ) return dropped, Job(job) + @override + def create_storage_integration( + self, + name: str, + storage_prefix: str, + *, + aws_role_arn: Optional[str] = None, + 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, + namespace_path: Optional[List[str]] = None, + ) -> StorageIntegrationInfo: + return LOOP.run( + self._conn.create_storage_integration( + name, + storage_prefix, + aws_role_arn=aws_role_arn, + aws_access_key_id_secret=aws_access_key_id_secret, + aws_secret_access_key_secret=aws_secret_access_key_secret, + aws_region=aws_region, + namespace_path=namespace_path, + ) + ) + + @override + def describe_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> StorageIntegrationInfo: + return LOOP.run( + self._conn.describe_storage_integration(name, namespace_path=namespace_path) + ) + + @override + def list_storage_integrations( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + return Listing( + self._conn.list_storage_integrations( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) + + @override + def drop_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + LOOP.run( + self._conn.drop_storage_integration(name, namespace_path=namespace_path) + ) + @override def create_secret( self, name: str, value: str, *, namespace_path: Optional[List[str]] = None @@ -3267,6 +3358,102 @@ class AsyncConnection(object): ) return dropped, AsyncJob(job) + async def create_storage_integration( + self, + name: str, + storage_prefix: str, + *, + aws_role_arn: Optional[str] = None, + 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, + namespace_path: Optional[List[str]] = None, + ) -> StorageIntegrationInfo: + """Register database-owned S3 access using a role or two Secret references. + + ``namespace_path`` follows Secret naming rules. A SecretReference can name + another database; a string names a root Secret in this database. + 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 + ) + 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" + ) + + def reference(value): + if value is None: + return None + if isinstance(value, str): + value = SecretReference(value) + if not isinstance(value, SecretReference): + raise TypeError( + "Credential secrets must be names or SecretReference objects" + ) + return asdict(value) + + request = { + "storage_provider": "S3", + "storage_prefix": storage_prefix, + "aws_role_arn": aws_role_arn, + "aws_region": aws_region, + "aws_access_key_id_secret": reference(aws_access_key_id_secret), + "aws_secret_access_key_secret": reference(aws_secret_access_key_secret), + } + info = await self._inner.create_storage_integration( + validate_secret_name(name), + json.dumps(request), + list(validate_namespace_path(namespace_path)), + ) + return StorageIntegrationInfo(**json.loads(info)) + + async def describe_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> StorageIntegrationInfo: + """Return configuration and Secret references without credential values.""" + return StorageIntegrationInfo( + **json.loads( + await self._inner.describe_storage_integration( + validate_secret_name(name), + list(validate_namespace_path(namespace_path)), + ) + ) + ) + + def list_storage_integrations( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> AsyncListing[str]: + """Lazily list one namespace of this database's storage integrations.""" + return AsyncListing( + self._inner.list_storage_integrations( + namespace_path=list(validate_namespace_path(namespace_path)), + page_token=page_token, + page_limit=page_limit, + ) + ) + + async def drop_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + """Drop a storage integration with ownership permission.""" + await self._inner.drop_storage_integration( + validate_secret_name(name), list(validate_namespace_path(namespace_path)) + ) + async def create_secret( self, name: str, value: str, *, namespace_path: Optional[List[str]] = None ) -> None: diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index d97e6526b..436f4f10e 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -43,6 +43,7 @@ from ..sql import Query as SqlQuery from ..sql import QueryDescription, QueryParameters from ..materialized_view import MaterializedView, MaterializedViewSource, SelectArg from ..secrets import EnvVarSecret, SecretInfo +from ..storage_integration import SecretReference, StorageIntegrationInfo from ..view import ViewDescription if TYPE_CHECKING: @@ -1036,6 +1037,62 @@ class RemoteDBConnection(DBConnection): ) return dropped, Job(job) + @override + def create_storage_integration( + self, + name: str, + storage_prefix: str, + *, + aws_role_arn: Optional[str] = None, + 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, + namespace_path: Optional[List[str]] = None, + ) -> StorageIntegrationInfo: + return LOOP.run( + self._conn.create_storage_integration( + name, + storage_prefix, + aws_role_arn=aws_role_arn, + aws_access_key_id_secret=aws_access_key_id_secret, + aws_secret_access_key_secret=aws_secret_access_key_secret, + aws_region=aws_region, + namespace_path=namespace_path, + ) + ) + + @override + def describe_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> StorageIntegrationInfo: + return LOOP.run( + self._conn.describe_storage_integration(name, namespace_path=namespace_path) + ) + + @override + def list_storage_integrations( + self, + *, + namespace_path: Optional[List[str]] = None, + page_token: Optional[str] = None, + page_limit: Optional[int] = None, + ) -> Listing[str]: + return Listing( + self._conn.list_storage_integrations( + namespace_path=namespace_path, + page_token=page_token, + page_limit=page_limit, + ) + ) + + @override + def drop_storage_integration( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + LOOP.run( + self._conn.drop_storage_integration(name, namespace_path=namespace_path) + ) + @override def create_secret( self, name: str, value: str, *, namespace_path: Optional[List[str]] = None diff --git a/python/python/lancedb/storage_integration.py b/python/python/lancedb/storage_integration.py new file mode 100644 index 000000000..f3ede6631 --- /dev/null +++ b/python/python/lancedb/storage_integration.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Database-owned storage integrations and references to credential Secrets.""" + +from dataclasses import dataclass +from typing import Any, Dict, Literal, Optional + + +@dataclass(frozen=True) +class SecretReference: + """A credential Secret, optionally qualified by database and namespace.""" + + name: str + database: Optional[str] = None + namespace_path: tuple[str, ...] = () + + +@dataclass(frozen=True) +class StorageIntegrationInfo: + """Registered S3 configuration and policies to give to your AWS administrator. + + ``configuration_instructions`` is a single sentence of setup guidance for + display; its wording may change. Use the structured identity and policy fields + instead of parsing that text. Creating this object does not verify AWS access + or configure a storage operation to use it. Static-key authentication stores + only Secret references. No field contains AWS credentials. + """ + + name: str + database: str + namespace_path: list[str] + storage_provider: Literal["S3"] + storage_prefix: str + aws_region: Optional[str] + aws_account_id: str + aws_iam_principal_arn: str + aws_external_id: str + created_at_millis: int + access_verified: bool + trust_policy: Optional[Dict[str, Any]] + permissions_policy: 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 + + def __post_init__(self): + for field in ("aws_access_key_id_secret", "aws_secret_access_key_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 new file mode 100644 index 000000000..064ab6fe1 --- /dev/null +++ b/python/python/tests/test_storage_integration.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import contextlib +import json +from pathlib import Path + +import pytest +import pyarrow as pa +import pyarrow.flight as flight + +import lancedb +from lancedb.authz import Object, Privilege +from lancedb.db import AsyncConnection, DBConnection +from lancedb.catalog import Catalog, AsyncCatalog +from lancedb.storage_integration import SecretReference, StorageIntegrationInfo + + +INFO = json.loads( + ( + Path(__file__).parents[3] + / "rust/lancedb/tests/fixtures/storage_integration.json" + ).read_text() +) +NAME = "raw" +ROLE = "arn:aws:iam::123456789012:role/access" + + +@contextlib.contextmanager +def remote_server(text_type=pa.string()): + calls = [] + info = dict(INFO) + + class Server(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") + + def do_get(self, context, ticket): + command = ticket.ticket + calls.append(command) + if b"SHOW STORAGE INTEGRATIONS" in command: + result = pa.table({"name": pa.array(["zeta", NAME], type=text_type)}) + elif b"DROP STORAGE INTEGRATION" in command: + result = pa.table({"dropped": [True]}) + else: + if b"STORAGE_AWS_ACCESS_KEY_ID_SECRET" in command: + info["aws_role_arn"] = None + info["trust_policy"] = None + info["aws_access_key_id_secret"] = { + "database": "shared", + "namespace_path": ["credentials"], + "name": "access_key", + } + info["aws_secret_access_key_secret"] = { + "database": "shared", + "namespace_path": ["credentials"], + "name": "secret_key", + } + row = dict(info) + for field in [ + "namespace_path", + "trust_policy", + "permissions_policy", + "aws_access_key_id_secret", + "aws_secret_access_key_secret", + ]: + if field in row and row[field] is not None: + row[field] = json.dumps(row[field]) + row.setdefault("aws_access_key_id_secret", None) + row.setdefault("aws_secret_access_key_secret", None) + result = pa.table( + { + field: pa.array( + [value], + type=( + pa.bool_() + if field == "access_verified" + else pa.int64() + if field == "created_at_millis" + else text_type + ), + ) + for field, value in row.items() + } + ) + return flight.RecordBatchStream(result) + + with Server() as server: + yield f"grpc://127.0.0.1:{server.port}", calls + + +def connect(host): + return lancedb.connect( + "db://default", + api_key="fake", + sql_host_override=host, + host_override="http://127.0.0.1:1", + ) + + +def test_storage_integration_methods_belong_to_databases(): + assert hasattr(DBConnection, "create_storage_integration") + assert hasattr(AsyncConnection, "create_storage_integration") + assert not hasattr(Catalog, "create_storage_integration") + assert not hasattr(AsyncCatalog, "create_storage_integration") + + +@pytest.mark.parametrize( + "text_type", [pa.string(), pa.large_string(), pa.string_view()] +) +def test_sync_storage_integration_lifecycle_and_aws_information(text_type): + with remote_server(text_type) as (host, calls): + db = connect(host) + info = db.create_storage_integration( + NAME, INFO["storage_prefix"], aws_role_arn=ROLE + ) + assert isinstance(info, lancedb.StorageIntegrationInfo) + assert info.aws_account_id == "001122334455" + assert not info.access_verified + assert ( + info.trust_policy["Statement"][0]["Condition"]["StringEquals"][ + "sts:ExternalId" + ] + == info.aws_external_id + ) + assert db.describe_storage_integration(NAME) == info + listing = db.list_storage_integrations(page_limit=1) + assert len(calls) == 2 + assert list(listing) == ["raw", "zeta"] + assert db.drop_storage_integration(NAME) is None + assert b'CREATE STORAGE INTEGRATION "raw"' in calls[0] + assert b"STORAGE_AWS_ROLE_ARN = '" + ROLE.encode() + b"'" in calls[0] + assert b'DROP STORAGE INTEGRATION "raw"' in calls[-1] + + +@pytest.mark.asyncio +async def test_async_storage_integration_lifecycle(): + 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( + NAME, INFO["storage_prefix"], aws_role_arn=ROLE + ) + assert await db.describe_storage_integration(NAME) == info + assert [name async for name in db.list_storage_integrations()] == [ + "raw", + "zeta", + ] + await db.drop_storage_integration(NAME) + with pytest.raises(ValueError): + await db.describe_storage_integration("a//b") + + +def test_integration_authorization_object_and_creation_privilege(): + value = Object.storage_integration(database="default", name=NAME) + assert str(value) == "storage_integration:default:public$raw" + assert Privilege.CREATE_STORAGE_INTEGRATION.value == "CREATE_STORAGE_INTEGRATION" + + +@pytest.mark.parametrize( + "text_type", [pa.string(), pa.large_string(), pa.string_view()] +) +def test_static_credentials_use_explicit_cross_database_secret_references(text_type): + with remote_server(text_type) as (host, calls): + db = connect(host) + info = db.create_storage_integration( + "raw", + INFO["storage_prefix"], + aws_access_key_id_secret=SecretReference( + "access_key", database="shared", namespace_path=("credentials",) + ), + aws_secret_access_key_secret=SecretReference( + "secret_key", database="shared", namespace_path=("credentials",) + ), + ) + assert info.aws_role_arn is None + assert db.describe_storage_integration("raw") == info + assert info.aws_access_key_id_secret.database == "shared" + assert b"STORAGE_AWS_ROLE_ARN" not in calls[0] + assert ( + b'STORAGE_AWS_ACCESS_KEY_ID_SECRET = "shared"."credentials"."access_key"' + in calls[0] + ) + with pytest.raises(ValueError): + db.create_storage_integration( + "bad", + INFO["storage_prefix"], + aws_role_arn=ROLE, + aws_access_key_id_secret="key", + ) + assert len(calls) == 2 + + +def test_metadata_defaults_an_omitted_role_to_none(): + data = dict(INFO) + data.pop("aws_role_arn") + assert StorageIntegrationInfo(**data).aws_role_arn is None + + +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) diff --git a/python/src/authz.rs b/python/src/authz.rs index 6ad3e69b1..8e5ae5b7c 100644 --- a/python/src/authz.rs +++ b/python/src/authz.rs @@ -262,6 +262,19 @@ impl AuthzObject { })) } + #[staticmethod] + fn storage_integration( + database: &str, + namespace: &Bound<'_, PyAny>, + name: &str, + ) -> PyResult { + Self::validated(Object::StorageIntegration(StorageIntegrationObject { + database: database.into(), + namespace: namespace_path(namespace)?, + storage_integration: name.into(), + })) + } + #[staticmethod] fn function(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult { Self::validated(Object::Function(FunctionObject { diff --git a/python/src/connection.rs b/python/src/connection.rs index 54ea615de..004456737 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -888,6 +888,72 @@ impl Connection { }) } + #[pyo3(signature = (name, request, namespace_path=None))] + pub fn create_storage_integration<'py>( + &self, + py: Python<'py>, + name: String, + request: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self.get_inner()?.clone(); + let request: lancedb::storage_integration::CreateStorageIntegrationRequest = + serde_json::from_str(&request) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + future_into_py(py, async move { + let info = inner + .create_storage_integration(name, request, &namespace_path.unwrap_or_default()) + .await + .infer_error()?; + serde_json::to_string(&info).map_err(|error| PyValueError::new_err(error.to_string())) + }) + } + #[pyo3(signature = (name, namespace_path=None))] + pub fn describe_storage_integration<'py>( + &self, + py: Python<'py>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self.get_inner()?.clone(); + future_into_py(py, async move { + let info = inner + .describe_storage_integration(name, &namespace_path.unwrap_or_default()) + .await + .infer_error()?; + serde_json::to_string(&info).map_err(|error| PyValueError::new_err(error.to_string())) + }) + } + #[pyo3(signature = (*, namespace_path=None, page_token=None, page_limit=None))] + pub fn list_storage_integrations( + &self, + namespace_path: Option>, + page_token: Option, + page_limit: Option, + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_storage_integrations( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, page_limit), + ), + )) + } + #[pyo3(signature = (name, namespace_path=None))] + pub fn drop_storage_integration<'py>( + &self, + py: Python<'py>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self.get_inner()?.clone(); + future_into_py(py, async move { + inner + .drop_storage_integration(name, &namespace_path.unwrap_or_default()) + .await + .infer_error() + }) + } + #[pyo3(signature = (name, value, namespace_path=None))] pub fn create_secret( self_: PyRef<'_, Self>, diff --git a/rust/lancedb/src/authz/mod.rs b/rust/lancedb/src/authz/mod.rs index 3191627e0..f53cd3fb1 100644 --- a/rust/lancedb/src/authz/mod.rs +++ b/rust/lancedb/src/authz/mod.rs @@ -42,7 +42,7 @@ pub mod subject; pub use api::*; pub use object::{ DatabaseObject, FunctionObject, NamespaceObject, NamespacePath, Object, SecretObject, - TableObject, ViewObject, + StorageIntegrationObject, TableObject, ViewObject, }; pub use privilege::*; pub use subject::{Subject, SubjectKind}; diff --git a/rust/lancedb/src/authz/object.rs b/rust/lancedb/src/authz/object.rs index fd324b748..13279e911 100644 --- a/rust/lancedb/src/authz/object.rs +++ b/rust/lancedb/src/authz/object.rs @@ -16,6 +16,7 @@ pub const NAMESPACE_TYPE_STRING: &str = "namespace"; pub const TABLE_TYPE_STRING: &str = "table"; pub const VIEW_TYPE_STRING: &str = "view"; pub const SECRET_TYPE_STRING: &str = "secret"; +pub const STORAGE_INTEGRATION_TYPE_STRING: &str = "storage_integration"; pub const FUNCTION_TYPE_STRING: &str = "function"; pub const UNKNOWN_TYPE_STRING: &str = "unknown"; @@ -102,6 +103,7 @@ pub enum Object { Table(TableObject), View(ViewObject), Secret(SecretObject), + StorageIntegration(StorageIntegrationObject), Function(FunctionObject), Unknown(String), } @@ -116,6 +118,7 @@ impl Object { Self::Table(object) => object.validate(), Self::View(object) => object.validate(), Self::Secret(object) => object.validate(), + Self::StorageIntegration(object) => object.validate(), Self::Function(object) => object.validate(), Self::Unknown(value) if value.is_empty() => Err("object must not be empty".into()), Self::Unknown(_) => Ok(()), @@ -131,6 +134,7 @@ impl Object { Self::Table(object) => Some(&object.database), Self::View(object) => Some(&object.database), Self::Secret(object) => Some(&object.database), + Self::StorageIntegration(object) => Some(&object.database), Self::Function(object) => Some(&object.database), Self::Unknown(_) => None, } @@ -144,6 +148,7 @@ impl Object { Self::Table(_) => TABLE_TYPE_STRING, Self::View(_) => VIEW_TYPE_STRING, Self::Secret(_) => SECRET_TYPE_STRING, + Self::StorageIntegration(_) => STORAGE_INTEGRATION_TYPE_STRING, Self::Function(_) => FUNCTION_TYPE_STRING, Self::Unknown(_) => UNKNOWN_TYPE_STRING, } @@ -170,6 +175,7 @@ impl std::fmt::Display for Object { Self::Table(object) => object.fmt(f), Self::View(object) => object.fmt(f), Self::Secret(object) => object.fmt(f), + Self::StorageIntegration(object) => object.fmt(f), Self::Function(object) => object.fmt(f), Self::Unknown(s) => f.write_str(s), } @@ -215,6 +221,12 @@ impl<'de> Deserialize<'de> for Object { let view = ViewObject::deserialize(serde::de::value::StringDeserializer::::new(s))?; Ok(Self::View(view)) + } else if kind == STORAGE_INTEGRATION_TYPE_STRING { + let object = + StorageIntegrationObject::deserialize(serde::de::value::StringDeserializer::< + D::Error, + >::new(s))?; + Ok(Self::StorageIntegration(object)) } else if kind == SECRET_TYPE_STRING { let secret = SecretObject::deserialize( serde::de::value::StringDeserializer::::new(s), @@ -503,6 +515,61 @@ impl<'de> Deserialize<'de> for SecretObject { } } +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct StorageIntegrationObject { + pub database: String, + pub namespace: NamespacePath, + pub storage_integration: String, +} + +impl StorageIntegrationObject { + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + validate_secret_name(&self.storage_integration).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for StorageIntegrationObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "storage_integration:{}:{}${}", + self.database, self.namespace, self.storage_integration + ) + } +} + +impl Serialize for StorageIntegrationObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for StorageIntegrationObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let invalid = || D::Error::custom(format!("invalid StorageIntegrationObject {s:?}")); + let rest = s.strip_prefix("storage_integration:").ok_or_else(invalid)?; + let (database, path) = rest.split_once(':').ok_or_else(invalid)?; + let (namespace, storage_integration) = path.rsplit_once('$').ok_or_else(invalid)?; + let object = Self { + database: database.to_string(), + namespace: NamespacePath::parse(namespace).map_err(D::Error::custom)?, + storage_integration: storage_integration.to_string(), + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} + /// A function resource, encoded as `function:DATABASE:NAMESPACE$FUNCTION`. /// /// ``` diff --git a/rust/lancedb/src/authz/privilege.rs b/rust/lancedb/src/authz/privilege.rs index 66b07643a..2782a17bd 100644 --- a/rust/lancedb/src/authz/privilege.rs +++ b/rust/lancedb/src/authz/privilege.rs @@ -7,7 +7,7 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _}; /// A LanceDB privilege, serialized by name (for example, `SELECT`). /// Numeric aliases use `PRIVILEGE` followed by decimal digits. Numbers 0 through -/// 13 map to the known privileges in declaration order. Unassigned numbers are +/// 14 map to the known privileges in declaration order. Unassigned numbers are /// preserved verbatim; their meaning and validity are determined by the server. /// /// ``` @@ -81,6 +81,9 @@ pub enum Privilege { /// - can read server configuration. Operate, + /// On a namespace: can register a storage integration. + CreateStorageIntegration, + /// A future privilege name, preserved verbatim. #[strum(to_string = "{0}")] Unknown(UnknownPrivilege), @@ -97,7 +100,7 @@ impl fmt::Display for UnknownPrivilege { } impl Privilege { - pub const KNOWN: [Self; 14] = [ + pub const KNOWN: [Self; 15] = [ Self::Ownership, Self::Usage, Self::Select, @@ -112,6 +115,7 @@ impl Privilege { Self::CreateSecret, Self::CreateFunction, Self::Operate, + Self::CreateStorageIntegration, ]; pub fn valid_for_system(&self) -> bool { @@ -138,6 +142,7 @@ impl Privilege { | Self::CreateMaterializedView | Self::CreateView | Self::CreateSecret + | Self::CreateStorageIntegration | Self::CreateFunction | Self::Unknown(_) ) @@ -170,6 +175,11 @@ impl Privilege { matches!(self, Self::Ownership | Self::Usage | Self::Unknown(_)) } + /// Privileges accepted on a storage integration. + pub fn valid_for_storage_integration(&self) -> bool { + matches!(self, Self::Ownership | Self::Usage | Self::Unknown(_)) + } + /// Whether this privilege can mutate database contents or schema. /// Unknown privileges are conservatively treated as writes. pub fn is_write(&self) -> bool { @@ -203,6 +213,7 @@ impl FromStr for Privilege { Ok(11) => Self::CreateSecret, Ok(12) => Self::CreateFunction, Ok(13) => Self::Operate, + Ok(14) => Self::CreateStorageIntegration, _ => Self::Unknown(UnknownPrivilege(name.to_owned())), }); } @@ -242,6 +253,7 @@ impl TryFrom<&Privilege> for u8 { Privilege::CreateSecret => 11, Privilege::CreateFunction => 12, Privilege::Operate => 13, + Privilege::CreateStorageIntegration => 14, Privilege::Unknown(unknown) => unknown.0["PRIVILEGE".len()..].parse()?, }) } @@ -288,6 +300,16 @@ impl<'de> Deserialize<'de> for Privilege { mod tests { use super::Privilege; + #[test] + fn storage_integration_creation_is_namespace_scoped() { + let privilege = Privilege::CreateStorageIntegration; + assert!(!privilege.valid_for_system()); + assert!(privilege.valid_for_namespace()); + assert!(!privilege.valid_for_database()); + assert!(!privilege.valid_for_storage_integration()); + assert!(privilege.is_write()); + } + #[test] fn display_matches_serde() { for privilege in Privilege::KNOWN { @@ -337,6 +359,7 @@ mod tests { (11, Privilege::CreateSecret), (12, Privilege::CreateFunction), (13, Privilege::Operate), + (14, Privilege::CreateStorageIntegration), ] { assert_eq!(u8::try_from(&privilege).unwrap(), number); assert_eq!(u8::try_from(privilege.clone()).unwrap(), number); @@ -364,7 +387,7 @@ mod tests { #[test] fn unknown_privileges_convert_to_u8() { - for number in 14..=u8::MAX { + for number in 15..=u8::MAX { for name in [ format!("PRIVILEGE{number}"), format!("PRIVILEGE{number:03}"), @@ -399,7 +422,7 @@ mod tests { #[test] fn future_privilege_names_round_trip() { for name in [ - "PRIVILEGE14", + "PRIVILEGE15", "PRIVILEGE20", "PRIVILEGE64", "PRIVILEGE256", @@ -442,7 +465,7 @@ mod tests { #[test] fn unknown_privileges_are_valid_for_all_objects_and_treated_as_writes() { - for name in ["PRIVILEGE14", "PRIVILEGE20", "PRIVILEGE64", "PRIVILEGE256"] { + for name in ["PRIVILEGE15", "PRIVILEGE20", "PRIVILEGE64", "PRIVILEGE256"] { let privilege: Privilege = name.parse().unwrap(); assert!(privilege.valid_for_system()); assert!(privilege.valid_for_database()); @@ -475,6 +498,7 @@ mod tests { | Privilege::CreateView | Privilege::CreateSecret | Privilege::CreateFunction + | Privilege::CreateStorageIntegration ), "unexpected write classification for {privilege}" ); diff --git a/rust/lancedb/src/catalog.rs b/rust/lancedb/src/catalog.rs index 18373ed5a..cdd363fe0 100644 --- a/rust/lancedb/src/catalog.rs +++ b/rust/lancedb/src/catalog.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! Catalogs manage databases. A remote catalog is the root namespace of a server. +//! Catalogs manage databases at the server root. Storage integrations belong to databases. //! //! ``` //! # #[cfg(feature = "remote")] @@ -200,7 +200,7 @@ impl Stream for DatabaseNames { } } -/// A backend that manages databases. Implementations own database lifecycle semantics. +/// A backend that manages databases. #[async_trait::async_trait] pub trait Catalog: Send + Sync + std::fmt::Debug + 'static { /// Obtain deployment-wide authorization, if supported by this backend. @@ -265,7 +265,6 @@ impl CatalogConnection { pub fn catalog(&self) -> &Arc { &self.catalog } - /// Create a database; pass a name or [`CreateDatabaseRequest`] for additional options. pub async fn create_database( &self, diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 20bb0b360..1b1920414 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -747,6 +747,69 @@ impl Connection { .await } + /// Register storage access in this database. Names follow Secret namespace rules. + /// + /// ``` + /// # async fn example(db: &lancedb::Connection) -> lancedb::Result<()> { + /// use lancedb::storage_integration::CreateStorageIntegrationRequest; + /// let info = db.create_storage_integration("raw", CreateStorageIntegrationRequest::s3("s3://customer-bucket/raw/", "arn:aws:iam::123456789012:role/access"), &[]).await?; + /// println!("{}", info.configuration_instructions); + /// # Ok(()) + /// # } + /// ``` + pub async fn create_storage_integration( + &self, + name: impl AsRef, + request: crate::storage_integration::CreateStorageIntegrationRequest, + namespace_path: &[String], + ) -> Result { + validate_secret_reference(name.as_ref(), namespace_path)?; + self.internal + .create_storage_integration(name.as_ref(), request, namespace_path) + .await + } + /// Describe a storage integration, returning secret references without credential values. + pub async fn describe_storage_integration( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result { + validate_secret_reference(name.as_ref(), namespace_path)?; + self.internal + .describe_storage_integration(name.as_ref(), namespace_path) + .await + } + /// Lazily list storage integrations in a namespace of this database. + pub fn list_storage_integrations( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + database + .list_storage_integrations(&namespace_path, options) + .await + } + }) + } + /// Drop an integration with ownership permission; integrations in use cannot be dropped. + pub async fn drop_storage_integration( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result<()> { + validate_secret_reference(name.as_ref(), namespace_path)?; + self.internal + .drop_storage_integration(name.as_ref(), namespace_path) + .await + } + /// Create a named Secret in this database. /// /// Fails if the name is taken, so a create can never silently become a diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 54d51c24f..4528257d1 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -473,6 +473,40 @@ pub trait Database: let dropped = self.drop_function(name, version, namespace_path).await?; Ok((dropped, crate::job::Job::new_done())) } + /// Create a database-owned storage integration in a namespace. + async fn create_storage_integration( + &self, + _name: &str, + _request: crate::storage_integration::CreateStorageIntegrationRequest, + _namespace_path: &[String], + ) -> Result { + secret_catalog_not_supported() + } + /// Describe a database-owned storage integration. + async fn describe_storage_integration( + &self, + _name: &str, + _namespace_path: &[String], + ) -> Result { + secret_catalog_not_supported() + } + /// List one namespace of storage integrations. + async fn list_storage_integrations( + &self, + _namespace_path: &[String], + _options: ListingOptions, + ) -> Result> { + secret_catalog_not_supported() + } + /// Drop a database-owned storage integration. + async fn drop_storage_integration( + &self, + _name: &str, + _namespace_path: &[String], + ) -> Result<()> { + secret_catalog_not_supported() + } + /// Create a named Secret in this database. Fails if the name is taken, so /// a create can never silently become a rotation. async fn create_secret( diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index fa6850f14..b1abf58af 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -215,6 +215,7 @@ pub mod remote; pub mod rerankers; pub mod secrets; pub mod sql; +pub mod storage_integration; pub mod table; #[cfg(test)] pub mod test_utils; diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 5abe4dcf7..23b320837 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -14,6 +14,7 @@ pub(crate) mod job; pub mod oauth; mod retry; pub(crate) mod sql; +mod storage_integration; pub(crate) mod table; pub(crate) mod token_cache; pub(crate) mod util; diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 5c85d0792..445961205 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use crate::storage_integration::{CreateStorageIntegrationRequest, StorageIntegrationInfo}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -48,9 +49,11 @@ use super::client::{ }; use super::sql::SqlClient; use super::sql::materialized_view::quote_identifier; +use super::storage_integration as integration_sql; use super::table::RemoteTable; use super::util::parse_server_version; use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id}; +use futures::TryStreamExt; // Request structure for the remote clone table API #[derive(serde::Serialize)] @@ -220,7 +223,7 @@ impl RemoteDatabaseOptionsBuilder { #[derive(Debug)] pub struct RemoteDatabase { - client: RestfulLanceDbClient, + pub(super) client: RestfulLanceDbClient, // Cache existence and server capabilities, not mutable per-handle table state. table_cache: Cache, uri: String, @@ -1103,6 +1106,83 @@ impl Database for RemoteDatabase { Ok((dropped.dropped, job)) } + async fn create_storage_integration( + &self, + name: &str, + request: CreateStorageIntegrationRequest, + namespace_path: &[String], + ) -> Result { + let sql = integration_sql::create_sql(name, namespace_path, &request)?; + let result = self + .execute_query_async(ExecuteQueryRequest { + query: sql, + default_namespace_path: namespace_path.to_vec(), + parameters: None, + }) + .await?; + integration_sql::decode_info(&result.reader().await?.try_collect::>().await?) + } + + async fn describe_storage_integration( + &self, + name: &str, + namespace_path: &[String], + ) -> Result { + let name = integration_sql::object_name(name, namespace_path)?; + let result = self + .execute_query_async(ExecuteQueryRequest { + query: format!("DESCRIBE STORAGE INTEGRATION {name}"), + default_namespace_path: namespace_path.to_vec(), + parameters: None, + }) + .await?; + integration_sql::decode_info(&result.reader().await?.try_collect::>().await?) + } + + async fn list_storage_integrations( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Result> { + validate_namespace(namespace_path)?; + let result = self + .execute_query_async(ExecuteQueryRequest { + query: "SHOW STORAGE INTEGRATIONS".into(), + default_namespace_path: namespace_path.to_vec(), + parameters: None, + }) + .await?; + let mut names = + integration_sql::decode_names(&result.reader().await?.try_collect::>().await?)?; + if let Some(token) = options.page_token { + names.retain(|name| name > &token); + } + let mut page_token = None; + if let Some(limit) = options.page_limit + && names.len() > limit as usize + { + names.truncate(limit as usize); + page_token = names.last().cloned(); + } + Ok(ListingPage { + items: names, + page_token, + }) + } + + async fn drop_storage_integration(&self, name: &str, namespace_path: &[String]) -> Result<()> { + let name = integration_sql::object_name(name, namespace_path)?; + let result = self + .execute_query_async(ExecuteQueryRequest { + query: format!("DROP STORAGE INTEGRATION {name}"), + default_namespace_path: namespace_path.to_vec(), + parameters: None, + }) + .await?; + result.reader().await?.try_collect::>().await?; + Ok(()) + } + async fn create_secret( &self, name: &str, diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 6a0421fa6..3418235bd 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -28,6 +28,7 @@ use crate::remote::client::HeaderProvider; use crate::remote::client::test_utils::client_with_handler; use crate::remote::db::RemoteDatabase; use crate::remote::table::RemoteTable; +use crate::storage_integration::{CreateStorageIntegrationRequest, SecretReference}; use crate::table::BaseTable; #[derive(Debug, Default)] @@ -2156,3 +2157,116 @@ async fn remote_describe_materialized_view_preserves_the_sql_definition() { shutdown.send(()).unwrap(); } } + +#[tokio::test] +async fn storage_integration_methods_use_sql_with_namespace_and_typed_metadata() { + let info: serde_json::Value = serde_json::from_str(include_str!( + "../../tests/fixtures/storage_integration.json" + )) + .unwrap(); + let columns: Vec<(String, ArrayRef)> = info + .as_object() + .unwrap() + .iter() + .map(|(name, value)| { + let column: ArrayRef = match value { + serde_json::Value::Bool(value) => { + Arc::new(arrow_array::BooleanArray::from(vec![*value])) + } + serde_json::Value::Number(value) => { + Arc::new(Int64Array::from(vec![value.as_i64().unwrap()])) + } + serde_json::Value::Null => Arc::new(StringArray::from(vec![None::])), + serde_json::Value::String(value) => { + Arc::new(StringArray::from(vec![value.clone()])) + } + value => Arc::new(StringArray::from(vec![value.to_string()])), + }; + (name.clone(), column) + }) + .collect(); + for defer in [false, true] { + let service = TestSqlService { + result: RecordBatch::try_from_iter(columns.clone()).unwrap(), + adaptive: true, + defer, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let database = RemoteDatabase::new_mock(|request| -> http::Response { + panic!( + "Storage integration operation used REST: {}", + request.url().path() + ); + }) + .with_sql_client(client); + let namespace = vec!["platform".into(), "aws".into()]; + let configuration = CreateStorageIntegrationRequest::s3( + "s3://customer-bucket/raw/", + "arn:aws:iam::123456789012:role/access", + ) + .with_aws_region("us-east-1"); + let created = database + .create_storage_integration("raw", configuration, &namespace) + .await + .unwrap(); + assert_eq!(created.database, "default"); + assert_eq!( + created.trust_policy["Statement"][0]["Condition"]["StringEquals"]["sts:ExternalId"], + created.aws_external_id + ); + assert_eq!( + database + .describe_storage_integration("raw", &namespace) + .await + .unwrap() + .aws_external_id, + created.aws_external_id + ); + assert_eq!( + database + .list_storage_integrations(&namespace, Default::default()) + .await + .unwrap() + .items, + ["raw"] + ); + database + .drop_storage_integration("raw", &namespace) + .await + .unwrap(); + let credential = SecretReference::new("key") + .with_database("platform/raw") + .with_namespace(vec!["credentials".into()]); + database + .create_storage_integration( + "raw", + CreateStorageIntegrationRequest::s3_with_secrets( + "s3://customer-bucket/raw/", + credential.clone(), + credential, + ), + &namespace, + ) + .await + .unwrap(); + let statements = statements.lock().unwrap(); + assert!(statements[0].contains("STORAGE_AWS_REGION = 'us-east-1'")); + assert_eq!(statements[1], "DESCRIBE STORAGE INTEGRATION \"raw\""); + assert_eq!(statements[2], "SHOW STORAGE INTEGRATIONS"); + assert_eq!(statements[3], "DROP STORAGE INTEGRATION \"raw\""); + assert!(statements[4].contains( + "STORAGE_AWS_ACCESS_KEY_ID_SECRET = \"platform/raw\".\"credentials\".\"key\"" + )); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|header| header.namespace_path == "platform$aws") + ); + let _ = shutdown.send(()); + } +} diff --git a/rust/lancedb/src/remote/storage_integration.rs b/rust/lancedb/src/remote/storage_integration.rs new file mode 100644 index 000000000..48ac51f27 --- /dev/null +++ b/rust/lancedb/src/remote/storage_integration.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use arrow_array::{ + Array, BooleanArray, Int64Array, LargeStringArray, RecordBatch, StringArray, StringViewArray, +}; +use serde_json::Value; + +use crate::storage_integration::{ + CreateStorageIntegrationRequest, SecretReference, StorageIntegrationInfo, +}; +use crate::utils::{validate_database_name, validate_namespace, validate_secret_name}; +use crate::{Error, Result}; + +fn invalid(message: impl Into) -> Error { + Error::InvalidInput { + message: message.into(), + } +} + +pub(super) fn identifier(name: &str) -> String { + format!("\"{}\"", name.replace('"', "\"\"")) +} + +fn literal(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +pub(super) fn object_name(name: &str, namespace: &[String]) -> Result { + validate_secret_name(name)?; + validate_namespace(namespace)?; + Ok(identifier(name)) +} + +fn secret_name(reference: &SecretReference) -> Result { + validate_secret_name(&reference.name)?; + validate_namespace(&reference.namespace_path)?; + let mut parts = Vec::new(); + if let Some(database) = &reference.database { + validate_database_name(database)?; + parts.push(identifier(database)); + } + if reference.namespace_path.is_empty() { + parts.push(identifier("public")); + } else { + parts.extend(reference.namespace_path.iter().map(|name| identifier(name))); + } + parts.push(identifier(&reference.name)); + Ok(parts.join(".")) +} + +pub(super) fn create_sql( + name: &str, + namespace: &[String], + request: &CreateStorageIntegrationRequest, +) -> Result { + let mut sql = format!( + "CREATE STORAGE INTEGRATION {} STORAGE_PROVIDER = 'S3' STORAGE_PREFIX = {}", + object_name(name, namespace)?, + literal(&request.storage_prefix) + ); + match ( + &request.aws_role_arn, + &request.aws_access_key_id_secret, + &request.aws_secret_access_key_secret, + ) { + (Some(role), None, None) => { + sql.push_str(&format!(" STORAGE_AWS_ROLE_ARN = {}", literal(role))) + } + (None, Some(key), Some(secret)) => sql.push_str(&format!( + " STORAGE_AWS_ACCESS_KEY_ID_SECRET = {} STORAGE_AWS_SECRET_ACCESS_KEY_SECRET = {}", + secret_name(key)?, + secret_name(secret)? + )), + _ => { + return Err(invalid( + "Specify either an AWS role or both credential Secret references", + )); + } + } + if let Some(region) = &request.aws_region { + sql.push_str(&format!(" STORAGE_AWS_REGION = {}", literal(region))); + } + Ok(sql) +} + +fn string_value(column: &dyn Array, row: usize) -> Option<&str> { + if let Some(array) = column.as_any().downcast_ref::() { + Some(array.value(row)) + } else if let Some(array) = column.as_any().downcast_ref::() { + Some(array.value(row)) + } else if let Some(array) = column.as_any().downcast_ref::() { + Some(array.value(row)) + } else { + None + } +} + +pub(super) fn decode_info(batches: &[RecordBatch]) -> Result { + let mut rows = batches.iter().filter(|batch| batch.num_rows() > 0); + let batch = rows + .next() + .ok_or_else(|| invalid("SQL returned no storage integration"))?; + if batch.num_rows() != 1 || rows.next().is_some() { + return Err(invalid("SQL must return exactly one storage integration")); + } + let mut values = serde_json::Map::new(); + for (field, column) in batch.schema().fields().iter().zip(batch.columns()) { + let mut value = if column.is_null(0) { + Value::Null + } else if let Some(value) = string_value(column.as_ref(), 0) { + Value::String(value.into()) + } else if let Some(array) = column.as_any().downcast_ref::() { + Value::Bool(array.value(0)) + } else if let Some(array) = column.as_any().downcast_ref::() { + Value::from(array.value(0)) + } else { + return Err(invalid(format!("Unexpected SQL type for {}", field.name()))); + }; + if matches!( + field.name().as_str(), + "namespace_path" + | "aws_access_key_id_secret" + | "aws_secret_access_key_secret" + | "trust_policy" + | "permissions_policy" + ) && let Some(json) = value.as_str() + { + value = serde_json::from_str(json) + .map_err(|error| invalid(format!("Invalid {} metadata: {error}", field.name())))?; + } + values.insert(field.name().clone(), value); + } + serde_json::from_value(Value::Object(values)) + .map_err(|error| invalid(format!("Invalid storage integration SQL result: {error}"))) +} + +pub(super) fn decode_names(batches: &[RecordBatch]) -> Result> { + let mut names = Vec::new(); + for batch in batches { + let column = batch + .column_by_name("name") + .ok_or_else(|| invalid("SQL storage integration listing is missing names"))?; + for row in 0..batch.num_rows() { + if column.is_null(row) { + return Err(invalid("SQL storage integration name cannot be null")); + } + names.push( + string_value(column.as_ref(), row) + .ok_or_else(|| { + invalid("SQL storage integration listing requires string names") + })? + .to_string(), + ); + } + } + names.sort(); + Ok(names) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + + #[test] + fn sql_quotes_values_and_keeps_secret_root_namespace_explicit() { + let role = CreateStorageIntegrationRequest::s3("s3://bucket/it's/", "role'value"); + let sql = create_sql("raw", &[], &role).unwrap(); + assert!(sql.contains("STORAGE_PREFIX = 's3://bucket/it''s/'")); + assert!(sql.contains("STORAGE_AWS_ROLE_ARN = 'role''value'")); + let reference = SecretReference::new("key").with_database("platform/raw"); + let static_config = CreateStorageIntegrationRequest::s3_with_secrets( + "s3://bucket/raw/", + reference.clone(), + reference, + ); + let sql = create_sql("raw", &["nested".into()], &static_config).unwrap(); + assert!( + sql.contains("STORAGE_AWS_ACCESS_KEY_ID_SECRET = \"platform/raw\".\"public\".\"key\"") + ); + assert!(create_sql("a/b", &[], &role).is_err()); + } + + #[test] + fn sql_metadata_rejects_missing_rows_and_invalid_json() { + assert!(decode_info(&[]).is_err()); + let batch = RecordBatch::try_from_iter([( + "namespace_path", + Arc::new(StringArray::from(vec!["invalid json"])) as arrow_array::ArrayRef, + )]) + .unwrap(); + assert!( + decode_info(&[batch]) + .unwrap_err() + .to_string() + .contains("namespace_path") + ); + } +} diff --git a/rust/lancedb/src/storage_integration.rs b/rust/lancedb/src/storage_integration.rs new file mode 100644 index 000000000..955cefe49 --- /dev/null +++ b/rust/lancedb/src/storage_integration.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Database-owned cloud-storage access configurations managed by the remote service. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// A Secret object, optionally qualified by its owning database and namespace. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SecretReference { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub database: Option, + #[serde(default)] + pub namespace_path: Vec, + pub name: String, +} +impl SecretReference { + /// Reference a Secret in the integration's database. + pub fn new(name: impl Into) -> Self { + Self { + database: None, + namespace_path: vec![], + name: name.into(), + } + } + /// Qualify the database that owns the Secret. + pub fn with_database(mut self, database: impl Into) -> Self { + self.database = Some(database.into()); + self + } + /// Qualify the namespace that owns the Secret. + pub fn with_namespace(mut self, namespace: Vec) -> Self { + self.namespace_path = namespace; + self + } +} + +/// Configuration for registering an AWS S3 storage integration. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct CreateStorageIntegrationRequest { + pub storage_provider: StorageProvider, + pub storage_prefix: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aws_role_arn: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aws_access_key_id_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aws_secret_access_key_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aws_region: Option, +} + +/// Cloud storage providers supported by storage integrations. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum StorageProvider { + #[serde(rename = "S3")] + S3, +} + +impl CreateStorageIntegrationRequest { + /// Register a prefix and the customer-owned IAM role authorized to access it. + pub fn s3(storage_prefix: impl Into, aws_role_arn: impl Into) -> Self { + Self { + storage_provider: StorageProvider::S3, + storage_prefix: storage_prefix.into(), + aws_role_arn: Some(aws_role_arn.into()), + aws_access_key_id_secret: None, + aws_secret_access_key_secret: None, + aws_region: None, + } + } + + /// Register static S3 credentials by Secret reference, without accepting key values. + pub fn s3_with_secrets( + storage_prefix: impl Into, + access_key: SecretReference, + secret_key: SecretReference, + ) -> Self { + Self { + storage_provider: StorageProvider::S3, + storage_prefix: storage_prefix.into(), + aws_role_arn: None, + aws_access_key_id_secret: Some(access_key), + aws_secret_access_key_secret: Some(secret_key), + aws_region: None, + } + } + + /// 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()); + self + } +} + +/// Persisted configuration and complete AWS 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 +/// AWS access credentials. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct StorageIntegrationInfo { + pub name: String, + pub database: String, + pub namespace_path: Vec, + pub storage_provider: StorageProvider, + pub storage_prefix: String, + #[serde(default)] + pub aws_role_arn: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aws_access_key_id_secret: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub aws_secret_access_key_secret: Option, + pub aws_region: Option, + pub aws_account_id: String, + pub aws_iam_principal_arn: String, + pub aws_external_id: String, + pub created_at_millis: i64, + pub access_verified: bool, + pub trust_policy: Value, + pub permissions_policy: Value, + pub configuration_instructions: String, +} diff --git a/rust/lancedb/tests/fixtures/storage_integration.json b/rust/lancedb/tests/fixtures/storage_integration.json new file mode 100644 index 000000000..4b131ebe1 --- /dev/null +++ b/rust/lancedb/tests/fixtures/storage_integration.json @@ -0,0 +1,36 @@ +{ + "name": "raw", + "storage_provider": "S3", + "storage_prefix": "s3://customer-bucket/raw/", + "aws_role_arn": "arn:aws:iam::123456789012:role/access", + "aws_region": null, + "aws_account_id": "001122334455", + "aws_iam_principal_arn": "arn:aws:iam::001122334455:role/lancedb-storage-access", + "aws_external_id": "55b54d68-4686-4d38-97da-9d77e54a643e", + "created_at_millis": 1, + "access_verified": false, + "trust_policy": { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "sts:AssumeRole", + "sts:TagSession" + ], + "Principal": { + "AWS": "arn:aws:iam::001122334455:role/lancedb-storage-access" + }, + "Condition": { + "StringEquals": { + "sts:ExternalId": "55b54d68-4686-4d38-97da-9d77e54a643e" + } + } + } + ] + }, + "permissions_policy": {}, + "configuration_instructions": "Update IAM role arn:aws:iam::123456789012:role/access to trust arn:aws:iam::001122334455:role/lancedb-storage-access with external ID 55b54d68-4686-4d38-97da-9d77e54a643e and attach the S3 permissions policy for s3://customer-bucket/raw/.", + "database": "default", + "namespace_path": [] +}