Skip to main content

systemprompt_database/services/postgres/
conversion.rs

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