use std::collections::BTreeMap;
use std::path::PathBuf;
use clap::{Args, ValueEnum};
pub const CREDENTIAL_ENV: &str = "BIFROST_CATALOG_CREDENTIAL";
pub const TOKEN_ENV: &str = "BIFROST_CATALOG_TOKEN";
pub const SECRET_ENV: [&str; 2] = [CREDENTIAL_ENV, TOKEN_ENV];
#[derive(ValueEnum, Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum CatalogKind {
#[default]
Rest,
Aistor,
Polaris,
Unity,
Gravitino,
}
impl CatalogKind {
pub fn as_str(self) -> &'static str {
match self {
CatalogKind::Rest => "rest",
CatalogKind::Aistor => "aistor",
CatalogKind::Polaris => "polaris",
CatalogKind::Unity => "unity",
CatalogKind::Gravitino => "gravitino",
}
}
}
#[derive(Args, Debug, Clone, Default, PartialEq, Eq)]
#[group(skip)]
pub struct CatalogArgs {
#[arg(long, value_enum, default_value_t = CatalogKind::Rest)]
pub catalog: CatalogKind,
#[arg(long, value_name = "SCOPE")]
pub rest_scope: Option<String>,
#[arg(long, value_name = "URL")]
pub rest_oauth2_uri: Option<String>,
#[arg(long = "rest-header", value_name = "NAME=VALUE")]
pub rest_headers: Vec<String>,
#[arg(long, value_name = "BOOL")]
pub vended_credentials: Option<bool>,
#[arg(long, value_name = "USER")]
pub gravitino_user: Option<String>,
#[arg(long, value_name = "URL")]
pub grpc_uri: Option<String>,
#[arg(long, value_name = "PEM")]
pub grpc_ca: Option<PathBuf>,
#[arg(long, value_name = "PEM", requires = "grpc_key")]
pub grpc_cert: Option<PathBuf>,
#[arg(long, value_name = "PEM", requires = "grpc_cert")]
pub grpc_key: Option<PathBuf>,
}
#[cfg(feature = "skade-rest")]
fn env(name: &str) -> Option<String> {
std::env::var(name).ok().filter(|v| !v.is_empty())
}
#[cfg(feature = "skade-rest")]
fn missing(kind: CatalogKind) -> String {
format!(
"--catalog {}: this build has no {0} connector - rebuild with --features catalog-{0}",
kind.as_str()
)
}
impl CatalogArgs {
pub fn headers(&self) -> Result<BTreeMap<String, String>, String> {
let mut out = BTreeMap::new();
for h in &self.rest_headers {
let (k, v) = h.split_once('=').ok_or_else(|| format!("--rest-header {h:?}: want Name=Value"))?;
let k = k.trim();
if k.is_empty() {
return Err(format!("--rest-header {h:?}: empty name"));
}
if k.eq_ignore_ascii_case("authorization") {
return Err("--rest-header Authorization: a secret never goes on the command line - set BIFROST_CATALOG_TOKEN".into());
}
out.insert(k.to_string(), v.trim().to_string());
}
Ok(out)
}
#[cfg(feature = "skade-rest")]
pub fn rest(
&self,
uri: &str,
warehouse: Option<String>,
sigv4: Option<String>,
s3: skade::S3Config,
read_only: bool,
) -> Result<skade::connect::IcebergRest, String> {
self.rest_with(uri, warehouse, sigv4, s3, read_only, env(CREDENTIAL_ENV), env(TOKEN_ENV))
}
#[cfg(feature = "skade-rest")]
#[allow(clippy::too_many_arguments)]
fn rest_with(
&self,
uri: &str,
warehouse: Option<String>,
sigv4: Option<String>,
s3: skade::S3Config,
read_only: bool,
credential: Option<String>,
token: Option<String>,
) -> Result<skade::connect::IcebergRest, String> {
use skade::connect::IcebergRest;
let headers = self.headers()?;
if self.gravitino_user.is_some() && self.catalog != CatalogKind::Gravitino {
return Err("--gravitino-user needs --catalog gravitino".into());
}
#[allow(unused_variables)]
let need = |what: &str| {
warehouse
.clone()
.filter(|w| !w.trim().is_empty())
.ok_or_else(|| format!("--catalog {} needs --rest-warehouse ({what})", self.catalog.as_str()))
};
let mut rest = match self.catalog {
CatalogKind::Rest => IcebergRest {
uri: uri.to_string(),
warehouse,
token,
credential,
scope: self.rest_scope.clone(),
oauth2_server_uri: self.rest_oauth2_uri.clone(),
tls: s3.tls.clone(),
vended_credentials: self.vended_credentials.unwrap_or(false),
s3,
read_only,
..IcebergRest::default()
},
#[cfg(feature = "catalog-aistor")]
CatalogKind::Aistor => {
let mut r = skade::connect::AIStor { s3, warehouse: need("the AIStor warehouse")?, read_only }.rest();
r.uri = format!("{}/_iceberg", uri.trim_end_matches('/').trim_end_matches("/_iceberg"));
r.token = token;
r.vended_credentials = self.vended_credentials.unwrap_or(false);
r
}
#[cfg(feature = "catalog-polaris")]
CatalogKind::Polaris => {
let mut p = skade::connect::Polaris::new(uri, need("the Polaris catalog name")?);
p.credential = credential;
p.scope = self.rest_scope.clone();
p.oauth2_server_uri = self.rest_oauth2_uri.clone();
p.token = token;
p.tls = s3.tls.clone();
p.s3 = s3;
p.vended_credentials = self.vended_credentials.unwrap_or(true);
p.read_only = read_only;
p.rest()
}
#[cfg(feature = "catalog-unity")]
CatalogKind::Unity => {
if credential.is_some() {
return Err("--catalog unity authenticates by token: set BIFROST_CATALOG_TOKEN, not BIFROST_CATALOG_CREDENTIAL".into());
}
let mut u = skade::connect::UnityCatalog::new(uri, need("the Unity catalog name")?);
u.token = token;
u.tls = s3.tls.clone();
u.s3 = s3;
u.vended_credentials = self.vended_credentials.unwrap_or(true);
u.read_only = read_only;
u.rest()
}
#[cfg(feature = "catalog-gravitino")]
CatalogKind::Gravitino => {
let mut g = skade::connect::Gravitino::new(uri);
g.warehouse = warehouse;
g.simple_user = self.gravitino_user.clone();
g.credential = credential;
g.scope = self.rest_scope.clone();
g.oauth2_server_uri = self.rest_oauth2_uri.clone();
g.token = token;
g.tls = s3.tls.clone();
g.s3 = s3;
g.vended_credentials = self.vended_credentials.unwrap_or(false);
g.read_only = read_only;
g.rest()
}
#[allow(unreachable_patterns)]
other => return Err(missing(other)),
};
rest.headers.extend(headers);
if let Some(service) = sigv4 {
rest.sigv4 = Some(skade::connect::SigV4 { service });
}
Ok(rest)
}
#[cfg(feature = "skade-grpc")]
pub fn grpc(&self, uri: &str, s3: Option<skade::S3Config>, read_only: bool) -> skade::connect::SkadeGrpc {
let tls = (self.grpc_ca.is_some() || self.grpc_cert.is_some()).then(|| skade::tls::ClientTls {
ca_pem: self.grpc_ca.clone(),
client_cert_pem: self.grpc_cert.clone(),
client_key_pem: self.grpc_key.clone(),
});
skade::connect::SkadeGrpc { uri: uri.to_string(), tls, s3, read_only }
}
pub fn check_unused(&self, has_rest: bool) -> Result<(), String> {
if !has_rest
&& (self.catalog != CatalogKind::Rest
|| self.rest_scope.is_some()
|| self.rest_oauth2_uri.is_some()
|| !self.rest_headers.is_empty()
|| self.vended_credentials.is_some()
|| self.gravitino_user.is_some())
{
return Err("--catalog / --rest-scope / --rest-oauth2-uri / --rest-header / --vended-credentials / --gravitino-user need --rest-uri".into());
}
if self.grpc_uri.is_none() && (self.grpc_ca.is_some() || self.grpc_cert.is_some()) {
return Err("--grpc-ca / --grpc-cert / --grpc-key need --grpc-uri".into());
}
Ok(())
}
pub fn to_args(&self) -> Vec<String> {
let mut a = Vec::new();
if self.catalog != CatalogKind::Rest {
a.extend(["--catalog".to_string(), self.catalog.as_str().to_string()]);
}
let mut push = |k: &str, v: Option<String>| {
if let Some(v) = v {
a.push(format!("--{k}"));
a.push(v);
}
};
push("rest-scope", self.rest_scope.clone());
push("rest-oauth2-uri", self.rest_oauth2_uri.clone());
for h in &self.rest_headers {
push("rest-header", Some(h.clone()));
}
push("vended-credentials", self.vended_credentials.map(|b| b.to_string()));
push("gravitino-user", self.gravitino_user.clone());
push("grpc-uri", self.grpc_uri.clone());
push("grpc-ca", self.grpc_ca.as_ref().map(|p| p.display().to_string()));
push("grpc-cert", self.grpc_cert.as_ref().map(|p| p.display().to_string()));
push("grpc-key", self.grpc_key.as_ref().map(|p| p.display().to_string()));
a
}
pub fn absolutized(&self, cwd: &std::path::Path) -> Self {
let abs = |p: &Option<PathBuf>| p.as_ref().map(|p| if p.is_absolute() { p.clone() } else { cwd.join(p) });
Self { grpc_ca: abs(&self.grpc_ca), grpc_cert: abs(&self.grpc_cert), grpc_key: abs(&self.grpc_key), ..self.clone() }
}
}
#[cfg(test)]
mod tests {
use super::*;
use clap::Parser;
#[derive(Parser)]
struct Cli {
#[command(flatten)]
c: CatalogArgs,
}
fn parse(args: &[&str]) -> CatalogArgs {
Cli::try_parse_from(std::iter::once("t").chain(args.iter().copied())).unwrap().c
}
#[cfg(feature = "skade-rest")]
fn s3() -> skade::S3Config {
skade::S3Config { endpoint: "https://minio:9000".into(), bucket: "lake".into(), ..Default::default() }
}
#[test]
fn to_args_round_trips_and_carries_no_secret() {
let c = parse(&[
"--catalog", "gravitino", "--rest-scope", "s", "--rest-oauth2-uri", "https://idp/token",
"--rest-header", "X-Tenant=a", "--vended-credentials", "true", "--gravitino-user", "ana",
"--grpc-uri", "https://k:50051", "--grpc-ca", "/ca.pem",
]);
let args = c.to_args();
let refs: Vec<&str> = args.iter().map(String::as_str).collect();
assert_eq!(parse(&refs), c);
let joined = args.join(" ");
assert!(!joined.contains(CREDENTIAL_ENV) && !joined.contains(TOKEN_ENV));
}
#[test]
fn the_default_is_plain_rest_and_emits_no_flag() {
assert_eq!(parse(&[]).to_args(), Vec::<String>::new());
}
#[test]
fn an_authorization_header_is_refused() {
let c = parse(&["--rest-header", "authorization=Bearer x"]);
assert!(c.headers().unwrap_err().contains("BIFROST_CATALOG_TOKEN"));
assert!(parse(&["--rest-header", "nope"]).headers().is_err());
}
#[cfg(feature = "skade-rest")]
#[test]
fn plain_rest_takes_secrets_from_the_caller_and_headers_from_flags() {
let c = parse(&["--rest-header", "X-Tenant=a", "--rest-scope", "catalog"]);
let r = c
.rest_with("https://cat", Some("wh".into()), Some("s3tables".into()), s3(), true, Some("id:sec".into()), None)
.unwrap();
assert_eq!(r.uri, "https://cat");
assert_eq!(r.credential.as_deref(), Some("id:sec"));
assert_eq!(r.headers.get("X-Tenant").map(String::as_str), Some("a"));
assert_eq!(r.sigv4.map(|s| s.service).as_deref(), Some("s3tables"));
assert!(r.read_only && !r.vended_credentials);
}
#[test]
fn rest_only_flags_without_a_rest_uri_are_refused() {
assert!(parse(&["--catalog", "polaris"]).check_unused(false).is_err());
assert!(parse(&["--catalog", "polaris"]).check_unused(true).is_ok());
assert!(parse(&["--grpc-ca", "/c"]).check_unused(false).is_err());
}
#[cfg(all(feature = "skade-rest", not(feature = "catalog-polaris")))]
#[test]
fn a_vendor_this_build_lacks_names_its_feature() {
let e = parse(&["--catalog", "polaris"]).rest_with("u", Some("c".into()), None, s3(), true, None, None);
assert!(e.unwrap_err().contains("--features catalog-polaris"));
}
#[cfg(feature = "catalog-polaris")]
#[test]
fn polaris_gets_its_path_scope_and_vended_credentials() {
let r = parse(&["--catalog", "polaris"])
.rest_with("https://polaris:8181", Some("cat".into()), None, s3(), true, Some("id:sec".into()), None)
.unwrap();
assert_eq!(r.uri, "https://polaris:8181/api/catalog");
assert_eq!(r.warehouse.as_deref(), Some("cat"));
assert_eq!(r.scope.as_deref(), Some("PRINCIPAL_ROLE:ALL"));
assert!(r.vended_credentials && r.read_only);
assert!(parse(&["--catalog", "polaris"]).rest_with("u", None, None, s3(), true, None, None).is_err());
}
#[cfg(feature = "catalog-aistor")]
#[test]
fn aistor_points_at_the_minio_iceberg_path() {
for uri in ["https://minio:9000", "https://minio:9000/_iceberg/"] {
let r = parse(&["--catalog", "aistor"])
.rest_with(uri, Some("analytics".into()), None, s3(), true, None, None)
.unwrap();
assert_eq!(r.uri, "https://minio:9000/_iceberg");
assert_eq!(r.sigv4.map(|s| s.service).as_deref(), Some("s3tables"));
}
}
#[cfg(feature = "skade-grpc")]
#[test]
fn grpc_tls_comes_from_the_pem_flags() {
let c = parse(&["--grpc-uri", "https://k:50051", "--grpc-ca", "/ca.pem"]);
let g = c.grpc("https://k:50051", None, true);
assert_eq!(g.tls.unwrap().ca_pem.as_deref(), Some(std::path::Path::new("/ca.pem")));
assert!(g.read_only && g.s3.is_none());
}
}