lancedb 0.41.0

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

//! Catalogs manage databases. A remote catalog is the root namespace of a server.
//!
//! ```
//! # #[cfg(feature = "remote")]
//! # async fn example() -> lancedb::Result<()> {
//! let catalog = lancedb::connect_catalog("https://my-server.example")
//!     .api_key("my-api-key")
//!     .execute().await?;
//! let database = catalog.create_database("analytics").await?;
//! # Ok(())
//! # }
//! ```

use std::fmt;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use futures::future::BoxFuture;
use futures::{FutureExt, Stream};

use crate::Result;
use crate::connection::Connection;
use crate::database::Database;
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};

/// Options for creating a database. By default an existing name is an error.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct CreateDatabaseRequest {
    /// Logical database name, including any literal slashes.
    pub name: String,
    /// Open the existing database if it is already registered.
    pub exist_ok: bool,
}

impl CreateDatabaseRequest {
    /// Initialize a request with the default behavior.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            exist_ok: false,
        }
    }

    /// Open an existing database instead of failing if its name already exists.
    pub fn exist_ok(mut self, value: bool) -> Self {
        self.exist_ok = value;
        self
    }
}

impl<T: Into<String>> From<T> for CreateDatabaseRequest {
    fn from(name: T) -> Self {
        Self::new(name)
    }
}

/// Options for database deletion. The database must be empty unless `cascade` is set.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct DropDatabaseRequest {
    /// Logical database name, including any literal slashes.
    pub name: String,
    /// Succeed if the database is absent.
    pub ignore_missing: bool,
    /// Drop everything in the database first.
    pub cascade: bool,
}

impl DropDatabaseRequest {
    /// Initialize a request with the default behavior.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            ignore_missing: false,
            cascade: false,
        }
    }

    /// Succeed when the database does not exist.
    pub fn ignore_missing(mut self, value: bool) -> Self {
        self.ignore_missing = value;
        self
    }

    /// Drop every table, view, function, secret and namespace in the database first. The
    /// server cleans up their storage in the background, and the name stays taken until then.
    pub fn cascade(mut self, value: bool) -> Self {
        self.cascade = value;
        self
    }
}

impl<T: Into<String>> From<T> for DropDatabaseRequest {
    fn from(name: T) -> Self {
        Self::new(name)
    }
}

/// Pagination options for listing databases.
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct ListDatabasesRequest {
    /// Maximum number of names to return. None uses the server default.
    pub limit: Option<u32>,
    /// Opaque continuation token from a previous response. None starts a listing.
    pub page_token: Option<String>,
}

impl ListDatabasesRequest {
    /// Set the maximum page size (1 through 2147483647).
    pub fn limit(mut self, limit: u32) -> Self {
        self.limit = Some(limit);
        self
    }
    /// Resume a listing from an opaque continuation token.
    pub fn page_token(mut self, token: impl Into<String>) -> Self {
        self.page_token = Some(token.into());
        self
    }
}

/// One page of database names, relative to the catalog.
#[derive(Clone, Debug, Default)]
pub struct ListDatabasesResponse {
    /// Logical database names on this page.
    pub databases: Vec<String>,
    /// None indicates the end of the listing.
    pub page_token: Option<String>,
}

/// A lazy stream of database names with inspectable pagination state.
///
/// The continuation token advances when a page is fetched, so resuming from it
/// skips any names still cached in this stream. Drain the cache before saving a
/// token if those names must be included. Errors terminate the stream; the token
/// for the failed request remains available for resuming with a new stream.
pub struct DatabaseNames {
    catalog: Arc<dyn Catalog>,
    names: std::vec::IntoIter<String>,
    page_token: Option<String>,
    page_limit: Option<u32>,
    done: bool,
    pending: Option<BoxFuture<'static, Result<ListDatabasesResponse>>>,
}

impl DatabaseNames {
    /// Number of names available without another REST request.
    pub fn num_page_results(&self) -> usize {
        self.names.len()
    }

    /// Token for the next REST request, or `None` after the last page is fetched.
    ///
    /// Before the first request, returns the supplied starting token (`None`
    /// means start from the beginning). A final page can still have cached names.
    pub fn page_token(&self) -> Option<&str> {
        self.page_token.as_deref()
    }
}

impl Stream for DatabaseNames {
    type Item = Result<String>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        loop {
            if let Some(name) = this.names.next() {
                return Poll::Ready(Some(Ok(name)));
            }
            if this.done {
                return Poll::Ready(None);
            }
            let pending = this.pending.get_or_insert_with(|| {
                let catalog = this.catalog.clone();
                let request = ListDatabasesRequest {
                    page_token: this.page_token.clone(),
                    limit: this.page_limit,
                };
                async move { catalog.list_databases(request).await }.boxed()
            });
            let response = futures::ready!(pending.as_mut().poll(cx));
            this.pending = None;
            match response {
                Ok(response) => {
                    this.page_token = response.page_token.filter(|token| !token.is_empty());
                    this.done = this.page_token.is_none();
                    this.names = response.databases.into_iter();
                }
                Err(error) => {
                    this.done = true;
                    return Poll::Ready(Some(Err(error)));
                }
            }
        }
    }
}

/// A backend that manages databases. Implementations own database lifecycle semantics.
#[async_trait::async_trait]
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
    /// Obtain deployment-wide authorization, if supported by this backend.
    fn authz(&self) -> Result<Arc<dyn crate::authz::Authorization>> {
        Err(crate::Error::NotSupported {
            message: "Authorization administration is not supported by this catalog".into(),
        })
    }

    /// Catalog endpoint or location.
    fn uri(&self) -> &str;
    /// Create a database, or open it when `exist_ok` permits.
    async fn create_database(&self, request: CreateDatabaseRequest) -> Result<Arc<dyn Database>>;
    /// Drop a database, which must be empty unless the request cascades.
    async fn drop_database(&self, request: DropDatabaseRequest) -> Result<()>;
    /// List a page of database names.
    async fn list_databases(&self, request: ListDatabasesRequest) -> Result<ListDatabasesResponse>;
    /// Connect to an existing database by its logical name.
    async fn connect_database(&self, name: &str) -> Result<Arc<dyn Database>>;
}

/// A catalog connection that returns ordinary LanceDB database connections.
#[derive(Clone)]
pub struct CatalogConnection {
    catalog: Arc<dyn Catalog>,
    embedding_registry: Arc<dyn EmbeddingRegistry>,
}

impl fmt::Debug for CatalogConnection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CatalogConnection")
            .field("uri", &self.uri())
            .finish_non_exhaustive()
    }
}

impl CatalogConnection {
    /// Wrap a catalog implementation using the default in-memory embedding registry.
    pub fn new(catalog: Arc<dyn Catalog>) -> Self {
        Self {
            catalog,
            embedding_registry: Arc::new(MemoryRegistry::new()),
        }
    }

    /// Provide the registry used by databases opened through this connection.
    pub fn with_embedding_registry(mut self, registry: Arc<dyn EmbeddingRegistry>) -> Self {
        self.embedding_registry = registry;
        self
    }

    /// The catalog endpoint or location.
    pub fn uri(&self) -> &str {
        self.catalog.uri()
    }
    /// Obtain a handle sharing this catalog's deployment scope and credentials.
    pub fn authz(&self) -> Result<Arc<dyn crate::authz::Authorization>> {
        self.catalog.authz()
    }

    /// Access the underlying backend.
    pub fn catalog(&self) -> &Arc<dyn Catalog> {
        &self.catalog
    }

    /// Create a database; pass a name or [`CreateDatabaseRequest`] for additional options.
    pub async fn create_database(
        &self,
        request: impl Into<CreateDatabaseRequest>,
    ) -> Result<Connection> {
        Ok(Connection::new(
            self.catalog.create_database(request.into()).await?,
            self.embedding_registry.clone(),
        ))
    }

    /// Drop a database; pass a name or [`DropDatabaseRequest`] for options such as `cascade`.
    pub async fn drop_database(&self, request: impl Into<DropDatabaseRequest>) -> Result<()> {
        self.catalog.drop_database(request.into()).await
    }

    /// Iterate lazily over all database names, fetching more results as needed.
    ///
    /// `page_token` starts at a saved continuation token; `None` starts at the
    /// beginning. `page_limit` limits each REST response (not the whole stream);
    /// `None` uses the server default. Errors terminate the listing.
    ///
    /// ```
    /// # async fn example(catalog: &lancedb::catalog::CatalogConnection) -> lancedb::Result<()> {
    /// use futures::TryStreamExt;
    /// let mut names = catalog.list_databases(None, None);
    /// while let Some(name) = names.try_next().await? {
    ///     println!("{name}");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn list_databases(
        &self,
        page_token: Option<String>,
        page_limit: Option<u32>,
    ) -> DatabaseNames {
        DatabaseNames {
            catalog: self.catalog.clone(),
            names: Vec::new().into_iter(),
            page_token,
            page_limit,
            done: false,
            pending: None,
        }
    }

    /// Connect to a database without creating it if it is missing.
    pub async fn connect_database(&self, name: impl AsRef<str>) -> Result<Connection> {
        Ok(Connection::new(
            self.catalog.connect_database(name.as_ref()).await?,
            self.embedding_registry.clone(),
        ))
    }
}

/// Configure a connection to a remote catalog.
#[cfg(feature = "remote")]
#[derive(Debug)]
pub struct ConnectCatalogBuilder {
    endpoint: String,
    options: crate::remote::RemoteCatalogOptions,
}

#[cfg(feature = "remote")]
impl ConnectCatalogBuilder {
    /// Start configuring a remote HTTP(S) catalog connection.
    pub fn new(endpoint: impl Into<String>) -> Self {
        Self {
            endpoint: endpoint.into(),
            options: Default::default(),
        }
    }
    /// Authenticate with an API key.
    pub fn api_key(mut self, key: impl Into<String>) -> Self {
        self.options.api_key = Some(key.into());
        self
    }
    /// Configure headers, TLS, timeouts, and other shared client settings.
    pub fn client_config(mut self, config: crate::remote::ClientConfig) -> Self {
        self.options.client_config = config;
        self
    }
    /// Set the SQL service endpoint inherited by database connections.
    ///
    /// Required to execute SQL when the catalog endpoint uses HTTPS. The SQL
    /// connection is initialized lazily, using the ordinary remote SQL client.
    pub fn sql_host_override(mut self, endpoint: impl Into<String>) -> Self {
        self.options.sql_host_override = Some(endpoint.into());
        self
    }
    /// Configure table read consistency for opened databases.
    pub fn read_consistency_interval(mut self, interval: std::time::Duration) -> Self {
        self.options.read_consistency_interval = Some(interval);
        self
    }
    /// Authenticate using OAuth; mutually exclusive with API keys and header providers.
    pub fn oauth_config(mut self, config: crate::remote::OAuthConfig) -> Self {
        self.options.oauth_config = Some(config);
        self
    }
    /// Connect to the server's root namespace. Database-scoped headers are omitted.
    pub async fn execute(self) -> Result<CatalogConnection> {
        Ok(CatalogConnection::new(Arc::new(
            crate::remote::RemoteCatalog::try_new(&self.endpoint, self.options)?,
        )))
    }
}