athena_rs 3.3.0

//! Scylla backend implementation.
use crate::client::backend::{
    BackendError, BackendResult, BackendType, DatabaseBackend, HealthStatus, QueryLanguage,
    QueryResult, TranslatedQuery,
};
use crate::drivers::scylla::client as scylla_client;
use crate::drivers::scylla::client::ScyllaConnectionInfo;
use crate::drivers::scylla::health;

use async_trait::async_trait;
use health::{HostHealthTracker, SystemClock};
use std::any::Any;

/// ## `ScyllaBackend`
/// A Scylla backend powered by scylla-rs.
/// It is used to execute queries on the Scylla database.
///
/// # Arguments
///
/// * `info` - The Scylla connection info.
///
/// # Returns
///
/// A `ScyllaBackend` instance.
///
pub struct ScyllaBackend {
    info: ScyllaConnectionInfo,
}

impl ScyllaBackend {
    /// ## `new`
    /// Create a new Scylla backend.
    ///
    /// # Arguments
    ///
    /// * `info` - The Scylla connection info.
    ///
    /// # Returns
    ///
    /// A `ScyllaBackend` instance.
    ///
    pub fn new(info: ScyllaConnectionInfo) -> Self {
        Self { info }
    }
}

#[async_trait]
impl DatabaseBackend for ScyllaBackend {
    /// ## `execute_query`
    /// Execute a query on the Scylla backend.
    ///
    /// # Arguments
    ///
    /// * `query` - The translated query to execute.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the query result.
    async fn execute_query(&self, query: TranslatedQuery) -> BackendResult<QueryResult> {
        if !matches!(query.language, QueryLanguage::Cql) {
            return Err(BackendError::Generic(
                "Scylla backend only supports CQL".to_string(),
            ));
        }

        let (rows, columns) = scylla_client::execute_query(query.sql)
            .await
            .map_err(|err| BackendError::Generic(err.to_string()))?;

        Ok(QueryResult::new(rows, columns, None, None, None))
    }

    /// ## `health_check`
    /// Check the health of the Scylla backend.
    ///
    /// # Arguments
    ///
    /// * `self` - The Scylla backend.
    ///
    /// # Returns
    ///
    /// A `BackendResult` containing the health status.
    async fn health_check(&self) -> BackendResult<HealthStatus> {
        let tracker: &HostHealthTracker<SystemClock> = health::global_tracker();

        if let Some(deadline) = tracker.offline_until(&self.info.host) {
            return Ok(HealthStatus::Offline(deadline));
        }

        Ok(HealthStatus::Healthy)
    }

    /// ## `backend_type`
    /// Get the backend type.
    ///
    /// # Arguments
    ///
    /// * `self` - The Scylla backend.
    ///
    /// # Returns
    ///
    /// A `BackendType` containing the backend type.
    fn backend_type(&self) -> BackendType {
        BackendType::Scylla
    }

    /// ## `supports_sql`
    /// Check if the backend supports SQL.
    ///
    /// # Arguments
    ///
    /// * `self` - The Scylla backend.
    ///
    /// # Returns
    ///
    /// A `bool` indicating if the backend supports SQL.
    fn supports_sql(&self) -> bool {
        false
    }

    /// ## `supports_cql`
    /// Check if the backend supports CQL.
    ///
    /// # Arguments
    ///
    /// * `self` - The Scylla backend.
    ///
    /// # Returns
    ///
    /// A `bool` indicating if the backend supports CQL.
    fn supports_cql(&self) -> bool {
        true
    }

    /// ## `as_any`
    /// Get the backend as any.
    ///
    /// # Arguments
    ///
    /// * `self` - The Scylla backend.
    ///
    /// # Returns
    ///
    /// A `&dyn Any` containing the backend.
    fn as_any(&self) -> &dyn Any {
        self
    }
}