use crate::error::{TViewError, TViewResult};
use pgrx::JsonB;
use pgrx::datum::DatumWithOid;
use pgrx::prelude::*;
#[allow(dead_code)] pub fn insert_array_element(
table_name: &str,
pk_column: &str,
pk_value: i64,
array_path: &[String],
new_element: JsonB,
sort_key: Option<String>,
) -> TViewResult<()> {
let path_str = array_path.join(",");
let path_array = format!("ARRAY['{path_str}']");
let sql = sort_key.map_or_else(
|| {
format!(
r"
UPDATE {table_name} SET
data = jsonb_array_insert_where(data, {path_array}, $1, NULL, NULL),
updated_at = now()
WHERE {pk_column} = $2
"
)
},
|key| {
format!(
r"
UPDATE {table_name} SET
data = jsonb_array_insert_where(data, {path_array}, $1, '{key}', 'ASC'),
updated_at = now()
WHERE {pk_column} = $2
"
)
},
);
let args = vec![
unsafe { DatumWithOid::new(new_element, PgOid::BuiltIn(PgBuiltInOids::JSONBOID).value()) },
unsafe { DatumWithOid::new(pk_value, PgOid::BuiltIn(PgBuiltInOids::INT8OID).value()) },
];
Spi::run_with_args(&sql, &args).map_err(|e| TViewError::SpiError {
query: sql,
error: e.to_string(),
})?;
Ok(())
}
#[allow(dead_code)] pub fn delete_array_element(
table_name: &str,
pk_column: &str,
pk_value: i64,
array_path: &[String],
match_key: &str,
match_value: JsonB,
) -> TViewResult<()> {
let path_str = array_path.join(",");
let path_array = format!("ARRAY['{path_str}']");
let sql = format!(
r"
UPDATE {table_name} SET
data = jsonb_array_delete_where(data, {path_array}, '{match_key}', $1),
updated_at = now()
WHERE {pk_column} = $2
"
);
let args = vec![
unsafe { DatumWithOid::new(match_value, PgOid::BuiltIn(PgBuiltInOids::JSONBOID).value()) },
unsafe { DatumWithOid::new(pk_value, PgOid::BuiltIn(PgBuiltInOids::INT8OID).value()) },
];
Spi::run_with_args(&sql, &args).map_err(|e| TViewError::SpiError {
query: sql,
error: e.to_string(),
})?;
Ok(())
}
#[allow(dead_code)] pub fn check_array_functions_available() -> TViewResult<bool> {
let sql = r"
SELECT EXISTS(
SELECT 1 FROM pg_proc
WHERE proname = 'jsonb_array_insert_where'
OR proname = 'jsonb_array_delete_where'
)
";
Spi::get_one::<bool>(sql)
.map_err(|e| TViewError::SpiError {
query: sql.to_string(),
error: e.to_string(),
})
.map(|opt| opt.unwrap_or(false))
}
#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
use super::*;
#[pg_test]
fn test_insert_array_element() {
Spi::run(
r#"
CREATE TABLE test_array_ops (
id BIGINT PRIMARY KEY,
data JSONB DEFAULT '{"items": []}'::jsonb
)
"#,
)
.unwrap();
Spi::run(
r#"
INSERT INTO test_array_ops VALUES (1, '{"items": []}'::jsonb)
"#,
)
.unwrap();
let new_element = JsonB(serde_json::json!({"id": 1, "name": "Test Item"}));
insert_array_element(
"test_array_ops",
"id",
1,
&["items".to_string()],
new_element,
None,
)
.unwrap();
let result = Spi::get_one::<JsonB>("SELECT data FROM test_array_ops WHERE id = 1")
.unwrap()
.unwrap();
let items = result.0["items"].as_array().unwrap();
assert_eq!(items.len(), 1);
assert_eq!(items[0]["name"], "Test Item");
}
#[pg_test]
fn test_delete_array_element() {
Spi::run(
r#"
CREATE TABLE test_array_ops (
id BIGINT PRIMARY KEY,
data JSONB DEFAULT '{"items": [{"id": 1, "name": "Test Item"}]}'::jsonb
)
"#,
)
.unwrap();
Spi::run(r#"
INSERT INTO test_array_ops VALUES (1, '{"items": [{"id": 1, "name": "Test Item"}]}'::jsonb)
"#).unwrap();
let match_value = JsonB(serde_json::json!(1));
delete_array_element(
"test_array_ops",
"id",
1,
&["items".to_string()],
"id",
match_value,
)
.unwrap();
let result = Spi::get_one::<JsonB>("SELECT data FROM test_array_ops WHERE id = 1")
.unwrap()
.unwrap();
let items = result.0["items"].as_array().unwrap();
assert_eq!(items.len(), 0);
}
#[pg_test]
fn test_check_array_functions_available() {
let available = check_array_functions_available().unwrap();
let _ = available;
}
}