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},
}
#[derive(Clone)]
pub struct ClientSql{
pub(crate) clients: Arc<HashMap<String, Mutex<DbClient>>>,
}
pub struct TempParquetFile(pub String);
impl Drop for TempParquetFile {
fn drop(&mut self) {
let _=std::fs::remove_file(&self.0);
}
}
#[derive(Clone)]
pub struct TabGoes{
pub lazy: polars::prelude::LazyFrame,
pub _files: Vec<Arc<TempParquetFile>>,
pub client: ClientSql,
}
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: Arc::new(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())?;
let lf = Self::standardize_columns(lf)?;
return Ok(TabGoes {
lazy: lf,
_files: vec![Arc::new(TempParquetFile(file_path))] ,
client: self.clone(),
});
}
};
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())?;
let lf = Self::standardize_columns(lf)?;
Ok(TabGoes {
lazy: lf,
_files: vec![Arc::new(TempParquetFile(file_path))],
client: self.clone(),
})
}
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>>),
DateTime(Vec<Option<i64>>),
}
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<data_goes::chrono::NaiveDateTime>" | "Option<data_goes::chrono::NaiveDate>" => ColBuilder::DateTime(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())),
ColBuilder::DateTime(vec) => {
if let Some(dt) = row.get::<crate::chrono::NaiveDateTime>(i).ok() {
vec.push(Some(dt.and_utc().timestamp_millis())); } else if let Some(d) = row.get::<crate::chrono::NaiveDate>(i).ok() {
vec.push(d.and_hms_opt(0, 0, 0).map(|dt| dt.and_utc().timestamp_millis())); } else {
vec.push(None);
}
},
}
}
}
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),
ColBuilder::DateTime(vec) => {
let temp_series = Series::new(col_name.into(), vec);
temp_series.cast(&polars::prelude::DataType::Datetime(polars::prelude::TimeUnit::Milliseconds, None)).unwrap_or(temp_series)
},
};
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>>),
DateTime(Vec<Option<i64>>),
}
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<data_goes::chrono::NaiveDateTime>" | "Option<data_goes::chrono::NaiveDate>" => ColBuilder::DateTime(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()),
ColBuilder::DateTime(vec) => {
if let Some(dt) = row.get::<Option<crate::chrono::NaiveDateTime>, _>(i).flatten() {
vec.push(Some(dt.and_utc().timestamp_millis())); } else if let Some(d) = row.get::<Option<crate::chrono::NaiveDate>, _>(i).flatten() {
vec.push(d.and_hms_opt(0, 0, 0).map(|dt| dt.and_utc().timestamp_millis())); } else {
vec.push(None);
}
},
}
}
}
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),
ColBuilder::DateTime(vec) => {
let temp_series = Series::new(col_name.into(), vec);
temp_series.cast(&polars::prelude::DataType::Datetime(polars::prelude::TimeUnit::Milliseconds, None)).unwrap_or(temp_series)
},
};
columns.push(s.into());
}
Ok(DataFrame::new(columns)?)
}
fn standardize_columns(mut lf: polars::prelude::LazyFrame) -> anyhow::Result<polars::prelude::LazyFrame> {
let schema = lf.collect_schema()?;
let old_names: Vec<String> = schema.iter_names().map(|s| s.to_string()).collect();
let new_names: Vec<String> = old_names.iter().map(|n| {
n.to_lowercase()
.replace(" ", "_")
.replace("-", "_")
}).collect();
Ok(lf.rename(old_names, new_names, true))
}
fn polars_to_ch_type(dtype: &polars::prelude::DataType) -> &'static str {
let type_str = format!("{:?}", dtype).to_lowercase();
if type_str.contains("int") && type_str.contains("64") { return "Int64"; }
if type_str.contains("int") { return "Int32"; }
if type_str.contains("float64") { return "Float64"; }
if type_str.contains("float") { return "Float32"; }
if type_str.contains("bool") { return "Bool"; }
if type_str.contains("datetime") { return "DateTime64(3)"; }
if type_str.contains("date") { return "Date"; }
"String"
}
pub async fn sink_parquet(&self, server_name: &str, db_name: &str, table: &str, path: &str, mode: crate::IfExists, schema: &polars::prelude::Schema) -> anyhow::Result<()> {
let client_mutex = self.clients.get(server_name).ok_or_else(|| anyhow::anyhow!("Servidor não achado"))?;
let mut client_lock = client_mutex.lock().await;
if let DbClient::ClickHouse{client, url, user, pass} = &mut *client_lock {
if matches!(mode, crate::IfExists::Replace) {
client.query(&format!("DROP TABLE IF EXISTS {}.{}", db_name, table)).execute().await?;
}
let mut col_defs = Vec::new();
for (name, dtype) in schema.iter() {
col_defs.push(format!("{} {}", name, Self::polars_to_ch_type(dtype)));
}
let create_sql = format!("CREATE TABLE IF NOT EXISTS {}.{} ({}) ENGINE = MergeTree() ORDER BY tuple()", db_name, table, col_defs.join(", "));
client.query(&create_sql).execute().await?;
if matches!(mode, crate::IfExists::Fail) {
let count = client.query(&format!("SELECT count() FROM {}.{}", db_name, table)).fetch_one::<u64>().await?;
if count > 0 { anyhow::bail!("Tabela já contém dados (Modo Fail)"); }
}
let file_bytes = tokio::fs::read(path).await?;
let insert_url = format!("{}/?query=INSERT+INTO+{}.{}+FORMAT+Parquet", url, db_name, table);
let resp = reqwest::Client::new().post(insert_url)
.basic_auth(user.as_str(), Some(pass.as_str()))
.body(file_bytes).send().await?;
if !resp.status().is_success() { anyhow::bail!("Erro ClickHouse: {}", resp.text().await?); }
}
Ok(())
}
}