mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-26 05:26:14 +00:00
feat!: return lazy iterators for table and namespace listings (#4599)
Make table and namespace listings consistent with other resource
listings across Rust, Python, and TypeScript. Callers can iterate over
individual names across all pages, with configurable page sizes and
continuation tokens.
In order to avoid breaking existing Python users, this change adds some
compatibility methods to the returned Python iterators. Code such as:
```
for t in db.list_tables().tables:
print(t)
```
will continue to work as expected because the iterator has a `tables`
field and a `page_token` field, which function as expected to reveal the
first page of results. In other words, it is never necessary to use the
iterator __as an iterator__. It may be used as if it were an old-style
response object. (As the iterator progresses, the `tables` list may
switch to the second, third, etc. page of results.) Similar
compatibility code allows the async python API to keep working as well.
This commit is contained in:
@@ -19,11 +19,9 @@ from .index import (
|
||||
FTS,
|
||||
)
|
||||
from lance_namespace import (
|
||||
ListNamespacesResponse,
|
||||
CreateNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
ListTablesResponse,
|
||||
)
|
||||
from .remote import ClientConfig
|
||||
from .types import BaseTokenizerType
|
||||
@@ -114,12 +112,12 @@ class Connection(object):
|
||||
uri: str
|
||||
async def is_open(self): ...
|
||||
async def close(self): ...
|
||||
async def list_namespaces(
|
||||
def list_namespaces(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse: ...
|
||||
) -> NameListing: ...
|
||||
async def create_namespace(
|
||||
self,
|
||||
namespace_path: List[str],
|
||||
@@ -136,12 +134,12 @@ class Connection(object):
|
||||
self,
|
||||
namespace_path: List[str],
|
||||
) -> DescribeNamespaceResponse: ...
|
||||
async def list_tables(
|
||||
def list_tables(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse: ...
|
||||
) -> NameListing: ...
|
||||
async def table_names(
|
||||
self,
|
||||
namespace_path: Optional[List[str]],
|
||||
@@ -1010,6 +1008,7 @@ async def connect_catalog(
|
||||
class NameListing:
|
||||
def num_page_results(self) -> int: ...
|
||||
def page_token(self) -> Optional[str]: ...
|
||||
async def next_page(self) -> Optional[dict[str, Any]]: ...
|
||||
def __aiter__(self) -> NameListing: ...
|
||||
async def __anext__(self) -> str: ...
|
||||
|
||||
|
||||
+40
-30
@@ -32,18 +32,24 @@ if sys.version_info >= (3, 12):
|
||||
else:
|
||||
from overrides import EnforceOverrides, override # type: ignore
|
||||
|
||||
from lancedb.listing import AsyncListing, Listing
|
||||
from lancedb.listing import (
|
||||
AsyncListing,
|
||||
AsyncNamespaceListing,
|
||||
AsyncTableListing,
|
||||
Listing,
|
||||
NamespaceListing,
|
||||
TableListing,
|
||||
validate_page_limit,
|
||||
)
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
|
||||
from lancedb.common import data_to_reader, sanitize_uri, validate_schema
|
||||
from lancedb.background_loop import LOOP
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
ListNamespacesResponse,
|
||||
CreateNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
ListTablesResponse,
|
||||
connect as namespace_connect,
|
||||
)
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
@@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse:
|
||||
) -> NamespaceListing:
|
||||
"""List immediate child namespace names in the given namespace.
|
||||
|
||||
Parameters
|
||||
@@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListNamespacesResponse
|
||||
NamespaceListing
|
||||
Response containing namespace names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return ListNamespacesResponse(namespaces=[], page_token=None)
|
||||
return NamespaceListing(AsyncNamespaceListing(None, empty=True))
|
||||
|
||||
def create_namespace(
|
||||
self,
|
||||
@@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List a page of tables in this database.
|
||||
) -> TableListing:
|
||||
"""List tables in this database, one page at a time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListTablesResponse
|
||||
TableListing
|
||||
Response containing table names and optional page_token for pagination.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
@@ -1436,7 +1442,7 @@ class LanceDBConnection(DBConnection):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse:
|
||||
) -> NamespaceListing:
|
||||
"""List immediate child namespace names in the given namespace.
|
||||
|
||||
Parameters
|
||||
@@ -1452,16 +1458,14 @@ class LanceDBConnection(DBConnection):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListNamespacesResponse
|
||||
NamespaceListing
|
||||
Response containing namespace names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
return NamespaceListing(
|
||||
self._conn.list_namespaces(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
limit=limit,
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1512,8 +1516,8 @@ class LanceDBConnection(DBConnection):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List a page of tables in this database.
|
||||
) -> TableListing:
|
||||
"""List tables in this database, one page at a time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -1529,12 +1533,12 @@ class LanceDBConnection(DBConnection):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListTablesResponse
|
||||
TableListing
|
||||
Response containing table names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
return TableListing(
|
||||
self._conn.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
@@ -2286,12 +2290,12 @@ class AsyncConnection(object):
|
||||
else:
|
||||
return None
|
||||
|
||||
async def list_namespaces(
|
||||
def list_namespaces(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse:
|
||||
) -> AsyncNamespaceListing:
|
||||
"""List immediate child namespace names in the given namespace.
|
||||
|
||||
Parameters
|
||||
@@ -2306,15 +2310,18 @@ class AsyncConnection(object):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListNamespacesResponse
|
||||
AsyncNamespaceListing
|
||||
Response containing namespace names and optional pagination token
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
result = await self._inner.list_namespaces(
|
||||
validate_page_limit(limit)
|
||||
inner = self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
return ListNamespacesResponse(**result)
|
||||
if limit == 0:
|
||||
return AsyncNamespaceListing(None, empty=True)
|
||||
return AsyncNamespaceListing(inner, page_token)
|
||||
|
||||
async def create_namespace(
|
||||
self,
|
||||
@@ -2393,13 +2400,13 @@ class AsyncConnection(object):
|
||||
result = await self._inner.describe_namespace(namespace_path)
|
||||
return DescribeNamespaceResponse(**result)
|
||||
|
||||
async def list_tables(
|
||||
def list_tables(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List a page of tables in this database.
|
||||
) -> AsyncTableListing:
|
||||
"""List tables in this database, one page at a time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -2415,15 +2422,18 @@ class AsyncConnection(object):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListTablesResponse
|
||||
AsyncTableListing
|
||||
Response containing table names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
result = await self._inner.list_tables(
|
||||
validate_page_limit(limit)
|
||||
inner = self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
return ListTablesResponse(**result)
|
||||
if limit == 0:
|
||||
return AsyncTableListing(None, empty=True)
|
||||
return AsyncTableListing(inner, page_token)
|
||||
|
||||
async def table_names(
|
||||
self,
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar
|
||||
|
||||
from .background_loop import LOOP
|
||||
@@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]):
|
||||
for initial, cached, final-page, and error behavior.
|
||||
"""
|
||||
return self._inner.page_token()
|
||||
|
||||
|
||||
def validate_page_limit(limit: Optional[int]) -> None:
|
||||
"""Validate the legacy limit argument before passing it to the bindings."""
|
||||
if limit is not None and (
|
||||
not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1
|
||||
):
|
||||
raise ValueError("limit must be between 0 and 2147483647")
|
||||
|
||||
|
||||
class _AsyncPagedNameListing(AsyncIterator[str]):
|
||||
"""An async listing with an inspectable copy of its current backend page."""
|
||||
|
||||
def __init__(
|
||||
self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False
|
||||
):
|
||||
self._inner = inner
|
||||
self._token = page_token
|
||||
self._page: Optional[list[str]] = [] if empty else None
|
||||
self._position = 0
|
||||
self._done = empty
|
||||
self._failure: Optional[BaseException] = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def page_token(self) -> Optional[str]:
|
||||
"""The continuation token for the next backend request."""
|
||||
return self._token
|
||||
|
||||
def num_page_results(self) -> int:
|
||||
"""Count the remaining names in the current page."""
|
||||
return 0 if self._page is None else len(self._page) - self._position
|
||||
|
||||
def _page_copy(self) -> list[str]:
|
||||
if self._page is None:
|
||||
raise RuntimeError("Await or iterate the listing before reading its page")
|
||||
return self._page.copy()
|
||||
|
||||
async def _fetch_page(self) -> None:
|
||||
try:
|
||||
page = await self._inner.next_page()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except BaseException as exc:
|
||||
self._failure = exc
|
||||
self._done = True
|
||||
raise
|
||||
if page is None:
|
||||
self._page = []
|
||||
self._token = None
|
||||
self._done = True
|
||||
else:
|
||||
self._page = page["items"]
|
||||
self._token = page["page_token"]
|
||||
self._done = self._token is None
|
||||
self._position = 0
|
||||
|
||||
async def _load_initial(self):
|
||||
async with self._lock:
|
||||
if self._failure is not None:
|
||||
raise self._failure
|
||||
if self._page is None:
|
||||
await self._fetch_page()
|
||||
return self
|
||||
|
||||
def __await__(self):
|
||||
return self._load_initial().__await__()
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> str:
|
||||
async with self._lock:
|
||||
if self._failure is not None:
|
||||
raise StopAsyncIteration
|
||||
if self._page is None:
|
||||
await self._fetch_page()
|
||||
assert self._page is not None
|
||||
while self._position >= len(self._page):
|
||||
if self._done:
|
||||
raise StopAsyncIteration
|
||||
await self._fetch_page()
|
||||
assert self._page is not None
|
||||
item = self._page[self._position]
|
||||
self._position += 1
|
||||
return item
|
||||
|
||||
|
||||
class AsyncTableListing(_AsyncPagedNameListing):
|
||||
"""Awaitable table iterator; awaiting loads one page.
|
||||
|
||||
For example::
|
||||
|
||||
page = await db.list_tables(limit=20)
|
||||
print(page.tables, page.page_token)
|
||||
async for name in db.list_tables(limit=20):
|
||||
print(name)
|
||||
"""
|
||||
|
||||
@property
|
||||
def tables(self) -> list[str]:
|
||||
"""A copy of every table name in the most recently fetched page."""
|
||||
return self._page_copy()
|
||||
|
||||
|
||||
class AsyncNamespaceListing(_AsyncPagedNameListing):
|
||||
"""Awaitable namespace iterator; awaiting loads one page.
|
||||
|
||||
For example::
|
||||
|
||||
page = await db.list_namespaces(limit=20)
|
||||
print(page.namespaces, page.page_token)
|
||||
"""
|
||||
|
||||
@property
|
||||
def namespaces(self) -> list[str]:
|
||||
"""A copy of every name in the most recently fetched page."""
|
||||
return self._page_copy()
|
||||
|
||||
|
||||
class _PagedNameListing(Iterator[str]):
|
||||
def __init__(self, inner: _AsyncPagedNameListing):
|
||||
self._inner = inner
|
||||
LOOP.run(inner._load_initial())
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> str:
|
||||
try:
|
||||
return LOOP.run(self._inner.__anext__())
|
||||
except StopAsyncIteration:
|
||||
raise StopIteration from None
|
||||
|
||||
@property
|
||||
def page_token(self) -> Optional[str]:
|
||||
return self._inner.page_token
|
||||
|
||||
def num_page_results(self) -> int:
|
||||
return self._inner.num_page_results()
|
||||
|
||||
|
||||
class TableListing(_PagedNameListing):
|
||||
"""Table iterator retaining the current page as a list.
|
||||
|
||||
For example::
|
||||
|
||||
page = db.list_tables(limit=20)
|
||||
print(page.tables, page.page_token)
|
||||
for name in db.list_tables(limit=20):
|
||||
print(name)
|
||||
"""
|
||||
|
||||
@property
|
||||
def tables(self) -> list[str]:
|
||||
return self._inner.tables
|
||||
|
||||
|
||||
class NamespaceListing(_PagedNameListing):
|
||||
"""Namespace iterator retaining the current page as a list.
|
||||
|
||||
For example::
|
||||
|
||||
page = db.list_namespaces(limit=20)
|
||||
print(page.namespaces, page.page_token)
|
||||
"""
|
||||
|
||||
@property
|
||||
def namespaces(self) -> list[str]:
|
||||
return self._inner.namespaces
|
||||
|
||||
@@ -52,7 +52,14 @@ from lancedb.background_loop import LOOP
|
||||
from lancedb.arrow import AsyncRecordBatchReader
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.job import AsyncJob, Job
|
||||
from lancedb.listing import AsyncListing, Listing
|
||||
from lancedb.listing import (
|
||||
AsyncListing,
|
||||
AsyncNamespaceListing,
|
||||
AsyncTableListing,
|
||||
Listing,
|
||||
NamespaceListing,
|
||||
TableListing,
|
||||
)
|
||||
from lancedb.sql import AsyncQuery as AsyncSqlQuery
|
||||
from lancedb.sql import QueryDescription, QueryParameters
|
||||
from lance_namespace import (
|
||||
@@ -61,8 +68,6 @@ from lance_namespace import (
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
NamespaceExistsRequest,
|
||||
TableExistsRequest,
|
||||
)
|
||||
@@ -854,7 +859,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse:
|
||||
) -> NamespaceListing:
|
||||
"""
|
||||
List child namespaces under the given namespace.
|
||||
|
||||
@@ -871,12 +876,12 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListNamespacesResponse
|
||||
NamespaceListing
|
||||
Response containing namespace names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
return NamespaceListing(
|
||||
self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
@@ -1023,7 +1028,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
) -> TableListing:
|
||||
"""
|
||||
List all tables in this database with pagination support.
|
||||
|
||||
@@ -1040,12 +1045,12 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListTablesResponse
|
||||
TableListing
|
||||
Response containing table names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
return TableListing(
|
||||
self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
@@ -1459,12 +1464,12 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace_path = []
|
||||
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
||||
|
||||
async def list_namespaces(
|
||||
def list_namespaces(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse:
|
||||
) -> AsyncNamespaceListing:
|
||||
"""
|
||||
List child namespaces under the given namespace.
|
||||
|
||||
@@ -1481,12 +1486,12 @@ class AsyncLanceNamespaceDBConnection:
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListNamespacesResponse
|
||||
NamespaceListing
|
||||
Response containing namespace names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return await self._inner.list_namespaces(
|
||||
return self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
@@ -1616,12 +1621,12 @@ class AsyncLanceNamespaceDBConnection:
|
||||
except TableNotFoundError:
|
||||
return False
|
||||
|
||||
async def list_tables(
|
||||
def list_tables(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
) -> AsyncTableListing:
|
||||
"""
|
||||
List all tables in this database with pagination support.
|
||||
|
||||
@@ -1638,12 +1643,12 @@ class AsyncLanceNamespaceDBConnection:
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListTablesResponse
|
||||
TableListing
|
||||
Response containing table names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return await self._inner.list_tables(
|
||||
return self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
|
||||
@@ -36,7 +36,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig
|
||||
import pyarrow as pa
|
||||
|
||||
from ..common import DATA
|
||||
from lancedb.listing import Listing
|
||||
from lancedb.listing import Listing, NamespaceListing, TableListing
|
||||
from ..db import DBConnection, LOOP
|
||||
from ..functions import FunctionVersion, UdfDefinition
|
||||
from ..job import AsyncJob, Job
|
||||
@@ -54,8 +54,6 @@ from lance_namespace import (
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
)
|
||||
from ..pydantic import LanceModel
|
||||
from ..table import Table
|
||||
@@ -268,7 +266,7 @@ class RemoteDBConnection(DBConnection):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListNamespacesResponse:
|
||||
) -> NamespaceListing:
|
||||
"""List immediate child namespace names in the given namespace.
|
||||
|
||||
Parameters
|
||||
@@ -284,12 +282,12 @@ class RemoteDBConnection(DBConnection):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListNamespacesResponse
|
||||
NamespaceListing
|
||||
Response containing namespace names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
return NamespaceListing(
|
||||
self._conn.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
@@ -379,7 +377,7 @@ class RemoteDBConnection(DBConnection):
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
) -> TableListing:
|
||||
"""List a page of tables in this database.
|
||||
|
||||
Parameters
|
||||
@@ -396,12 +394,12 @@ class RemoteDBConnection(DBConnection):
|
||||
|
||||
Returns
|
||||
-------
|
||||
ListTablesResponse
|
||||
TableListing
|
||||
Response containing table names and optional page_token for pagination.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
return TableListing(
|
||||
self._conn.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
@@ -896,6 +896,8 @@ def test_transient_embedding_error_is_retried_in_table_add(tmp_path):
|
||||
assert len(calls) == 3
|
||||
assert mock_sleep.call_count == 2
|
||||
assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list)
|
||||
|
||||
|
||||
@patch("time.sleep")
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
@@ -941,6 +943,8 @@ def test_retry_on_wrapped_provider_error(mock_sleep):
|
||||
assert wrapped_func() == [0.1]
|
||||
assert mock_func.call_count == 2
|
||||
assert mock_sleep.call_count == 1
|
||||
|
||||
|
||||
def test_openai_bad_request_partial_batch_recovery():
|
||||
# Regression test for issue #1677
|
||||
|
||||
|
||||
@@ -88,6 +88,143 @@ async def advance(items, asynchronous):
|
||||
KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_compatible_name_listing_pages(listing_server, kind, asynchronous):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, response_page(kind, ["a", "b"], "next")),
|
||||
(200, response_page(kind, ["c"], None)),
|
||||
]
|
||||
)
|
||||
db = await connection(endpoint, asynchronous)
|
||||
listing = getattr(db, "list_" + kind)(
|
||||
namespace_path=["team"], page_token="start", limit=2
|
||||
)
|
||||
field = kind
|
||||
assert len(requests) == (0 if asynchronous else 1)
|
||||
if asynchronous:
|
||||
with pytest.raises(RuntimeError, match="Await or iterate"):
|
||||
getattr(listing, field)
|
||||
assert listing.page_token == "start"
|
||||
assert await listing is listing
|
||||
assert len(requests) == 1
|
||||
page_names = getattr(listing, field)
|
||||
assert type(page_names) is list
|
||||
assert page_names == ["a", "b"]
|
||||
page_names.append("changed")
|
||||
assert getattr(listing, field) == ["a", "b"]
|
||||
assert listing.page_token == "next"
|
||||
assert listing.num_page_results() == 2
|
||||
assert await advance(listing, asynchronous) == "a"
|
||||
assert getattr(listing, field) == ["a", "b"]
|
||||
assert await advance(listing, asynchronous) == "b"
|
||||
assert len(requests) == 1
|
||||
assert await advance(listing, asynchronous) == "c"
|
||||
assert getattr(listing, field) == ["c"]
|
||||
assert listing.page_token is None
|
||||
assert len(requests) == 2
|
||||
assert ([item async for item in listing] if asynchronous else list(listing)) == []
|
||||
if asynchronous:
|
||||
assert await listing is listing
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
|
||||
async def test_name_listing_skips_empty_page_only_during_iteration(
|
||||
listing_server, kind
|
||||
):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, response_page(kind, [], "next")),
|
||||
(200, response_page(kind, ["a"], None)),
|
||||
]
|
||||
)
|
||||
db = await connection(endpoint, True)
|
||||
listing = getattr(db, "list_" + kind)()
|
||||
await listing
|
||||
assert getattr(listing, kind) == []
|
||||
assert listing.page_token == "next"
|
||||
assert len(requests) == 1
|
||||
assert [name async for name in listing] == ["a"]
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_name_listing_zero_limit_uses_no_request(
|
||||
listing_server, kind, asynchronous
|
||||
):
|
||||
endpoint, requests, _ = listing_server
|
||||
db = await connection(endpoint, asynchronous)
|
||||
listing = getattr(db, "list_" + kind)(limit=0)
|
||||
if asynchronous:
|
||||
assert await listing is listing
|
||||
assert [name async for name in listing] == []
|
||||
else:
|
||||
assert list(listing) == []
|
||||
assert getattr(listing, kind) == []
|
||||
assert listing.page_token is None
|
||||
assert requests == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
|
||||
async def test_name_listing_failure_retains_token(listing_server, kind):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.append((401, {}))
|
||||
db = await connection(endpoint, True)
|
||||
listing = getattr(db, "list_" + kind)(page_token="retry")
|
||||
with pytest.raises(HttpError):
|
||||
await listing
|
||||
assert listing.page_token == "retry"
|
||||
with pytest.raises(HttpError):
|
||||
await listing
|
||||
assert [name async for name in listing] == []
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
|
||||
async def test_name_listing_direct_async_iteration(listing_server, kind):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, response_page(kind, ["a"], "next")),
|
||||
(200, response_page(kind, ["b"], None)),
|
||||
]
|
||||
)
|
||||
db = await connection(endpoint, True)
|
||||
listing = getattr(db, "list_" + kind)(limit=1)
|
||||
assert requests == []
|
||||
assert [name async for name in listing] == ["a", "b"]
|
||||
assert getattr(listing, kind) == ["b"]
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["tables", "namespaces"])
|
||||
async def test_name_listing_serializes_concurrent_advances(listing_server, kind):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.append((200, response_page(kind, ["a", "b"], None)))
|
||||
db = await connection(endpoint, True)
|
||||
listing = getattr(db, "list_" + kind)()
|
||||
results = await asyncio.gather(
|
||||
listing.__anext__(),
|
||||
listing.__anext__(),
|
||||
listing.__anext__(),
|
||||
return_exceptions=True,
|
||||
)
|
||||
assert results[:2] == ["a", "b"]
|
||||
assert isinstance(results[2], StopAsyncIteration)
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", KINDS)
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
@@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_embeddings():
|
||||
httpx = pytest.importorskip("httpx")
|
||||
pytest.importorskip("ollama")
|
||||
requests = []
|
||||
|
||||
|
||||
+16
-40
@@ -590,29 +590,17 @@ impl Connection {
|
||||
|
||||
#[pyo3(signature = (namespace_path=None, page_token=None, limit=None))]
|
||||
pub fn list_namespaces(
|
||||
self_: PyRef<'_, Self>,
|
||||
&self,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
page_token: Option<String>,
|
||||
limit: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let py = self_.py();
|
||||
future_into_py(py, async move {
|
||||
use lance_namespace::models::ListNamespacesRequest;
|
||||
let request = ListNamespacesRequest {
|
||||
id: namespace_path,
|
||||
page_token,
|
||||
limit: limit.map(|l| l as i32),
|
||||
..Default::default()
|
||||
};
|
||||
let response = inner.list_namespaces(request).await.infer_error()?;
|
||||
Python::attach(|py| -> PyResult<Py<PyDict>> {
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("namespaces", response.namespaces)?;
|
||||
dict.set_item("page_token", response.page_token)?;
|
||||
Ok(dict.unbind())
|
||||
})
|
||||
})
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
Ok(crate::listing::NameListing::new(
|
||||
self.get_inner()?.list_namespaces(
|
||||
&namespace_path.unwrap_or_default(),
|
||||
crate::listing::options(page_token, limit),
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (namespace_path, mode=None, properties=None))]
|
||||
@@ -725,29 +713,17 @@ impl Connection {
|
||||
|
||||
#[pyo3(signature = (namespace_path=None, page_token=None, limit=None))]
|
||||
pub fn list_tables(
|
||||
self_: PyRef<'_, Self>,
|
||||
&self,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
page_token: Option<String>,
|
||||
limit: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let py = self_.py();
|
||||
future_into_py(py, async move {
|
||||
use lance_namespace::models::ListTablesRequest;
|
||||
let request = ListTablesRequest {
|
||||
id: namespace_path,
|
||||
page_token,
|
||||
limit: limit.map(|l| l as i32),
|
||||
..Default::default()
|
||||
};
|
||||
let response = inner.list_tables(request).await.infer_error()?;
|
||||
Python::attach(|py| -> PyResult<Py<PyDict>> {
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("tables", response.tables)?;
|
||||
dict.set_item("page_token", response.page_token)?;
|
||||
Ok(dict.unbind())
|
||||
})
|
||||
})
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
Ok(crate::listing::NameListing::new(
|
||||
self.get_inner()?.list_tables(
|
||||
&namespace_path.unwrap_or_default(),
|
||||
crate::listing::options(page_token, limit),
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
/// Get the configuration for constructing an equivalent namespace client.
|
||||
|
||||
+28
-2
@@ -3,10 +3,11 @@
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
use futures::{StreamExt, future::poll_fn};
|
||||
use futures::{Future, StreamExt, future::poll_fn};
|
||||
use lancedb::listing::{Listing, ListingOptions};
|
||||
use pyo3::exceptions::PyStopAsyncIteration;
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3::types::{PyDict, PyDictMethods};
|
||||
use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
pub fn options(page_token: Option<String>, page_limit: Option<u32>) -> ListingOptions {
|
||||
@@ -52,6 +53,31 @@ impl NameListing {
|
||||
Ok(value)
|
||||
})
|
||||
}
|
||||
|
||||
fn next_page(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _guard = next_lock.lock().await;
|
||||
let page = futures::future::poll_fn(|cx| {
|
||||
let mut inner = inner.lock().unwrap();
|
||||
let future = inner.next_page();
|
||||
futures::pin_mut!(future);
|
||||
future.poll(cx)
|
||||
})
|
||||
.await
|
||||
.infer_error()?;
|
||||
Python::attach(|py| -> PyResult<Option<Py<PyDict>>> {
|
||||
page.map(|page| {
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("items", page.items)?;
|
||||
dict.set_item("page_token", page.page_token)?;
|
||||
Ok(dict.unbind())
|
||||
})
|
||||
.transpose()
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
|
||||
Reference in New Issue
Block a user