echoes/backend/src/storage/sql/postgresql.rs

118 lines
3.9 KiB
Rust
Raw Normal View History

use super::{
builder::{self, SafeValue},
schema, DatabaseTrait,
};
use crate::common::error::{CustomError, CustomErrorInto, CustomResult};
use crate::config;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde_json::Value;
use sqlx::{Column, Executor, PgPool, Row, TypeInfo};
use std::collections::HashMap;
2024-11-11 01:38:58 +08:00
#[derive(Clone)]
pub struct Postgresql {
2024-11-11 01:38:58 +08:00
pool: PgPool,
}
#[async_trait]
impl DatabaseTrait for Postgresql {
async fn initialization(db_config: config::SqlConfig) -> CustomResult<()> {
let db_prefix = SafeValue::Text(
format!("{}", db_config.db_prefix),
builder::ValidationLevel::Strict,
);
let grammar = schema::generate_schema(schema::DatabaseType::PostgreSQL, db_prefix)?;
let connection_str = format!(
"postgres://{}:{}@{}:{}",
db_config.user, db_config.password, db_config.address, db_config.port
);
let pool = PgPool::connect(&connection_str).await?;
pool.execute(format!("CREATE DATABASE {}", db_config.db_name).as_str())
.await?;
let new_connection_str = format!(
"postgres://{}:{}@{}:{}/{}",
db_config.user,
db_config.password,
db_config.address,
db_config.port,
db_config.db_name
);
let new_pool = PgPool::connect(&new_connection_str).await?;
new_pool.execute(grammar.as_str()).await?;
Ok(())
}
async fn connect(db_config: &config::SqlConfig) -> CustomResult<Self> {
2024-11-11 01:38:58 +08:00
let connection_str = format!(
"postgres://{}:{}@{}:{}/{}",
db_config.user,
db_config.password,
db_config.address,
db_config.port,
db_config.db_name
2024-11-11 01:38:58 +08:00
);
let pool = PgPool::connect(&connection_str).await?;
2024-11-11 01:38:58 +08:00
Ok(Postgresql { pool })
}
async fn execute_query<'a>(
&'a self,
builder: &builder::QueryBuilder,
) -> CustomResult<Vec<HashMap<String, Value>>> {
let (query, values) = builder.build()?;
let mut sqlx_query = sqlx::query(&query);
for value in values {
match value {
SafeValue::Null => sqlx_query = sqlx_query.bind(None::<String>),
SafeValue::Bool(b) => sqlx_query = sqlx_query.bind(b),
SafeValue::Integer(i) => sqlx_query = sqlx_query.bind(i),
SafeValue::Float(f) => sqlx_query = sqlx_query.bind(f),
SafeValue::Text(s, _) => sqlx_query = sqlx_query.bind(s),
SafeValue::DateTime(dt) => sqlx_query = sqlx_query.bind(dt.to_rfc3339()),
}
}
let rows = sqlx_query.fetch_all(&self.pool).await?;
2024-11-11 01:38:58 +08:00
Ok(rows
.into_iter()
.map(|row| {
row.columns()
.iter()
.map(|col| {
let value = match col.type_info().name() {
"INT4" | "INT8" => Value::Number(
row.try_get::<i64, _>(col.name()).unwrap_or_default().into(),
),
"FLOAT4" | "FLOAT8" => Value::Number(
serde_json::Number::from_f64(
row.try_get::<f64, _>(col.name()).unwrap_or(0.0),
)
.unwrap_or_else(|| 0.into()),
),
"BOOL" => Value::Bool(row.try_get(col.name()).unwrap_or_default()),
_ => Value::String(row.try_get(col.name()).unwrap_or_default()),
};
(col.name().to_string(), value)
})
.collect()
})
.collect())
}
}
impl Postgresql {
fn get_sdb(&self){
let a=self.pool;
}
}