use crate::cli::args::{DbCommand, GlobalArgs, ExportFormat as CliExportFormat};
use crate::cli::output::OutputWriter;
use crate::db::duckdb::{AnalyticsDb, ExportFormat};
use crate::error::{ExitStatus, GdeltError, Result};
use serde_json::json;
pub async fn handle_db(cmd: DbCommand, global: &GlobalArgs) -> Result<ExitStatus> {
match cmd {
DbCommand::Stats(args) => {
let db = AnalyticsDb::open()?;
let output = OutputWriter::new(global);
let stats = db.stats()?;
if args.detailed {
let events_by_year = db.query(
"SELECT year, COUNT(*) as count FROM events GROUP BY year ORDER BY year DESC LIMIT 10"
)?;
let top_countries = db.query(
"SELECT action_geo_country_code as country, COUNT(*) as count
FROM events WHERE action_geo_country_code IS NOT NULL
GROUP BY action_geo_country_code ORDER BY count DESC LIMIT 10"
)?;
output.write_value(&json!({
"summary": stats,
"events_by_year": events_by_year,
"top_countries": top_countries,
}))?;
} else {
output.write_value(&stats)?;
}
Ok(ExitStatus::Success)
}
DbCommand::Vacuum(args) => {
let db = AnalyticsDb::open()?;
let output = OutputWriter::new(global);
if args.analyze {
db.analyze()?;
}
db.vacuum()?;
let stats = db.stats()?;
output.write_value(&json!({
"status": "success",
"message": "Database optimized",
"file_size_bytes": stats.file_size_bytes,
}))?;
Ok(ExitStatus::Success)
}
DbCommand::Export(args) => {
let db = AnalyticsDb::open()?;
let output = OutputWriter::new(global);
let format = match args.format {
CliExportFormat::Parquet => ExportFormat::Parquet,
CliExportFormat::Csv => ExportFormat::Csv,
CliExportFormat::Json => ExportFormat::Json,
};
let table = if let Some(ref filter) = args.filter {
let temp_view = format!(
"CREATE OR REPLACE TEMP VIEW export_view AS SELECT * FROM {} WHERE {}",
args.table, filter
);
db.query(&temp_view)?;
"export_view"
} else {
&args.table
};
db.export(table, &args.output, format)?;
output.write_value(&json!({
"status": "success",
"table": args.table,
"output": args.output.display().to_string(),
"format": format!("{:?}", args.format).to_lowercase(),
}))?;
Ok(ExitStatus::Success)
}
DbCommand::Import(args) => {
let db = AnalyticsDb::open()?;
let output = OutputWriter::new(global);
let ext = args.input.extension()
.and_then(|e| e.to_str())
.unwrap_or("");
let format = args.format.unwrap_or_else(|| {
match ext {
"parquet" => CliExportFormat::Parquet,
"csv" => CliExportFormat::Csv,
"json" | "jsonl" => CliExportFormat::Json,
_ => CliExportFormat::Csv,
}
});
let path_str = args.input.to_string_lossy();
let sql = match format {
CliExportFormat::Parquet => format!(
"INSERT INTO {} SELECT * FROM read_parquet('{}')",
args.table, path_str
),
CliExportFormat::Csv => format!(
"INSERT INTO {} SELECT * FROM read_csv_auto('{}', header=true)",
args.table, path_str
),
CliExportFormat::Json => format!(
"INSERT INTO {} SELECT * FROM read_json_auto('{}')",
args.table, path_str
),
};
db.query(&sql)?;
let stats = db.stats()?;
output.write_value(&json!({
"status": "success",
"input": args.input.display().to_string(),
"table": args.table,
"current_counts": {
"events": stats.events_count,
"gkg": stats.gkg_count,
"mentions": stats.mentions_count,
}
}))?;
Ok(ExitStatus::Success)
}
DbCommand::Query(args) => {
let db = AnalyticsDb::open()?;
let output = OutputWriter::new(global);
let query_lower = args.query.to_lowercase();
if query_lower.contains("drop ") || query_lower.contains("truncate ") ||
query_lower.contains("delete ") || query_lower.contains("alter ") {
return Err(GdeltError::Validation(
"Destructive queries not allowed. Use data management commands instead.".into()
));
}
let result = db.query(&args.query)?;
output.write_value(&result)?;
Ok(ExitStatus::Success)
}
}
}