Skip to main content

athena_driver/postgresql/
json_row_expr.rs

1//! SQL fragments for JSON row projections in gateway fetch paths.
2
3/// PostgreSQL caps function arguments at 100; `jsonb_build_object` uses two per key/value pair.
4const MAX_JSONB_BUILD_OBJECT_PAIRS: usize = 50;
5
6/// SQL expression that builds one JSON object from pair fragments (key literal + column ref each).
7///
8/// Each `pairs` entry must be one comma-separated fragment such as `"'id', t.id"`.
9/// Callers must pass a non-empty slice; an empty slice yields `jsonb_build_object()` as a safe fallback.
10pub fn merge_jsonb_build_object_expr(pairs: &[String]) -> String {
11    if pairs.is_empty() {
12        return "jsonb_build_object()".to_string();
13    }
14    let chunks: Vec<String> = pairs
15        .chunks(MAX_JSONB_BUILD_OBJECT_PAIRS)
16        .map(|chunk| format!("jsonb_build_object({})", chunk.join(", ")))
17        .collect();
18    match chunks.len() {
19        1 => chunks[0].clone(),
20        _ => format!("({})", chunks.join(" || ")),
21    }
22}
23
24#[cfg(test)]
25mod tests {
26    use super::merge_jsonb_build_object_expr;
27
28    fn sample_pairs(n: usize) -> Vec<String> {
29        (0..n).map(|i| format!("'c{i}', t.c{i}")).collect()
30    }
31
32    #[test]
33    fn single_chunk_no_merge_operator() {
34        let pairs = sample_pairs(50);
35        let expr = merge_jsonb_build_object_expr(&pairs);
36        assert!(expr.starts_with("jsonb_build_object("));
37        assert!(!expr.contains("||"));
38        assert_eq!(expr.matches("jsonb_build_object").count(), 1);
39    }
40
41    #[test]
42    fn fifty_one_pairs_merge_two_objects() {
43        let pairs = sample_pairs(51);
44        let expr = merge_jsonb_build_object_expr(&pairs);
45        assert!(expr.starts_with('('));
46        assert!(expr.contains("||"));
47        assert_eq!(expr.matches("jsonb_build_object").count(), 2);
48    }
49
50    #[test]
51    fn one_pair_single_call() {
52        let pairs = sample_pairs(1);
53        let expr = merge_jsonb_build_object_expr(&pairs);
54        assert_eq!(expr, "jsonb_build_object('c0', t.c0)");
55    }
56}