gdelt 0.1.0

CLI for GDELT Project - optimized for agentic usage with local data caching
//! Database management command handlers.

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;

/// Handle DB commands
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 {
                // Add more detailed stats
                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,
            };

            // If filter specified, create a temp view
            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);

            // Detect format from extension if not specified
            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);

            // Basic SQL injection prevention
            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)
        }
    }
}