lancedb/docs/web-tests/py/test_connection.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors

from pathlib import Path
import shutil

import pytest


def test_connection():
    # --8<-- [start:connect]
    import lancedb

    uri = "ex_lancedb"
    db = lancedb.connect(uri)
    # --8<-- [end:connect]
    assert db is not None
    shutil.rmtree(uri, ignore_errors=True)
    assert not Path(uri).exists()


async def connect_async_example():
    # --8<-- [start:connect_async]
    import lancedb

    uri = "ex_lancedb"
    async_db = await lancedb.connect_async(uri)
    # --8<-- [end:connect_async]

    return async_db


def connect_enterprise_quickstart_config():
    # --8<-- [start:connect_enterprise_quickstart]
    import os

    import lancedb

    db = lancedb.connect(
        uri=os.environ["LANCEDB_URI"],  # db://your-database
        api_key=os.environ["LANCEDB_API_KEY"],
        region=os.environ.get("LANCEDB_REGION", "us-east-1"),
        host_override=os.environ["LANCEDB_HOST_OVERRIDE"],  # https://your-endpoint
        sql_host_override=os.environ["LANCEDB_SQL_HOST_OVERRIDE"],
    )
    # --8<-- [end:connect_enterprise_quickstart]
    return db


def test_connect_enterprise_quickstart(monkeypatch):
    import lancedb

    captured = {}

    def fake_connect(**kwargs):
        captured.update(kwargs)
        return object()

    monkeypatch.setattr(lancedb, "connect", fake_connect)
    monkeypatch.setenv("LANCEDB_URI", "db://your-database")
    monkeypatch.setenv("LANCEDB_API_KEY", "your-api-key")
    monkeypatch.delenv("LANCEDB_REGION", raising=False)
    monkeypatch.setenv("LANCEDB_HOST_OVERRIDE", "https://your-endpoint")
    monkeypatch.setenv(
        "LANCEDB_SQL_HOST_OVERRIDE", "grpc+tls://your-sql-endpoint:10026"
    )

    assert connect_enterprise_quickstart_config() is not None
    assert captured == {
        "uri": "db://your-database",
        "api_key": "your-api-key",
        "region": "us-east-1",
        "host_override": "https://your-endpoint",
        "sql_host_override": "grpc+tls://your-sql-endpoint:10026",
    }


def test_connect_enterprise_quickstart_requires_an_endpoint(monkeypatch):
    monkeypatch.setenv("LANCEDB_URI", "db://your-database")
    monkeypatch.setenv("LANCEDB_API_KEY", "your-api-key")
    monkeypatch.delenv("LANCEDB_HOST_OVERRIDE", raising=False)

    with pytest.raises(KeyError, match="LANCEDB_HOST_OVERRIDE"):
        connect_enterprise_quickstart_config()


def test_connect_enterprise_quickstart_requires_a_sql_endpoint(monkeypatch):
    monkeypatch.setenv("LANCEDB_URI", "db://your-database")
    monkeypatch.setenv("LANCEDB_API_KEY", "your-api-key")
    monkeypatch.setenv("LANCEDB_HOST_OVERRIDE", "https://your-endpoint")
    monkeypatch.delenv("LANCEDB_SQL_HOST_OVERRIDE", raising=False)

    with pytest.raises(KeyError, match="LANCEDB_SQL_HOST_OVERRIDE"):
        connect_enterprise_quickstart_config()


def connect_object_storage_config():
    # --8<-- [start:connect_object_storage]
    import lancedb

    uri = "s3://your-bucket/path"
    # You can also use "gs://your-bucket/path" or "az://your-container/path".
    db = lancedb.connect(uri)
    # --8<-- [end:connect_object_storage]

    return db


def namespace_table_ops_example():
    # --8<-- [start:namespace_table_ops]
    import lancedb

    db = lancedb.connect_namespace("dir", {"root": "./local_lancedb"})

    # Create namespace tree: prod/search
    db.create_namespace(["prod"], mode="exist_ok")
    db.create_namespace(["prod", "search"], mode="exist_ok")
    db.create_namespace(["prod", "recommendations"], mode="exist_ok")

    db.create_table(
        "user",
        data=[{"id": 1, "vector": [0.1, 0.2], "name": "alice"}],
        namespace_path=["prod", "search"],
        mode="create",  # use "overwrite" only if you want to replace existing table
    )

    db.create_table(
        "user",
        data=[{"id": 2, "vector": [0.3, 0.4], "name": "bob"}],
        namespace_path=["prod", "recommendations"],
        mode="create",  # use "overwrite" only if you want to replace existing table
    )

    # Verify: listings return a response with the names and a page token
    print(db.list_namespaces().namespaces)  # ['prod']
    print(db.list_namespaces(namespace_path=["prod"]).namespaces)
    # ['recommendations', 'search']
    print(db.list_tables(namespace_path=["prod", "search"]).tables)  # ['user']
    print(db.list_tables(namespace_path=["prod", "recommendations"]).tables)  # ['user']
    # --8<-- [end:namespace_table_ops]


def test_namespace_table_ops(tmp_path, monkeypatch, capsys):
    monkeypatch.chdir(tmp_path)
    namespace_table_ops_example()
    assert capsys.readouterr().out.splitlines() == [
        "['prod']",
        "['recommendations', 'search']",
        "['user']",
        "['user']",
    ]


def namespace_admin_ops_example():
    # --8<-- [start:namespace_admin_ops]
    import lancedb

    db = lancedb.connect_namespace("dir", {"root": "./local_lancedb"})
    namespace = ["prod", "search"]

    db.create_namespace(["prod"])
    db.create_namespace(["prod", "search"])

    child_namespaces = db.list_namespaces(namespace_path=["prod"]).namespaces
    print(f"Child namespaces under ['prod']: {child_namespaces}")
    # Child namespaces under ['prod']: ['search']

    metadata = db.describe_namespace(["prod", "search"])
    print(f"Metadata for namespace {namespace}: {metadata}")
    # Metadata for namespace ['prod', 'search']: properties=None

    db.drop_namespace(["prod", "search"], mode="skip")
    db.drop_namespace(["prod"], mode="skip")
    # --8<-- [end:namespace_admin_ops]
    return db, child_namespaces, metadata


def test_namespace_admin_ops(tmp_path, monkeypatch, capsys):
    monkeypatch.chdir(tmp_path)
    db, child_namespaces, metadata = namespace_admin_ops_example()
    assert child_namespaces == ["search"]
    assert capsys.readouterr().out.splitlines()[0] == (
        "Child namespaces under ['prod']: ['search']"
    )
    assert db.list_namespaces().namespaces == []


async def connect_object_storage_config_async():
    # --8<-- [start:connect_object_storage_async]
    import lancedb

    uri = "s3://your-bucket/path"
    # You can also use "gs://your-bucket/path" or "az://your-container/path".
    async_db = await lancedb.connect_async(uri)
    # --8<-- [end:connect_object_storage_async]

    return async_db


def test_connect_enterprise_placeholders(monkeypatch):
    import lancedb

    captured = {}

    def fake_connect(**kwargs):
        captured.update(kwargs)
        return object()

    monkeypatch.setattr(lancedb, "connect", fake_connect)
    # --8<-- [start:connect_enterprise_placeholders]
    import lancedb

    db = lancedb.connect(
        uri="db://your-database",
        api_key="your-api-key",
        region="us-east-1",
        host_override="https://your-endpoint",
        sql_host_override="grpc+tls://your-sql-endpoint:10026",
    )
    # --8<-- [end:connect_enterprise_placeholders]
    assert db is not None
    assert captured == {
        "uri": "db://your-database",
        "api_key": "your-api-key",
        "region": "us-east-1",
        "host_override": "https://your-endpoint",
        "sql_host_override": "grpc+tls://your-sql-endpoint:10026",
    }