athena-gateway 3.18.0

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::{
    GatewayDeleteRequest, GatewayFetchRequest, GatewayInsertRequest, GatewayRpcRequest,
    GatewayUpdateRequest,
};

/// Portable execution modes accepted by gateway SQL request payloads.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GatewaySqlExecutionMode {
    SingleTransaction,
    PerStatement,
}

impl Default for GatewaySqlExecutionMode {
    fn default() -> Self {
        Self::SingleTransaction
    }
}

/// Canonical `/gateway/query` request shape.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GatewaySqlRequest {
    pub query: String,
    #[serde(default, alias = "schema")]
    pub schema_name: Option<String>,
    #[serde(default)]
    pub execution_mode: Option<GatewaySqlExecutionMode>,
    #[serde(default)]
    pub params: Vec<Value>,
}

/// Canonical `/query/sql` + `/gateway/sql` request shape.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GatewaySqlExecutionRequest {
    pub query: String,
    pub driver: String,
    pub db_name: String,
    #[serde(default, alias = "schema")]
    pub schema_name: Option<String>,
    #[serde(default)]
    pub execution_mode: Option<GatewaySqlExecutionMode>,
    #[serde(default)]
    pub params: Vec<Value>,
}

/// API-level operation variants for normalized gateway requests.
#[derive(Debug, Clone)]
pub enum GatewayApiRequestPayload {
    Fetch(GatewayFetchRequest),
    Insert(GatewayInsertRequest),
    Update(GatewayUpdateRequest),
    Delete(GatewayDeleteRequest),
    Sql(GatewaySqlRequest),
    Rpc(GatewayRpcRequest),
}

/// Unified API-level gateway request wrapper.
#[derive(Debug, Clone)]
pub struct GatewayApiRequest {
    payload: GatewayApiRequestPayload,
}

impl GatewayApiRequest {
    /// Returns a borrowed view of the normalized gateway payload.
    pub fn payload(&self) -> &GatewayApiRequestPayload {
        &self.payload
    }

    /// Consumes the wrapper and returns the normalized gateway payload.
    pub fn into_payload(self) -> GatewayApiRequestPayload {
        self.payload
    }
}

impl From<GatewayFetchRequest> for GatewayApiRequest {
    fn from(value: GatewayFetchRequest) -> Self {
        Self {
            payload: GatewayApiRequestPayload::Fetch(value),
        }
    }
}

impl From<GatewayInsertRequest> for GatewayApiRequest {
    fn from(value: GatewayInsertRequest) -> Self {
        Self {
            payload: GatewayApiRequestPayload::Insert(value),
        }
    }
}

impl From<GatewayUpdateRequest> for GatewayApiRequest {
    fn from(value: GatewayUpdateRequest) -> Self {
        Self {
            payload: GatewayApiRequestPayload::Update(value),
        }
    }
}

impl From<GatewayDeleteRequest> for GatewayApiRequest {
    fn from(value: GatewayDeleteRequest) -> Self {
        Self {
            payload: GatewayApiRequestPayload::Delete(value),
        }
    }
}

impl From<GatewaySqlRequest> for GatewayApiRequest {
    fn from(value: GatewaySqlRequest) -> Self {
        Self {
            payload: GatewayApiRequestPayload::Sql(value),
        }
    }
}

impl From<GatewayRpcRequest> for GatewayApiRequest {
    fn from(value: GatewayRpcRequest) -> Self {
        Self {
            payload: GatewayApiRequestPayload::Rpc(value),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn gateway_sql_request_supports_execution_mode() {
        let parsed: GatewaySqlRequest = serde_json::from_value(json!({
            "query": "SELECT 1",
            "execution_mode": "per_statement",
            "params": [1, "two"]
        }))
        .expect("request should parse");

        assert_eq!(
            parsed.execution_mode,
            Some(GatewaySqlExecutionMode::PerStatement)
        );
        assert_eq!(parsed.params, vec![json!(1), json!("two")]);
    }
}