athena_rs 3.0.1

//! 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 gateway_api;
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, QueryResult,
};
use backends::{
    gateway::GatewayBackend, postgres::PostgresBackend, scylla::ScyllaBackend,
    supabase::SupabaseBackend,
};
use builder::AthenaClientBuilder;
use config::ClientConfig;
pub use gateway_api::{
    Gateway, GatewayDeleteRequest, GatewayDriverRequest, GatewayFetchRequest, GatewayInsertRequest,
    GatewayOperation, GatewayPath, GatewayQueryResult, GatewayRequest, GatewayRequestFactory,
    GatewayRequestPayload, GatewayRoutes, GatewaySqlRequest, GatewayUpdateRequest,
    build_gateway_endpoint, request,
};
use query_builder::{DeleteBuilder, InsertBuilder, SelectBuilder, UpdateBuilder};
use translator::{CqlTranslator, PostgrestTranslator, QueryTranslator, SqlTranslator};

pub use backend::{QueryLanguage, TranslatedQuery};
pub use query_builder::Condition;
pub use query_builder::ConditionOperator;
pub use query_builder::OrderDirection;

/// Backward-compatible route alias. Prefer `Gateway::FETCH_PATH`.
pub const GATEWAY_FETCH_PATH: &str = Gateway::FETCH_PATH;
/// Backward-compatible route alias. Prefer `Gateway::INSERT_PATH`.
pub const GATEWAY_INSERT_PATH: &str = Gateway::INSERT_PATH;
/// Backward-compatible route alias. Prefer `Gateway::UPDATE_PATH`.
pub const GATEWAY_UPDATE_PATH: &str = Gateway::UPDATE_PATH;
/// Backward-compatible route alias. Prefer `Gateway::DELETE_PATH`.
pub const GATEWAY_DELETE_PATH: &str = Gateway::DELETE_PATH;
/// Backward-compatible route alias. Prefer `Gateway::QUERY_PATH`.
pub const GATEWAY_QUERY_PATH: &str = Gateway::QUERY_PATH;
/// Backward-compatible route alias. Prefer `Gateway::SQL_PATH`.
pub const GATEWAY_SQL_PATH: &str = Gateway::SQL_PATH;
/// Backward-compatible route alias. Prefer `Gateway::LEGACY_SQL_PATH`.
pub const LEGACY_SQL_PATH: &str = Gateway::LEGACY_SQL_PATH;

/// ## `AthenaClient`
/// The main client for Athena.
///
/// # Arguments
///
/// * `backend` - The backend.
/// * `config` - The client configuration.
///
/// # Returns
///
/// A `AthenaClient` containing the client.
///
pub struct AthenaClient {
    backend: Box<dyn DatabaseBackend>,
    config: ClientConfig,
}

impl AthenaClient {
    /// ## `builder`
    /// Create a new Athena client builder.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `AthenaClientBuilder` containing the builder.
    ///
    pub fn builder() -> AthenaClientBuilder {
        AthenaClientBuilder::new()
    }

    /// ## `build`
    /// Build a new Athena client.
    ///
    /// # Arguments
    ///
    /// * `builder` - The Athena client builder.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the client.
    ///
    pub async fn build(builder: AthenaClientBuilder) -> BackendResult<Self> {
        let config: ClientConfig = builder.build_config()?;
        Self::from_config(config).await
    }

    /// ## `new`
    /// Create a new Athena client.
    ///
    /// # Arguments
    ///
    /// * `url` - The connection URL.
    /// * `key` - The connection key.
    /// * `client` - The client name.
    /// * `backend` - The backend type.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the client.
    ///
    pub async fn new(
        url: impl Into<String>,
        key: impl Into<String>,
        client: impl Into<String>,
    ) -> BackendResult<Self> {
        Self::new_with_backend(url, key, client, BackendType::Native).await
    }

    /// ## `new_with_backend`
    /// Create a new Athena client with a specific backend.
    ///
    /// # Arguments
    ///
    /// * `url` - The connection URL.
    /// * `key` - The connection key.
    /// * `client` - The client name.
    /// * `backend` - The backend type.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the client.
    ///
    pub async fn new_with_backend(
        url: impl Into<String>,
        key: impl Into<String>,
        client: impl Into<String>,
        backend: BackendType,
    ) -> BackendResult<Self> {
        let builder: AthenaClientBuilder = Self::builder()
            .backend(backend)
            .url(url)
            .key(key)
            .client(client);
        Self::build(builder).await
    }

    /// ## `new_with_backend_name`
    /// Create a new Athena client with a specific backend name.
    ///
    /// # Arguments
    ///
    /// * `url` - The connection URL.
    /// * `key` - The connection key.
    /// * `client` - The client name.
    /// * `backend_name` - The backend name.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the client.
    ///
    pub async fn new_with_backend_name(
        url: impl Into<String>,
        key: impl Into<String>,
        client: impl Into<String>,
        backend_name: &str,
    ) -> BackendResult<Self> {
        let backend: BackendType = parse_backend_name(backend_name);
        Self::new_with_backend(url, key, client, backend).await
    }

    /// ## `new_direct`
    /// Create a new Athena client directly.
    ///
    /// # Arguments
    ///
    /// * `url` - The connection URL.
    /// * `key` - The connection key.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the client.
    ///
    pub async fn new_direct(url: impl Into<String>, key: impl Into<String>) -> BackendResult<Self> {
        let builder: AthenaClientBuilder = Self::builder().url(url).key(key);
        Self::build(builder).await
    }

    /// Return fully-qualified gateway endpoints for this client base URL.
    pub fn gateway_routes(&self) -> GatewayRoutes {
        GatewayRoutes::for_base_url(&self.config.connection.url)
    }

    /// Build a fully-qualified endpoint for the provided path using this client URL.
    pub fn gateway_endpoint(&self, path: &str) -> String {
        build_gateway_endpoint(&self.config.connection.url, path)
    }

    /// ## `select`
    /// Create a new select builder.
    ///
    /// # Arguments
    ///
    /// * `table` - The table name.
    ///
    /// # Returns
    ///
    /// A `SelectBuilder` containing the builder.
    ///
    pub fn select(&self, table: &str) -> SelectBuilder<'_> {
        SelectBuilder::new(self, table)
    }

    /// Alias for `select` that mirrors gateway `/gateway/fetch` naming.
    pub fn fetch(&self, table: &str) -> SelectBuilder<'_> {
        self.select(table)
    }

    /// ## `insert`
    /// Create a new insert builder.
    ///
    /// # Arguments
    ///
    /// * `table` - The table name.
    ///
    /// # Returns
    ///
    /// A `InsertBuilder` containing the builder.
    ///
    pub fn insert(&self, table: &str) -> InsertBuilder<'_> {
        InsertBuilder::new(self, table)
    }

    /// ## `update`
    /// Create a new update builder.
    ///
    /// # Arguments
    ///
    /// * `table` - The table name.
    /// # Returns
    ///
    /// A `UpdateBuilder` containing the builder.
    ///
    pub fn update(&self, table: &str) -> UpdateBuilder<'_> {
        UpdateBuilder::new(self, table, None)
    }

    /// ## `update_by_id`
    /// Create a new update builder pre-filtered by row id (`id = <row_id>`).
    pub fn update_by_id(&self, table: &str, row_id: impl Into<String>) -> UpdateBuilder<'_> {
        UpdateBuilder::new(self, table, Some(row_id.into()))
    }

    /// ## `delete`
    /// Create a new delete builder.
    ///
    /// # Arguments
    ///
    /// * `table` - The table name.
    ///
    /// # Returns
    ///
    /// A `DeleteBuilder` containing the builder.
    ///
    pub fn delete(&self, table: &str) -> DeleteBuilder<'_> {
        DeleteBuilder::new(self, table, None)
    }

    /// ## `delete_by_id`
    /// Create a new delete builder pre-filtered by row id (`id = <row_id>`).
    pub fn delete_by_id(&self, table: &str, row_id: impl Into<String>) -> DeleteBuilder<'_> {
        DeleteBuilder::new(self, table, Some(row_id.into()))
    }

    /// Execute a typed gateway request.
    pub async fn gateway_request<R>(&self, request: R) -> BackendResult<GatewayQueryResult>
    where
        R: Into<GatewayRequest>,
    {
        let request: GatewayRequest = request.into();
        match request.into_payload() {
            GatewayRequestPayload::Fetch(request) => {
                let mut builder = self.select(&request.table);
                if !request.columns.is_empty() {
                    builder = builder.columns(request.columns);
                }
                if let Some(raw_select) = request.raw_select {
                    builder = builder.raw_select(raw_select);
                }
                builder = builder.where_conditions(request.conditions);
                for (column, direction) in request.order_by {
                    builder = builder.order_by(&column, direction);
                }
                if let Some(limit) = request.limit {
                    builder = builder.limit(limit);
                }
                if let Some(offset) = request.offset {
                    builder = builder.offset(offset);
                }
                builder.execute().await
            }
            GatewayRequestPayload::Insert(request) => {
                self.insert(&request.table)
                    .payload(request.payload)
                    .execute()
                    .await
            }
            GatewayRequestPayload::Update(request) => {
                let update_scope = request.scope;
                let mut builder = if let Some(row_id) = update_scope.row_id {
                    self.update_by_id(&request.table, row_id)
                } else {
                    self.update(&request.table)
                };
                builder = builder.where_conditions(update_scope.conditions);
                if update_scope.allow_unfiltered {
                    builder = builder.unsafe_unfiltered();
                }
                builder.payload(request.payload).execute().await
            }
            GatewayRequestPayload::Delete(request) => {
                let mut builder = if let Some(row_id) = request.row_id {
                    self.delete_by_id(&request.table, row_id)
                } else {
                    self.delete(&request.table)
                };
                builder = builder.where_conditions(request.conditions);
                if request.allow_unfiltered {
                    builder = builder.unsafe_unfiltered();
                }
                builder.execute().await
            }
            GatewayRequestPayload::Sql(request) => self.sql(&request.query).await,
        }
    }

    /// Execute a typed fetch/select request using shared condition primitives.
    pub async fn fetch_request<R>(&self, request: R) -> BackendResult<GatewayQueryResult>
    where
        R: Into<GatewayRequest>,
    {
        self.gateway_request(request).await
    }

    /// Execute a typed insert request.
    pub async fn insert_request<R>(&self, request: R) -> BackendResult<GatewayQueryResult>
    where
        R: Into<GatewayRequest>,
    {
        self.gateway_request(request).await
    }

    /// Execute a typed update request.
    pub async fn update_request<R>(&self, request: R) -> BackendResult<GatewayQueryResult>
    where
        R: Into<GatewayRequest>,
    {
        self.gateway_request(request).await
    }

    /// Execute a typed delete request.
    pub async fn delete_request<R>(&self, request: R) -> BackendResult<GatewayQueryResult>
    where
        R: Into<GatewayRequest>,
    {
        self.gateway_request(request).await
    }

    /// ## `execute_sql`
    /// Execute a SQL query.
    ///
    /// # Arguments
    ///
    /// * `sql` - The SQL query.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub async fn execute_sql(&self, sql: &str) -> BackendResult<QueryResult> {
        let translated = TranslatedQuery::sql(sql, Vec::new(), None);
        self.backend.execute_query(translated).await
    }

    /// Alias for `execute_sql` matching athena-js style.
    pub async fn sql(&self, sql: &str) -> BackendResult<QueryResult> {
        self.execute_sql(sql).await
    }

    /// Execute a typed SQL request.
    pub async fn sql_request<R>(&self, request: R) -> BackendResult<GatewayQueryResult>
    where
        R: Into<GatewayRequest>,
    {
        self.gateway_request(request).await
    }

    /// ## `execute_cql`
    /// Execute a CQL query.
    ///
    /// # Arguments
    ///
    /// * `cql` - The CQL query.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub async fn execute_cql(&self, cql: &str) -> BackendResult<QueryResult> {
        let translated = TranslatedQuery::cql(cql, Vec::new(), None);
        self.backend.execute_query(translated).await
    }

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

    /// ## `config`
    /// Get the client configuration.
    ///
    /// # Arguments
    ///
    /// # Returns
    ///
    /// A `ClientConfig` containing the client configuration.
    ///
    pub fn config(&self) -> &ClientConfig {
        &self.config
    }

    /// Returns true when the client is configured to route through gateway CRUD APIs.
    fn uses_gateway_crud_routing(&self) -> bool {
        self.config.client_name.is_some() && self.backend.supports_sql()
    }

    /// Translate a typed gateway request to the appropriate backend-specific query.
    fn translate_request(&self, request: &GatewayRequest) -> BackendResult<TranslatedQuery> {
        if self.uses_gateway_crud_routing() {
            return PostgrestTranslator.translate_request(request);
        }

        match self.backend.backend_type() {
            BackendType::Supabase | BackendType::Postgrest => {
                PostgrestTranslator.translate_request(request)
            }
            BackendType::Scylla => CqlTranslator.translate_request(request),
            BackendType::PostgreSQL | BackendType::Native | BackendType::Neon => {
                SqlTranslator.translate_request(request)
            }
        }
    }
    /// ## `execute_select`
    /// Execute a select query.
    ///
    /// # Arguments
    ///
    /// * `builder` - The select builder.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub(crate) async fn execute_select(
        &self,
        builder: SelectBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let request: GatewayRequest = builder.into_request().into();
        let translated: TranslatedQuery = self.translate_request(&request)?;
        self.backend.execute_query(translated).await
    }

    /// ## `execute_insert`
    /// Execute an insert query.
    ///
    /// # Arguments
    ///
    /// * `builder` - The insert builder.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub(crate) async fn execute_insert(
        &self,
        builder: InsertBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let request: GatewayRequest = builder.into_request().into();
        let translated: TranslatedQuery = self.translate_request(&request)?;
        self.backend.execute_query(translated).await
    }

    /// ## `execute_update`
    /// Execute an update query.
    ///
    /// # Arguments
    ///
    /// * `builder` - The update builder.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    ///
    pub(crate) async fn execute_update(
        &self,
        builder: UpdateBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let request: GatewayUpdateRequest = builder.into_request();
        if !query_has_scope(request.scope.row_id.as_deref(), &request.scope.conditions)
            && !request.scope.allow_unfiltered
        {
            return Err(BackendError::Generic(
                "Refusing unfiltered update; add `where_*`/`row_id(...)` or call `unsafe_unfiltered()`"
                    .to_string(),
            ));
        }
        let request: GatewayRequest = request.into();
        let translated: TranslatedQuery = self.translate_request(&request)?;
        self.backend.execute_query(translated).await
    }

    /// ## `execute_delete`
    /// Execute a delete query.
    ///
    /// # Arguments
    ///
    /// * `builder` - The delete builder.
    ///
    /// # Returns   
    pub(crate) async fn execute_delete(
        &self,
        builder: DeleteBuilder<'_>,
    ) -> BackendResult<QueryResult> {
        let request: GatewayDeleteRequest = builder.into_request();
        if !query_has_scope(request.row_id.as_deref(), &request.conditions)
            && !request.allow_unfiltered
        {
            return Err(BackendError::Generic(
                "Refusing unfiltered delete; add `where_*`/`row_id(...)` or call `unsafe_unfiltered()`"
                    .to_string(),
            ));
        }
        let request: GatewayRequest = request.into();
        let translated: TranslatedQuery = self.translate_request(&request)?;
        self.backend.execute_query(translated).await
    }

    /// ## `from_config`
    /// Create a new Athena client from a configuration.
    ///
    /// # Arguments
    ///
    /// * `config` - The client configuration.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the client.
    ///
    async fn from_config(config: ClientConfig) -> BackendResult<Self> {
        if let Some(client_name) = config.client_name.clone() {
            let key: String = config.connection.key.clone().ok_or_else(|| {
                BackendError::Generic(
                    "Athena key is required when using client-routed gateway mode".to_string(),
                )
            })?;
            let backend: GatewayBackend = GatewayBackend::new(
                config.connection.url.clone(),
                key,
                client_name,
                config.backend_type,
            );
            return Ok(Self {
                backend: Box::new(backend),
                config,
            });
        }

        let backend: Box<dyn DatabaseBackend> = match config.backend_type {
            BackendType::Supabase => {
                let key: String =
                    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 = 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 =
                    PostgresBackend::from_connection_string(&config.connection.url).await?;
                Box::new(backend)
            }
        };

        Ok(Self { backend, config })
    }
}

/// ## `parse_backend_name`
/// Parse a backend name to a backend type.
///
/// # Arguments
///
/// * `backend_name` - The backend name.
///
/// # Returns
///
/// A `BackendType` containing the backend type.
///
fn parse_backend_name(backend_name: &str) -> BackendType {
    match backend_name.to_ascii_lowercase().as_str() {
        "supabase" => BackendType::Supabase,
        "postgrest" => BackendType::Postgrest,
        "scylla" => BackendType::Scylla,
        "neon" => BackendType::Neon,
        "postgresql" | "postgres" => BackendType::PostgreSQL,
        _ => BackendType::Native,
    }
}

fn query_has_scope(row_id: Option<&str>, conditions: &[Condition]) -> bool {
    row_id.is_some() || !conditions.is_empty()
}

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

    #[test]
    fn parse_backend_name_maps_supported_aliases() {
        assert_eq!(parse_backend_name("supabase"), BackendType::Supabase);
        assert_eq!(parse_backend_name("postgrest"), BackendType::Postgrest);
        assert_eq!(parse_backend_name("scylla"), BackendType::Scylla);
        assert_eq!(parse_backend_name("neon"), BackendType::Neon);
        assert_eq!(parse_backend_name("postgresql"), BackendType::PostgreSQL);
        assert_eq!(parse_backend_name("postgres"), BackendType::PostgreSQL);
        assert_eq!(parse_backend_name("unknown"), BackendType::Native);
    }

    #[tokio::test]
    async fn new_defaults_to_native_and_sets_client_routing() {
        let client = AthenaClient::new("http://localhost:4052", "secret", "reporting")
            .await
            .expect("client should build without network");

        assert_eq!(client.config().backend_type, BackendType::Native);
        assert_eq!(client.config().client_name.as_deref(), Some("reporting"));
        assert_eq!(client.config().connection.url, "http://localhost:4052");
    }

    #[tokio::test]
    async fn new_with_backend_name_resolves_backend_case_insensitive() {
        let client = AthenaClient::new_with_backend_name(
            "http://localhost:4052",
            "secret",
            "reporting",
            "NeOn",
        )
        .await
        .expect("client should build without network");

        assert_eq!(client.config().backend_type, BackendType::Neon);
    }

    #[tokio::test]
    async fn gateway_routed_clients_use_postgrest_translation_for_crud() {
        let client = AthenaClient::new("http://localhost:4052", "secret", "reporting")
            .await
            .expect("client should build without network");

        let translated = client
            .translate_request(&GatewayRequest::fetch("users"))
            .expect("translation should succeed");

        assert_eq!(translated.language, QueryLanguage::Postgrest);
    }

    #[tokio::test]
    async fn gateway_routes_expose_expected_endpoints() {
        let client = AthenaClient::new("http://localhost:4052", "secret", "reporting")
            .await
            .expect("client should build without network");

        let routes = client.gateway_routes();
        assert_eq!(routes.fetch, "http://localhost:4052/gateway/fetch");
        assert_eq!(routes.insert, "http://localhost:4052/gateway/insert");
        assert_eq!(routes.update, "http://localhost:4052/gateway/update");
        assert_eq!(routes.delete, "http://localhost:4052/gateway/delete");
        assert_eq!(routes.query, "http://localhost:4052/gateway/query");
        assert_eq!(routes.sql, "http://localhost:4052/gateway/sql");
    }

    #[test]
    fn query_has_scope_requires_row_id_or_conditions() {
        assert!(!query_has_scope(None, &[]));
        assert!(query_has_scope(Some("row_1"), &[]));
        assert!(query_has_scope(
            None,
            &[Condition::new(
                "status",
                ConditionOperator::Eq,
                vec![serde_json::json!("active")]
            )]
        ));
    }
}