crates/knut-popviz/src/lake.rs
//! **Where the Iceberg target's tables are** — a local skade warehouse directory, or
//! an Iceberg REST catalog: MinIO AIStor's `/_iceberg` (SigV4 service `s3tables`),
//! Nessie, Polaris. Generate ([`crate::run_load`]) and permutate
//! ([`crate::permutate::permutate_lake`]) write through it, and bifrost reads the same
//! tables through the same catalog — the population, Spark-shaped, in the customer's
//! own lake.
//!
//! Table names are bare (`Person`, `rel_LIVES_AT`) in ONE namespace: `main` in a local
//! warehouse (skade's default), [`RestLake::namespace`] in a catalog (default `main`
//! too, so one mapping file serves both lakes).
//!
//! The S3 keys come from `MINIO_ACCESS_KEY` / `MINIO_SECRET_KEY` only — a flag lands in
//! `/proc/<pid>/cmdline`. They sign the data files and, with SigV4, every catalog call.
use std::path::PathBuf;
/// An Iceberg REST catalog and the S3 store its tables' files are in.
#[derive(Debug, Clone, PartialEq)]
pub struct RestLake {
/// The catalog root, e.g. `https://minio:9000/_iceberg`.
pub uri: String,
/// The REST `warehouse` parameter — for AIStor the warehouse (= bucket) name.
pub warehouse: Option<String>,
/// The namespace every table is in (default `main`).
pub namespace: String,
/// The S3 endpoint of the data files, e.g. `https://minio:9000`.
pub s3_endpoint: String,
pub s3_region: String,
/// The data bucket; `None` takes [`RestLake::warehouse`].
pub s3_bucket: Option<String>,
/// A PEM CA for an `https://` catalog and S3 endpoint with a private CA.
pub s3_ca: Option<PathBuf>,
/// SigV4 service to sign catalog calls with — `s3tables` for AIStor.
pub sigv4: Option<String>,
}
/// Where the tables are.
#[derive(Debug, Clone, PartialEq)]
pub enum Lake {
/// A local skade warehouse directory.
Dir(PathBuf),
/// An Iceberg REST catalog (feature `iceberg-rest`).
Rest(RestLake),
}
impl Lake {
/// For logs and the `graph` field of a report. Never a secret.
pub fn describe(&self) -> String {
match self {
Lake::Dir(d) => format!("iceberg:{}", d.display()),
Lake::Rest(r) => format!(
"iceberg-rest:{}#{}{}",
r.uri,
r.warehouse.as_deref().map(|w| format!("{w}/")).unwrap_or_default(),
r.namespace
),
}
}
}
#[cfg(feature = "iceberg")]
mod open {
use std::sync::Arc;
use arrow_array::RecordBatch;
use skade::iceberg::{Catalog, NamespaceIdent, TableIdent};
use super::Lake;
/// An opened [`Lake`]: the warehouse, or the catalog and the namespace.
pub(crate) enum Open {
Dir(skade::Warehouse),
Rest { catalog: Arc<dyn Catalog>, ns: NamespaceIdent },
}
/// The format a table the target creates in a catalog gets: v2, which every REST
/// catalog (AIStor included) and every reader takes. A local warehouse keeps skade's
/// own default.
const REST_FORMAT: skade::iceberg::spec::FormatVersion = skade::iceberg::spec::FormatVersion::V2;
impl Lake {
pub(crate) async fn open(&self) -> Result<Open, String> {
match self {
Lake::Dir(dir) => Ok(Open::Dir(skade::Warehouse::open(dir).await.map_err(|e| format!("{}: {e}", dir.display()))?)),
Lake::Rest(r) => {
let ns = NamespaceIdent::from_strs(r.namespace.split('.')).map_err(|e| format!("namespace {:?}: {e}", r.namespace))?;
Ok(Open::Rest { catalog: rest_catalog(r).await?, ns })
}
}
}
}
#[cfg(feature = "iceberg-rest")]
async fn rest_catalog(r: &super::RestLake) -> Result<Arc<dyn Catalog>, String> {
use skade::connect::CatalogConnector as _;
let key = |k: &str| std::env::var(k).ok().filter(|v| !v.is_empty()).ok_or_else(|| format!("{k} is not set - the lake's keys come from the environment only"));
let bucket = r.s3_bucket.clone().or_else(|| r.warehouse.clone()).ok_or("--s3-bucket is required (no --rest-warehouse to take it from)")?;
let s3 = skade::S3Config {
endpoint: r.s3_endpoint.clone(),
region: r.s3_region.clone(),
bucket,
access_key_id: key("MINIO_ACCESS_KEY")?,
secret_access_key: key("MINIO_SECRET_KEY")?,
path_style: true,
tls: r.s3_ca.clone().map(|ca| skade::tls::ClientTls { ca_pem: Some(ca), ..Default::default() }),
};
let rest = skade::connect::IcebergRest {
uri: r.uri.clone(),
warehouse: r.warehouse.clone(),
prefix_or_ref: None,
token: None,
tls: s3.tls.clone(),
s3,
read_only: false,
sigv4: r.sigv4.clone().map(|service| skade::connect::SigV4 { service }),
..Default::default()
};
let d = rest.describe();
rest.open().await.map_err(|e| format!("{d}: {e}"))
}
#[cfg(not(feature = "iceberg-rest"))]
async fn rest_catalog(r: &super::RestLake) -> Result<Arc<dyn Catalog>, String> {
Err(format!("{}: this build has no REST catalog client - rebuild with --features iceberg-rest", r.uri))
}
impl Open {
fn ident(&self, table: &str) -> Result<TableIdent, String> {
match self {
Open::Dir(wh) => wh.table_ident(table).map_err(|e| e.to_string()),
Open::Rest { ns, .. } => Ok(TableIdent::new(ns.clone(), table.to_string())),
}
}
async fn load(&self, table: &str) -> Result<skade::iceberg::table::Table, String> {
let e = |e: &dyn std::fmt::Display| format!("{table}: {e}");
match self {
Open::Dir(wh) => Ok(wh.table(table).await.map_err(|x| e(&x))?.inner().clone()),
Open::Rest { catalog, .. } => catalog.load_table(&self.ident(table)?).await.map_err(|x| e(&x)),
}
}
/// Append `batch` to `table`, creating the table from the batch's schema first if
/// absent. One commit. Returns the rows appended.
pub(crate) async fn append(&self, table: &str, batch: RecordBatch) -> Result<usize, String> {
let rows = batch.num_rows();
if rows == 0 {
return Ok(0);
}
let e = |e: &dyn std::fmt::Display| format!("{table}: {e}");
match self {
Open::Dir(wh) => {
let mut t = wh.table_or_create(table, &batch.schema()).await.map_err(|x| e(&x))?;
t.append(&[batch]).await.map_err(|x| e(&x))?;
}
Open::Rest { catalog, .. } => {
let t = skade::table_or_create_in(catalog.as_ref(), &self.ident(table)?, &batch.schema(), REST_FORMAT).await.map_err(|x| e(&x))?;
skade::append(catalog.as_ref(), &t, &[batch]).await.map_err(|x| e(&x))?;
}
}
Ok(rows)
}
/// Every live row of `table`.
pub(crate) async fn read(&self, table: &str) -> Result<Vec<RecordBatch>, String> {
let t = self.load(table).await?;
skade::read_all(&t).await.map_err(|e| format!("{table}: {e}"))
}
/// The named columns of every live row of `table`.
pub(crate) async fn read_columns(&self, table: &str, columns: &[&str]) -> Result<Vec<RecordBatch>, String> {
let t = self.load(table).await?;
skade::read_columns(&t, columns).await.map_err(|e| format!("{table}: {e}"))
}
/// Replace every live data file of `table` with `batches`, in ONE `overwrite`
/// commit — Spark's copy-on-write UPDATE/DELETE shape.
pub(crate) async fn overwrite_all(&self, table: &str, batches: Vec<RecordBatch>) -> Result<(), String> {
let batches: Vec<RecordBatch> = batches.into_iter().filter(|b| b.num_rows() > 0).collect();
let e = |e: &dyn std::fmt::Display| format!("{table}: {e}");
match self {
Open::Dir(wh) => {
let mut t = wh.table(table).await.map_err(|x| e(&x))?;
let live = t.live_data_files().await.map_err(|x| e(&x))?;
t.overwrite(&live, &batches).await.map_err(|x| e(&x))
}
Open::Rest { catalog, .. } => {
let t = self.load(table).await?;
let live = skade::live_data_files(&t).await.map_err(|x| e(&x))?;
skade::overwrite(catalog.as_ref(), &t, &live, &batches).await.map(|_| ()).map_err(|x| e(&x))
}
}
}
/// Drop `tables` from a catalog lake's namespace, those that exist — Generate's
/// `--drop`, which names exactly the tables it writes, so a namespace shared with
/// anything else keeps everything else. Returns the tables dropped. A local
/// warehouse is wiped as a directory by the caller before it is opened, so this is
/// a no-op there.
pub(crate) async fn drop_tables(&self, tables: &[String]) -> Result<usize, String> {
let Open::Rest { catalog, .. } = self else { return Ok(0) };
let mut n = 0;
for t in tables {
let ident = self.ident(t)?;
if catalog.table_exists(&ident).await.map_err(|e| format!("{t}: {e}"))? {
catalog.drop_table(&ident).await.map_err(|e| format!("drop {t}: {e}"))?;
n += 1;
}
}
Ok(n)
}
}
}