seaweedfs/seaweed-worker/crates/lance/src/dataset.rs

//! Opening a Lance dataset with credentials the namespace vended.
//!
//! The worker holds no storage credentials of its own: it asks the namespace to
//! describe a table with `vend_credentials`, and the `storage_options` that come
//! back are handed to lance as-is. They expire, so a job that runs longer than
//! their lifetime re-describes rather than carrying one set throughout.

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::{Context, Result};
use lance::dataset::Dataset;
use lance::dataset::builder::DatasetBuilder;
use lance::io::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor};

use crate::catalog::{NamespaceClient, TableDescription};

/// A table the worker is about to work on.
pub struct OpenTable {
    pub id: Vec<String>,
    pub location: String,
    pub dataset: Dataset,
}

/// What detection needs to decide whether a table is worth a job. Reading it
/// opens the dataset but touches no data files.
pub struct TableStats {
    pub fragments: usize,
    pub version: u64,
    pub total_versions: usize,
    pub rows: usize,
    /// The Arrow schema as JSON, which is the only description of this table
    /// anything outside the format can produce.
    pub schema: Option<String>,
}

/// Storage options an operator supplies for deployments that vend none.
pub type FallbackOptions = HashMap<String, String>;

/// Resolve a table through the namespace and open it.
pub async fn open(
    client: &NamespaceClient,
    id: &[String],
    fallback: &FallbackOptions,
) -> Result<OpenTable> {
    let description = client.describe_table(id).await?;
    let dataset = open_at(&description, fallback).await?;
    Ok(OpenTable {
        id: id.to_vec(),
        location: description.location,
        dataset,
    })
}

/// The namespace vends object_store's own option names, so they pass straight
/// through. What it vends always wins: the fallback exists because a deployment
/// without STS vends no credentials at all, and then the worker has no other way
/// to reach the data.
async fn open_at(description: &TableDescription, fallback: &FallbackOptions) -> Result<Dataset> {
    let mut options: HashMap<String, String> = fallback.clone();
    options.extend(description.storage_options.clone());
    DatasetBuilder::from_uri(&description.location)
        .with_storage_options(options)
        .load()
        .await
        .with_context(|| format!("open lance dataset at {}", description.location))
}

/// Remove every object under a table's location. A table the catalog can no
/// longer describe is already dropped, so the call stays idempotent.
pub async fn drop(
    client: &NamespaceClient,
    id: &[String],
    fallback: &FallbackOptions,
) -> Result<()> {
    let description = match client.describe_table(id).await {
        Ok(description) => description,
        Err(err) if describe_not_found(&err) => return Ok(()),
        Err(err) => return Err(err),
    };
    let mut options: HashMap<String, String> = fallback.clone();
    options.extend(description.storage_options.clone());
    let (store, path) = ObjectStore::from_uri_and_params(
        Arc::new(ObjectStoreRegistry::default()),
        &description.location,
        &ObjectStoreParams {
            storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
                options,
            ))),
            ..Default::default()
        },
    )
    .await?;
    store.remove_dir_all(path).await?;
    Ok(())
}

fn describe_not_found(err: &anyhow::Error) -> bool {
    err.chain().any(|cause| {
        cause
            .downcast_ref::<reqwest::Error>()
            .and_then(|e| e.status())
            == Some(reqwest::StatusCode::NOT_FOUND)
    })
}

impl OpenTable {
    pub async fn stats(&self) -> Result<TableStats> {
        let versions = self.dataset.versions().await?;
        Ok(TableStats {
            fragments: self.dataset.get_fragments().len(),
            version: self.dataset.version().version,
            total_versions: versions.len(),
            rows: self.dataset.count_rows(None).await.unwrap_or(0),
            schema: schema_json(&self.dataset),
        })
    }
}

/// Renders the dataset's schema as JSON. Best effort: a schema that will not
/// serialise is not a reason to fail a maintenance sweep.
fn schema_json(dataset: &Dataset) -> Option<String> {
    let arrow: arrow_schema::Schema = dataset.schema().into();
    serde_json::to_string(
        &arrow
            .fields()
            .iter()
            .map(|f| {
                serde_json::json!({
                    "name": f.name(),
                    "type": f.data_type().to_string(),
                    "nullable": f.is_nullable(),
                })
            })
            .collect::<Vec<_>>(),
    )
    .ok()
}