use std::{collections::HashMap};
use mysql_async::prelude::Queryable;
use tokio::sync::Mutex;
pub use crate::config::ConfigClient;
use mssql_client::{Ready};
use anyhow::{Ok};
use polars::{prelude::*};
use std::fs;
use std::sync::Arc;
pub enum DbClient {
SqlServer(mssql_client::Client<Ready>),
MySql(mysql_async::Conn),
ClickHouse{client:clickhouse::Client, url: String, user: String, pass: String},
}
pub struct ClientSql{
pub(crate) clients: HashMap<String, Mutex<DbClient>>,
}
pub struct TempParquetFile(String);
impl Drop for TempParquetFile {
fn drop(&mut self) {
let _=std::fs::remove_file(&self.0);
}
}
#[derive(Clone)]
pub struct TabGoes{
pub(crate) lazy: polars::prelude::LazyFrame,
pub(crate) _files: Vec<Arc<TempParquetFile>>
}
impl ClientSql {
pub async fn new() -> anyhow::Result<Self>{
let toml_str = std::fs::read_to_string("database.toml").expect("Arquivo database.toml");
let toml_str = toml_str.replace("\r\n", "\n");
let app_config: crate::config::AppConfig = toml::from_str(&toml_str)?;
let mut clients = std::collections::HashMap::new();
for server_cfg in app_config.servers{
let server_key = server_cfg.alias.clone().unwrap_or_else(|| server_cfg.host.clone());
let conn_mssql = format!("Server={};User Id={};Password={};TrustServerCertificate=true", server_cfg.host, server_cfg.user, server_cfg.password);
if let std::result::Result::Ok(config) = mssql_client::Config::from_connection_string(&conn_mssql){
if let std::result::Result::Ok(client) = mssql_client::Client::connect(config).await{
clients.insert(server_key.clone(), Mutex::new(DbClient::SqlServer(client)));
continue;
}
}
let conn_mysql = format!("mysql://{}:{}@{}:3306", server_cfg.user, server_cfg.password, server_cfg.host);
if let std::result::Result::Ok(opts)= mysql_async::Opts::from_url(&conn_mysql){
if let std::result::Result::Ok(client)= mysql_async::Conn::new(opts).await{
clients.insert(server_key.clone(), Mutex::new(DbClient::MySql(client)));
continue;
}
}
let ch_url = format!("http://{}:8123", server_cfg.host);
let client_ch = clickhouse::Client::default()
.with_url(&ch_url)
.with_user(server_cfg.user.to_string())
.with_password(server_cfg.password.to_string());
if let std::result::Result::Ok(_) = client_ch.query("SELECT 1").fetch_one::<u8>().await{
clients.insert(server_key.clone(), Mutex::new(DbClient::ClickHouse{
client: client_ch,
url: ch_url,
user: server_cfg.user.clone(),
pass: server_cfg.password.clone()
}));
continue;
}
}
Ok(Self{clients})
}
pub async fn query_to_parquet(&self,server_name: &str, sql_gerado: &str, expected_cols: &[&str], expected_types: &[&str])-> anyhow::Result<TabGoes>{
let client_mutex =self.clients.get(server_name)
.ok_or_else(||anyhow::anyhow!("Falha de conexão com o servidor: {}",server_name))?;
let mut client_lock = client_mutex.lock().await;
let file_path = format!("temp_query_{}.parquet", crate::uuid::Uuid::new_v4());
let mut df = match &mut *client_lock{
DbClient::SqlServer(client)=>{
Self::fetch_mssql_to_df(client, sql_gerado, expected_cols, expected_types).await?
},
DbClient::MySql(client)=>{
Self::fetch_mysql_to_df(client, sql_gerado, expected_cols, expected_types).await?
},
DbClient::ClickHouse{url,user,pass, ..}=>{
let req_client = reqwest::Client::new();
let sql_parquet = format!("{} FORMAT Parquet", sql_gerado);
let resp = req_client.post(url.as_str())
.basic_auth(user.as_str(), Some(pass.as_str()))
.body(sql_parquet)
.send()
.await?;
if !resp.status().is_success(){
anyhow::bail!("Erro ClickHouse: {}", resp.text().await?);
}
std::fs::write(&file_path, resp.bytes().await?)?;
let lf = polars::prelude::LazyFrame::scan_parquet(&file_path, polars::prelude::ScanArgsParquet::default())?;
return Ok(TabGoes {
lazy: lf,
_files: vec![Arc::new(TempParquetFile(file_path))]
});
}
};
let mut file = std::fs::File::create(&file_path)?;
polars::prelude::ParquetWriter::new(&mut file).finish(&mut df)?;
drop(df);
let lf = polars::prelude::LazyFrame::scan_parquet(&file_path, polars::prelude::ScanArgsParquet::default())?;
Ok(TabGoes {
lazy: lf,
_files: vec![Arc::new(TempParquetFile(file_path))],
})
}
async fn fetch_mssql_to_df(
client: &mut mssql_client::Client<Ready>,
sql: &str,
expected_cols: &[&str],
expected_types: &[&str],
)->anyhow::Result<DataFrame>{
enum ColBuilder {
I16(Vec<Option<i16>>),
I32(Vec<Option<i32>>),
I64(Vec<Option<i64>>),
F64(Vec<Option<f64>>),
Bool(Vec<Option<bool>>),
String(Vec<Option<String>>),
}
let mut builders:Vec<ColBuilder> = expected_types.iter().map(|&t|{
match t{
"Option<i16>" => ColBuilder::I16(Vec::new()),
"Option<i32>" => ColBuilder::I32(Vec::new()),
"Option<i64>" => ColBuilder::I64(Vec::new()),
"Option<f64>" | "Option<f32>" | "Option<data_goes::rust_decimal::Decimal>" => ColBuilder::F64(Vec::new()),
"Option<bool>" => ColBuilder::Bool(Vec::new()),
_ => ColBuilder::String(Vec::new()),
}
}).collect();
let rows = client.query(sql, &[]).await?;
for result in rows{
let row = result?;
for (i, builder) in builders.iter_mut().enumerate(){
match builder {
ColBuilder::I16(vec) => vec.push(row.get::<i16>(i).ok()),
ColBuilder::I32(vec) => vec.push(row.get::<i32>(i).ok()),
ColBuilder::I64(vec) => vec.push(row.get::<i64>(i).ok()),
ColBuilder::F64(vec) => vec.push(row.get::<f64>(i).ok()),
ColBuilder::Bool(vec) => vec.push(row.get::<bool>(i).ok()),
ColBuilder::String(vec) => vec.push(row.get::<String>(i).ok().map(|s| s.to_string())),
}
}
}
let mut columns = Vec::new();
for (i, builder) in builders.into_iter().enumerate(){
let col_name = expected_cols.get(i).unwrap_or(&"").split('.').last().unwrap_or("col");
let s = match builder {
ColBuilder::I16(vec) => Series::new(col_name.into(), vec),
ColBuilder::I32(vec) => Series::new(col_name.into(), vec),
ColBuilder::I64(vec) => Series::new(col_name.into(), vec),
ColBuilder::F64(vec) => Series::new(col_name.into(), vec),
ColBuilder::Bool(vec) => Series::new(col_name.into(), vec),
ColBuilder::String(vec) => Series::new(col_name.into(), vec),
};
columns.push(s.into());
}
Ok(DataFrame::new(columns)?)
}
async fn fetch_mysql_to_df(
client: &mut mysql_async::Conn,
sql: &str,
expected_cols: &[&str],
expected_types: &[&str],
)->anyhow::Result<DataFrame>{
enum ColBuilder{
I16(Vec<Option<i16>>),
I32(Vec<Option<i32>>),
I64(Vec<Option<i64>>),
F64(Vec<Option<f64>>),
Bool(Vec<Option<bool>>),
String(Vec<Option<String>>),
}
let mut builders: Vec<ColBuilder>= expected_types.iter().map(|&t|{
match t {
"Option<i16>" => ColBuilder::I16(Vec::new()),
"Option<i32>" => ColBuilder::I32(Vec::new()),
"Option<i64>" => ColBuilder::I64(Vec::new()),
"Option<f64>" | "Option<f32>" | "Option<data_goes::rust_decimal::Decimal>" => ColBuilder::F64(Vec::new()),
"Option<bool>" => ColBuilder::Bool(Vec::new()),
_ => ColBuilder::String(Vec::new()),
}
}).collect();
let mut result = client.query_iter(sql).await?;
while let Some(row)= result.next().await? {
for (i, builder) in builders.iter_mut().enumerate(){
match builder{
ColBuilder::I16(vec) => vec.push(row.get::<Option<i16>, _>(i).flatten()),
ColBuilder::I32(vec) => vec.push(row.get::<Option<i32>, _>(i).flatten()),
ColBuilder::I64(vec) => vec.push(row.get::<Option<i64>, _>(i).flatten()),
ColBuilder::F64(vec) => vec.push(row.get::<Option<f64>, _>(i).flatten()),
ColBuilder::Bool(vec) => vec.push(row.get::<Option<bool>, _>(i).flatten()),
ColBuilder::String(vec) => vec.push(row.get::<Option<String>, _>(i).flatten()),
}
}
}
let mut columns = Vec::new();
for (i,builder) in builders.into_iter().enumerate(){
let col_name = expected_cols.get(i).unwrap_or(&"").split('.').last().unwrap_or("col");
let s = match builder {
ColBuilder::I16(vec) => Series::new(col_name.into(), vec),
ColBuilder::I32(vec) => Series::new(col_name.into(), vec),
ColBuilder::I64(vec) => Series::new(col_name.into(), vec),
ColBuilder::F64(vec) => Series::new(col_name.into(), vec),
ColBuilder::Bool(vec) => Series::new(col_name.into(), vec),
ColBuilder::String(vec) => Series::new(col_name.into(), vec),
};
columns.push(s.into());
}
Ok(DataFrame::new(columns)?)
}
}