athena_rs 2.0.2

//! Backup and restore API endpoints.
//!
//! Provides HTTP handlers for creating database backups via `pg_dump` (directory
//! format), uploading them to S3, listing available backups, and restoring via
//! `pg_restore`.
//!
//! ## Environment variables
//!
//! | Variable | Description |
//! |---|---|
//! | `ATHENA_BACKUP_S3_BUCKET` | S3 bucket name (required) |
//! | `ATHENA_BACKUP_S3_REGION` | AWS region (default: `us-east-1`) |
//! | `ATHENA_BACKUP_S3_ACCESS_KEY` | AWS access key ID |
//! | `ATHENA_BACKUP_S3_SECRET_KEY` | AWS secret access key |
//! | `ATHENA_BACKUP_S3_ENDPOINT` | Custom S3-compatible endpoint URL (optional) |
//! | `ATHENA_BACKUP_S3_PREFIX` | Key prefix for stored objects (default: `backups`) |
//! | `ATHENA_PG_DUMP_PATH` / `ATHENA_PG_RESTORE_PATH` | Optional absolute paths to pg_dump / pg_restore. When unset the server uses PATH, then auto-downloads a portable Linux x86_64 bundle. |
//!
//! ## Endpoints
//!
//! - `POST /admin/backups` – Run `pg_dump` and upload to S3.
//! - `GET /admin/backups` – List backup objects stored in S3.
//! - `GET /admin/backups/{key}/download` – Download a backup archive from S3.
//! - `POST /admin/backups/{key}/restore` – Download from S3 and run `pg_restore`.
//! - `DELETE /admin/backups/{key}` – Delete a backup from S3.

use std::env;
use std::path::PathBuf;

use actix_web::{
    HttpRequest, HttpResponse, delete, get, post,
    web::{self, Data, Json, Path},
};
use aws_config::BehaviorVersion;
use aws_sdk_s3::Client as S3Client;
use aws_sdk_s3::config::{Credentials, Region};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use serde_json::json;
use tokio::process::Command;
use uuid::Uuid;

use crate::AppState;
use crate::api::auth::authorize_static_admin_key;
use crate::api::response::{api_success, bad_request, internal_error, service_unavailable};
use crate::parser::resolve_postgres_uri;
use crate::utils::pg_tools::{PgToolsPaths, ensure_pg_tools};

// ── configuration helpers ────────────────────────────────────────────────────

/// Read a required S3 environment variable.
fn s3_env(key: &str) -> Option<String> {
    env::var(key).ok().filter(|v| !v.trim().is_empty())
}

struct S3Config {
    bucket: String,
    region: String,
    prefix: String,
    access_key: Option<String>,
    secret_key: Option<String>,
    endpoint: Option<String>,
}

impl S3Config {
    fn from_env() -> Option<Self> {
        let bucket = s3_env("ATHENA_BACKUP_S3_BUCKET")?;
        Some(Self {
            bucket,
            region: s3_env("ATHENA_BACKUP_S3_REGION").unwrap_or_else(|| "us-east-1".to_string()),
            prefix: s3_env("ATHENA_BACKUP_S3_PREFIX").unwrap_or_else(|| "backups".to_string()),
            access_key: s3_env("ATHENA_BACKUP_S3_ACCESS_KEY"),
            secret_key: s3_env("ATHENA_BACKUP_S3_SECRET_KEY"),
            endpoint: s3_env("ATHENA_BACKUP_S3_ENDPOINT"),
        })
    }
}

async fn build_s3_client(cfg: &S3Config) -> S3Client {
    let region: Region = Region::new(cfg.region.clone());

    let aws_config: aws_config::SdkConfig = if let (Some(ak), Some(sk)) =
        (&cfg.access_key, &cfg.secret_key)
    {
        let creds: Credentials = Credentials::new(ak, sk, None, None, "athena-env");
        let mut builder: aws_config::ConfigLoader = aws_config::defaults(BehaviorVersion::latest())
            .region(region)
            .credentials_provider(creds);
        if let Some(ep) = &cfg.endpoint {
            builder = builder.endpoint_url(ep);
        }
        builder.load().await
    } else {
        let mut builder: aws_config::ConfigLoader =
            aws_config::defaults(BehaviorVersion::latest()).region(region);
        if let Some(ep) = &cfg.endpoint {
            builder = builder.endpoint_url(ep);
        }
        builder.load().await
    };

    let mut s3_cfg_builder: aws_sdk_s3::config::Builder =
        aws_sdk_s3::config::Builder::from(&aws_config);
    if cfg.endpoint.is_some() {
        // Required for MinIO / path-style custom S3 endpoints.
        s3_cfg_builder = s3_cfg_builder.force_path_style(true);
    }

    S3Client::from_conf(s3_cfg_builder.build())
}

// ── request / response types ─────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
struct CreateBackupRequest {
    /// Logical client name whose Postgres URI will be used for `pg_dump`.
    client_name: String,
    /// Optional human-readable label stored in the S3 object metadata.
    #[serde(default)]
    label: Option<String>,
}

#[derive(Debug, Deserialize)]
struct RestoreBackupRequest {
    /// Logical client name whose Postgres URI is the restore target.
    client_name: String,
}

#[derive(Debug, Serialize)]
struct BackupObject {
    id: String,
    key: String,
    client_name: String,
    label: Option<String>,
    size_bytes: i64,
    created_at: String,
}

// ── helpers ───────────────────────────────────────────────────────────────────

/// Retrieve the Postgres connection URI for `client_name` from the registry.
fn resolve_pg_uri(state: &AppState, client_name: &str) -> Result<String, HttpResponse> {
    let registered: crate::drivers::postgresql::sqlx_driver::RegisteredClient = state
        .pg_registry
        .registered_client(client_name)
        .ok_or_else(|| {
            bad_request(
                "Unknown client",
                format!("No Postgres client named '{}' is registered.", client_name),
            )
        })?;

    // Resolve the URI: prefer config_uri_template (handles env var references),
    // then fall back to pg_uri if it was set directly.
    let uri = registered
        .config_uri_template
        .as_deref()
        .map(resolve_postgres_uri)
        .or(registered.pg_uri)
        .ok_or_else(|| {
            bad_request(
                "Client URI unavailable",
                format!("No Postgres URI is available for client '{}'.", client_name),
            )
        })?;

    Ok(uri)
}

/// Extract the password from a Postgres URI and return
/// `(sanitized_uri_without_password, Option<password>)`.
///
/// Keeps the password out of process command-line arguments (which are visible
/// in `ps` listings) by passing it through the `PGPASSWORD` environment
/// variable instead.
fn extract_pg_password(pg_uri: &str) -> (String, Option<String>) {
    // URI format: postgresql://[user[:password]@]host[:port][/dbname][?...]
    let prefix = if pg_uri.starts_with("postgresql://") {
        "postgresql://"
    } else if pg_uri.starts_with("postgres://") {
        "postgres://"
    } else {
        // Not a URI format we can parse; return as-is.
        return (pg_uri.to_string(), None);
    };

    let after_scheme = &pg_uri[prefix.len()..];

    // Find the last '@' which separates user-info from host.
    if let Some(at_pos) = after_scheme.rfind('@') {
        let userinfo = &after_scheme[..at_pos];
        let after_at = &after_scheme[at_pos..]; // includes the '@'

        if let Some(colon_pos) = userinfo.find(':') {
            let user = &userinfo[..colon_pos];
            let password = userinfo[colon_pos + 1..].to_string();
            let sanitized = format!("{}{}{}", prefix, user, after_at);
            return (sanitized, Some(password));
        }
    }

    (pg_uri.to_string(), None)
}

/// Create a temporary directory, run `pg_dump`, archive the result, return the
/// archive path.  All temporary intermediate directories are cleaned up on
/// error; the caller is responsible for cleaning up the returned path.
async fn run_pg_dump(pg_uri: &str) -> Result<PathBuf, String> {
    let tmp_root = env::temp_dir().join(format!("athena_backup_{}", Uuid::new_v4()));
    let dump_dir = tmp_root.join("dump");
    let archive_path = tmp_root.join("backup.tar.gz");

    let pg_tools: PgToolsPaths = ensure_pg_tools()
        .await
        .map_err(|e| format!("pg_dump resolution failed: {e}"))?;

    tokio::fs::create_dir_all(&dump_dir)
        .await
        .map_err(|e| format!("Could not create temp directory: {e}"))?;

    // Strip password from URI and pass it via PGPASSWORD to avoid exposing
    // credentials in process command-line listings.
    let (pg_uri_safe, pg_password) = extract_pg_password(pg_uri);

    // Run pg_dump with directory format.
    let mut cmd = Command::new(&pg_tools.pg_dump);
    if let Some(pass) = pg_password {
        cmd.env("PGPASSWORD", pass);
    }
    let status = cmd
        .args(["--format=directory", "--file"])
        .arg(&dump_dir)
        .arg(&pg_uri_safe)
        .status()
        .await
        .map_err(|e| {
            if e.kind() == std::io::ErrorKind::NotFound {
                "pg_dump binary could not be resolved (PATH/env/download). Set ATHENA_PG_DUMP_PATH to override."
                    .to_string()
            } else {
                format!("Failed to start pg_dump: {e}")
            }
        })?;

    if !status.success() {
        let _ = tokio::fs::remove_dir_all(&tmp_root).await;
        return Err(format!("pg_dump exited with status {status}"));
    }

    // Archive the dump directory into a tar.gz.
    let dump_dir_clone = dump_dir.clone();
    let archive_path_clone = archive_path.clone();
    tokio::task::spawn_blocking(move || -> Result<(), String> {
        let file = std::fs::File::create(&archive_path_clone)
            .map_err(|e| format!("Cannot create archive: {e}"))?;
        let enc = flate2::write::GzEncoder::new(file, flate2::Compression::default());
        let mut builder = tar::Builder::new(enc);
        builder
            .append_dir_all("dump", &dump_dir_clone)
            .map_err(|e| format!("Cannot archive dump directory: {e}"))?;
        builder
            .finish()
            .map_err(|e| format!("Cannot finalize archive: {e}"))?;
        Ok(())
    })
    .await
    .map_err(|e| format!("Archive task panicked: {e}"))??;

    // Remove uncompressed dump dir, keep only the archive.
    let _ = tokio::fs::remove_dir_all(&dump_dir).await;

    Ok(archive_path)
}

/// Download `key` from S3, extract to a temp directory, run `pg_restore`.
async fn run_pg_restore(
    s3_client: &S3Client,
    bucket: &str,
    key: &str,
    pg_uri: &str,
) -> Result<(), String> {
    // Download the backup archive.
    let resp: aws_sdk_s3::operation::get_object::GetObjectOutput = s3_client
        .get_object()
        .bucket(bucket)
        .key(key)
        .send()
        .await
        .map_err(|e| format!("S3 download failed: {e}"))?;

    let bytes: web::Bytes = resp
        .body
        .collect()
        .await
        .map_err(|e| format!("S3 read failed: {e}"))?
        .into_bytes();

    let tmp_root: PathBuf = env::temp_dir().join(format!("athena_restore_{}", Uuid::new_v4()));
    tokio::fs::create_dir_all(&tmp_root)
        .await
        .map_err(|e| format!("Could not create temp dir: {e}"))?;

    let archive_path: PathBuf = tmp_root.join("backup.tar.gz");
    let restore_dir: PathBuf = tmp_root.join("dump");

    tokio::fs::write(&archive_path, &bytes)
        .await
        .map_err(|e| format!("Could not write archive: {e}"))?;

    // Extract the tar.gz.
    let archive_path_clone: PathBuf = archive_path.clone();
    let restore_dir_clone: PathBuf = restore_dir.clone();
    tokio::task::spawn_blocking(move || -> Result<(), String> {
        let file = std::fs::File::open(&archive_path_clone)
            .map_err(|e| format!("Cannot open archive: {e}"))?;
        let dec = flate2::read::GzDecoder::new(file);
        let mut archive: tar::Archive<flate2::read::GzDecoder<std::fs::File>> =
            tar::Archive::new(dec);
        archive
            .unpack(&restore_dir_clone)
            .map_err(|e| format!("Cannot extract archive: {e}"))?;
        Ok(())
    })
    .await
    .map_err(|e| format!("Extract task panicked: {e}"))??;

    // The pg_dump directory format stores files under the name we passed to
    // append_dir_all ("dump"), so the actual dump content is at restore_dir/dump.
    let inner_dump_dir: PathBuf = restore_dir.join("dump");

    // Strip password from URI and pass via PGPASSWORD.
    let (pg_uri_safe, pg_password) = extract_pg_password(pg_uri);

    let pg_tools: PgToolsPaths = ensure_pg_tools()
        .await
        .map_err(|e| format!("pg_restore resolution failed: {e}"))?;

    let mut cmd: Command = Command::new(&pg_tools.pg_restore);
    if let Some(pass) = pg_password {
        cmd.env("PGPASSWORD", pass);
    }
    let status: std::process::ExitStatus = cmd
        .args(["--format=directory", "--clean", "--if-exists", "--dbname"])
        .arg(&pg_uri_safe)
        .arg(&inner_dump_dir)
        .status()
        .await
        .map_err(|e| {
            if e.kind() == std::io::ErrorKind::NotFound {
                "pg_restore binary not found in PATH — ensure PostgreSQL client tools are installed"
                    .to_string()
            } else {
                format!("Failed to start pg_restore: {e}")
            }
        })?;

    let _ = tokio::fs::remove_dir_all(&tmp_root).await;

    if !status.success() {
        return Err(format!("pg_restore exited with status {status}"));
    }

    Ok(())
}

/// Upload a file to S3 and return its object key.
async fn upload_to_s3(
    s3_client: &S3Client,
    cfg: &S3Config,
    local_path: &PathBuf,
    client_name: &str,
    label: Option<&str>,
) -> Result<String, String> {
    let backup_id: String = Uuid::new_v4().to_string();
    let key: String = format!("{}/{}/{}.tar.gz", cfg.prefix, client_name, backup_id);

    let data: Vec<u8> = tokio::fs::read(local_path)
        .await
        .map_err(|e| format!("Cannot read archive file: {e}"))?;

    let mut req: aws_sdk_s3::operation::put_object::builders::PutObjectFluentBuilder = s3_client
        .put_object()
        .bucket(&cfg.bucket)
        .key(&key)
        .body(data.into())
        .content_type("application/gzip")
        .metadata("client_name", client_name)
        .metadata("backup_id", &backup_id)
        .metadata("created_at", Utc::now().to_rfc3339());

    if let Some(lbl) = label {
        req = req.metadata("label", lbl);
    }

    req.send()
        .await
        .map_err(|e| format!("S3 upload failed: {e}"))?;

    Ok(key)
}

// ── handlers ─────────────────────────────────────────────────────────────────

/// Create a database backup via `pg_dump` and upload it to S3.
///
/// # Request body
/// ```json
/// { "client_name": "my_db", "label": "pre-deploy" }
/// ```
#[post("/admin/backups")]
pub async fn admin_create_backup(
    req: HttpRequest,
    state: Data<AppState>,
    body: Json<CreateBackupRequest>,
) -> HttpResponse {
    if let Err(resp) = authorize_static_admin_key(&req) {
        return resp;
    }

    let Some(s3_cfg) = S3Config::from_env() else {
        return service_unavailable(
            "S3 not configured",
            "Set ATHENA_BACKUP_S3_BUCKET and related environment variables to enable backups.",
        );
    };

    let pg_uri = match resolve_pg_uri(&state, &body.client_name) {
        Ok(uri) => uri,
        Err(resp) => return resp,
    };

    // Run pg_dump.
    let archive_path = match run_pg_dump(&pg_uri).await {
        Ok(p) => p,
        Err(err) => return internal_error("pg_dump failed", err),
    };

    // Upload to S3.
    let s3_client = build_s3_client(&s3_cfg).await;
    let key = match upload_to_s3(
        &s3_client,
        &s3_cfg,
        &archive_path,
        &body.client_name,
        body.label.as_deref(),
    )
    .await
    {
        Ok(k) => k,
        Err(err) => {
            let _ = tokio::fs::remove_file(&archive_path).await;
            return internal_error("S3 upload failed", err);
        }
    };

    // Clean up local archive.
    if let Some(parent) = archive_path.parent() {
        let _ = tokio::fs::remove_dir_all(parent).await;
    }

    api_success(
        "Backup created",
        json!({
            "key": key,
            "client_name": body.client_name,
            "label": body.label,
        }),
    )
}

/// List backup objects stored in S3.
///
/// Optionally filter by `?client_name=…`.
#[get("/admin/backups")]
pub async fn admin_list_backups(
    req: HttpRequest,
    _state: Data<AppState>,
    query: web::Query<std::collections::HashMap<String, String>>,
) -> HttpResponse {
    if let Err(resp) = authorize_static_admin_key(&req) {
        return resp;
    }

    let Some(s3_cfg) = S3Config::from_env() else {
        return service_unavailable(
            "S3 not configured",
            "Set ATHENA_BACKUP_S3_BUCKET and related environment variables to enable backups.",
        );
    };

    let s3_client = build_s3_client(&s3_cfg).await;

    let filter_client = query.get("client_name").cloned();
    let prefix = match &filter_client {
        Some(cn) => format!("{}/{}/", s3_cfg.prefix, cn),
        None => format!("{}/", s3_cfg.prefix),
    };

    let resp = match s3_client
        .list_objects_v2()
        .bucket(&s3_cfg.bucket)
        .prefix(&prefix)
        .send()
        .await
    {
        Ok(r) => r,
        Err(err) => return internal_error("Failed to list S3 objects", err.to_string()),
    };

    let mut backups: Vec<BackupObject> = Vec::new();

    for obj in resp.contents() {
        let key = obj.key().unwrap_or_default().to_string();

        // Expected key format: {prefix}/{client_name}/{backup_id}.tar.gz
        let parts: Vec<&str> = key.split('/').collect();
        let (client_name, id) = if parts.len() >= 3 {
            let cn = parts[parts.len() - 2].to_string();
            let id = parts
                .last()
                .and_then(|s| s.strip_suffix(".tar.gz"))
                .unwrap_or_else(|| parts.last().copied().unwrap_or(""))
                .to_string();
            (cn, id)
        } else {
            // Unexpected key format; still include the object but log a warning.
            tracing::warn!(key = %key, "S3 backup key does not match expected format <prefix>/<client_name>/<id>.tar.gz");
            (
                "unknown".to_string(),
                parts
                    .last()
                    .and_then(|s| s.strip_suffix(".tar.gz"))
                    .unwrap_or_else(|| parts.last().copied().unwrap_or(&key))
                    .to_string(),
            )
        };

        // Fetch object metadata for label.
        let label = match s3_client
            .head_object()
            .bucket(&s3_cfg.bucket)
            .key(&key)
            .send()
            .await
        {
            Ok(head) => head.metadata().and_then(|m| m.get("label")).cloned(),
            Err(_) => None,
        };

        let size_bytes = obj.size().unwrap_or(0);
        let created_at = obj
            .last_modified()
            .map(|t| t.to_string())
            .unwrap_or_default();

        backups.push(BackupObject {
            id,
            key,
            client_name,
            label,
            size_bytes,
            created_at,
        });
    }

    // Most-recent first.
    backups.sort_by(|a, b| b.created_at.cmp(&a.created_at));

    api_success("Listed backups", json!({ "backups": backups }))
}

/// Restore a database from a backup stored in S3 via `pg_restore`.
///
/// `{key:.*}` is the S3 object key (may contain slashes).
///
/// # Request body
/// ```json
/// { "client_name": "my_db" }
/// ```
#[post("/admin/backups/{key:.*}/restore")]
pub async fn admin_restore_backup(
    req: HttpRequest,
    state: Data<AppState>,
    key_param: Path<String>,
    body: Json<RestoreBackupRequest>,
) -> HttpResponse {
    if let Err(resp) = authorize_static_admin_key(&req) {
        return resp;
    }

    let Some(s3_cfg) = S3Config::from_env() else {
        return service_unavailable(
            "S3 not configured",
            "Set ATHENA_BACKUP_S3_BUCKET and related environment variables to enable backups.",
        );
    };

    let pg_uri = match resolve_pg_uri(&state, &body.client_name) {
        Ok(uri) => uri,
        Err(resp) => return resp,
    };

    // The key_param path segment contains the S3 object key (URL-decoded by Actix).
    let key = key_param.into_inner();
    if key.is_empty() {
        return bad_request(
            "Missing backup key",
            "Provide the S3 object key as the path segment.",
        );
    }

    let s3_client = build_s3_client(&s3_cfg).await;

    match run_pg_restore(&s3_client, &s3_cfg.bucket, &key, &pg_uri).await {
        Ok(()) => api_success(
            "Restore completed",
            json!({ "key": key, "client_name": body.client_name }),
        ),
        Err(err) => internal_error("pg_restore failed", err),
    }
}

/// Download a backup archive directly from S3.
///
/// `{key:.*}` is the S3 object key (may contain slashes).  Returns the raw
/// `application/gzip` bytes with a `Content-Disposition: attachment` header.
#[get("/admin/backups/{key:.*}/download")]
pub async fn admin_download_backup(
    req: HttpRequest,
    _state: Data<AppState>,
    key_param: Path<String>,
) -> HttpResponse {
    if let Err(resp) = authorize_static_admin_key(&req) {
        return resp;
    }

    let Some(s3_cfg) = S3Config::from_env() else {
        return service_unavailable(
            "S3 not configured",
            "Set ATHENA_BACKUP_S3_BUCKET and related environment variables to enable backups.",
        );
    };

    let key = key_param.into_inner();
    if key.is_empty() {
        return bad_request(
            "Missing backup key",
            "Provide the S3 object key as the path segment.",
        );
    }

    let s3_client = build_s3_client(&s3_cfg).await;

    let resp = match s3_client
        .get_object()
        .bucket(&s3_cfg.bucket)
        .key(&key)
        .send()
        .await
    {
        Ok(r) => r,
        Err(err) => return internal_error("S3 download failed", err.to_string()),
    };

    let bytes = match resp.body.collect().await {
        Ok(b) => b.into_bytes(),
        Err(err) => return internal_error("S3 read failed", err.to_string()),
    };

    let filename = key
        .rsplit('/')
        .next()
        .unwrap_or("backup.tar.gz")
        .to_string();

    HttpResponse::Ok()
        .content_type("application/gzip")
        .insert_header((
            "Content-Disposition",
            format!("attachment; filename=\"{}\"", filename),
        ))
        .body(bytes.to_vec())
}

/// Delete a backup from S3.
///
/// `{key:.*}` is the S3 object key (may contain slashes).
#[delete("/admin/backups/{key:.*}")]
pub async fn admin_delete_backup(
    req: HttpRequest,
    _state: Data<AppState>,
    key_param: Path<String>,
) -> HttpResponse {
    if let Err(resp) = authorize_static_admin_key(&req) {
        return resp;
    }

    let Some(s3_cfg) = S3Config::from_env() else {
        return service_unavailable(
            "S3 not configured",
            "Set ATHENA_BACKUP_S3_BUCKET and related environment variables to enable backups.",
        );
    };

    // The key_param path segment contains the S3 object key (URL-decoded by Actix).
    let key = key_param.into_inner();
    if key.is_empty() {
        return bad_request(
            "Missing backup key",
            "Provide the S3 object key as the path segment.",
        );
    }

    let s3_client: S3Client = build_s3_client(&s3_cfg).await;

    match s3_client
        .delete_object()
        .bucket(&s3_cfg.bucket)
        .key(&key)
        .send()
        .await
    {
        Ok(_) => api_success("Backup deleted", json!({ "key": key })),
        Err(err) => internal_error("S3 delete failed", err.to_string()),
    }
}

pub fn services(cfg: &mut web::ServiceConfig) {
    cfg.service(admin_create_backup)
        .service(admin_list_backups)
        .service(admin_download_backup)
        .service(admin_restore_backup)
        .service(admin_delete_backup);
}