chadcatlett / pg-jq Goto Github PK

use pgx::*;
use rustc_hash::FxHashMap;

pg_module_magic!();

extension_sql!(r#"CREATE DOMAIN JsonString AS json;"#);

#[inline]
fn get_jq_program(q: &str) -> &'static mut jq_rs::JqProgram {
    unsafe {
        static mut JQ_PROGRAMS: Option<FxHashMap<String, jq_rs::JqProgram>> = None;

        let map = JQ_PROGRAMS.get_or_insert_with(|| Default::default());
        map.entry(q.to_string())
            .or_insert_with(|| jq_rs::compile(q).expect("failed to compile jq program"))
    }
}

#[pg_operator(parallel_safe, immutable)]
#[opname(~>)]
fn jqeval_json(input: JsonString, query: &str) -> JsonString {
    let program = get_jq_program(query);
    let results = program.run(&input.0).expect("failed to run jq program");
    JsonString(results)
}

#[pg_operator(parallel_safe, immutable)]
#[opname(~>)]
fn jqeval_jsonb(input: JsonB, query: &str) -> JsonB {
    let program = get_jq_program(query);
    let input = serde_json::to_string(&input.0).unwrap();
    let results = program.run(&input).expect("failed to run jq program");
    JsonB(serde_json::from_str(&results).expect("results not valid json"))
}

#[cfg(any(test, feature = "pg_test"))]
mod tests {
    use pgx::*;

    #[pg_test]
    fn test_hello_jq_pgx() {
        assert_eq!("Hello, jq_pgx", crate::hello_jq_pgx());
    }
}

#[cfg(test)]
pub mod pg_test {
    pub fn setup(_options: Vec<&str>) {
        // perform one-off initialization when the pg_test framework starts
    }

    pub fn postgresql_conf_options() -> Vec<&'static str> {
        // return any postgresql.conf settings that are required for your tests
        vec![]
    }
}