athena_rs 0.75.4

WIP Database API gateway
Documentation
//! Cache Check Module
//!
//! This module provides functionality to check and retrieve cached responses based on cache control headers.

use crate::AppState;
use crate::utils::redis_client::GLOBAL_REDIS;
use actix_web::{HttpRequest, HttpResponse, web::Data};
use serde_json::{Value, json};
use tracing::info;

/// ## Get Cached Response
///
/// This function retrieves a cached response for a given cache key.
///
/// ### Parameters
///
/// - `app_state`: A `Data<AppState>` instance containing the shared cache.
/// - `cache_key`: A string slice representing the key to look up in the cache.
///
/// ### Returns
///
/// - `Option<HttpResponse>`: Returns an `HttpResponse` containing the cached data if found, otherwise `None`.
///
/// ### Example
///
/// ```rust,no_run
/// # use actix_web::web::Data;
/// # use athena_rs::api::cache::check::get_cached_response;
/// # use athena_rs::AppState;
/// # use athena_rs::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
/// # use moka::future::Cache;
/// # use reqwest::Client;
/// # use serde_json::json;
/// # use std::sync::Arc;
/// # async fn doc_example() {
/// #     let cache = Cache::builder().build();
/// #     let immortal = Cache::builder().build();
/// #     let app_state = AppState {
/// #         cache: Arc::new(cache),
/// #         immortal_cache: Arc::new(immortal),
/// #         client: Client::new(),
/// #         pg_registry: Arc::new(PostgresClientRegistry::empty()),
/// #         gateway_force_camel_case_to_snake_case: false,
/// #         pipeline_registry: None,
/// #     };
/// #     let cache = Data::new(app_state);
/// let response = get_cached_response(cache, "some_cache_key").await;
/// #     let _ = response;
/// # }
/// ```
pub async fn get_cached_response(
    app_state: Data<AppState>,
    cache_key: &str,
) -> Option<HttpResponse> {
    if let Some(cached_response) = app_state.cache.get(cache_key).await {
        return Some(HttpResponse::Ok().json(cached_response.clone()));
    }
    // Try Redis on local miss
    if let Some(redis) = GLOBAL_REDIS.get() {
        if let Ok(value) = redis.lock().await.get(cache_key).await {
            if !value.is_null() {
                app_state
                    .cache
                    .insert(cache_key.to_string(), value.clone())
                    .await;
                return Some(HttpResponse::Ok().json(value));
            }
        }
    }
    None
}

/// ## Check Cache Control and Get Response
///
/// This function checks if the "Cache-Control" header is set to "no-cache" and retrieves the cached response if not.
///
/// ### Parameters
///
/// - `req`: A reference to the `HttpRequest` object.
/// - `app_state`: A `Data<AppState>` instance containing the shared cache.
/// - `cache_key`: A string slice representing the key to look up in the cache.
///
/// ### Returns
///
/// - `Option<HttpResponse>`: Returns an `HttpResponse` containing the cached data if found and cache control is not "no-cache", otherwise `None`.
///
/// ### Example
///
/// ```rust,no_run
/// # use actix_web::http::header::CACHE_CONTROL;
/// # use actix_web::test::TestRequest;
/// # use actix_web::web::Data;
/// # use athena_rs::api::cache::check::check_cache_control_and_get_response;
/// # use athena_rs::AppState;
/// # use athena_rs::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
/// # use moka::future::Cache;
/// # use reqwest::Client;
/// # use serde_json::json;
/// # use std::sync::Arc;
/// # async fn doc_example() {
/// #     let cache = Cache::builder().build();
/// #     let immortal = Cache::builder().build();
/// #     let app_state = AppState {
/// #         cache: Arc::new(cache),
/// #         immortal_cache: Arc::new(immortal),
/// #         client: Client::new(),
/// #         pg_registry: Arc::new(PostgresClientRegistry::empty()),
/// #         gateway_force_camel_case_to_snake_case: false,
/// #         pipeline_registry: None,
/// #     };
/// #     let cache = Data::new(app_state);
/// #     let req = TestRequest::default()
/// #         .insert_header((CACHE_CONTROL, "max-age=0"))
/// #         .to_http_request();
/// let response = check_cache_control_and_get_response(&req, cache, "some_cache_key").await;
/// #     let _ = response;
/// # }
/// ```
pub async fn check_cache_control_and_get_response(
    req: &HttpRequest,
    app_state: Data<AppState>,
    cache_key: &str,
) -> Option<HttpResponse> {
    if let Some(cache_control_header) = req.headers().get("Cache-Control") {
        if let Ok(cache_control_value) = cache_control_header.to_str() {
            if cache_control_value.contains("no-cache") {
                return None;
            }
        }
    }
    if let Some(cached_response) = app_state.cache.get(cache_key).await {
        let mut response_data: Value = cached_response.clone();

        // Check for X-Strip-Nulls header
        if let Some(strip_nulls_header) = req.headers().get("X-Strip-Nulls") {
            if let Ok(strip_nulls_value) = strip_nulls_header.to_str() {
                if strip_nulls_value.to_lowercase() == "true" {
                    if let Value::Object(map) = &mut response_data {
                        map.retain(|_, v| !v.is_null());
                        if let Some(data_array) = map.get_mut("data") {
                            if let Value::Array(array) = data_array {
                                for item in array.iter_mut() {
                                    if let Value::Object(obj) = item {
                                        obj.retain(|_, v| !v.is_null());
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        return Some(HttpResponse::Ok().json(json!({
            "success": true,
            "data": response_data
        })));
    }
    None
}

/// ## Check Cache Control and Get Response V2
///
/// This function checks if the "Cache-Control" header is set to "no-cache" and retrieves the cached response if not.
///
/// ### Parameters
///
pub async fn check_cache_control_and_get_response_v2(
    req: &HttpRequest,
    app_state: Data<AppState>,
    cache_key: &str,
) -> Option<HttpResponse> {
    if let Some(cache_control_header) = req.headers().get("Cache-Control") {
        if let Ok(cache_control_value) = cache_control_header.to_str() {
            if cache_control_value.contains("no-cache") {
                return None;
            }
        }
    }
    // Prefer local cache; on miss, try Redis
    if let Some(cached_response) = app_state.cache.get(cache_key).await {
        // Removed verbose logging of entire cached response
        let response_data: Value = cached_response.clone();

        // If the response is an array with a single item, return just that item
        if let Value::Array(arr) = &response_data {
            if arr.len() == 1 {
                return Some(HttpResponse::Ok().json(&arr[0]));
            }
        }

        return Some(HttpResponse::Ok().json(response_data));
    }
    // local miss: try Redis
    if let Some(redis) = GLOBAL_REDIS.get() {
        if let Ok(value) = redis.lock().await.get(cache_key).await {
            if !value.is_null() {
                // backfill local cache and return
                app_state
                    .cache
                    .insert(cache_key.to_string(), value.clone())
                    .await;

                if let Value::Array(arr) = &value {
                    if arr.len() == 1 {
                        return Some(HttpResponse::Ok().json(&arr[0]));
                    }
                }
                return Some(HttpResponse::Ok().json(value));
            }
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
    use actix_web::http::header::CACHE_CONTROL;
    use actix_web::test::TestRequest;
    use actix_web::web::Data;
    use moka::future::Cache;
    use reqwest::Client;
    use serde_json::json;
    use std::sync::Arc;

    #[tokio::test]
    async fn test_get_cached_response_returns_value() {
        let cache: Cache<String, Value> = Cache::builder().build();
        let immortal: Cache<String, Value> = Cache::builder().build();
        let app: AppState = AppState {
            cache: Arc::new(cache),
            immortal_cache: Arc::new(immortal),
            client: Client::new(),
            pg_registry: Arc::new(PostgresClientRegistry::empty()),
            gateway_force_camel_case_to_snake_case: false,
            pipeline_registry: None,
        };
        let data: Data<AppState> = Data::new(app);
        let key = "k".to_string();
        data.cache.insert(key.clone(), json!({"x": 1})).await;
        let maybe = get_cached_response(data, &key).await;
        assert!(maybe.is_some());
    }

    #[tokio::test]
    async fn test_check_cache_control_no_cache_skips() {
        let cache: Cache<String, Value> = Cache::builder().build();
        let immortal: Cache<String, Value> = Cache::builder().build();
        let app: AppState = AppState {
            cache: Arc::new(cache),
            immortal_cache: Arc::new(immortal),
            client: Client::new(),
            pg_registry: Arc::new(PostgresClientRegistry::empty()),
            gateway_force_camel_case_to_snake_case: false,
            pipeline_registry: None,
        };
        let data: Data<AppState> = Data::new(app);
        let key = "k2".to_string();
        data.cache.insert(key.clone(), json!({"y": 2})).await;

        let req = TestRequest::default()
            .insert_header((CACHE_CONTROL, "no-cache"))
            .to_http_request();

        let got = check_cache_control_and_get_response(&req, data, &key).await;
        assert!(got.is_none());
    }

    #[tokio::test]
    async fn test_v2_flattens_single_item_array() {
        let cache: Cache<String, Value> = Cache::builder().build();
        let immortal: Cache<String, Value> = Cache::builder().build();
        let app: AppState = AppState {
            cache: Arc::new(cache),
            immortal_cache: Arc::new(immortal),
            client: Client::new(),
            pg_registry: Arc::new(PostgresClientRegistry::empty()),
            gateway_force_camel_case_to_snake_case: false,
            pipeline_registry: None,
        };
        let data: Data<AppState> = Data::new(app);
        let key = "k3".to_string();
        data.cache.insert(key.clone(), json!([{"z": 3}])).await;

        let req = TestRequest::default().to_http_request();
        let resp = check_cache_control_and_get_response_v2(&req, data, &key).await;
        assert!(resp.is_some());
    }

    #[tokio::test]
    async fn test_v2_cache_miss_returns_none() {
        let cache: Cache<String, Value> = Cache::builder().build();
        let immortal: Cache<String, Value> = Cache::builder().build();
        let app: AppState = AppState {
            cache: Arc::new(cache),
            immortal_cache: Arc::new(immortal),
            client: Client::new(),
            pg_registry: Arc::new(PostgresClientRegistry::empty()),
            gateway_force_camel_case_to_snake_case: false,
            pipeline_registry: None,
        };
        let data: Data<AppState> = Data::new(app);
        let key = "missing".to_string();
        let req = TestRequest::default().to_http_request();
        let resp = check_cache_control_and_get_response_v2(&req, data, &key).await;
        assert!(resp.is_none());
    }

    #[tokio::test]
    async fn test_check_strips_nulls_when_header_set() {
        let cache: Cache<String, Value> = Cache::builder().build();
        let immortal: Cache<String, Value> = Cache::builder().build();
        let app = AppState {
            cache: Arc::new(cache),
            immortal_cache: Arc::new(immortal),
            client: Client::new(),
            pg_registry: Arc::new(PostgresClientRegistry::empty()),
            gateway_force_camel_case_to_snake_case: false,
            pipeline_registry: None,
        };
        let data: Data<AppState> = Data::new(app);
        let key: String = "strip".to_string();
        data.cache
            .insert(key.clone(), json!({"data": [{"a": 1, "b": null}]}))
            .await;
        let req = TestRequest::default()
            .insert_header(("X-Strip-Nulls", "true"))
            .to_http_request();
        let resp = check_cache_control_and_get_response(&req, data, &key).await;
        assert!(resp.is_some());
    }
}