exarrow-rs 0.28.0

//! Transport protocol abstraction trait.
//!
//! This module defines the `TransportProtocol` trait that abstracts the underlying
//! communication mechanism for connecting to Exasol. This allows for different
//! protocol implementations (WebSocket in Phase 1, gRPC in Phase 2).

use crate::error::TransportError;
use crate::{DRIVER_NAME, DRIVER_VERSION};
use async_trait::async_trait;

use super::messages::{ColumnInfo, DataType, ResultData, ResultSetHandle, SessionInfo};

/// Connection parameters for establishing a transport connection.
#[derive(Debug, Clone)]
pub struct ConnectionParams {
    /// Database host
    pub host: String,
    /// Database port
    pub port: u16,
    /// Use TLS/SSL
    pub use_tls: bool,
    /// Validate server certificate (default: true)
    pub validate_server_certificate: bool,
    /// Expected SHA-256 hex fingerprint of the server's DER certificate
    pub certificate_fingerprint: Option<String>,
    /// Connection timeout in milliseconds
    pub timeout_ms: u64,
}

impl ConnectionParams {
    /// Create new connection parameters.
    pub fn new(host: String, port: u16) -> Self {
        Self {
            host,
            port,
            use_tls: true,
            validate_server_certificate: true,
            certificate_fingerprint: None,
            timeout_ms: 30_000, // 30 seconds default
        }
    }

    /// Set whether to use TLS.
    pub fn with_tls(mut self, use_tls: bool) -> Self {
        self.use_tls = use_tls;
        self
    }

    /// Set whether to validate the server certificate.
    ///
    /// # Security Warning
    ///
    /// Disabling certificate validation makes the connection vulnerable to
    /// man-in-the-middle attacks. Only disable in development environments
    /// with self-signed certificates.
    pub fn with_validate_server_certificate(mut self, validate: bool) -> Self {
        self.validate_server_certificate = validate;
        self
    }

    /// Pin TLS connection to a specific certificate fingerprint (SHA-256 hex of DER cert).
    pub fn with_certificate_fingerprint(mut self, fingerprint: String) -> Self {
        self.certificate_fingerprint = Some(fingerprint);
        self
    }

    /// Set connection timeout.
    pub fn with_timeout(mut self, timeout_ms: u64) -> Self {
        self.timeout_ms = timeout_ms;
        self
    }

    /// Build the WebSocket URL from parameters.
    pub fn to_websocket_url(&self) -> String {
        let scheme = if self.use_tls { "wss" } else { "ws" };
        format!("{}://{}:{}", scheme, self.host, self.port)
    }
}

/// Login identity: the user, the password, and the client the server reports.
///
/// The client name and version appear in `EXA_DBA_SESSIONS.CLIENT` as
/// `<client_name> <client_version>`. They default to the driver name and the
/// crate version; the login message always reports the driver name separately.
#[derive(Debug, Clone)]
pub struct Credentials {
    /// Username
    pub username: String,
    /// Password (will be cleared after use)
    pub password: String,
    /// Client name the server reports for the session
    pub client_name: String,
    /// Client version the server reports for the session
    pub client_version: String,
}

impl Credentials {
    /// Create new credentials that identify the client as this driver.
    pub fn new(username: String, password: String) -> Self {
        Self {
            username,
            password,
            client_name: DRIVER_NAME.to_string(),
            client_version: DRIVER_VERSION.to_string(),
        }
    }

    /// Replace the client name and version the server reports for the session.
    pub fn with_client(mut self, client_name: String, client_version: String) -> Self {
        self.client_name = client_name;
        self.client_version = client_version;
        self
    }
}

impl Drop for Credentials {
    fn drop(&mut self) {
        self.password.clear();
    }
}

/// Handle to a prepared statement on the server.
#[derive(Debug, Clone)]
pub struct PreparedStatementHandle {
    /// Server-side statement identifier
    pub handle: i32,
    /// Number of parameters expected
    pub num_params: i32,
    /// Parameter type information (if available)
    pub parameter_types: Vec<DataType>,
    /// Parameter names from Exasol metadata (if available)
    pub parameter_names: Vec<Option<String>>,
    /// Result-set column metadata, empty for row-count-producing statements
    pub result_columns: Vec<ColumnInfo>,
}

impl PreparedStatementHandle {
    /// Create a new prepared statement handle.
    pub fn new(
        handle: i32,
        num_params: i32,
        parameter_types: Vec<DataType>,
        parameter_names: Vec<Option<String>>,
    ) -> Self {
        Self {
            handle,
            num_params,
            parameter_types,
            parameter_names,
            result_columns: Vec::new(),
        }
    }

    /// Attach result-set column metadata decoded from the same reply as the parameters.
    pub(crate) fn with_result_columns(mut self, result_columns: Vec<ColumnInfo>) -> Self {
        self.result_columns = result_columns;
        self
    }
}

/// Transport protocol trait for database communication.
///
/// This trait abstracts the underlying transport mechanism, allowing for
/// different implementations (WebSocket, gRPC, etc.).
#[async_trait]
pub trait TransportProtocol: Send + Sync {
    /// Connect to the database server.
    ///
    /// Starts the connection-timeout deadline (`params.timeout_ms`). That one
    /// deadline bounds the TCP connect, the TLS handshake, and the WebSocket
    /// upgrade together, and [`authenticate`](Self::authenticate) gets only
    /// what remains of it. A step that runs out of time fails with
    /// `Connection timeout after <ms>ms (<step>)`, naming the step.
    ///
    /// # Arguments
    ///
    /// * `params` - Connection parameters
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if connection fails or the deadline elapses.
    async fn connect(&mut self, params: &ConnectionParams) -> Result<(), TransportError>;

    /// Authenticate with the database.
    ///
    /// The login is bounded by what remains of the deadline that
    /// [`connect`](Self::connect) started, so connect and login together never
    /// take longer than the connection timeout. A login that runs out of time
    /// fails with `Connection timeout after <ms>ms (login)` and closes the
    /// connection, because the abandoned response would corrupt a retry.
    ///
    /// # Arguments
    ///
    /// * `credentials` - User credentials
    ///
    /// # Returns
    ///
    /// Session information on successful authentication.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if authentication fails.
    async fn authenticate(
        &mut self,
        credentials: &Credentials,
    ) -> Result<SessionInfo, TransportError>;

    /// Execute a SQL query.
    ///
    /// # Arguments
    ///
    /// * `sql` - SQL statement to execute
    ///
    /// # Returns
    ///
    /// Result set handle for SELECT queries, or result data for other statements.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if execution fails.
    async fn execute_query(&mut self, sql: &str) -> Result<QueryResult, TransportError>;

    /// Fetch result data from a result set.
    ///
    /// # Arguments
    ///
    /// * `handle` - Result set handle from execute_query
    ///
    /// # Returns
    ///
    /// Result data containing rows and metadata.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if fetch fails.
    async fn fetch_results(
        &mut self,
        handle: ResultSetHandle,
    ) -> Result<ResultData, TransportError>;

    /// Close a result set.
    ///
    /// # Arguments
    ///
    /// * `handle` - Result set handle to close
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if close fails.
    async fn close_result_set(&mut self, handle: ResultSetHandle) -> Result<(), TransportError>;

    /// Create a prepared statement from SQL text.
    ///
    /// # Arguments
    ///
    /// * `sql` - SQL statement with parameter placeholders (?)
    ///
    /// # Returns
    ///
    /// A statement handle for later execution.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if statement creation fails.
    async fn create_prepared_statement(
        &mut self,
        sql: &str,
    ) -> Result<PreparedStatementHandle, TransportError>;

    /// Execute a prepared statement with parameters.
    ///
    /// # Arguments
    ///
    /// * `handle` - Prepared statement handle from create_prepared_statement
    /// * `parameters` - Parameters in column-major format (each inner Vec is a column)
    ///
    /// # Returns
    ///
    /// Query result (result set or row count).
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if execution fails.
    async fn execute_prepared_statement(
        &mut self,
        handle: &PreparedStatementHandle,
        parameters: Option<Vec<Vec<serde_json::Value>>>,
    ) -> Result<QueryResult, TransportError>;

    /// Close a prepared statement and release server-side resources.
    ///
    /// # Arguments
    ///
    /// * `handle` - Prepared statement handle to close
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if close fails.
    async fn close_prepared_statement(
        &mut self,
        handle: &PreparedStatementHandle,
    ) -> Result<(), TransportError>;

    /// Close the connection.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if disconnect fails.
    async fn close(&mut self) -> Result<(), TransportError>;

    /// Drop the connection's socket without any protocol round-trip.
    ///
    /// This exists for the case where the driver gives up on an in-flight
    /// request and can no longer match a response to the request that produced
    /// it. A graceful [`close`](Self::close) is unusable there: it sends a
    /// disconnect command and waits for a reply that would arrive behind the
    /// abandoned response, so it blocks for as long as the abandoned request
    /// keeps running on the server. Terminating instead abandons the socket and
    /// leaves the server to reap the session when that socket closes.
    ///
    /// The non-async signature is what forbids awaited I/O here, so no
    /// implementation can reintroduce a round-trip.
    ///
    /// Afterwards [`is_connected`](Self::is_connected) reports `false`,
    /// [`close`](Self::close) succeeds without I/O, and every other operation
    /// fails with the terminated-transport error, which tells the caller to
    /// reconnect, instead of reading the abandoned response.
    fn terminate(&mut self);

    /// Check if the connection is still active.
    fn is_connected(&self) -> bool;

    /// Set the autocommit mode on the server.
    ///
    /// # Arguments
    ///
    /// * `enabled` - Whether to enable autocommit
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if the operation fails.
    async fn set_autocommit(&mut self, enabled: bool) -> Result<(), TransportError>;

    /// Set the query timeout on the server (seconds; `0` means unlimited).
    ///
    /// # Arguments
    ///
    /// * `timeout_secs` - The query timeout in whole seconds
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if the operation fails.
    async fn set_query_timeout(&mut self, timeout_secs: u64) -> Result<(), TransportError>;

    /// Set the session's current schema with the protocol's set-attributes
    /// command, then record the schema the server opened through
    /// [`refresh_current_schema`](Self::refresh_current_schema).
    ///
    /// The name travels as a plain attribute value, never as an `OPEN SCHEMA`
    /// statement, so the server resolves it by exact name first and by its
    /// upper-case form second. The set-attributes response does not echo the
    /// schema, which is why the get-attributes request follows.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` with the server's message when the server
    /// rejects the schema. The rejection returns before the get-attributes
    /// request, so the recorded current schema stays unchanged.
    async fn set_current_schema(&mut self, schema: &str) -> Result<(), TransportError>;

    /// Ask the server for the session's current schema with one get-attributes
    /// request, record the answer, and return it.
    ///
    /// Returns `None` when the session has no current schema. Unlike
    /// [`current_schema`](Self::current_schema), this also sees a change that
    /// no response of this session reported, such as another session dropping
    /// the schema.
    ///
    /// # Errors
    ///
    /// Returns `TransportError` if the request fails.
    async fn refresh_current_schema(&mut self) -> Result<Option<String>, TransportError>;

    /// The session's current schema as the server last reported it, read
    /// without a server request; `None` means no current schema.
    ///
    /// The transport is the single owner of this value because it is the only
    /// module that sees every response: it records the current schema from
    /// each response that carries the attribute, including the responses to
    /// `OPEN SCHEMA`, `CREATE SCHEMA`, `CLOSE SCHEMA`, and `DROP SCHEMA` run as
    /// SQL, and from every get-attributes response.
    fn current_schema(&self) -> Option<String>;
}

/// Result of a query execution.
#[derive(Debug, Clone)]
pub enum QueryResult {
    /// Result set from a SELECT query
    ResultSet {
        /// Handle for fetching more data (None if all data is in the response)
        handle: Option<ResultSetHandle>,
        /// Result data (may include first batch of rows)
        data: ResultData,
    },
    /// Row count from an INSERT/UPDATE/DELETE query
    RowCount {
        /// Number of affected rows
        count: i64,
    },
}

impl QueryResult {
    /// Create a result set query result.
    pub fn result_set(handle: Option<ResultSetHandle>, data: ResultData) -> Self {
        Self::ResultSet { handle, data }
    }

    /// Create a row count query result.
    pub fn row_count(count: i64) -> Self {
        Self::RowCount { count }
    }

    /// Check if this is a result set.
    pub fn is_result_set(&self) -> bool {
        matches!(self, Self::ResultSet { .. })
    }

    /// Check if this is a row count.
    pub fn is_row_count(&self) -> bool {
        matches!(self, Self::RowCount { .. })
    }

    /// Get the result set handle if this is a result set.
    /// Returns None if all data was included in the initial response.
    pub fn handle(&self) -> Option<ResultSetHandle> {
        match self {
            Self::ResultSet { handle, .. } => *handle,
            _ => None,
        }
    }

    /// Get the row count if this is a row count result.
    pub fn get_row_count(&self) -> Option<i64> {
        match self {
            Self::RowCount { count } => Some(*count),
            _ => None,
        }
    }

    /// Check if this result has more data to fetch.
    pub fn has_more_data(&self) -> bool {
        match self {
            Self::ResultSet { handle, data } => {
                handle.is_some() && (data.data.num_rows() as i64) < data.total_rows
            }
            _ => false,
        }
    }
}

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

    #[test]
    fn test_connection_params_default() {
        let params = ConnectionParams::new("localhost".to_string(), 8563);
        assert_eq!(params.host, "localhost");
        assert_eq!(params.port, 8563);
        assert!(params.use_tls);
        assert!(params.validate_server_certificate);
        assert_eq!(params.timeout_ms, 30_000);
    }

    #[test]
    fn test_connection_params_builder() {
        let params = ConnectionParams::new("db.example.com".to_string(), 9000)
            .with_tls(false)
            .with_timeout(60_000);

        assert_eq!(params.host, "db.example.com");
        assert_eq!(params.port, 9000);
        assert!(!params.use_tls);
        assert!(params.validate_server_certificate);
        assert_eq!(params.timeout_ms, 60_000);
    }

    #[test]
    fn test_connection_params_validate_certificate_disabled() {
        let params = ConnectionParams::new("localhost".to_string(), 8563)
            .with_tls(true)
            .with_validate_server_certificate(false);

        assert!(params.use_tls);
        assert!(!params.validate_server_certificate);
    }

    #[test]
    fn test_websocket_url_with_tls() {
        let params = ConnectionParams::new("localhost".to_string(), 8563).with_tls(true);
        assert_eq!(params.to_websocket_url(), "wss://localhost:8563");
    }

    #[test]
    fn test_websocket_url_without_tls() {
        let params = ConnectionParams::new("localhost".to_string(), 8563).with_tls(false);
        assert_eq!(params.to_websocket_url(), "ws://localhost:8563");
    }

    #[test]
    fn test_credentials_creation() {
        let creds = Credentials::new("user".to_string(), "pass".to_string());
        assert_eq!(creds.username, "user");
        assert_eq!(creds.password, "pass");
    }

    #[test]
    fn credentials_default_to_the_driver_identity() {
        let creds = Credentials::new("user".to_string(), "pass".to_string());

        assert_eq!(creds.client_name, "exarrow-rs");
        assert_eq!(creds.client_version, env!("CARGO_PKG_VERSION"));
    }

    #[test]
    fn with_client_replaces_the_client_name_and_version() {
        let creds = Credentials::new("user".to_string(), "pass".to_string())
            .with_client("dbt".to_string(), "9.9".to_string());

        assert_eq!(creds.client_name, "dbt");
        assert_eq!(creds.client_version, "9.9");
        assert_eq!(creds.username, "user");
    }

    #[test]
    fn test_credentials_drop_clears_password() {
        let creds = Credentials::new("user".to_string(), "secret".to_string());
        assert_eq!(creds.password, "secret");
        drop(creds);
        // Password should be cleared (can't test directly after drop)
    }

    #[test]
    fn test_prepared_statement_handle_creation() {
        let param_types = vec![
            DataType {
                type_name: "DECIMAL".to_string(),
                precision: Some(18),
                scale: Some(0),
                size: None,
                character_set: None,
                with_local_time_zone: None,
                fraction: None,
                srid: None,
            },
            DataType {
                type_name: "VARCHAR".to_string(),
                precision: None,
                scale: None,
                size: Some(100),
                character_set: Some("UTF8".to_string()),
                with_local_time_zone: None,
                fraction: None,
                srid: None,
            },
        ];

        let handle = PreparedStatementHandle::new(42, 2, param_types, vec![]);
        assert_eq!(handle.handle, 42);
        assert_eq!(handle.num_params, 2);
        assert_eq!(handle.parameter_types.len(), 2);
        assert_eq!(handle.parameter_types[0].type_name, "DECIMAL");
        assert_eq!(handle.parameter_types[1].type_name, "VARCHAR");
    }

    #[test]
    fn test_prepared_statement_handle_no_params() {
        let handle = PreparedStatementHandle::new(1, 0, vec![], vec![]);
        assert_eq!(handle.handle, 1);
        assert_eq!(handle.num_params, 0);
        assert!(handle.parameter_types.is_empty());
    }

    #[test]
    fn prepared_statement_handle_result_columns_default_empty() {
        let handle = PreparedStatementHandle::new(1, 0, vec![], vec![]);
        assert!(handle.result_columns.is_empty());
    }

    #[test]
    fn prepared_statement_handle_with_result_columns_sets_them() {
        let columns = vec![ColumnInfo {
            name: "ID".to_string(),
            data_type: DataType {
                type_name: "DECIMAL".to_string(),
                precision: Some(18),
                scale: Some(0),
                size: None,
                character_set: None,
                with_local_time_zone: None,
                fraction: None,
                srid: None,
            },
        }];

        let handle =
            PreparedStatementHandle::new(1, 0, vec![], vec![]).with_result_columns(columns.clone());

        assert_eq!(handle.result_columns.len(), 1);
        assert_eq!(handle.result_columns[0].name, "ID");
        assert_eq!(handle.result_columns[0].data_type.type_name, "DECIMAL");
    }

    #[test]
    fn test_query_result_result_set() {
        use super::super::messages::{ColumnInfo, DataType, ResultData, ResultPayload};

        let data = ResultData {
            columns: vec![ColumnInfo {
                name: "id".to_string(),
                data_type: DataType {
                    type_name: "DECIMAL".to_string(),
                    precision: Some(18),
                    scale: Some(0),
                    size: None,
                    character_set: None,
                    with_local_time_zone: None,
                    fraction: None,
                    srid: None,
                },
            }],
            data: ResultPayload::Json(vec![]),
            total_rows: 0,
        };

        let result = QueryResult::result_set(Some(ResultSetHandle::new(1)), data);
        assert!(result.is_result_set());
        assert!(!result.is_row_count());
        assert_eq!(result.handle().unwrap().as_i32(), 1);
        assert!(result.get_row_count().is_none());
    }

    #[test]
    fn test_query_result_row_count() {
        let result = QueryResult::row_count(42);
        assert!(!result.is_result_set());
        assert!(result.is_row_count());
        assert_eq!(result.get_row_count().unwrap(), 42);
        assert!(result.handle().is_none());
    }
}