ytsaurus-client 0.2.0

//! `launch` — runs a whole operation with no Python on the machine.
//!
//! Does end to end what `tests/e2e/run_e2e.sh` does through the `yt` CLI:
//! creates tables, uploads the worker, writes input, starts a map, waits for
//! it, and reads the result back. The point is that the only executable
//! involved is this one and the worker it uploads.
//!
//! ```sh
//! export YT_PROXY=http://localhost:8000
//! scripts/build-worker.sh cat
//! cargo run -p ytsaurus-examples --bin launch
//! ```
//!
//! Note this binary is a *launcher*, not a worker: it runs on your machine and
//! is built for the host, not for musl.

use std::process::ExitCode;

use ytsaurus_client::{Client, MapSpec};
use ytsaurus_yson::{Scan, YsonFormat, YsonNode, YsonValue, from_slice, scan::scan_value};

/// Where the demo keeps its tables.
const BASE: &str = "//tmp/ytsaurus_rs_launch";

/// The worker this launches, as produced by `scripts/build-worker.sh cat`.
const WORKER: &str = "target/x86_64-unknown-linux-musl/release-worker/cat";

fn main() -> ExitCode {
    match run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(e) => {
            eprintln!("\nlaunch failed: {e}");
            // The chain is where the cluster's actual complaint lives.
            let mut source = std::error::Error::source(&e);
            while let Some(cause) = source {
                eprintln!("  caused by: {cause}");
                source = cause.source();
            }
            ExitCode::FAILURE
        }
    }
}

fn run() -> Result<(), ytsaurus_client::ClientError> {
    let client = Client::from_env()?;

    if !std::path::Path::new(WORKER).exists() {
        eprintln!("worker not found at {WORKER}");
        eprintln!("build it first:  scripts/build-worker.sh cat");
        return Err(ytsaurus_client::ClientError::Config(
            "the worker binary has not been built".to_owned(),
        ));
    }

    step("Preparing Cypress");
    client.remove(BASE)?;
    client.create("map_node", BASE)?;
    client.create("table", &format!("{BASE}/input"))?;
    client.create("table", &format!("{BASE}/output"))?;
    done(BASE);

    step("Uploading the worker");
    client.upload_worker(WORKER, &format!("{BASE}/cat"))?;
    done(&format!("{BASE}/cat (executable)"));

    step("Writing input rows");
    let rows = sample_rows();
    client.write_table(&format!("{BASE}/input"), &rows)?;
    done(&format!(
        "{} rows",
        client.row_count(&format!("{BASE}/input"))?
    ));

    step("Starting the map operation");
    let spec = MapSpec::new(
        "./cat",
        [format!("{BASE}/input")],
        [format!("{BASE}/output")],
    )
    .with_local_file(format!("{BASE}/cat"))
    .with_memory_limit(512 * 1024 * 1024);

    let id = client.start_map(&spec)?;
    done(&format!("operation {id}"));

    step("Waiting for it to finish");
    client.wait_for_operation(&id)?;
    done("completed");

    step("Checking the result");
    // `cat` is the identity, so reading both tables back through the same path
    // must give identical bytes. Comparing against what we uploaded would fail
    // for an unrelated reason: the cluster re-encodes rows on ingest.
    let before = client.read_table(&format!("{BASE}/input"))?;
    let after = client.read_table(&format!("{BASE}/output"))?;

    if before != after {
        eprintln!(
            "   output differs from input: {} vs {} bytes",
            before.len(),
            after.len()
        );
        return Err(ytsaurus_client::ClientError::Config(
            "the identity map changed its input".to_owned(),
        ));
    }
    done(&format!("identical ({} bytes)", after.len()));

    step("Decoding a row, to prove it is real data");
    if let Some(first) = first_record(&after) {
        let value: YsonValue = from_slice(first, YsonFormat::Binary).map_err(|e| {
            ytsaurus_client::ClientError::Decode {
                command: "read_table".to_owned(),
                reason: e.to_string(),
            }
        })?;
        if let YsonNode::Map(m) = &value.node {
            let columns: Vec<String> = m
                .keys()
                .map(|k| String::from_utf8_lossy(k).into_owned())
                .collect();
            done(&format!("first row has columns {columns:?}"));
        }
    }

    println!("\nAll done — no Python was involved.");
    println!("Tables left at {BASE}");
    Ok(())
}

fn step(what: &str) {
    println!("\n== {what}");
}

fn done(what: &str) {
    println!("   ok {what}");
}

/// Returns the first record of a binary YSON list fragment.
fn first_record(data: &[u8]) -> Option<&[u8]> {
    let trimmed = data.strip_prefix(b";").unwrap_or(data);
    match scan_value(trimmed, YsonFormat::Binary).ok()? {
        Scan::Complete { len } => Some(&trimmed[..len]),
        Scan::Incomplete => None,
    }
}

/// A handful of rows, built with the codec rather than by hand.
fn sample_rows() -> Vec<u8> {
    use serde::Serialize;
    use ytsaurus_yson::to_vec;

    #[derive(Serialize)]
    struct Row<'a> {
        key: &'a str,
        count: i64,
        #[serde(with = "serde_bytes")]
        blob: &'a [u8],
        ratio: f64,
        flag: bool,
    }

    let rows = [
        Row {
            key: "alpha",
            count: 1,
            blob: &[0xDE, 0xAD],
            ratio: 0.5,
            flag: true,
        },
        Row {
            key: "beta",
            count: -2,
            blob: &[0x00, 0xFF],
            ratio: -1.25,
            flag: false,
        },
        Row {
            key: "",
            count: i64::MAX,
            blob: &[],
            ratio: 0.0,
            flag: true,
        },
    ];

    let mut out = Vec::new();
    for row in &rows {
        out.extend_from_slice(&to_vec(row, YsonFormat::Binary).expect("encodes"));
        out.push(b';');
    }
    out
}