systemprompt_database/services/postgres/
conversion.rs

1//! Conversions between `SQLx` `PostgreSQL` rows, [`QueryResult`] /
2//! [`crate::models::JsonRow`] values, and parameter binders.
3//!
4//! Part of the documented sqlx allowlist: the binder operates on the dynamic
5//! `sqlx::query::Query` value passed in by the trait implementation.
6//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10use sqlx::{Column, Row};
11use std::collections::HashMap;
12
13use crate::models::{DbValue, QueryResult, ToDbValue};
14
15pub fn rows_to_result(rows: Vec<sqlx::postgres::PgRow>, start: std::time::Instant) -> QueryResult {
16    let mut columns = Vec::new();
17    let mut result_rows = Vec::new();
18
19    if let Some(first_row) = rows.first() {
20        columns = first_row
21            .columns()
22            .iter()
23            .map(|c| c.name().to_owned())
24            .collect();
25    }
26
27    for row in rows {
28        result_rows.push(row_to_json(&row));
29    }
30
31    let row_count = result_rows.len();
32    let execution_time_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
33
34    QueryResult {
35        columns,
36        rows: result_rows,
37        row_count,
38        execution_time_ms,
39    }
40}
41
42pub fn row_to_json(row: &sqlx::postgres::PgRow) -> HashMap<String, serde_json::Value> {
43    row.columns()
44        .iter()
45        .map(|col| (col.name().to_owned(), column_to_json(row, col.ordinal())))
46        .collect()
47}
48
49fn column_to_json(row: &sqlx::postgres::PgRow, ordinal: usize) -> serde_json::Value {
50    if let Ok(val) = row.try_get::<Option<chrono::DateTime<chrono::Utc>>, _>(ordinal) {
51        return val.map_or(serde_json::Value::Null, |v| {
52            serde_json::Value::String(v.to_rfc3339())
53        });
54    }
55    if let Ok(val) = row.try_get::<Option<uuid::Uuid>, _>(ordinal) {
56        return val.map_or(serde_json::Value::Null, |v| {
57            serde_json::Value::String(v.to_string())
58        });
59    }
60    if let Ok(val) = row.try_get::<Option<String>, _>(ordinal) {
61        return val.map_or(serde_json::Value::Null, serde_json::Value::String);
62    }
63    if let Ok(val) = row.try_get::<Option<i64>, _>(ordinal) {
64        return val.map_or(serde_json::Value::Null, |v| {
65            serde_json::Value::Number(v.into())
66        });
67    }
68    if let Ok(val) = row.try_get::<Option<i32>, _>(ordinal) {
69        return val.map_or(serde_json::Value::Null, |v| {
70            serde_json::Value::Number(i64::from(v).into())
71        });
72    }
73    if let Ok(val) = row.try_get::<Option<f64>, _>(ordinal) {
74        return val.map_or(serde_json::Value::Null, |v| serde_json::json!(v));
75    }
76    if let Ok(val) = row.try_get::<Option<rust_decimal::Decimal>, _>(ordinal) {
77        return val.map_or(serde_json::Value::Null, |v| {
78            v.to_string().parse::<f64>().map_or_else(
79                |_| serde_json::Value::String(v.to_string()),
80                |f| serde_json::json!(f),
81            )
82        });
83    }
84    if let Ok(val) = row.try_get::<Option<bool>, _>(ordinal) {
85        return val.map_or(serde_json::Value::Null, serde_json::Value::Bool);
86    }
87    if let Ok(val) = row.try_get::<Option<Vec<String>>, _>(ordinal) {
88        return val.map_or(serde_json::Value::Null, |v| {
89            serde_json::Value::Array(v.into_iter().map(serde_json::Value::String).collect())
90        });
91    }
92    if let Ok(val) = row.try_get::<Option<serde_json::Value>, _>(ordinal) {
93        return val.unwrap_or(serde_json::Value::Null);
94    }
95    if let Ok(val) = row.try_get::<Option<Vec<u8>>, _>(ordinal) {
96        return val.map_or(serde_json::Value::Null, |bytes| {
97            use base64::Engine;
98            use base64::engine::general_purpose::STANDARD;
99            serde_json::Value::String(STANDARD.encode(&bytes))
100        });
101    }
102    serde_json::Value::Null
103}
104
105pub fn bind_params<'q>(
106    mut query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
107    params: &[&dyn ToDbValue],
108) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
109    for param in params {
110        let value = param.to_db_value();
111        query = match value {
112            DbValue::String(s) => query.bind(s),
113            DbValue::Int(i) => query.bind(i),
114            DbValue::Float(f) => query.bind(f),
115            DbValue::Bool(b) => query.bind(b),
116            DbValue::Bytes(b) => query.bind(b),
117            DbValue::Timestamp(dt) => query.bind(dt),
118            DbValue::StringArray(arr) => query.bind(arr),
119            DbValue::NullString => query.bind(None::<String>),
120            DbValue::NullInt => query.bind(None::<i64>),
121            DbValue::NullFloat => query.bind(None::<f64>),
122            DbValue::NullBool => query.bind(None::<bool>),
123            DbValue::NullBytes => query.bind(None::<Vec<u8>>),
124            DbValue::NullTimestamp => query.bind(None::<chrono::DateTime<chrono::Utc>>),
125            DbValue::NullStringArray => query.bind(None::<Vec<String>>),
126        };
127    }
128    query
129}
systemprompt_database/services/postgres/conversion.rs

systemprompt_database/services/postgres/
conversion.rs
