Working with SQLx and Age

TL;DR: I created two functions for working with Apache Age (accessed by SQLx) that returns JSON results directly from the database. The code could be improved and I’m unsure what the performance cost of the regex is but this works well enough to be used in osintbuddy.. We end up automatically casting the agtype results to JSON.

// External resources: https://age.apache.org/age-manual/master/intro/types.html#map
// Expected return must include: serde_json::Value
pub async fn with_cypher(
    query: &str,
    tx: &mut PgConnection,
) -> Result<Vec<serde_json::Value>, AppError> {
    let objs: Vec<PgRow> = sqlx::raw_sql(query).fetch_all(tx).await.map_err(|err| {
        error!("{err}");
        AppError {
            message: "We ran into an OpenCypher select transaction error!",
            kind: ErrorKind::Critical,
        }
    })?;
    let re = Regex::new(r"(::vertex|::edge)").unwrap();
    let json_results = objs
        .iter()
        .map(|row| {
            let mut json_objs: Vec<serde_json::Value> = Vec::new();
            for (i, col) in row.columns().iter().enumerate() {
                let value = row.get::<String, _>(i);
                re.replace_all(value, "");
                let v = serde_json::Value::from(value);
                json_objs.push(v);
            }
            return json_objs;
        })
        .flatten()
        .collect();
    Ok(json_results)
}

Usage looks like:

// pool: actix_web::web::Data<PgPool>
let mut tx = age_tx(pool.as_ref()).await?;
    let vertices = with_cypher(
        &graph_name,
        tx.as_mut(),
        "MATCH (v) WITH {id: id(v)} AS v RETURN v",
        "v agtype",
    )
    .await?;
    let edges = with_cypher(
        &graph_name,
        tx.as_mut(),
        "MATCH (v)-[e]->() WITH {id: id(e)} AS e RETURN e",
        "e agtype",
    )
    .await?;
    let second_degrees = with_cypher(
        &graph_name,
        tx.as_mut(),
        "MATCH ()-[]->(a)-[]->(v) WITH {id: id(v)} AS v RETURN v",
        "v agtype",
    )
    .await?;
    tx.commit().await.map_err(|err| {
        error!("{err}");
        AppError {
            message: "We ran into an error commiting the age transaction!",
            kind: ErrorKind::Critical,
        }
    })?;

I’ll slowly refine this over time but it works well enough for now.

When looking at the extension code or the issues it seems you’re supposed to be able to cast agtype directly to JSON however I can’t get this to work. I run into errors such as:

ERROR:  cannot cast type agtype to json
# or else I get errors like:
ERROR:  type "agtype_as_json" does not exist

I’ll try troubleshooting this further later, being able to return JSON from Apache Age with less casting would be really nice.