mirror of
https://github.com/lancedb/lancedb.git
synced 2026-06-03 20:30:42 +00:00
refactor!: consolidate namespace related naming and enterprise integration (#3205)
1. Refactored every client (Rust core, Python, Node/TypeScript) so “namespace” usage is explicit: code now keeps namespace paths (namespace_path) separate from namespace clients (namespace_client). Connections propagate the client, table creation routes through it, and managed versioning defaults are resolved from namespace metadata. Python gained LanceNamespaceDBConnection/async counterparts, and the namespace-focused tests were rewritten to match the clarified API surface. 2. Synchronized the workspace with Lance 5.0.0-beta.3 (see https://github.com/lance-format/lance/pull/6186 for the upstream namespace refactor), updating Cargo/uv lockfiles and ensuring all bindings align with the new namespace semantics. 3. Added a namespace-backed code path to lancedb.connect() via new keyword arguments (namespace_client_impl, namespace_client_properties, plus the existing pushdown-ops flag). When those kwargs are supplied, connect() delegates to connect_namespace, so users can opt into namespace clients without changing APIs. (The async helper will gain parity in a later change)
This commit is contained in:
Jack Ye
GitHub
@@ -183,8 +183,8 @@ def test_table_names(tmp_db: lancedb.DBConnection):
|
||||
result = list(tmp_db.table_names("test2", limit=2))
|
||||
assert result == ["test3"], f"Expected ['test3'], got {result}"
|
||||
|
||||
# Test that namespace parameter can be passed as keyword
|
||||
result = list(tmp_db.table_names(namespace=[]))
|
||||
# Test that namespace_path parameter can be passed as keyword
|
||||
result = list(tmp_db.table_names(namespace_path=[]))
|
||||
assert len(result) == 3
|
||||
|
||||
|
||||
@@ -909,7 +909,7 @@ def test_local_namespace_operations(tmp_path):
|
||||
NotImplementedError,
|
||||
match="Namespace operations are not supported for listing database",
|
||||
):
|
||||
db.list_namespaces(namespace=["test"])
|
||||
db.list_namespaces(namespace_path=["test"])
|
||||
|
||||
|
||||
def test_local_create_namespace_not_supported(tmp_path):
|
||||
|
||||
@@ -33,6 +33,16 @@ class TestNamespaceConnection:
|
||||
# Initially no tables in root
|
||||
assert len(list(db.table_names())) == 0
|
||||
|
||||
def test_connect_via_connect_helper(self):
|
||||
"""Connecting via lancedb.connect should delegate to namespace connection."""
|
||||
db = lancedb.connect(
|
||||
namespace_client_impl="dir",
|
||||
namespace_client_properties={"root": self.temp_dir},
|
||||
)
|
||||
|
||||
assert isinstance(db, lancedb.LanceNamespaceDBConnection)
|
||||
assert len(list(db.table_names())) == 0
|
||||
|
||||
def test_create_table_through_namespace(self):
|
||||
"""Test creating a table through namespace."""
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
@@ -50,14 +60,14 @@ class TestNamespaceConnection:
|
||||
)
|
||||
|
||||
# Create empty table in child namespace
|
||||
table = db.create_table("test_table", schema=schema, namespace=["test_ns"])
|
||||
table = db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
assert table is not None
|
||||
assert table.name == "test_table"
|
||||
assert table.namespace == ["test_ns"]
|
||||
assert table.id == "test_ns$test_table"
|
||||
|
||||
# Table should appear in child namespace
|
||||
table_names = list(db.table_names(namespace=["test_ns"]))
|
||||
table_names = list(db.table_names(namespace_path=["test_ns"]))
|
||||
assert "test_table" in table_names
|
||||
assert len(table_names) == 1
|
||||
|
||||
@@ -80,10 +90,10 @@ class TestNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
db.create_table("test_table", schema=schema, namespace=["test_ns"])
|
||||
db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Open the table
|
||||
table = db.open_table("test_table", namespace=["test_ns"])
|
||||
table = db.open_table("test_table", namespace_path=["test_ns"])
|
||||
assert table is not None
|
||||
assert table.name == "test_table"
|
||||
assert table.namespace == ["test_ns"]
|
||||
@@ -108,31 +118,31 @@ class TestNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
db.create_table("table1", schema=schema, namespace=["test_ns"])
|
||||
db.create_table("table2", schema=schema, namespace=["test_ns"])
|
||||
db.create_table("table1", schema=schema, namespace_path=["test_ns"])
|
||||
db.create_table("table2", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Verify both tables exist in child namespace
|
||||
table_names = list(db.table_names(namespace=["test_ns"]))
|
||||
table_names = list(db.table_names(namespace_path=["test_ns"]))
|
||||
assert "table1" in table_names
|
||||
assert "table2" in table_names
|
||||
assert len(table_names) == 2
|
||||
|
||||
# Drop one table
|
||||
db.drop_table("table1", namespace=["test_ns"])
|
||||
db.drop_table("table1", namespace_path=["test_ns"])
|
||||
|
||||
# Verify only table2 remains
|
||||
table_names = list(db.table_names(namespace=["test_ns"]))
|
||||
table_names = list(db.table_names(namespace_path=["test_ns"]))
|
||||
assert "table1" not in table_names
|
||||
assert "table2" in table_names
|
||||
assert len(table_names) == 1
|
||||
|
||||
# Drop the second table
|
||||
db.drop_table("table2", namespace=["test_ns"])
|
||||
assert len(list(db.table_names(namespace=["test_ns"]))) == 0
|
||||
db.drop_table("table2", namespace_path=["test_ns"])
|
||||
assert len(list(db.table_names(namespace_path=["test_ns"]))) == 0
|
||||
|
||||
# Should not be able to open dropped table
|
||||
with pytest.raises(TableNotFoundError):
|
||||
db.open_table("table1", namespace=["test_ns"])
|
||||
db.open_table("table1", namespace_path=["test_ns"])
|
||||
|
||||
def test_create_table_with_schema(self):
|
||||
"""Test creating a table with explicit schema through namespace."""
|
||||
@@ -151,7 +161,7 @@ class TestNamespaceConnection:
|
||||
)
|
||||
|
||||
# Create table with schema in child namespace
|
||||
table = db.create_table("test_table", schema=schema, namespace=["test_ns"])
|
||||
table = db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
assert table is not None
|
||||
assert table.namespace == ["test_ns"]
|
||||
|
||||
@@ -175,7 +185,7 @@ class TestNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
db.create_table("old_name", schema=schema, namespace=["test_ns"])
|
||||
db.create_table("old_name", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Rename should raise NotImplementedError
|
||||
with pytest.raises(NotImplementedError, match="rename_table is not supported"):
|
||||
@@ -196,20 +206,20 @@ class TestNamespaceConnection:
|
||||
]
|
||||
)
|
||||
for i in range(3):
|
||||
db.create_table(f"table{i}", schema=schema, namespace=["test_ns"])
|
||||
db.create_table(f"table{i}", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Verify tables exist in child namespace
|
||||
assert len(list(db.table_names(namespace=["test_ns"]))) == 3
|
||||
assert len(list(db.table_names(namespace_path=["test_ns"]))) == 3
|
||||
|
||||
# Drop all tables in child namespace
|
||||
db.drop_all_tables(namespace=["test_ns"])
|
||||
db.drop_all_tables(namespace_path=["test_ns"])
|
||||
|
||||
# Verify all tables are gone from child namespace
|
||||
assert len(list(db.table_names(namespace=["test_ns"]))) == 0
|
||||
assert len(list(db.table_names(namespace_path=["test_ns"]))) == 0
|
||||
|
||||
# Test that table_names works with keyword-only namespace parameter
|
||||
db.create_table("test_table", schema=schema, namespace=["test_ns"])
|
||||
result = list(db.table_names(namespace=["test_ns"]))
|
||||
db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
result = list(db.table_names(namespace_path=["test_ns"]))
|
||||
assert "test_table" in result
|
||||
|
||||
def test_table_operations(self):
|
||||
@@ -227,7 +237,7 @@ class TestNamespaceConnection:
|
||||
pa.field("text", pa.string()),
|
||||
]
|
||||
)
|
||||
table = db.create_table("test_table", schema=schema, namespace=["test_ns"])
|
||||
table = db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Verify empty table was created
|
||||
result = table.to_pandas()
|
||||
@@ -298,25 +308,25 @@ class TestNamespaceConnection:
|
||||
]
|
||||
)
|
||||
table = db.create_table(
|
||||
"test_table", schema=schema, namespace=["test_namespace"]
|
||||
"test_table", schema=schema, namespace_path=["test_namespace"]
|
||||
)
|
||||
assert table is not None
|
||||
|
||||
# Verify table exists in namespace
|
||||
tables_in_namespace = list(db.table_names(namespace=["test_namespace"]))
|
||||
tables_in_namespace = list(db.table_names(namespace_path=["test_namespace"]))
|
||||
assert "test_table" in tables_in_namespace
|
||||
assert len(tables_in_namespace) == 1
|
||||
|
||||
# Open table from namespace
|
||||
table = db.open_table("test_table", namespace=["test_namespace"])
|
||||
table = db.open_table("test_table", namespace_path=["test_namespace"])
|
||||
assert table is not None
|
||||
assert table.name == "test_table"
|
||||
|
||||
# Drop table from namespace
|
||||
db.drop_table("test_table", namespace=["test_namespace"])
|
||||
db.drop_table("test_table", namespace_path=["test_namespace"])
|
||||
|
||||
# Verify table no longer exists in namespace
|
||||
tables_in_namespace = list(db.table_names(namespace=["test_namespace"]))
|
||||
tables_in_namespace = list(db.table_names(namespace_path=["test_namespace"]))
|
||||
assert len(tables_in_namespace) == 0
|
||||
|
||||
# Drop namespace
|
||||
@@ -338,14 +348,14 @@ class TestNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
db.create_table("test_table", schema=schema, namespace=["test_namespace"])
|
||||
db.create_table("test_table", schema=schema, namespace_path=["test_namespace"])
|
||||
|
||||
# Try to drop namespace with tables - should fail
|
||||
with pytest.raises(NamespaceNotEmptyError):
|
||||
db.drop_namespace(["test_namespace"])
|
||||
|
||||
# Drop table first
|
||||
db.drop_table("test_table", namespace=["test_namespace"])
|
||||
db.drop_table("test_table", namespace_path=["test_namespace"])
|
||||
|
||||
# Now dropping namespace should work
|
||||
db.drop_namespace(["test_namespace"])
|
||||
@@ -368,10 +378,10 @@ class TestNamespaceConnection:
|
||||
|
||||
# Create table with same name in both namespaces
|
||||
table_a = db.create_table(
|
||||
"same_name_table", schema=schema, namespace=["namespace_a"]
|
||||
"same_name_table", schema=schema, namespace_path=["namespace_a"]
|
||||
)
|
||||
table_b = db.create_table(
|
||||
"same_name_table", schema=schema, namespace=["namespace_b"]
|
||||
"same_name_table", schema=schema, namespace_path=["namespace_b"]
|
||||
)
|
||||
|
||||
# Add different data to each table
|
||||
@@ -389,7 +399,9 @@ class TestNamespaceConnection:
|
||||
table_b.add(data_b)
|
||||
|
||||
# Verify data in namespace_a table
|
||||
opened_table_a = db.open_table("same_name_table", namespace=["namespace_a"])
|
||||
opened_table_a = db.open_table(
|
||||
"same_name_table", namespace_path=["namespace_a"]
|
||||
)
|
||||
result_a = opened_table_a.to_pandas().sort_values("id").reset_index(drop=True)
|
||||
assert len(result_a) == 2
|
||||
assert result_a["id"].tolist() == [1, 2]
|
||||
@@ -400,7 +412,9 @@ class TestNamespaceConnection:
|
||||
assert [v.tolist() for v in result_a["vector"]] == [[1.0, 2.0], [3.0, 4.0]]
|
||||
|
||||
# Verify data in namespace_b table
|
||||
opened_table_b = db.open_table("same_name_table", namespace=["namespace_b"])
|
||||
opened_table_b = db.open_table(
|
||||
"same_name_table", namespace_path=["namespace_b"]
|
||||
)
|
||||
result_b = opened_table_b.to_pandas().sort_values("id").reset_index(drop=True)
|
||||
assert len(result_b) == 3
|
||||
assert result_b["id"].tolist() == [10, 20, 30]
|
||||
@@ -420,8 +434,8 @@ class TestNamespaceConnection:
|
||||
assert "same_name_table" not in root_tables
|
||||
|
||||
# Clean up
|
||||
db.drop_table("same_name_table", namespace=["namespace_a"])
|
||||
db.drop_table("same_name_table", namespace=["namespace_b"])
|
||||
db.drop_table("same_name_table", namespace_path=["namespace_a"])
|
||||
db.drop_table("same_name_table", namespace_path=["namespace_b"])
|
||||
db.drop_namespace(["namespace_a"])
|
||||
db.drop_namespace(["namespace_b"])
|
||||
|
||||
@@ -449,6 +463,8 @@ class TestAsyncNamespaceConnection:
|
||||
table_names = await db.table_names()
|
||||
assert len(list(table_names)) == 0
|
||||
|
||||
# Async connect via namespace helper is not enabled yet.
|
||||
|
||||
async def test_create_table_async(self):
|
||||
"""Test creating a table asynchronously through namespace."""
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
@@ -467,13 +483,13 @@ class TestAsyncNamespaceConnection:
|
||||
|
||||
# Create empty table in child namespace
|
||||
table = await db.create_table(
|
||||
"test_table", schema=schema, namespace=["test_ns"]
|
||||
"test_table", schema=schema, namespace_path=["test_ns"]
|
||||
)
|
||||
assert table is not None
|
||||
assert isinstance(table, lancedb.AsyncTable)
|
||||
|
||||
# Table should appear in child namespace
|
||||
table_names = await db.table_names(namespace=["test_ns"])
|
||||
table_names = await db.table_names(namespace_path=["test_ns"])
|
||||
assert "test_table" in list(table_names)
|
||||
|
||||
async def test_open_table_async(self):
|
||||
@@ -490,10 +506,10 @@ class TestAsyncNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
await db.create_table("test_table", schema=schema, namespace=["test_ns"])
|
||||
await db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Open the table
|
||||
table = await db.open_table("test_table", namespace=["test_ns"])
|
||||
table = await db.open_table("test_table", namespace_path=["test_ns"])
|
||||
assert table is not None
|
||||
assert isinstance(table, lancedb.AsyncTable)
|
||||
|
||||
@@ -547,20 +563,20 @@ class TestAsyncNamespaceConnection:
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
]
|
||||
)
|
||||
await db.create_table("table1", schema=schema, namespace=["test_ns"])
|
||||
await db.create_table("table2", schema=schema, namespace=["test_ns"])
|
||||
await db.create_table("table1", schema=schema, namespace_path=["test_ns"])
|
||||
await db.create_table("table2", schema=schema, namespace_path=["test_ns"])
|
||||
|
||||
# Verify both tables exist in child namespace
|
||||
table_names = list(await db.table_names(namespace=["test_ns"]))
|
||||
table_names = list(await db.table_names(namespace_path=["test_ns"]))
|
||||
assert "table1" in table_names
|
||||
assert "table2" in table_names
|
||||
assert len(table_names) == 2
|
||||
|
||||
# Drop one table
|
||||
await db.drop_table("table1", namespace=["test_ns"])
|
||||
await db.drop_table("table1", namespace_path=["test_ns"])
|
||||
|
||||
# Verify only table2 remains
|
||||
table_names = list(await db.table_names(namespace=["test_ns"]))
|
||||
table_names = list(await db.table_names(namespace_path=["test_ns"]))
|
||||
assert "table1" not in table_names
|
||||
assert "table2" in table_names
|
||||
assert len(table_names) == 1
|
||||
@@ -589,20 +605,24 @@ class TestAsyncNamespaceConnection:
|
||||
]
|
||||
)
|
||||
table = await db.create_table(
|
||||
"test_table", schema=schema, namespace=["test_namespace"]
|
||||
"test_table", schema=schema, namespace_path=["test_namespace"]
|
||||
)
|
||||
assert table is not None
|
||||
|
||||
# Verify table exists in namespace
|
||||
tables_in_namespace = list(await db.table_names(namespace=["test_namespace"]))
|
||||
tables_in_namespace = list(
|
||||
await db.table_names(namespace_path=["test_namespace"])
|
||||
)
|
||||
assert "test_table" in tables_in_namespace
|
||||
assert len(tables_in_namespace) == 1
|
||||
|
||||
# Drop table from namespace
|
||||
await db.drop_table("test_table", namespace=["test_namespace"])
|
||||
await db.drop_table("test_table", namespace_path=["test_namespace"])
|
||||
|
||||
# Verify table no longer exists in namespace
|
||||
tables_in_namespace = list(await db.table_names(namespace=["test_namespace"]))
|
||||
tables_in_namespace = list(
|
||||
await db.table_names(namespace_path=["test_namespace"])
|
||||
)
|
||||
assert len(tables_in_namespace) == 0
|
||||
|
||||
# Drop namespace
|
||||
@@ -627,15 +647,98 @@ class TestAsyncNamespaceConnection:
|
||||
]
|
||||
)
|
||||
for i in range(3):
|
||||
await db.create_table(f"table{i}", schema=schema, namespace=["test_ns"])
|
||||
await db.create_table(
|
||||
f"table{i}", schema=schema, namespace_path=["test_ns"]
|
||||
)
|
||||
|
||||
# Verify tables exist in child namespace
|
||||
table_names = await db.table_names(namespace=["test_ns"])
|
||||
table_names = await db.table_names(namespace_path=["test_ns"])
|
||||
assert len(list(table_names)) == 3
|
||||
|
||||
# Drop all tables in child namespace
|
||||
await db.drop_all_tables(namespace=["test_ns"])
|
||||
await db.drop_all_tables(namespace_path=["test_ns"])
|
||||
|
||||
# Verify all tables are gone from child namespace
|
||||
table_names = await db.table_names(namespace=["test_ns"])
|
||||
table_names = await db.table_names(namespace_path=["test_ns"])
|
||||
assert len(list(table_names)) == 0
|
||||
|
||||
|
||||
class TestPushdownOperations:
|
||||
"""Test pushdown operations on namespace connections."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up test fixtures."""
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_query_table_pushdown_stored(self):
|
||||
"""Test that QueryTable pushdown is stored on sync connection."""
|
||||
db = lancedb.connect_namespace(
|
||||
"dir",
|
||||
{"root": self.temp_dir},
|
||||
namespace_client_pushdown_operations=["QueryTable"],
|
||||
)
|
||||
assert "QueryTable" in db._pushdown_operations
|
||||
|
||||
def test_create_table_pushdown_stored(self):
|
||||
"""Test that CreateTable pushdown is stored on sync connection."""
|
||||
db = lancedb.connect_namespace(
|
||||
"dir",
|
||||
{"root": self.temp_dir},
|
||||
namespace_client_pushdown_operations=["CreateTable"],
|
||||
)
|
||||
assert "CreateTable" in db._pushdown_operations
|
||||
|
||||
def test_both_pushdowns_stored(self):
|
||||
"""Test that both pushdown operations can be set together."""
|
||||
db = lancedb.connect_namespace(
|
||||
"dir",
|
||||
{"root": self.temp_dir},
|
||||
namespace_client_pushdown_operations=["QueryTable", "CreateTable"],
|
||||
)
|
||||
assert "QueryTable" in db._pushdown_operations
|
||||
assert "CreateTable" in db._pushdown_operations
|
||||
|
||||
def test_pushdown_defaults_to_empty(self):
|
||||
"""Test that pushdown operations default to empty."""
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
assert len(db._pushdown_operations) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestAsyncPushdownOperations:
|
||||
"""Test pushdown operations on async namespace connections."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
def teardown_method(self):
|
||||
"""Clean up test fixtures."""
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
async def test_async_query_table_pushdown_stored(self):
|
||||
"""Test that QueryTable pushdown is stored on async connection."""
|
||||
db = lancedb.connect_namespace_async(
|
||||
"dir",
|
||||
{"root": self.temp_dir},
|
||||
namespace_client_pushdown_operations=["QueryTable"],
|
||||
)
|
||||
assert "QueryTable" in db._pushdown_operations
|
||||
|
||||
async def test_async_create_table_pushdown_stored(self):
|
||||
"""Test that CreateTable pushdown is stored on async connection."""
|
||||
db = lancedb.connect_namespace_async(
|
||||
"dir",
|
||||
{"root": self.temp_dir},
|
||||
namespace_client_pushdown_operations=["CreateTable"],
|
||||
)
|
||||
assert "CreateTable" in db._pushdown_operations
|
||||
|
||||
async def test_async_pushdown_defaults_to_empty(self):
|
||||
"""Test that pushdown operations default to empty on async connection."""
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
assert len(db._pushdown_operations) == 0
|
||||
|
||||
@@ -4,9 +4,11 @@
|
||||
"""
|
||||
Integration tests for LanceDB Namespace with S3 and credential refresh.
|
||||
|
||||
This test simulates a namespace server that returns incrementing credentials
|
||||
and verifies that the credential refresh mechanism works correctly for both
|
||||
create_table and open_table operations.
|
||||
This test uses DirectoryNamespace with native ops_metrics and vend_input_storage_options
|
||||
features to track API calls and test credential refresh mechanisms.
|
||||
|
||||
Tests are parameterized to run with both DirectoryNamespace and a CustomNamespace
|
||||
wrapper to verify Python-Rust binding works correctly for custom implementations.
|
||||
|
||||
Tests verify:
|
||||
- Storage options provider is auto-created and used
|
||||
@@ -18,22 +20,136 @@ Tests verify:
|
||||
import copy
|
||||
import time
|
||||
import uuid
|
||||
from threading import Lock
|
||||
from typing import Dict
|
||||
from typing import Dict, Optional
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lance_namespace import (
|
||||
CreateEmptyTableRequest,
|
||||
CreateEmptyTableResponse,
|
||||
from lance.namespace import (
|
||||
DeclareTableRequest,
|
||||
DeclareTableResponse,
|
||||
DescribeTableRequest,
|
||||
DescribeTableResponse,
|
||||
DirectoryNamespace,
|
||||
LanceNamespace,
|
||||
)
|
||||
from lance_namespace import (
|
||||
CreateNamespaceRequest,
|
||||
CreateNamespaceResponse,
|
||||
CreateTableRequest,
|
||||
CreateTableResponse,
|
||||
CreateTableVersionRequest,
|
||||
CreateTableVersionResponse,
|
||||
DeregisterTableRequest,
|
||||
DeregisterTableResponse,
|
||||
DescribeNamespaceRequest,
|
||||
DescribeNamespaceResponse,
|
||||
DescribeTableVersionRequest,
|
||||
DescribeTableVersionResponse,
|
||||
DropNamespaceRequest,
|
||||
DropNamespaceResponse,
|
||||
DropTableRequest,
|
||||
DropTableResponse,
|
||||
ListNamespacesRequest,
|
||||
ListNamespacesResponse,
|
||||
ListTablesRequest,
|
||||
ListTablesResponse,
|
||||
ListTableVersionsRequest,
|
||||
ListTableVersionsResponse,
|
||||
NamespaceExistsRequest,
|
||||
RegisterTableRequest,
|
||||
RegisterTableResponse,
|
||||
TableExistsRequest,
|
||||
)
|
||||
from lancedb.namespace import LanceNamespaceDBConnection
|
||||
|
||||
|
||||
class CustomNamespace(LanceNamespace):
|
||||
"""A custom namespace wrapper that delegates to DirectoryNamespace.
|
||||
|
||||
This class verifies that the Python-Rust binding works correctly for
|
||||
custom namespace implementations that wrap the native DirectoryNamespace.
|
||||
All methods simply delegate to the underlying DirectoryNamespace instance.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: DirectoryNamespace):
|
||||
self._inner = inner
|
||||
|
||||
def namespace_id(self) -> str:
|
||||
return f"CustomNamespace[{self._inner.namespace_id()}]"
|
||||
|
||||
def create_namespace(
|
||||
self, request: CreateNamespaceRequest
|
||||
) -> CreateNamespaceResponse:
|
||||
return self._inner.create_namespace(request)
|
||||
|
||||
def describe_namespace(
|
||||
self, request: DescribeNamespaceRequest
|
||||
) -> DescribeNamespaceResponse:
|
||||
return self._inner.describe_namespace(request)
|
||||
|
||||
def namespace_exists(self, request: NamespaceExistsRequest) -> None:
|
||||
return self._inner.namespace_exists(request)
|
||||
|
||||
def drop_namespace(self, request: DropNamespaceRequest) -> DropNamespaceResponse:
|
||||
return self._inner.drop_namespace(request)
|
||||
|
||||
def list_namespaces(self, request: ListNamespacesRequest) -> ListNamespacesResponse:
|
||||
return self._inner.list_namespaces(request)
|
||||
|
||||
def create_table(
|
||||
self, request: CreateTableRequest, data: bytes
|
||||
) -> CreateTableResponse:
|
||||
return self._inner.create_table(request, data)
|
||||
|
||||
def declare_table(self, request: DeclareTableRequest) -> DeclareTableResponse:
|
||||
return self._inner.declare_table(request)
|
||||
|
||||
def describe_table(self, request: DescribeTableRequest) -> DescribeTableResponse:
|
||||
return self._inner.describe_table(request)
|
||||
|
||||
def table_exists(self, request: TableExistsRequest) -> None:
|
||||
return self._inner.table_exists(request)
|
||||
|
||||
def drop_table(self, request: DropTableRequest) -> DropTableResponse:
|
||||
return self._inner.drop_table(request)
|
||||
|
||||
def list_tables(self, request: ListTablesRequest) -> ListTablesResponse:
|
||||
return self._inner.list_tables(request)
|
||||
|
||||
def register_table(self, request: RegisterTableRequest) -> RegisterTableResponse:
|
||||
return self._inner.register_table(request)
|
||||
|
||||
def deregister_table(
|
||||
self, request: DeregisterTableRequest
|
||||
) -> DeregisterTableResponse:
|
||||
return self._inner.deregister_table(request)
|
||||
|
||||
def list_table_versions(
|
||||
self, request: ListTableVersionsRequest
|
||||
) -> ListTableVersionsResponse:
|
||||
return self._inner.list_table_versions(request)
|
||||
|
||||
def describe_table_version(
|
||||
self, request: DescribeTableVersionRequest
|
||||
) -> DescribeTableVersionResponse:
|
||||
return self._inner.describe_table_version(request)
|
||||
|
||||
def create_table_version(
|
||||
self, request: CreateTableVersionRequest
|
||||
) -> CreateTableVersionResponse:
|
||||
return self._inner.create_table_version(request)
|
||||
|
||||
def retrieve_ops_metrics(self) -> Optional[Dict[str, int]]:
|
||||
return self._inner.retrieve_ops_metrics()
|
||||
|
||||
|
||||
def _wrap_if_custom(ns_client: DirectoryNamespace, use_custom: bool):
|
||||
"""Wrap namespace client in CustomNamespace if use_custom is True."""
|
||||
if use_custom:
|
||||
return CustomNamespace(ns_client)
|
||||
return ns_client
|
||||
|
||||
|
||||
# LocalStack S3 configuration
|
||||
CONFIG = {
|
||||
"allow_http": "true",
|
||||
@@ -89,162 +205,88 @@ def delete_bucket(s3, bucket_name):
|
||||
pass
|
||||
|
||||
|
||||
class TrackingNamespace(LanceNamespace):
|
||||
def create_tracking_namespace(
|
||||
bucket_name: str,
|
||||
storage_options: dict,
|
||||
credential_expires_in_seconds: int = 60,
|
||||
use_custom: bool = False,
|
||||
):
|
||||
"""Create a DirectoryNamespace with ops metrics and credential vending enabled.
|
||||
|
||||
Uses native DirectoryNamespace features:
|
||||
- ops_metrics_enabled=true: Tracks API call counts via retrieve_ops_metrics()
|
||||
- vend_input_storage_options=true: Returns input storage options in responses
|
||||
- vend_input_storage_options_refresh_interval_millis: Adds expires_at_millis
|
||||
|
||||
Args:
|
||||
bucket_name: S3 bucket name or local path
|
||||
storage_options: Storage options to pass through (credentials, endpoint, etc.)
|
||||
credential_expires_in_seconds: Interval in seconds for credential expiration
|
||||
use_custom: If True, wrap in CustomNamespace for testing custom implementations
|
||||
|
||||
Returns:
|
||||
Tuple of (namespace_client, inner_namespace_client) where inner is always
|
||||
the DirectoryNamespace (used for metrics retrieval)
|
||||
"""
|
||||
Mock namespace that wraps DirectoryNamespace and tracks API calls.
|
||||
# Add refresh_offset_millis to storage options so that credentials are not
|
||||
# considered expired immediately. Set to 1 second (1000ms) so that refresh
|
||||
# checks work correctly with short-lived credentials in tests.
|
||||
storage_options_with_refresh = dict(storage_options)
|
||||
storage_options_with_refresh["refresh_offset_millis"] = "1000"
|
||||
|
||||
This namespace returns incrementing credentials with each API call to simulate
|
||||
credential rotation. It also tracks the number of times each API is called
|
||||
to verify caching behavior.
|
||||
"""
|
||||
dir_props = {f"storage.{k}": v for k, v in storage_options_with_refresh.items()}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket_name: str,
|
||||
storage_options: Dict[str, str],
|
||||
credential_expires_in_seconds: int = 60,
|
||||
):
|
||||
from lance.namespace import DirectoryNamespace
|
||||
if bucket_name.startswith("/") or bucket_name.startswith("file://"):
|
||||
dir_props["root"] = f"{bucket_name}/namespace_root"
|
||||
else:
|
||||
dir_props["root"] = f"s3://{bucket_name}/namespace_root"
|
||||
|
||||
self.bucket_name = bucket_name
|
||||
self.base_storage_options = storage_options
|
||||
self.credential_expires_in_seconds = credential_expires_in_seconds
|
||||
self.describe_call_count = 0
|
||||
self.create_call_count = 0
|
||||
self.lock = Lock()
|
||||
# Enable ops metrics tracking
|
||||
dir_props["ops_metrics_enabled"] = "true"
|
||||
# Enable storage options vending
|
||||
dir_props["vend_input_storage_options"] = "true"
|
||||
# Set refresh interval in milliseconds
|
||||
dir_props["vend_input_storage_options_refresh_interval_millis"] = str(
|
||||
credential_expires_in_seconds * 1000
|
||||
)
|
||||
|
||||
# Create underlying DirectoryNamespace with storage options
|
||||
dir_props = {f"storage.{k}": v for k, v in storage_options.items()}
|
||||
inner_ns_client = DirectoryNamespace(**dir_props)
|
||||
ns_client = _wrap_if_custom(inner_ns_client, use_custom)
|
||||
return ns_client, inner_ns_client
|
||||
|
||||
# Use S3 path for bucket name, local path for file paths
|
||||
if bucket_name.startswith("/") or bucket_name.startswith("file://"):
|
||||
dir_props["root"] = f"{bucket_name}/namespace_root"
|
||||
else:
|
||||
dir_props["root"] = f"s3://{bucket_name}/namespace_root"
|
||||
|
||||
self.inner = DirectoryNamespace(**dir_props)
|
||||
def get_describe_call_count(namespace_client) -> int:
|
||||
"""Get the number of describe_table calls made to the namespace client."""
|
||||
return namespace_client.retrieve_ops_metrics().get("describe_table", 0)
|
||||
|
||||
def get_describe_call_count(self) -> int:
|
||||
"""Thread-safe getter for describe call count."""
|
||||
with self.lock:
|
||||
return self.describe_call_count
|
||||
|
||||
def get_create_call_count(self) -> int:
|
||||
"""Thread-safe getter for create call count."""
|
||||
with self.lock:
|
||||
return self.create_call_count
|
||||
|
||||
def namespace_id(self) -> str:
|
||||
"""Return namespace identifier."""
|
||||
return f"TrackingNamespace {{ inner: {self.inner.namespace_id()} }}"
|
||||
|
||||
def _modify_storage_options(
|
||||
self, storage_options: Dict[str, str], count: int
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Add incrementing credentials with expiration timestamp.
|
||||
|
||||
This simulates a credential rotation system where each call returns
|
||||
new credentials that expire after credential_expires_in_seconds.
|
||||
"""
|
||||
# Start from base storage options (endpoint, region, allow_http, etc.)
|
||||
# because DirectoryNamespace returns None for storage_options from
|
||||
# describe_table/declare_table when no credential vendor is configured.
|
||||
modified = copy.deepcopy(self.base_storage_options)
|
||||
if storage_options:
|
||||
modified.update(storage_options)
|
||||
|
||||
# Increment credentials to simulate rotation
|
||||
modified["aws_access_key_id"] = f"AKID_{count}"
|
||||
modified["aws_secret_access_key"] = f"SECRET_{count}"
|
||||
modified["aws_session_token"] = f"TOKEN_{count}"
|
||||
|
||||
# Set expiration time
|
||||
expires_at_millis = int(
|
||||
(time.time() + self.credential_expires_in_seconds) * 1000
|
||||
)
|
||||
modified["expires_at_millis"] = str(expires_at_millis)
|
||||
|
||||
return modified
|
||||
|
||||
def declare_table(self, request: DeclareTableRequest) -> DeclareTableResponse:
|
||||
"""Track declare_table calls and inject rotating credentials."""
|
||||
with self.lock:
|
||||
self.create_call_count += 1
|
||||
count = self.create_call_count
|
||||
|
||||
response = self.inner.declare_table(request)
|
||||
response.storage_options = self._modify_storage_options(
|
||||
response.storage_options, count
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def create_empty_table(
|
||||
self, request: CreateEmptyTableRequest
|
||||
) -> CreateEmptyTableResponse:
|
||||
"""Track create_empty_table calls and inject rotating credentials."""
|
||||
with self.lock:
|
||||
self.create_call_count += 1
|
||||
count = self.create_call_count
|
||||
|
||||
response = self.inner.create_empty_table(request)
|
||||
response.storage_options = self._modify_storage_options(
|
||||
response.storage_options, count
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def describe_table(self, request: DescribeTableRequest) -> DescribeTableResponse:
|
||||
"""Track describe_table calls and inject rotating credentials."""
|
||||
with self.lock:
|
||||
self.describe_call_count += 1
|
||||
count = self.describe_call_count
|
||||
|
||||
response = self.inner.describe_table(request)
|
||||
response.storage_options = self._modify_storage_options(
|
||||
response.storage_options, count
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
# Pass through other methods to inner namespace
|
||||
def list_tables(self, request):
|
||||
return self.inner.list_tables(request)
|
||||
|
||||
def drop_table(self, request):
|
||||
return self.inner.drop_table(request)
|
||||
|
||||
def list_namespaces(self, request):
|
||||
return self.inner.list_namespaces(request)
|
||||
|
||||
def create_namespace(self, request):
|
||||
return self.inner.create_namespace(request)
|
||||
|
||||
def drop_namespace(self, request):
|
||||
return self.inner.drop_namespace(request)
|
||||
def get_declare_call_count(namespace_client) -> int:
|
||||
"""Get the number of declare_table calls made to the namespace client."""
|
||||
return namespace_client.retrieve_ops_metrics().get("declare_table", 0)
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_create_table_with_provider(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_create_table_with_provider(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test creating a table through namespace with storage options provider.
|
||||
|
||||
Verifies:
|
||||
- create_empty_table is called once to reserve location
|
||||
- declare_table is called once to reserve location
|
||||
- Storage options provider is auto-created
|
||||
- Table can be written successfully
|
||||
- Credentials are cached during write operations
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3600, # 1 hour
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -254,8 +296,8 @@ def test_namespace_create_table_with_provider(s3_bucket: str):
|
||||
namespace_path = [namespace_name]
|
||||
|
||||
# Verify initial state
|
||||
assert namespace.get_create_call_count() == 0
|
||||
assert namespace.get_describe_call_count() == 0
|
||||
assert get_declare_call_count(inner_ns_client) == 0
|
||||
assert get_describe_call_count(inner_ns_client) == 0
|
||||
|
||||
# Create table with data
|
||||
data = pa.table(
|
||||
@@ -266,12 +308,12 @@ def test_namespace_create_table_with_provider(s3_bucket: str):
|
||||
}
|
||||
)
|
||||
|
||||
table = db.create_table(table_name, data, namespace=namespace_path)
|
||||
table = db.create_table(table_name, data, namespace_path=namespace_path)
|
||||
|
||||
# Verify create_empty_table was called exactly once
|
||||
assert namespace.get_create_call_count() == 1
|
||||
# Verify declare_table was called exactly once
|
||||
assert get_declare_call_count(inner_ns_client) == 1
|
||||
# describe_table should NOT be called during create in create mode
|
||||
assert namespace.get_describe_call_count() == 0
|
||||
assert get_describe_call_count(inner_ns_client) == 0
|
||||
|
||||
# Verify table was created successfully
|
||||
assert table.name == table_name
|
||||
@@ -281,7 +323,8 @@ def test_namespace_create_table_with_provider(s3_bucket: str):
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_open_table_with_provider(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_open_table_with_provider(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test opening a table through namespace with storage options provider.
|
||||
|
||||
@@ -293,13 +336,14 @@ def test_namespace_open_table_with_provider(s3_bucket: str):
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3600,
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -317,21 +361,21 @@ def test_namespace_open_table_with_provider(s3_bucket: str):
|
||||
}
|
||||
)
|
||||
|
||||
db.create_table(table_name, data, namespace=namespace_path)
|
||||
db.create_table(table_name, data, namespace_path=namespace_path)
|
||||
|
||||
initial_create_count = namespace.get_create_call_count()
|
||||
assert initial_create_count == 1
|
||||
initial_declare_count = get_declare_call_count(inner_ns_client)
|
||||
assert initial_declare_count == 1
|
||||
|
||||
# Open the table
|
||||
opened_table = db.open_table(table_name, namespace=namespace_path)
|
||||
opened_table = db.open_table(table_name, namespace_path=namespace_path)
|
||||
|
||||
# Verify describe_table was called exactly once
|
||||
assert namespace.get_describe_call_count() == 1
|
||||
# create_empty_table should not be called again
|
||||
assert namespace.get_create_call_count() == initial_create_count
|
||||
assert get_describe_call_count(inner_ns_client) == 1
|
||||
# declare_table should not be called again
|
||||
assert get_declare_call_count(inner_ns_client) == initial_declare_count
|
||||
|
||||
# Perform multiple read operations
|
||||
describe_count_after_open = namespace.get_describe_call_count()
|
||||
describe_count_after_open = get_describe_call_count(inner_ns_client)
|
||||
|
||||
for _ in range(3):
|
||||
result = opened_table.to_pandas()
|
||||
@@ -340,11 +384,12 @@ def test_namespace_open_table_with_provider(s3_bucket: str):
|
||||
assert count == 5
|
||||
|
||||
# Verify credentials were cached (no additional describe_table calls)
|
||||
assert namespace.get_describe_call_count() == describe_count_after_open
|
||||
assert get_describe_call_count(inner_ns_client) == describe_count_after_open
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_credential_refresh_on_read(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_credential_refresh_on_read(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test credential refresh when credentials expire during read operations.
|
||||
|
||||
@@ -355,13 +400,14 @@ def test_namespace_credential_refresh_on_read(s3_bucket: str):
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3, # Short expiration for testing
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -378,16 +424,16 @@ def test_namespace_credential_refresh_on_read(s3_bucket: str):
|
||||
}
|
||||
)
|
||||
|
||||
db.create_table(table_name, data, namespace=namespace_path)
|
||||
db.create_table(table_name, data, namespace_path=namespace_path)
|
||||
|
||||
# Open table (triggers describe_table)
|
||||
opened_table = db.open_table(table_name, namespace=namespace_path)
|
||||
opened_table = db.open_table(table_name, namespace_path=namespace_path)
|
||||
|
||||
# Perform an immediate read (should use credentials from open)
|
||||
result = opened_table.to_pandas()
|
||||
assert len(result) == 3
|
||||
|
||||
describe_count_after_first_read = namespace.get_describe_call_count()
|
||||
describe_count_after_first_read = get_describe_call_count(inner_ns_client)
|
||||
|
||||
# Wait for credentials to expire (3 seconds + buffer)
|
||||
time.sleep(5)
|
||||
@@ -396,7 +442,7 @@ def test_namespace_credential_refresh_on_read(s3_bucket: str):
|
||||
result = opened_table.to_pandas()
|
||||
assert len(result) == 3
|
||||
|
||||
describe_count_after_refresh = namespace.get_describe_call_count()
|
||||
describe_count_after_refresh = get_describe_call_count(inner_ns_client)
|
||||
# Verify describe_table was called again (credential refresh)
|
||||
refresh_delta = describe_count_after_refresh - describe_count_after_first_read
|
||||
|
||||
@@ -409,7 +455,8 @@ def test_namespace_credential_refresh_on_read(s3_bucket: str):
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_credential_refresh_on_write(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_credential_refresh_on_write(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test credential refresh when credentials expire during write operations.
|
||||
|
||||
@@ -420,13 +467,14 @@ def test_namespace_credential_refresh_on_write(s3_bucket: str):
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3, # Short expiration
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -443,7 +491,7 @@ def test_namespace_credential_refresh_on_write(s3_bucket: str):
|
||||
}
|
||||
)
|
||||
|
||||
table = db.create_table(table_name, initial_data, namespace=namespace_path)
|
||||
table = db.create_table(table_name, initial_data, namespace_path=namespace_path)
|
||||
|
||||
# Add more data (should use cached credentials)
|
||||
new_data = pa.table(
|
||||
@@ -471,24 +519,26 @@ def test_namespace_credential_refresh_on_write(s3_bucket: str):
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_overwrite_mode(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_overwrite_mode(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test creating table in overwrite mode with credential tracking.
|
||||
|
||||
Verifies:
|
||||
- First create calls create_empty_table exactly once
|
||||
- First create calls declare_table exactly once
|
||||
- Overwrite mode calls describe_table exactly once to check existence
|
||||
- Storage options provider works in overwrite mode
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3600,
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -505,11 +555,11 @@ def test_namespace_overwrite_mode(s3_bucket: str):
|
||||
}
|
||||
)
|
||||
|
||||
table = db.create_table(table_name, data1, namespace=namespace_path)
|
||||
# Exactly one create_empty_table call for initial create
|
||||
assert namespace.get_create_call_count() == 1
|
||||
table = db.create_table(table_name, data1, namespace_path=namespace_path)
|
||||
# Exactly one declare_table call for initial create
|
||||
assert get_declare_call_count(inner_ns_client) == 1
|
||||
# No describe_table calls in create mode
|
||||
assert namespace.get_describe_call_count() == 0
|
||||
assert get_describe_call_count(inner_ns_client) == 0
|
||||
assert table.count_rows() == 2
|
||||
|
||||
# Overwrite the table
|
||||
@@ -521,14 +571,14 @@ def test_namespace_overwrite_mode(s3_bucket: str):
|
||||
)
|
||||
|
||||
table2 = db.create_table(
|
||||
table_name, data2, namespace=namespace_path, mode="overwrite"
|
||||
table_name, data2, namespace_path=namespace_path, mode="overwrite"
|
||||
)
|
||||
|
||||
# Should still have only 1 create_empty_table call
|
||||
# Should still have only 1 declare_table call
|
||||
# (overwrite reuses location from describe_table)
|
||||
assert namespace.get_create_call_count() == 1
|
||||
assert get_declare_call_count(inner_ns_client) == 1
|
||||
# Should have called describe_table exactly once to get existing table location
|
||||
assert namespace.get_describe_call_count() == 1
|
||||
assert get_describe_call_count(inner_ns_client) == 1
|
||||
|
||||
# Verify new data
|
||||
assert table2.count_rows() == 3
|
||||
@@ -537,7 +587,8 @@ def test_namespace_overwrite_mode(s3_bucket: str):
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_multiple_tables(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_multiple_tables(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test creating and opening multiple tables in the same namespace.
|
||||
|
||||
@@ -548,13 +599,14 @@ def test_namespace_multiple_tables(s3_bucket: str):
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3600,
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -564,22 +616,22 @@ def test_namespace_multiple_tables(s3_bucket: str):
|
||||
# Create first table
|
||||
table1_name = f"table1_{uuid.uuid4().hex}"
|
||||
data1 = pa.table({"id": [1, 2], "value": [10, 20]})
|
||||
db.create_table(table1_name, data1, namespace=namespace_path)
|
||||
db.create_table(table1_name, data1, namespace_path=namespace_path)
|
||||
|
||||
# Create second table
|
||||
table2_name = f"table2_{uuid.uuid4().hex}"
|
||||
data2 = pa.table({"id": [3, 4], "value": [30, 40]})
|
||||
db.create_table(table2_name, data2, namespace=namespace_path)
|
||||
db.create_table(table2_name, data2, namespace_path=namespace_path)
|
||||
|
||||
# Should have 2 create calls (one per table)
|
||||
assert namespace.get_create_call_count() == 2
|
||||
# Should have 2 declare calls (one per table)
|
||||
assert get_declare_call_count(inner_ns_client) == 2
|
||||
|
||||
# Open both tables
|
||||
opened1 = db.open_table(table1_name, namespace=namespace_path)
|
||||
opened2 = db.open_table(table2_name, namespace=namespace_path)
|
||||
opened1 = db.open_table(table1_name, namespace_path=namespace_path)
|
||||
opened2 = db.open_table(table2_name, namespace_path=namespace_path)
|
||||
|
||||
# Should have 2 describe calls (one per open)
|
||||
assert namespace.get_describe_call_count() == 2
|
||||
assert get_describe_call_count(inner_ns_client) == 2
|
||||
|
||||
# Verify both tables work independently
|
||||
assert opened1.count_rows() == 2
|
||||
@@ -593,7 +645,8 @@ def test_namespace_multiple_tables(s3_bucket: str):
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_namespace_with_schema_only(s3_bucket: str):
|
||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
||||
def test_namespace_with_schema_only(s3_bucket: str, use_custom: bool):
|
||||
"""
|
||||
Test creating empty table with schema only (no data).
|
||||
|
||||
@@ -604,13 +657,14 @@ def test_namespace_with_schema_only(s3_bucket: str):
|
||||
"""
|
||||
storage_options = copy.deepcopy(CONFIG)
|
||||
|
||||
namespace = TrackingNamespace(
|
||||
ns_client, inner_ns_client = create_tracking_namespace(
|
||||
bucket_name=s3_bucket,
|
||||
storage_options=storage_options,
|
||||
credential_expires_in_seconds=3600,
|
||||
use_custom=use_custom,
|
||||
)
|
||||
|
||||
db = LanceNamespaceDBConnection(namespace)
|
||||
db = LanceNamespaceDBConnection(ns_client)
|
||||
|
||||
# Create unique namespace for this test
|
||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
||||
@@ -628,12 +682,12 @@ def test_namespace_with_schema_only(s3_bucket: str):
|
||||
]
|
||||
)
|
||||
|
||||
table = db.create_table(table_name, schema=schema, namespace=namespace_path)
|
||||
table = db.create_table(table_name, schema=schema, namespace_path=namespace_path)
|
||||
|
||||
# Should have called create_empty_table once
|
||||
assert namespace.get_create_call_count() == 1
|
||||
# Should have called declare_table once
|
||||
assert get_declare_call_count(inner_ns_client) == 1
|
||||
# Should NOT have called describe_table in create mode
|
||||
assert namespace.get_describe_call_count() == 0
|
||||
assert get_describe_call_count(inner_ns_client) == 0
|
||||
|
||||
# Verify empty table
|
||||
assert table.count_rows() == 0
|
||||
|
||||
@@ -2108,7 +2108,7 @@ def test_stats(mem_db: DBConnection):
|
||||
stats = table.stats()
|
||||
print(f"{stats=}")
|
||||
assert stats == {
|
||||
"total_bytes": 38,
|
||||
"total_bytes": 60,
|
||||
"num_rows": 2,
|
||||
"num_indices": 0,
|
||||
"fragment_stats": {
|
||||
|
||||
Reference in New Issue
Block a user