From d46bc5dd6e1485b4e98014039075ccc34cabecf8 Mon Sep 17 00:00:00 2001
From: QianZhu <qian@lancedb.com>
Date: Mon, 16 Oct 2023 21:09:20 -0700
Subject: [PATCH] list table pagination draft (#574)

---
 python/lancedb/remote/client.py | 11 ++++++++---
 python/lancedb/remote/db.py     | 32 +++++++++++++++++++++++++++-----
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/python/lancedb/remote/client.py b/python/lancedb/remote/client.py
index 375b7894..24aa90df 100644
--- a/python/lancedb/remote/client.py
+++ b/python/lancedb/remote/client.py
@@ -151,10 +151,15 @@ class RestfulLanceDBClient:
             return await deserialize(resp)
 
     @_check_not_closed
-    async def list_tables(self):
+    async def list_tables(self, limit: int, page_token: str):
         """List all tables in the database."""
-        json = await self.get("/v1/table/", {})
-        return json["tables"]
+        try:
+            json = await self.get(
+                "/v1/table/", {"limit": limit, "page_token": page_token}
+            )
+            return json["tables"]
+        except StopAsyncIteration:
+            return []
 
     @_check_not_closed
     async def query(self, table_name: str, query: VectorQuery) -> VectorQueryResult:
diff --git a/python/lancedb/remote/db.py b/python/lancedb/remote/db.py
index 89d7a3f5..f087ddc5 100644
--- a/python/lancedb/remote/db.py
+++ b/python/lancedb/remote/db.py
@@ -13,7 +13,7 @@
 
 import asyncio
 import uuid
-from typing import List, Optional
+from typing import Iterator, Optional
 from urllib.parse import urlparse
 
 import pyarrow as pa
@@ -52,10 +52,27 @@ class RemoteDBConnection(DBConnection):
     def __repr__(self) -> str:
         return f"RemoveConnect(name={self.db_name})"
 
-    def table_names(self) -> List[str]:
-        """List the names of all tables in the database."""
-        result = self._loop.run_until_complete(self._client.list_tables())
-        return result
+    def table_names(self, last_token: str, limit=10) -> Iterator[str]:
+        """List the names of all tables in the database.
+        Parameters
+        ----------
+        last_token: str
+            The last token to start the new page.
+
+        Returns
+        -------
+        An iterator of table names.
+        """
+        while True:
+            result = self._loop.run_until_complete(
+                self._client.list_tables(limit, last_token)
+            )
+            if len(result) > 0:
+                last_token = result[len(result) - 1]
+            else:
+                break
+            for item in result:
+                yield result
 
     def open_table(self, name: str) -> Table:
         """Open a Lance Table in the database.
@@ -122,3 +139,8 @@ class RemoteDBConnection(DBConnection):
                 f"/v1/table/{name}/drop/",
             )
         )
+
+    async def close(self):
+        """Close the connection to the database."""
+        self._loop.close()
+        await self._client.close()