tellaro-query-language 1.3.8

//! OpenSearch client connection management.
//!
//! This module provides connection management for OpenSearch with configuration
//! from environment variables.

use super::error::{OpenSearchError, Result};
use opensearch::{
    auth::Credentials,
    cert::CertificateValidation,
    http::{
        transport::{SingleNodeConnectionPool, TransportBuilder},
        Url,
    },
    OpenSearch,
};
use std::env;

/// OpenSearch client configuration
#[derive(Debug, Clone)]
pub struct OpenSearchConfig {
    /// List of OpenSearch host URLs
    pub hosts: Vec<String>,
    /// Optional username for basic authentication
    pub username: Option<String>,
    /// Optional password for basic authentication
    pub password: Option<String>,
    /// Whether to verify SSL certificates
    pub verify_certs: bool,
    /// Whether to use SSL/TLS
    pub use_ssl: bool,
    /// Connection timeout in seconds
    pub timeout_secs: u64,
    /// Maximum number of retries
    pub max_retries: u32,
}

impl Default for OpenSearchConfig {
    fn default() -> Self {
        Self {
            hosts: vec!["http://localhost:9200".to_string()],
            username: None,
            password: None,
            verify_certs: false,
            use_ssl: false,
            timeout_secs: 30,
            max_retries: 3,
        }
    }
}

impl OpenSearchConfig {
    /// Load configuration from environment variables
    ///
    /// Environment variables:
    /// - `OPENSEARCH_HOSTS`: Comma-separated list of host URLs (default: http://localhost:9200)
    /// - `OPENSEARCH_USERNAME`: Username for authentication (optional)
    /// - `OPENSEARCH_PASSWORD`: Password for authentication (optional)
    /// - `OPENSEARCH_VERIFY_SSL`: Whether to verify SSL certs (default: false)
    /// - `OPENSEARCH_USE_SSL`: Whether to use SSL (default: false)
    /// - `OPENSEARCH_TIMEOUT`: Connection timeout in seconds (default: 30)
    /// - `OPENSEARCH_MAX_RETRIES`: Maximum retries (default: 3)
    ///
    /// # Example
    ///
    /// ```no_run
    /// use tql::opensearch::OpenSearchConfig;
    ///
    /// let config = OpenSearchConfig::from_env().unwrap();
    /// ```
    pub fn from_env() -> Result<Self> {
        let hosts = env::var("OPENSEARCH_HOSTS")
            .unwrap_or_else(|_| "http://localhost:9200".to_string())
            .split(',')
            .map(|s| s.trim().to_string())
            .collect();

        let username = env::var("OPENSEARCH_USERNAME").ok();
        let password = env::var("OPENSEARCH_PASSWORD").ok();

        let verify_certs = env::var("OPENSEARCH_VERIFY_SSL")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(false);

        let use_ssl = env::var("OPENSEARCH_USE_SSL")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(false);

        let timeout_secs = env::var("OPENSEARCH_TIMEOUT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(30);

        let max_retries = env::var("OPENSEARCH_MAX_RETRIES")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(3);

        Ok(Self {
            hosts,
            username,
            password,
            verify_certs,
            use_ssl,
            timeout_secs,
            max_retries,
        })
    }

    /// Create an OpenSearch client from this configuration
    ///
    /// # Example
    ///
    /// ```no_run
    /// use tql::opensearch::OpenSearchConfig;
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = OpenSearchConfig::from_env()?;
    /// let client = config.create_client()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn create_client(&self) -> Result<OpenSearch> {
        // Parse the first host URL
        let url = self
            .hosts
            .first()
            .ok_or_else(|| OpenSearchError::ConfigError("No hosts configured".to_string()))?
            .parse::<Url>()
            .map_err(|e| OpenSearchError::ConfigError(format!("Invalid URL: {}", e)))?;

        // Create connection pool
        let conn_pool = SingleNodeConnectionPool::new(url);

        // Build transport
        let mut transport_builder = TransportBuilder::new(conn_pool);

        // Add authentication if provided
        if let (Some(username), Some(password)) = (&self.username, &self.password) {
            let credentials = Credentials::Basic(username.clone(), password.clone());
            transport_builder = transport_builder.auth(credentials);
        }

        // Configure certificate validation
        if !self.verify_certs {
            transport_builder = transport_builder.cert_validation(CertificateValidation::None);
        }

        // Set timeout
        transport_builder =
            transport_builder.timeout(std::time::Duration::from_secs(self.timeout_secs));

        // Build transport
        let transport = transport_builder
            .build()
            .map_err(|e| OpenSearchError::ConnectionError(e.to_string()))?;

        // Create client
        Ok(OpenSearch::new(transport))
    }
}

/// OpenSearch client wrapper
pub struct OpenSearchClient {
    client: OpenSearch,
}

impl OpenSearchClient {
    /// Create a new OpenSearch client from configuration
    ///
    /// # Example
    ///
    /// ```no_run
    /// use tql::opensearch::{OpenSearchConfig, OpenSearchClient};
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = OpenSearchConfig::from_env()?;
    /// let client = OpenSearchClient::new(config)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(config: OpenSearchConfig) -> Result<Self> {
        let client = config.create_client()?;
        Ok(Self { client })
    }

    /// Get the underlying OpenSearch client
    pub fn client(&self) -> &OpenSearch {
        &self.client
    }

    /// Check if the OpenSearch cluster is available
    ///
    /// # Example
    ///
    /// ```no_run
    /// use tql::opensearch::{OpenSearchConfig, OpenSearchClient};
    ///
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = OpenSearchConfig::from_env()?;
    /// let client = OpenSearchClient::new(config)?;
    /// let available = client.is_available().await?;
    /// println!("OpenSearch available: {}", available);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn is_available(&self) -> Result<bool> {
        match self.client.ping().send().await {
            Ok(_) => Ok(true),
            Err(_) => Ok(false),
        }
    }
}

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

    #[test]
    fn test_default_config() {
        let config = OpenSearchConfig::default();
        assert_eq!(config.hosts, vec!["http://localhost:9200"]);
        assert_eq!(config.username, None);
        assert_eq!(config.password, None);
        assert!(!config.verify_certs);
        assert!(!config.use_ssl);
        assert_eq!(config.timeout_secs, 30);
        assert_eq!(config.max_retries, 3);
    }

    #[test]
    fn test_from_env_defaults() {
        // Clear any existing env vars
        env::remove_var("OPENSEARCH_HOSTS");
        env::remove_var("OPENSEARCH_USERNAME");
        env::remove_var("OPENSEARCH_PASSWORD");

        let config = OpenSearchConfig::from_env().unwrap();
        assert_eq!(config.hosts, vec!["http://localhost:9200"]);
        assert_eq!(config.username, None);
        assert_eq!(config.password, None);
    }

    #[test]
    fn test_from_env_with_values() {
        env::set_var(
            "OPENSEARCH_HOSTS",
            "https://opensearch1:9200,https://opensearch2:9200",
        );
        env::set_var("OPENSEARCH_USERNAME", "admin");
        env::set_var("OPENSEARCH_PASSWORD", "admin123");
        env::set_var("OPENSEARCH_VERIFY_SSL", "true");
        env::set_var("OPENSEARCH_USE_SSL", "true");
        env::set_var("OPENSEARCH_TIMEOUT", "60");
        env::set_var("OPENSEARCH_MAX_RETRIES", "5");

        let config = OpenSearchConfig::from_env().unwrap();
        assert_eq!(
            config.hosts,
            vec!["https://opensearch1:9200", "https://opensearch2:9200"]
        );
        assert_eq!(config.username, Some("admin".to_string()));
        assert_eq!(config.password, Some("admin123".to_string()));
        assert!(config.verify_certs);
        assert!(config.use_ssl);
        assert_eq!(config.timeout_secs, 60);
        assert_eq!(config.max_retries, 5);

        // Clean up
        env::remove_var("OPENSEARCH_HOSTS");
        env::remove_var("OPENSEARCH_USERNAME");
        env::remove_var("OPENSEARCH_PASSWORD");
        env::remove_var("OPENSEARCH_VERIFY_SSL");
        env::remove_var("OPENSEARCH_USE_SSL");
        env::remove_var("OPENSEARCH_TIMEOUT");
        env::remove_var("OPENSEARCH_MAX_RETRIES");
    }
}