athena_rs 1.1.0

//! Unified Athena client surface used by downstream crates.
pub mod backend;
pub mod backends;
pub mod builder;
pub mod config;
pub mod error;
pub mod query_builder;
pub mod translator;

use crate::drivers::scylla::client::ScyllaConnectionInfo;
use crate::drivers::supabase::client::SupabaseConnectionInfo;
use backend::{
    BackendError, BackendResult, BackendType, DatabaseBackend, HealthStatus, QueryLanguage,
    QueryResult, TranslatedQuery,
};
use backends::{postgres::PostgresBackend, scylla::ScyllaBackend, supabase::SupabaseBackend};
use builder::AthenaClientBuilder;
use config::ClientConfig;
use query_builder::{DeleteBuilder, InsertBuilder, SelectBuilder, UpdateBuilder};
use translator::{CqlTranslator, QueryTranslator, SqlTranslator};

pub struct AthenaClient {
    backend: Box<dyn DatabaseBackend>,
    config: ClientConfig,
}

impl AthenaClient {
    pub fn builder() -> AthenaClientBuilder {
        AthenaClientBuilder::new()
    }

    pub async fn build(builder: AthenaClientBuilder) -> BackendResult<Self> {
        let config = builder.build_config()?;
        Self::from_config(config).await
    }

    pub async fn new(url: impl Into<String>, key: impl Into<String>) -> BackendResult<Self> {
        Self::new_with_backend(url, key, BackendType::Native).await
    }

    pub async fn new_with_backend(
        url: impl Into<String>,
        key: impl Into<String>,
        backend: BackendType,
    ) -> BackendResult<Self> {
        let builder = Self::builder().backend(backend).url(url).key(key);
        Self::build(builder).await
    }

    pub async fn new_with_backend_name(
        url: impl Into<String>,
        key: impl Into<String>,
        backend_name: &str,
    ) -> BackendResult<Self> {
        let backend = match backend_name.to_lowercase().as_str() {
            "supabase" => BackendType::Supabase,
            "postgrest" => BackendType::Postgrest,
            "scylla" => BackendType::Scylla,
            "neon" => BackendType::Neon,
            _ => BackendType::Native,
        };
        Self::new_with_backend(url, key, backend).await
    }

    pub fn select(&self, table: &str) -> SelectBuilder<'_> {
        SelectBuilder::new(self, table)
    }

    pub fn insert(&self, table: &str) -> InsertBuilder<'_> {
        InsertBuilder::new(self, table)
    }

    pub fn update(&self, table: &str, row_id: Option<String>) -> UpdateBuilder<'_> {
        UpdateBuilder::new(self, table, row_id)
    }

    pub fn delete(&self, table: &str, row_id: Option<String>) -> DeleteBuilder<'_> {
        DeleteBuilder::new(self, table, row_id)
    }

    pub async fn execute_sql(&self, sql: &str) -> BackendResult<QueryResult> {
        let translated = TranslatedQuery::new(sql, QueryLanguage::Sql, Vec::new(), None);
        self.backend.execute_query(translated).await
    }

    pub async fn execute_cql(&self, cql: &str) -> BackendResult<QueryResult> {
        let translated = TranslatedQuery::new(cql, QueryLanguage::Cql, Vec::new(), None);
        self.backend.execute_query(translated).await
    }

    pub async fn health_check(&self) -> BackendResult<HealthStatus> {
        self.backend.health_check().await
    }

    pub fn config(&self) -> &ClientConfig {
        &self.config
    }

    pub(crate) async fn execute_select(
        &self,
        builder: SelectBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let query = builder.into_parts();
        let translated = if self.backend.supports_sql() {
            let translator = SqlTranslator;
            translator.translate_select(&query)
        } else {
            let translator = CqlTranslator;
            translator.translate_select(&query)
        };
        self.backend.execute_query(translated).await
    }

    pub(crate) async fn execute_insert(
        &self,
        builder: InsertBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let query = builder.into_parts();
        let translated = if self.backend.supports_sql() {
            let translator = SqlTranslator;
            translator.translate_insert(&query)
        } else {
            let translator = CqlTranslator;
            translator.translate_insert(&query)
        };
        self.backend.execute_query(translated).await
    }

    pub(crate) async fn execute_update(
        &self,
        builder: UpdateBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let query = builder.into_parts();
        let translated = if self.backend.supports_sql() {
            let translator = SqlTranslator;
            translator.translate_update(&query)
        } else {
            let translator = CqlTranslator;
            translator.translate_update(&query)
        };
        self.backend.execute_query(translated).await
    }

    pub(crate) async fn execute_delete(
        &self,
        builder: DeleteBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let query = builder.into_parts();
        let translated = if self.backend.supports_sql() {
            let translator = SqlTranslator;
            translator.translate_delete(&query)
        } else {
            let translator = CqlTranslator;
            translator.translate_delete(&query)
        };
        self.backend.execute_query(translated).await
    }

    async fn from_config(config: ClientConfig) -> BackendResult<Self> {
        let backend: Box<dyn DatabaseBackend> = match config.backend_type {
            BackendType::Supabase => {
                let key =
                    config.connection.key.clone().ok_or_else(|| {
                        BackendError::Generic("Supabase key is required".to_string())
                    })?;
                let info = SupabaseConnectionInfo::new(config.connection.url.clone(), key);
                Box::new(SupabaseBackend::new(info)?)
            }
            BackendType::Scylla => {
                let info = ScyllaConnectionInfo {
                    host: config.connection.url.clone(),
                    username: config.connection.database.clone().unwrap_or_default(),
                    password: config.connection.key.clone().unwrap_or_default(),
                };
                Box::new(ScyllaBackend::new(info))
            }
            BackendType::PostgreSQL
            | BackendType::Postgrest
            | BackendType::Native
            | BackendType::Neon => {
                let backend =
                    PostgresBackend::from_connection_string(&config.connection.url).await?;
                Box::new(backend)
            }
        };

        Ok(Self { backend, config })
    }
}