netsky 0.2.0

netsky CLI: the viable system launcher and subcommand dispatcher
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

//! DataFusion SQL over netsky observability tables.

use chrono::Utc;
use netsky_db::Db;
use serde_json::{Value, json};

use super::db_diag::{wrap_open_retry_error, wrap_query_retry_error};

pub fn run(sql: &str, json_out: bool) -> netsky_core::Result<()> {
    let db = super::db_diag::with_lock_retry(Db::open_read_only).map_err(wrap_open_retry_error)?;
    if !json_out {
        let output = super::db_diag::with_lock_retry(|| db.query(sql))
            .map_err(|e| wrap_query_retry_error(&db, e))?;
        println!("{output}");
        return Ok(());
    }
    let batches = super::db_diag::with_lock_retry(|| db.query_batches(sql))
        .map_err(|e| wrap_query_retry_error(&db, e))?;
    let rows = batches_to_rows(&batches)?;
    let envelope = json!({
        "command": "query",
        "status": "green",
        "summary": format!("{} row(s)", rows.len()),
        "generated_at": Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(),
        "data": {
            "sql": sql,
            "row_count": rows.len(),
            "rows": rows,
        },
    });
    println!("{}", serde_json::to_string_pretty(&envelope)?);
    Ok(())
}

fn batches_to_rows(batches: &[netsky_db::ArrowRecordBatch]) -> netsky_core::Result<Vec<Value>> {
    use netsky_db::arrow_array::ArrayRef;
    let mut rows = Vec::new();
    for batch in batches {
        let schema = batch.schema();
        let columns: Vec<(String, ArrayRef)> = schema
            .fields()
            .iter()
            .enumerate()
            .map(|(i, f)| (f.name().clone(), batch.column(i).clone()))
            .collect();
        for row_idx in 0..batch.num_rows() {
            let mut obj = serde_json::Map::with_capacity(columns.len());
            for (name, col) in &columns {
                obj.insert(name.clone(), array_value(col.as_ref(), row_idx));
            }
            rows.push(Value::Object(obj));
        }
    }
    Ok(rows)
}

fn array_value(arr: &dyn netsky_db::arrow_array::Array, idx: usize) -> Value {
    use netsky_db::arrow_array::{BooleanArray, Float64Array, Int32Array, Int64Array, StringArray};
    if arr.is_null(idx) {
        return Value::Null;
    }
    if let Some(a) = arr.as_any().downcast_ref::<Int64Array>() {
        return json!(a.value(idx));
    }
    if let Some(a) = arr.as_any().downcast_ref::<Int32Array>() {
        return json!(a.value(idx));
    }
    if let Some(a) = arr.as_any().downcast_ref::<Float64Array>() {
        return json!(a.value(idx));
    }
    if let Some(a) = arr.as_any().downcast_ref::<BooleanArray>() {
        return json!(a.value(idx));
    }
    if let Some(a) = arr.as_any().downcast_ref::<StringArray>() {
        return json!(a.value(idx));
    }
    // Fallback: stringify whatever DataFusion handed us so the row is still
    // emittable. Lossy on exotic types (timestamps with tz, lists, structs)
    // but never panics, which is the contract for `query --json` over
    // arbitrary user SQL.
    Value::String(format!("{:?}", arr.slice(idx, 1)))
}