lancedb/python/python/tests/test_storage_integration.py

# 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)