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: add database storage integration APIs (#4618)

Browse Source 
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.
This commit is contained in:
Jack Ye
2026-10-30 00:26:42 -07:00
committed by GitHub
parent a39a31807f
commit 6d126a3601
24 changed files with 1453 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/python/python.md
+81
View File
@@ -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.
python/python/lancedb/__init__.py
+1
View File
@@ -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,
python/python/lancedb/_lancedb.pyi
+20
View File
@@ -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: ...
python/python/lancedb/authz.py
+12
View File
@@ -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)
python/python/lancedb/catalog.py
+1 -1
View File
@@ -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
python/python/lancedb/db.py
+187
View File
@@ -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:
python/python/lancedb/remote/db.py
+57
View File
@@ -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
python/python/lancedb/storage_integration.py
+52
View File
@@ -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))
python/python/tests/test_storage_integration.py
+207
View File
@@ -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)
python/src/authz.rs
+13
View File
@@ -262,6 +262,19 @@ impl AuthzObject {
}))
}
#[staticmethod]
fn storage_integration(
database: &str,
namespace: &Bound<'_, PyAny>,
name: &str,
) -> PyResult<Self> {
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> {
Self::validated(Object::Function(FunctionObject {
python/src/connection.rs
+66
View File
@@ -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<Vec<String>>,
) -> PyResult<Bound<'py, PyAny>> {
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<Vec<String>>,
) -> PyResult<Bound<'py, PyAny>> {
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<Vec<String>>,
page_token: Option<String>,
page_limit: Option<u32>,
) -> PyResult<crate::listing::NameListing> {
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<Vec<String>>,
) -> PyResult<Bound<'py, PyAny>> {
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>,
rust/lancedb/src/authz/mod.rs
+1 -1
View File
@@ -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};
rust/lancedb/src/authz/object.rs
+67
View File
@@ -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::<D::Error>::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::<D::Error>::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<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
self.validate().map_err(S::Error::custom)?;
self.to_string().serialize(serializer)
}
}
impl<'de> Deserialize<'de> for StorageIntegrationObject {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
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`.
///
/// ```
rust/lancedb/src/authz/privilege.rs
+29 -5
View File
@@ -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}"
);
rust/lancedb/src/catalog.rs
+2 -3
View File
@@ -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<dyn Catalog> {
&self.catalog
}
/// Create a database; pass a name or [`CreateDatabaseRequest`] for additional options.
pub async fn create_database(
&self,
rust/lancedb/src/connection.rs
+63
View File
@@ -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<str>,
request: crate::storage_integration::CreateStorageIntegrationRequest,
namespace_path: &[String],
) -> Result<crate::storage_integration::StorageIntegrationInfo> {
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<str>,
namespace_path: &[String],
) -> Result<crate::storage_integration::StorageIntegrationInfo> {
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<String> {
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<str>,
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
rust/lancedb/src/database.rs
+34
View File
@@ -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<crate::storage_integration::StorageIntegrationInfo> {
secret_catalog_not_supported()
}
/// Describe a database-owned storage integration.
async fn describe_storage_integration(
&self,
_name: &str,
_namespace_path: &[String],
) -> Result<crate::storage_integration::StorageIntegrationInfo> {
secret_catalog_not_supported()
}
/// List one namespace of storage integrations.
async fn list_storage_integrations(
&self,
_namespace_path: &[String],
_options: ListingOptions,
) -> Result<ListingPage<String>> {
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(
rust/lancedb/src/lib.rs
+1
View File
@@ -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;
rust/lancedb/src/remote.rs
+1
View File
@@ -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;
rust/lancedb/src/remote/db.rs
+81 -1
View File
@@ -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<S: HttpSend = Sender> {
client: RestfulLanceDbClient<S>,
pub(super) client: RestfulLanceDbClient<S>,
// Cache existence and server capabilities, not mutable per-handle table state.
table_cache: Cache<String, ServerVersion>,
uri: String,
@@ -1103,6 +1106,83 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
Ok((dropped.dropped, job))
}
async fn create_storage_integration(
&self,
name: &str,
request: CreateStorageIntegrationRequest,
namespace_path: &[String],
) -> Result<StorageIntegrationInfo> {
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::<Vec<_>>().await?)
}
async fn describe_storage_integration(
&self,
name: &str,
namespace_path: &[String],
) -> Result<StorageIntegrationInfo> {
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::<Vec<_>>().await?)
}
async fn list_storage_integrations(
&self,
namespace_path: &[String],
options: ListingOptions,
) -> Result<ListingPage<String>> {
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::<Vec<_>>().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::<Vec<_>>().await?;
Ok(())
}
async fn create_secret(
&self,
name: &str,
rust/lancedb/src/remote/sql_test.rs
+114
View File
@@ -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::<String>])),
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<String> {
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(());
}
}
rust/lancedb/src/remote/storage_integration.rs
+201
View File
@@ -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<String>) -> 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<String> {
validate_secret_name(name)?;
validate_namespace(namespace)?;
Ok(identifier(name))
}
fn secret_name(reference: &SecretReference) -> Result<String> {
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<String> {
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::<StringArray>() {
Some(array.value(row))
} else if let Some(array) = column.as_any().downcast_ref::<LargeStringArray>() {
Some(array.value(row))
} else if let Some(array) = column.as_any().downcast_ref::<StringViewArray>() {
Some(array.value(row))
} else {
None
}
}
pub(super) fn decode_info(batches: &[RecordBatch]) -> Result<StorageIntegrationInfo> {
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::<BooleanArray>() {
Value::Bool(array.value(0))
} else if let Some(array) = column.as_any().downcast_ref::<Int64Array>() {
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<Vec<String>> {
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")
);
}
}
rust/lancedb/src/storage_integration.rs
+126
View File
@@ -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<String>,
#[serde(default)]
pub namespace_path: Vec<String>,
pub name: String,
}
impl SecretReference {
/// Reference a Secret in the integration's database.
pub fn new(name: impl Into<String>) -> 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<String>) -> Self {
self.database = Some(database.into());
self
}
/// Qualify the namespace that owns the Secret.
pub fn with_namespace(mut self, namespace: Vec<String>) -> 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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub aws_access_key_id_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub aws_secret_access_key_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub aws_region: Option<String>,
}
/// 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<String>, aws_role_arn: impl Into<String>) -> 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<String>,
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<String>) -> 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<String>,
pub storage_provider: StorageProvider,
pub storage_prefix: String,
#[serde(default)]
pub aws_role_arn: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub aws_access_key_id_secret: Option<SecretReference>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub aws_secret_access_key_secret: Option<SecretReference>,
pub aws_region: Option<String>,
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,
}
rust/lancedb/tests/fixtures/storage_integration.json
Vendored
+36
View File
@@ -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": []
}