refactor!: implement lazy iterators for several APIs (#4459)

Implement lazy iterators for:
- secrets
- views
- jobs
- functions
- materialized views

python/src/connection.rs

@@ -464,12 +464,16 @@ impl Connection {
         })
     }
 
-    pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.get_inner()?.clone();
-        future_into_py(self_.py(), async move {
-            let views = inner.list_materialized_views().await.infer_error()?;
-            Ok(views)
-        })
+    #[pyo3(signature = (*, page_token=None, page_limit=None))]
+    pub fn list_materialized_views(
+        &self,
+        page_token: Option<String>,
+        page_limit: Option<u32>,
+    ) -> PyResult<crate::listing::NameListing> {
+        Ok(crate::listing::NameListing::new(
+            self.get_inner()?
+                .list_materialized_views(crate::listing::options(page_token, page_limit)),
+        ))
     }
 
     #[pyo3(signature = (name, namespace_path=None))]
@@ -776,22 +780,19 @@ impl Connection {
         })
     }
 
-    #[pyo3(signature = (namespace_path=None))]
+    #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
     pub fn list_functions(
-        self_: PyRef<'_, Self>,
+        &self,
         namespace_path: Option<Vec<String>>,
-    ) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.get_inner()?.clone();
-        let namespace_path = namespace_path.unwrap_or_default();
-        future_into_py(self_.py(), async move {
-            inner
-                .list_functions(&namespace_path)
-                .await
-                .infer_error()?
-                .into_iter()
-                .map(|function| function.to_canonical_json().infer_error())
-                .collect::<PyResult<Vec<_>>>()
-        })
+        page_token: Option<String>,
+        page_limit: Option<u32>,
+    ) -> PyResult<crate::listing::FunctionListing> {
+        Ok(crate::listing::FunctionListing::new(
+            self.get_inner()?.list_functions(
+                &namespace_path.unwrap_or_default(),
+                crate::listing::options(page_token, page_limit),
+            ),
+        ))
     }
 
     #[pyo3(signature = (name, version, namespace_path=None))]
@@ -863,16 +864,19 @@ impl Connection {
         })
     }
 
-    #[pyo3(signature = (namespace_path=None))]
+    #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
     pub fn list_secrets(
-        self_: PyRef<'_, Self>,
+        &self,
         namespace_path: Option<Vec<String>>,
-    ) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.get_inner()?.clone();
-        let namespace_path = namespace_path.unwrap_or_default();
-        future_into_py(self_.py(), async move {
-            inner.list_secrets(&namespace_path).await.infer_error()
-        })
+        page_token: Option<String>,
+        page_limit: Option<u32>,
+    ) -> PyResult<crate::listing::NameListing> {
+        Ok(crate::listing::NameListing::new(
+            self.get_inner()?.list_secrets(
+                &namespace_path.unwrap_or_default(),
+                crate::listing::options(page_token, page_limit),
+            ),
+        ))
     }
 
     #[pyo3(signature = (name, namespace_path=None))]
@@ -973,27 +977,31 @@ impl Connection {
         })
     }
 
-    #[pyo3(signature = (namespace_path=None))]
+    #[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
     pub fn list_views(
-        self_: PyRef<'_, Self>,
+        &self,
         namespace_path: Option<Vec<String>>,
-    ) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.get_inner()?.clone();
-        let namespace_path = namespace_path.unwrap_or_default();
-        future_into_py(self_.py(), async move {
-            inner.list_views(&namespace_path).await.infer_error()
-        })
+        page_token: Option<String>,
+        page_limit: Option<u32>,
+    ) -> PyResult<crate::listing::NameListing> {
+        Ok(crate::listing::NameListing::new(
+            self.get_inner()?.list_views(
+                &namespace_path.unwrap_or_default(),
+                crate::listing::options(page_token, page_limit),
+            ),
+        ))
     }
 
-    pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
-        let inner = self_.get_inner()?.clone();
-        future_into_py(self_.py(), async move {
-            let jobs = inner.list_jobs().await.infer_error()?;
-            Ok(jobs
-                .into_iter()
-                .map(crate::job::JobInfo::from)
-                .collect::<Vec<_>>())
-        })
+    #[pyo3(signature = (*, page_token=None, page_limit=None))]
+    pub fn list_jobs(
+        &self,
+        page_token: Option<String>,
+        page_limit: Option<u32>,
+    ) -> PyResult<crate::listing::JobListing> {
+        Ok(crate::listing::JobListing::new(
+            self.get_inner()?
+                .list_jobs(crate::listing::options(page_token, page_limit)),
+        ))
     }
 
     pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {

python/src/lib.rs

@@ -29,6 +29,7 @@ pub mod expr;
 pub mod header;
 pub mod index;
 pub mod job;
+mod listing;
 pub mod namespace;
 pub mod oauth;
 pub mod otel;
@@ -65,6 +66,10 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<Connection>()?;
     m.add_class::<catalog::Catalog>()?;
     m.add_class::<catalog::DatabaseNames>()?;
+    m.add_class::<listing::NameListing>()?;
+    m.add_class::<listing::FunctionListing>()?;
+    m.add_class::<listing::JobListing>()?;
+
     m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
     m.add_class::<Session>()?;
     m.add_class::<Table>()?;

python/src/listing.rs

@@ -0,0 +1,131 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+use crate::error::PythonErrorExt;
+use crate::runtime::future_into_py;
+use futures::{StreamExt, future::poll_fn};
+use lancedb::listing::{Listing, ListingOptions};
+use pyo3::exceptions::PyStopAsyncIteration;
+use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
+use std::sync::{Arc, Mutex};
+
+pub fn options(page_token: Option<String>, page_limit: Option<u32>) -> ListingOptions {
+    let mut options = ListingOptions::default();
+    options.page_token = page_token;
+    options.page_limit = page_limit;
+    options
+}
+
+#[pyclass]
+pub struct NameListing {
+    inner: Arc<Mutex<Listing<String>>>,
+    next_lock: Arc<tokio::sync::Mutex<()>>,
+}
+impl NameListing {
+    pub fn new(inner: Listing<String>) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(inner)),
+            next_lock: Arc::new(tokio::sync::Mutex::new(())),
+        }
+    }
+}
+#[pymethods]
+impl NameListing {
+    fn num_page_results(&self, py: Python<'_>) -> usize {
+        py.detach(|| self.inner.lock().unwrap().num_page_results())
+    }
+    fn page_token(&self, py: Python<'_>) -> Option<String> {
+        py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
+    }
+    fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        self_
+    }
+    fn __anext__(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 value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
+                .await
+                .ok_or_else(|| PyStopAsyncIteration::new_err(""))?
+                .infer_error()?;
+            Ok(value)
+        })
+    }
+}
+
+#[pyclass]
+pub struct FunctionListing {
+    inner: Arc<Mutex<Listing<lancedb::function::FunctionVersion>>>,
+    next_lock: Arc<tokio::sync::Mutex<()>>,
+}
+impl FunctionListing {
+    pub fn new(inner: Listing<lancedb::function::FunctionVersion>) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(inner)),
+            next_lock: Arc::new(tokio::sync::Mutex::new(())),
+        }
+    }
+}
+#[pymethods]
+impl FunctionListing {
+    fn num_page_results(&self, py: Python<'_>) -> usize {
+        py.detach(|| self.inner.lock().unwrap().num_page_results())
+    }
+    fn page_token(&self, py: Python<'_>) -> Option<String> {
+        py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
+    }
+    fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        self_
+    }
+    fn __anext__(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 value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
+                .await
+                .ok_or_else(|| PyStopAsyncIteration::new_err(""))?
+                .infer_error()?;
+            value.to_canonical_json().infer_error()
+        })
+    }
+}
+
+#[pyclass]
+pub struct JobListing {
+    inner: Arc<Mutex<Listing<lancedb::database::JobInfo>>>,
+    next_lock: Arc<tokio::sync::Mutex<()>>,
+}
+impl JobListing {
+    pub fn new(inner: Listing<lancedb::database::JobInfo>) -> Self {
+        Self {
+            inner: Arc::new(Mutex::new(inner)),
+            next_lock: Arc::new(tokio::sync::Mutex::new(())),
+        }
+    }
+}
+#[pymethods]
+impl JobListing {
+    fn num_page_results(&self, py: Python<'_>) -> usize {
+        py.detach(|| self.inner.lock().unwrap().num_page_results())
+    }
+    fn page_token(&self, py: Python<'_>) -> Option<String> {
+        py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
+    }
+    fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        self_
+    }
+    fn __anext__(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 value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
+                .await
+                .ok_or_else(|| PyStopAsyncIteration::new_err(""))?
+                .infer_error()?;
+            Ok(crate::job::JobInfo::from(value))
+        })
+    }
+}