athena_rs 3.26.1

//! Backend trait definitions for Athena client.
use async_trait::async_trait;
use serde_json::Value;
use std::any::Any;
use std::fmt::{Debug, Formatter, Result as StdResult};
use std::result::Result;
use std::time::Instant;

/// ## `BackendType`
/// Supported database backends.
///
/// # Arguments
///
/// * `Native` - The native backend.
/// * `Supabase` - The Supabase backend.
/// * `Postgrest` - The Postgrest backend.
/// * `Scylla` - The Scylla backend.
/// * `Neon` - The Neon backend.
/// * `PostgreSQL` - The PostgreSQL backend.
/// * `D1` - The Cloudflare D1 backend.
///
/// # Returns
///
/// A `BackendType` containing the backend type.
///
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendType {
    Native,
    Supabase,
    Postgrest,
    Scylla,
    Neon,
    PostgreSQL,
    D1,
}

/// ## `QueryLanguage`
/// The type of query language that will be executed.
///
/// # Arguments
///
/// * `Sql` - The SQL query language.
/// * `Cql` - The CQL query language.
/// * `Postgrest` - The Postgrest query language.
///
/// # Returns
///
/// A `QueryLanguage` containing the query language.
///
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QueryLanguage {
    Sql,
    Cql,
    Postgrest,
}

/// HTTP method for PostgREST request execution.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PostgrestMethod {
    Get,
    Post,
    Patch,
    Delete,
}

/// ## `TranslatedQuery`
/// A query that has already been translated into backend-specific SQL/CQL.
///
/// # Arguments
///
/// * `sql` - The SQL query.
/// * `language` - The query language.
/// * `params` - The parameters.
/// * `table` - The table.
///
/// # Returns
///
/// A `TranslatedQuery` containing the translated query.
///
#[derive(Debug, Clone)]
pub struct TranslatedQuery {
    pub sql: String,
    pub language: QueryLanguage,
    pub params: Vec<Value>,
    pub table: Option<String>,
    pub schema_name: Option<String>,
    pub postgrest_method: Option<PostgrestMethod>,
    pub payload: Option<Value>,
}

impl TranslatedQuery {
    /// ## `new`
    /// Create a new translated query.
    ///
    /// # Arguments
    ///
    /// * `sql` - The SQL query.
    /// * `language` - The query language.
    /// * `params` - The parameters.
    /// * `table` - The table.
    ///
    /// # Returns
    ///
    /// A `TranslatedQuery` containing the translated query.
    ///
    pub fn new(
        sql: impl Into<String>,
        language: QueryLanguage,
        params: Vec<Value>,
        table: Option<String>,
    ) -> Self {
        Self {
            sql: sql.into(),
            language,
            params,
            table,
            schema_name: None,
            postgrest_method: None,
            payload: None,
        }
    }

    /// Create a SQL translated query.
    ///
    /// # Arguments
    ///
    /// * `sql` - The SQL query.
    /// * `params` - The parameters.
    /// * `table` - The table.
    ///
    /// # Returns
    ///
    /// A `TranslatedQuery` containing the translated query.
    ///
    pub fn sql(sql: impl Into<String>, params: Vec<Value>, table: Option<String>) -> Self {
        Self::new(sql, QueryLanguage::Sql, params, table)
    }

    /// Create a CQL translated query.
    ///
    /// # Arguments
    ///
    /// * `sql` - The CQL query.
    /// * `params` - The parameters.
    /// * `table` - The table.
    ///
    /// # Returns
    ///
    /// A `TranslatedQuery` containing the translated query.
    ///
    pub fn cql(sql: impl Into<String>, params: Vec<Value>, table: Option<String>) -> Self {
        Self::new(sql, QueryLanguage::Cql, params, table)
    }

    /// Create a PostgREST translated query.
    ///
    /// # Arguments
    ///
    /// * `sql` - The PostgREST query.
    /// * `params` - The parameters.
    /// * `table` - The table.
    ///
    /// # Returns
    ///
    /// A `TranslatedQuery` containing the translated query.
    ///
    pub fn postgrest(sql: impl Into<String>, params: Vec<Value>, table: Option<String>) -> Self {
        Self::new(sql, QueryLanguage::Postgrest, params, table)
    }

    /// ## `with_postgrest_method`
    /// Attach an explicit PostgREST HTTP method to this translated query.
    ///
    /// # Arguments
    ///
    /// * `method` - The PostgREST HTTP method.
    ///
    /// # Returns
    ///
    /// A `TranslatedQuery` containing the translated query.
    ///
    pub fn with_postgrest_method(mut self, method: PostgrestMethod) -> Self {
        self.postgrest_method = Some(method);
        self
    }

    /// ## `with_payload`
    /// Attach a JSON payload to this translated query.
    ///
    /// # Arguments
    ///
    /// * `payload` - The JSON payload.
    ///
    /// # Returns
    ///
    /// A `TranslatedQuery` containing the translated query.
    ///
    pub fn with_payload(mut self, payload: Value) -> Self {
        self.payload = Some(payload);
        self
    }

    pub fn with_schema_name(mut self, schema_name: impl Into<String>) -> Self {
        self.schema_name = Some(schema_name.into());
        self
    }
}

/// ## `QueryResult`
/// The normalized result of executing a query.
///
/// # Arguments
///
/// * `rows` - The rows.
/// * `columns` - The columns.
///
/// # Returns
///
/// A `QueryResult` containing the query result.
///
#[derive(Clone)]
pub struct QueryResult {
    pub rows: Vec<Value>,
    pub columns: Vec<String>,
    pub duration: Option<u64>,
    pub message: Option<String>,
    pub status: Option<String>,
    pub count: Option<u64>,
    pub metadata: Option<Value>,
}

/// Debug representation for query results.
impl Debug for QueryResult {
    /// ## `fmt`
    /// Format the query result.
    ///
    /// # Arguments
    ///
    /// * `f` - The formatter.
    ///
    /// # Returns
    ///
    /// A `StdResult` containing the formatted query result.
    ///
    fn fmt(&self, f: &mut Formatter<'_>) -> StdResult {
        f.debug_struct("QueryResult")
            .field("rows", &self.rows)
            .field("columns", &self.columns)
            .field("count", &self.count)
            .finish()
    }
}

impl QueryResult {
    /// ## `new`
    /// Create a new query result.
    ///
    /// # Arguments
    ///
    /// * `rows` - The rows.
    /// * `columns` - The columns.
    /// * `duration` - The duration.
    /// * `message` - The message.
    /// * `status` - The status (should be derived from an ENUM)
    ///
    /// # Returns
    ///
    /// A `QueryResult` containing the query result.
    ///
    pub fn new(
        rows: Vec<Value>,
        columns: Vec<String>,
        duration: Option<u64>,
        message: Option<String>,
        status: Option<String>,
    ) -> Self {
        Self {
            rows,
            columns,
            duration,
            message,
            status,
            count: None,
            metadata: None,
        }
    }

    pub fn with_count(mut self, count: Option<u64>) -> Self {
        self.count = count;
        self
    }

    pub fn with_metadata(mut self, metadata: Option<Value>) -> Self {
        self.metadata = metadata;
        self
    }

    pub fn bookmark(&self) -> Option<&str> {
        self.metadata
            .as_ref()
            .and_then(|value| value.get("bookmark"))
            .and_then(Value::as_str)
    }
}

/// ## `HealthStatus`
/// Health status returned by a backend health check.
///
/// # Arguments
///
/// * `Healthy` - The backend is healthy.
/// * `Offline` - The backend is offline.
///
/// # Returns
///
/// A `HealthStatus` containing the health status.
///
pub enum HealthStatus {
    Healthy,
    Offline(Instant),
}

/// ## `BackendResult`
/// Common result type for backend operations.
///
/// # Arguments
///
/// * `T` - The type of the result.
///
/// # Returns
///
pub type BackendResult<T> = Result<T, BackendError>;

/// Unified error type for backend implementations.
#[derive(thiserror::Error, Debug)]
pub enum BackendError {
    #[error("backend error: {0}")]
    Generic(String),
}

/// Trait that every backend must implement.
#[async_trait]
pub trait DatabaseBackend: Send + Sync {
    /// ## `execute_query`
    /// Execute a query on the backend.
    ///
    /// # Arguments
    ///
    /// * `query` - The translated query to execute.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    async fn execute_query(&self, query: TranslatedQuery) -> BackendResult<QueryResult>;

    /// ## `health_check`
    /// Check the health of the backend.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the health status.
    ///
    async fn health_check(&self) -> BackendResult<HealthStatus>;

    /// ## `backend_type`
    /// Get the backend type.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `BackendType` containing the backend type.
    ///
    fn backend_type(&self) -> BackendType;

    /// ## `supports_sql`
    /// Check if the backend supports SQL.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `bool` indicating if the backend supports SQL.
    ///
    fn supports_sql(&self) -> bool;

    /// ## `supports_cql`
    /// Check if the backend supports CQL.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `bool` indicating if the backend supports CQL.
    ///
    fn supports_cql(&self) -> bool;

    /// ## `as_any`
    /// Expose the concrete backend for downcasting when deeper capabilities are required
    /// (for example, schema introspection in SQL translators).
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `&dyn Any` containing the backend.
    fn as_any(&self) -> &dyn Any;
}