use std::{collections::HashMap};
use mysql_async::prelude::Queryable;
use tokio::sync::Mutex;
pub use crate::config::ConfigClient;
use mssql_client::{Ready};
use polars::{prelude::*};
use std::fs;
use polars::prelude::SerReader;
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(crate) origins: Arc<HashMap<String, (String, Option<String>, Option<String>)>>,
}
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;Command Timeout=0", 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;
}
}
let mut origins_map = std::collections::HashMap::new();
for origin in app_config.origins{
let path_obj = std::path::Path::new(&origin.path);
let folder_name = path_obj.file_name().and_then(|n| n.to_str()).unwrap_or("datalake");
let safe_name = folder_name.to_lowercase().replace(" ", "_").replace("-", "_").replace(".", "");
origins_map.insert(safe_name, (origin.path.clone(), origin.alias_db.clone(), origin.alias_table.clone()));
}
Ok(Self { clients: Arc::new(clients), origins: Arc::new(origins_map) })
}
pub async fn query_to_parquet(&self,server_name: &str, generate_sql: &str, expected_cols: &[&str], expected_types: &[&str])-> anyhow::Result<TabGoes>{
if let Some((origin_path, alias_db,alias_table)) = self.origins.get(server_name) {
return self.execute_sql_in_file(origin_path,alias_db.as_deref(),alias_table.as_deref(), generate_sql).await;
}
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, generate_sql, expected_cols, expected_types).await?
},
DbClient::MySql(client)=>{
Self::fetch_mysql_to_df(client, generate_sql, expected_cols, expected_types).await?
},
DbClient::ClickHouse{url,user,pass, ..}=>{
let req_client = reqwest::Client::new();
let sql_parquet = format!("{} FORMAT Parquet", generate_sql);
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(())
}
async fn execute_sql_in_file(&self, origin_path: &str,origin_alias_db:Option<&str>, origin_alias_table:Option<&str>, generate_sql: &str) -> anyhow::Result<TabGoes> {
let mut sql_clone = generate_sql.to_string();
let path_clone = origin_path.to_string();
let alias_db_clone= origin_alias_db.map(|s| s.to_string());
let alias_table_clone= origin_alias_table.map(|s| s.to_string());
let mut lazy_result = tokio::task::spawn_blocking(move || -> anyhow::Result<polars::prelude::LazyFrame> {
use polars::sql::SQLContext;
use calamine::{open_workbook_auto, Reader};
use std::io::Cursor;
let mut ctx = SQLContext::new();
let mut dir:Vec<_> = std::fs::read_dir(&path_clone)?.filter_map(Result::ok).collect();
dir.sort_by_key(|e|e.path());
let mut db_counter = 0;
for entry in dir {
let path = entry.path();
if path.is_file() && path.extension().unwrap_or_default() == "xlsx" {
if let Some(file_stem) = path.file_stem().and_then(|n| n.to_str()) {
let db_name = if let Some(ref alias)= alias_db_clone{
let name = if db_counter==0 {alias.clone()}else{format!("{}{}",alias,db_counter)};
db_counter+=1;
name
} else{if let Some(file_stem)=path.file_stem().and_then(|n|n.to_str()){
file_stem.to_lowercase().replace(" ", "_").replace("-", "_").replace(".", "")
}else{continue;}
};
if sql_clone.to_lowercase().contains(&db_name) {
let mut workbook = open_workbook_auto(&path)?;
let sheet_names = workbook.sheet_names().to_owned();
let mut tb_counter = 0;
for sheet_name in sheet_names {
let table_name = if let Some(ref alias) = alias_table_clone{
let name = if tb_counter ==0 {alias.clone()}else{format!("{}{}", alias,tb_counter)};
tb_counter += 1;
name
} else{
sheet_name.to_lowercase().replace(" ", "_").replace("-", "_").replace(".", "")
};
let table_identifier = format!("{}.{}", db_name, table_name);
if sql_clone.to_lowercase().contains(&table_identifier) {
let safe_identifier = table_identifier.replace(".", "_");
sql_clone = sql_clone.replace(&table_identifier, &safe_identifier);
if let Ok(range) = workbook.worksheet_range(&sheet_name) {
if !range.is_empty() {
let mut csv_buffer = csv::Writer::from_writer(Vec::new());
for row in range.rows() {
let row_strings: Vec<String> = row.iter().map(|cell| cell.to_string()).collect();
csv_buffer.write_record(&row_strings)?;
}
let cursor = Cursor::new(csv_buffer.into_inner()?);
let df = polars::prelude::CsvReadOptions::default()
.with_has_header(true)
.into_reader_with_file_handle(cursor)
.finish()?;
ctx.register(&safe_identifier, df.lazy());
}
}
}
}
}
}
}
}
let lf = ctx.execute(&sql_clone)?;
Ok(lf)
}).await??;
lazy_result = Self::standardize_columns(lazy_result)?;
let mut df_final = tokio::task::spawn_blocking(move || {
lazy_result.collect()
}).await??;
let file_path = format!("temp_query_{}.parquet", crate::uuid::Uuid::new_v4());
let mut file = std::fs::File::create(&file_path)?;
polars::prelude::ParquetWriter::new(&mut file).finish(&mut df_final)?;
drop(df_final);
let lf = polars::prelude::LazyFrame::scan_parquet(&file_path, polars::prelude::ScanArgsParquet::default())?;
Ok(TabGoes {
lazy: lf,
_files: vec![Arc::new(TempParquetFile(file_path))],
client: self.clone(),
})
}
pub async fn import_csv(&self, dir: &str, separator: u8)-> anyhow::Result<TabGoes>{
let dir_clone = dir.to_string();
let mut lazy_result = tokio::task::spawn_blocking(move || -> anyhow::Result<polars::prelude::LazyFrame>{
let lf = polars::prelude::LazyCsvReader::new(&dir_clone)
.with_has_header(true)
.with_separator(separator)
.finish()?;
Ok(lf)
}).await??;
let lazy_result = Self::standardize_columns(lazy_result)?;
Ok(TabGoes { lazy: lazy_result, _files: vec![], client: self.clone() })
}
}