use std::{collections::HashMap, fs::read_to_string, io::Write};
use anyhow::Result;
use rbatis::Rbatis;
use rbs::{to_value, Value};
use tokio_nsq::NSQProducer;
use crate::{dbs, files, nsqs, sqls, strs};
pub async fn db2file(
rb: &Rbatis,
sql: &str,
id_field: &str,
batch_num: u64,
mut args: HashMap<String, Value>,
spl: &str,
file_path: &str,
schema: &str,
) {
let mut scm = vec![];
let schema_file_name = format!("{}.schema", file_path);
if schema.is_empty() {
if files::exists(&schema_file_name) {
let sc = read_to_string(&schema_file_name).unwrap();
scm = strs::str_2_vec::<String>(&sc.trim(), ",");
}
} else {
scm = strs::str_2_vec::<String>(schema.trim(), ",");
}
if scm.len() == 0 {
let schema = sqls::get_select_fields(sql);
for ele in schema {
scm.push(ele.to_string());
}
}
if !scm.contains(&id_field.to_string()) {
println!("id字段[{id_field}]不存在");
return;
}
let mut id_idx = 0;
let mut idx = 0;
for f in &scm {
if f.eq(id_field) {
id_idx = idx;
break;
}
idx += 1;
}
let last_line = files::read_last_line(file_path, 500);
let mut exist_id = 0u64;
match last_line {
Some(line) => {
let sc: Vec<&str> = line.split(spl).collect();
let get = sc.get(id_idx);
exist_id = get.unwrap().parse().unwrap_or_default();
}
None => {}
}
args.insert("exist_id".to_string(), to_value!(exist_id));
args.insert("limit".to_string(), to_value!(batch_num));
let sql = &format!(
"select * from ({sql}) as _t where _t.{id_field}>@exist_id limit @limit offset @offset "
);
let mut offset = 0;
loop {
args.insert("offset".to_string(), to_value!(offset));
let list_ret: Result<Vec<HashMap<String, Value>>> =
dbs::named_fetch(rb, sql, args.clone()).await;
match list_ret {
Ok(list) => {
let len = list.len();
offset += batch_num;
if scm.len() == 1 && scm.get(0).unwrap().eq("*") {
scm.clear();
}
let mut str = String::new();
for map in list {
if scm.len() == 0 {
for (k, v) in map {
let value = k.to_string();
scm.push(value);
if v.is_str() {
str.push_str(v.as_str().unwrap());
} else {
str.push_str(v.to_string().as_str());
}
str.push_str(spl);
}
str.pop();
str.push('\n');
continue;
}
for k in &scm {
match map.get(k) {
Some(v) => {
if v.is_str() {
str.push_str(v.as_str().unwrap());
} else {
str.push_str(v.to_string().as_str());
}
}
None => {}
};
str.push_str(spl);
}
str.pop();
str.push('\n');
}
if !str.is_empty() {
files::open_file(file_path)
.write_all(str.as_bytes())
.unwrap();
let join = scm.join(",");
files::create_file(&schema_file_name)
.unwrap()
.write_all(join.as_bytes())
.unwrap();
}
print!("\nsync {len} records to {file_path}");
if len < batch_num as usize {
break;
}
}
Err(err) => {
println!("{err}");
break;
}
}
}
}
pub async fn db2file_with_nsq(
rb: &Rbatis,
sql: &str,
id_field: &str,
batch_num: u64,
mut args: HashMap<String, Value>,
spl: &str,
file_path: &str,
nsq_topic: &str,
schema: &str,
mut nsq_producer: NSQProducer,
) {
let mut scm = vec![];
let schema_file_name = format!("{}.schema", file_path);
if schema.is_empty() {
if files::exists(&schema_file_name) {
let sc = read_to_string(&schema_file_name).unwrap();
scm = strs::str_2_vec::<String>(&sc.trim(), ",");
}
} else {
scm = strs::str_2_vec::<String>(schema.trim(), ",");
}
if scm.len() == 0 {
let schema = sqls::get_select_fields(sql);
for ele in schema {
scm.push(ele.to_string());
}
}
if !scm.contains(&id_field.to_string()) {
println!("id字段[{id_field}]不存在");
return;
}
let mut id_idx = 0;
let mut idx = 0;
for f in &scm {
if f.eq(id_field) {
id_idx = idx;
break;
}
idx += 1;
}
let last_line = files::read_last_line(file_path, 500);
let mut exist_id = 0u64;
match last_line {
Some(line) => {
let sc: Vec<&str> = line.split(spl).collect();
let get = sc.get(id_idx);
exist_id = get.unwrap().parse().unwrap_or_default();
}
None => {}
}
args.insert("exist_id".to_string(), to_value!(exist_id));
args.insert("limit".to_string(), to_value!(batch_num));
let sql = &format!(
"select * from ({sql}) as _t where _t.{id_field}>@exist_id limit @limit offset @offset "
);
let mut offset = 0;
let topic = nsqs::topic(nsq_topic);
loop {
args.insert("offset".to_string(), to_value!(offset));
let list_ret: Result<Vec<HashMap<String, Value>>> =
dbs::named_fetch(rb, sql, args.clone()).await;
match list_ret {
Ok(list) => {
let len = list.len();
offset += batch_num;
if scm.len() == 1 && scm.get(0).unwrap().eq("*") {
scm.clear();
}
let mut str = String::new();
for map in list {
let mut line = String::new();
if scm.len() == 0 {
for (k, v) in map {
let value = k.to_string();
scm.push(value);
if v.is_str() {
line.push_str(v.as_str().unwrap());
} else {
line.push_str(v.to_string().as_str());
}
line.push_str(spl);
}
} else {
for k in &scm {
match map.get(k) {
Some(v) => {
if v.is_str() {
line.push_str(v.as_str().unwrap());
} else {
line.push_str(v.to_string().as_str());
}
}
None => {}
};
line.push_str(spl);
}
}
line.pop();
nsq_producer
.publish(&topic, line.clone().as_bytes().to_vec())
.await
.unwrap();
line.push('\n');
str.push_str(&line);
}
if !str.is_empty() {
let evt = nsq_producer.consume().await.unwrap();
if !nsqs::is_ok(evt) {
break;
}
files::open_file(file_path)
.write_all(str.as_bytes())
.unwrap();
let join = scm.join(",");
files::create_file(&schema_file_name)
.unwrap()
.write_all(join.as_bytes())
.unwrap();
}
print!("\nsync {len} records to {file_path}");
if len < batch_num as usize {
break;
}
}
Err(err) => {
println!("{err}");
break;
}
}
}
}