exarrow-rs 0.29.0

//! Session management for Exasol database connections.
//!
//! This module handles session lifecycle, state tracking, and connection pooling support.

use crate::connection::auth::AuthResponseData;
use crate::connection::version::parse_release_version;
use crate::error::ConnectionError;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tokio::sync::RwLock;

/// Session configuration.
#[derive(Debug, Clone)]
pub struct SessionConfig {
    /// Maximum number of retries for failed operations
    pub max_retries: u32,

    /// Enable transaction auto-commit mode
    pub auto_commit: bool,

    /// Default fetch size for queries
    pub default_fetch_size: usize,

    /// Query timeout currently applied to the session's `queryTimeout`
    /// attribute. `None` means no attribute is set (server `QUERY_TIMEOUT`
    /// governs). This is the reconcile baseline: `connect()` seeds it and
    /// `execute_statement` reads and rewrites it after each reconcile.
    pub query_timeout: Option<Duration>,
}

impl Default for SessionConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            auto_commit: true,
            default_fetch_size: 1000,
            query_timeout: None,
        }
    }
}

/// Session state.
///
/// Only the session lifecycle changes it: executing a statement never does, so
/// no failed or abandoned operation can leave a state that blocks the next one.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionState {
    /// Session is connected and ready
    Ready,

    /// Session is in a transaction
    InTransaction,

    /// Session is being closed
    Closing,

    /// Session is closed
    Closed,
}

impl SessionState {
    /// Check if the session can execute queries.
    pub fn can_execute(&self) -> bool {
        matches!(self, SessionState::Ready | SessionState::InTransaction)
    }
}

/// Database session information and state tracking.
pub struct Session {
    /// Session ID from the server
    session_id: String,

    /// Server information from authentication
    server_info: AuthResponseData,

    /// Session configuration
    config: SessionConfig,

    /// Current session state
    state: Arc<RwLock<SessionState>>,

    /// Query execution counter
    query_count: AtomicU64,

    /// Transaction active flag
    in_transaction: AtomicBool,

    /// Memoized parse of `server_info.release_version`; `None` when it does not parse.
    release_version: OnceLock<Option<(u32, u32, u32)>>,
}

impl Session {
    /// Create a new session.
    pub fn new(session_id: String, server_info: AuthResponseData, config: SessionConfig) -> Self {
        Self {
            session_id,
            server_info,
            config,
            state: Arc::new(RwLock::new(SessionState::Ready)),
            query_count: AtomicU64::new(0),
            in_transaction: AtomicBool::new(false),
            release_version: OnceLock::new(),
        }
    }

    /// Whether the connected Exasol server supports native Parquet IMPORT
    /// (added in Exasol 2025.1.11). An unparsable release version counts as
    /// unsupported.
    pub fn supports_native_parquet_import(&self) -> bool {
        self.parsed_release_version()
            .is_some_and(crate::connection::version::supports_native_parquet_import)
    }

    /// Whether the server release is at least `(major, minor, patch)`, compared
    /// lexicographically. An unparsable release version is below every threshold,
    /// so callers fall back to the legacy path. Issues no SQL: the release version
    /// comes from login and is parsed once per session.
    pub fn supports_at_least(&self, major: u32, minor: u32, patch: u32) -> bool {
        self.parsed_release_version()
            .is_some_and(|version| version >= (major, minor, patch))
    }

    fn parsed_release_version(&self) -> Option<(u32, u32, u32)> {
        *self
            .release_version
            .get_or_init(|| parse_release_version(&self.server_info.release_version))
    }

    /// Get the session ID.
    pub fn session_id(&self) -> &str {
        &self.session_id
    }

    /// Get server information.
    pub fn server_info(&self) -> &AuthResponseData {
        &self.server_info
    }

    /// Get session configuration.
    pub fn config(&self) -> &SessionConfig {
        &self.config
    }

    /// Get mutable access to the session configuration.
    ///
    /// Used by `execute_statement` to read and write back the applied
    /// `query_timeout` baseline as part of the queryTimeout reconcile.
    pub(crate) fn config_mut(&mut self) -> &mut SessionConfig {
        &mut self.config
    }

    /// Get current session state.
    pub async fn state(&self) -> SessionState {
        *self.state.read().await
    }

    /// Set session state.
    pub async fn set_state(&self, new_state: SessionState) {
        let mut state = self.state.write().await;
        *state = new_state;
    }

    /// Increment query counter.
    pub fn increment_query_count(&self) -> u64 {
        self.query_count.fetch_add(1, Ordering::SeqCst) + 1
    }

    /// Get total query count.
    pub fn query_count(&self) -> u64 {
        self.query_count.load(Ordering::SeqCst)
    }

    /// Check if in transaction.
    pub fn in_transaction(&self) -> bool {
        self.in_transaction.load(Ordering::SeqCst)
    }

    /// Check that a transaction can start, without changing any state.
    ///
    /// Fails with `ConnectionClosed` for a `Closing` or `Closed` session and
    /// with "Transaction already active" while a transaction is open.
    pub(crate) async fn ensure_can_begin_transaction(&self) -> Result<(), ConnectionError> {
        self.validate_ready().await?;

        if self.in_transaction() {
            return Err(ConnectionError::InvalidParameter {
                parameter: "transaction".to_string(),
                message: "Transaction already active".to_string(),
            });
        }

        Ok(())
    }

    /// Begin a transaction.
    pub async fn begin_transaction(&self) -> Result<(), ConnectionError> {
        self.ensure_can_begin_transaction().await?;

        self.in_transaction.store(true, Ordering::SeqCst);
        self.set_state(SessionState::InTransaction).await;

        Ok(())
    }

    /// Commit the current transaction.
    pub async fn commit_transaction(&self) -> Result<(), ConnectionError> {
        if !self.in_transaction() {
            return Ok(());
        }

        self.in_transaction.store(false, Ordering::SeqCst);
        self.set_state(SessionState::Ready).await;

        Ok(())
    }

    /// Rollback the current transaction.
    pub async fn rollback_transaction(&self) -> Result<(), ConnectionError> {
        if !self.in_transaction() {
            return Ok(());
        }

        self.in_transaction.store(false, Ordering::SeqCst);
        self.set_state(SessionState::Ready).await;

        Ok(())
    }

    /// Close the session.
    pub async fn close(&self) -> Result<(), ConnectionError> {
        self.set_state(SessionState::Closing).await;

        // Clean up resources
        if self.in_transaction() {
            // Force rollback if in transaction
            self.in_transaction.store(false, Ordering::SeqCst);
        }

        self.set_state(SessionState::Closed).await;

        Ok(())
    }

    /// Check if session is closed.
    pub async fn is_closed(&self) -> bool {
        matches!(self.state().await, SessionState::Closed)
    }

    /// Validate session is ready for operations.
    ///
    /// A `Closing` or `Closed` session reports `ConnectionError::ConnectionClosed`.
    pub async fn validate_ready(&self) -> Result<(), ConnectionError> {
        if self.state().await.can_execute() {
            Ok(())
        } else {
            Err(ConnectionError::ConnectionClosed)
        }
    }
}

impl std::fmt::Debug for Session {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Session")
            .field("session_id", &self.session_id)
            .field("config", &self.config)
            .field("query_count", &self.query_count())
            .field("in_transaction", &self.in_transaction())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::connection::auth::AuthResponseData;

    fn mock_server_info() -> AuthResponseData {
        mock_server_info_with_version("7.1.0")
    }

    fn mock_server_info_with_version(release_version: &str) -> AuthResponseData {
        AuthResponseData {
            session_id: "test_session".to_string(),
            protocol_version: 3,
            release_version: release_version.to_string(),
            database_name: "EXA".to_string(),
            product_name: "Exasol".to_string(),
            max_data_message_size: 4_194_304,
            max_identifier_length: 128,
            max_varchar_length: 2_000_000,
            identifier_quote_string: "\"".to_string(),
            time_zone: "UTC".to_string(),
            time_zone_behavior: "INVALID TIMESTAMP TO DOUBLE".to_string(),
        }
    }

    #[tokio::test]
    async fn test_session_creation() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert_eq!(session.session_id(), "sess123");
        assert_eq!(session.state().await, SessionState::Ready);
        assert_eq!(session.query_count(), 0);
        assert!(!session.in_transaction());
    }

    #[tokio::test]
    async fn test_session_state_transitions() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert_eq!(session.state().await, SessionState::Ready);

        session.set_state(SessionState::InTransaction).await;
        assert_eq!(session.state().await, SessionState::InTransaction);

        session.set_state(SessionState::Closing).await;
        assert_eq!(session.state().await, SessionState::Closing);

        session.set_state(SessionState::Closed).await;
        assert_eq!(session.state().await, SessionState::Closed);
    }

    #[tokio::test]
    async fn test_session_query_counter() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert_eq!(session.increment_query_count(), 1);
        assert_eq!(session.increment_query_count(), 2);
        assert_eq!(session.query_count(), 2);
    }

    #[tokio::test]
    async fn test_session_transaction() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert!(!session.in_transaction());

        // Begin transaction
        session.begin_transaction().await.unwrap();
        assert!(session.in_transaction());
        assert_eq!(session.state().await, SessionState::InTransaction);

        // Cannot begin another transaction
        match session.begin_transaction().await {
            Err(ConnectionError::InvalidParameter { message, .. }) => {
                assert_eq!(message, "Transaction already active")
            }
            other => panic!("expected an active-transaction error, got {:?}", other),
        }

        // Commit transaction
        session.commit_transaction().await.unwrap();
        assert!(!session.in_transaction());
        assert_eq!(session.state().await, SessionState::Ready);
    }

    #[tokio::test]
    async fn test_session_rollback() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        session.begin_transaction().await.unwrap();
        assert!(session.in_transaction());

        session.rollback_transaction().await.unwrap();
        assert!(!session.in_transaction());
        assert_eq!(session.state().await, SessionState::Ready);
    }

    #[tokio::test]
    async fn test_commit_without_transaction_is_noop() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert!(!session.in_transaction());
        session.commit_transaction().await.unwrap();
        assert!(!session.in_transaction());
        assert_eq!(session.state().await, SessionState::Ready);
    }

    #[tokio::test]
    async fn test_rollback_without_transaction_is_noop() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert!(!session.in_transaction());
        session.rollback_transaction().await.unwrap();
        assert!(!session.in_transaction());
        assert_eq!(session.state().await, SessionState::Ready);
    }

    #[tokio::test]
    async fn test_session_close() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert!(!session.is_closed().await);

        session.close().await.unwrap();
        assert!(session.is_closed().await);
        assert_eq!(session.state().await, SessionState::Closed);
    }

    /// Scenario: A closed session rejects operations as closed
    #[tokio::test]
    async fn test_session_validate_ready() {
        let session = Session::new(
            "sess123".to_string(),
            mock_server_info(),
            SessionConfig::default(),
        );

        assert!(session.validate_ready().await.is_ok());

        session.set_state(SessionState::InTransaction).await;
        assert!(session.validate_ready().await.is_ok());

        for state in [SessionState::Closing, SessionState::Closed] {
            session.set_state(state).await;
            assert!(
                matches!(
                    session.validate_ready().await,
                    Err(ConnectionError::ConnectionClosed)
                ),
                "{:?} must report the connection as closed",
                state
            );
        }
    }

    #[tokio::test]
    async fn test_begin_transaction_reports_a_closed_session_as_closed() {
        for state in [SessionState::Closing, SessionState::Closed] {
            let session = Session::new(
                "sess123".to_string(),
                mock_server_info(),
                SessionConfig::default(),
            );
            session.set_state(state).await;

            assert!(
                matches!(
                    session.begin_transaction().await,
                    Err(ConnectionError::ConnectionClosed)
                ),
                "{:?} must report the connection as closed",
                state
            );
            assert!(!session.in_transaction());
        }
    }

    #[test]
    fn test_session_state_checks() {
        assert!(SessionState::Ready.can_execute());
        assert!(SessionState::InTransaction.can_execute());
        assert!(!SessionState::Closing.can_execute());
        assert!(!SessionState::Closed.can_execute());
    }

    /// Scenario: Native Parquet import threshold
    #[test]
    fn test_supports_native_parquet_import_is_pure_and_memoized() {
        let session_71 = Session::new(
            "s1".to_string(),
            mock_server_info_with_version("7.1.0"),
            SessionConfig::default(),
        );
        assert!(!session_71.supports_native_parquet_import());

        let session_2025_1_11 = Session::new(
            "s2".to_string(),
            mock_server_info_with_version("2025.1.11"),
            SessionConfig::default(),
        );
        assert!(session_2025_1_11.supports_native_parquet_import());

        let session_2025_2_1 = Session::new(
            "s3".to_string(),
            mock_server_info_with_version("2025.2.1"),
            SessionConfig::default(),
        );
        assert!(session_2025_2_1.supports_native_parquet_import());

        let session_garbage = Session::new(
            "s4".to_string(),
            mock_server_info_with_version("garbage"),
            SessionConfig::default(),
        );
        assert!(!session_garbage.supports_native_parquet_import());

        // Verify memoization: calling twice returns the same value without panic
        let first = session_2025_1_11.supports_native_parquet_import();
        let second = session_2025_1_11.supports_native_parquet_import();
        assert_eq!(first, second);
    }

    /// Scenario: Compare the server version against a caller-supplied threshold
    #[test]
    fn supports_at_least_compares_major_minor_and_patch() {
        let session = Session::new(
            "s".to_string(),
            mock_server_info_with_version("8.27.1"),
            SessionConfig::default(),
        );

        let cases = [
            ((8, 27, 1), true),
            ((8, 27, 2), false),
            ((8, 28, 0), false),
            ((8, 26, 99), true),
            ((7, 99, 99), true),
            ((9, 0, 0), false),
        ];
        for ((major, minor, patch), expected) in cases {
            assert_eq!(
                session.supports_at_least(major, minor, patch),
                expected,
                "8.27.1 against ({major}, {minor}, {patch})"
            );
        }
    }

    /// Scenario: Reject an unparsable release_version
    #[test]
    fn supports_at_least_is_false_for_an_unparsable_release_version() {
        for release in ["", "unknown", "v1"] {
            let session = Session::new(
                "s".to_string(),
                mock_server_info_with_version(release),
                SessionConfig::default(),
            );

            for (major, minor, patch) in [(0, 0, 0), (9, 0, 0)] {
                assert!(
                    !session.supports_at_least(major, minor, patch),
                    "release {release:?} must count as below ({major}, {minor}, {patch})"
                );
            }
        }
    }
}