solidb 1.2.1

//! Connection handler for native driver protocol
//!
//! Processes incoming commands and executes them against the storage engine.

use crate::driver::protocol::{
    decode_message, encode_response, Command, DriverError, Response, MAX_MESSAGE_SIZE,
};
use crate::storage::StorageEngine;
use crate::transaction::TransactionId;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

pub mod admin;
pub mod auth;
pub mod database;
pub mod document;
pub mod graph;
pub mod index;
pub mod query;
pub mod scheduler;
pub mod transaction;

/// Handler for a single driver connection
pub struct DriverHandler {
    pub(crate) storage: Arc<StorageEngine>,
    /// Replication log; driver mutations must be logged here or replicas
    /// silently diverge from data written over the binary protocol.
    pub(crate) replication: Option<Arc<crate::sync::log::SyncLog>>,
    /// Active transactions for this connection
    pub(crate) transactions: HashMap<String, TransactionId>,
    /// Authenticated database (None = not authenticated)
    pub(crate) authenticated_db: Option<String>,
    /// Principal identity (username or API key id) for audit logs
    pub(crate) session_subject: String,
    /// Permissions resolved from the principal's roles at auth time.
    /// Connection-lifetime snapshot: a revoked role applies on reconnect.
    pub(crate) session_permissions: std::collections::HashSet<crate::server::Permission>,
    /// The role names those permissions came from. Kept so the binary protocol
    /// can hand the query executor the same principal the HTTP handlers do —
    /// `CURRENT_USER` / `CURRENT_ROLES`, row policies, and the write-side query
    /// paths must not depend on which protocol the query arrived over.
    pub(crate) session_roles: Vec<String>,
    /// Database restriction for scoped API keys
    pub(crate) session_scoped_databases: Option<Vec<String>>,
    /// `_key` of the API key the session authenticated with, if any. Kept so
    /// the session can be re-validated against `_api_keys` (Audit L2).
    pub(crate) session_api_key_id: Option<String>,
    /// When the session's credential was last checked against storage.
    pub(crate) session_validated_at: Option<std::time::Instant>,
    /// Peer IP, the IP half of the login limiter bucket (same `ip|username`
    /// format as `/_api/auth/login`).
    pub(crate) peer_ip: String,
}

/// How long a client has to deliver a payload once it has sent its length.
const PAYLOAD_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);

/// Frame ceiling before the connection has authenticated. An `Auth` command
/// is a few hundred bytes; the 16 MB post-auth cap let an anonymous peer make
/// the server allocate and decode 16 MB per frame (Audit H5).
pub(crate) const MAX_PREAUTH_MESSAGE_SIZE: usize = 64 * 1024;

/// How often an authenticated session's credential is re-checked: a deleted
/// user or API key, an expired key, or a revoked role used to keep working
/// for the lifetime of the TCP connection (Audit L2).
const SESSION_REVALIDATE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);

/// The IP half of `ip:port`, matching `ConnectInfo<SocketAddr>::ip()` on the
/// HTTP login path so both protocols share one limiter bucket per user.
pub(crate) fn peer_ip_of(addr: &str) -> String {
    addr.parse::<std::net::SocketAddr>()
        .map(|a| a.ip().to_string())
        .unwrap_or_else(|_| addr.to_string())
}

impl DriverHandler {
    /// The session's identity as the query executor sees it.
    ///
    /// Resolved with the same check the command dispatch uses, against the
    /// database the query actually names, so a Read-only session cannot reach
    /// the write-side query paths (auto-index creation) and a row policy
    /// applies here exactly as it does over HTTP.
    pub(crate) fn query_principal(&self, database: &str) -> crate::sdbql::QueryPrincipal {
        let allows = |action| {
            crate::server::AuthorizationService::check_permission_raw(
                &self.session_permissions,
                action,
                Some(database),
                self.session_scoped_databases.as_deref(),
            )
            .is_ok()
        };
        crate::sdbql::QueryPrincipal {
            user: self.session_subject.clone(),
            roles: self.session_roles.clone(),
            can_read: allows(crate::server::PermissionAction::Read),
            can_write: allows(crate::server::PermissionAction::Write),
            can_admin: allows(crate::server::PermissionAction::Admin),
        }
    }

    /// Create a new handler
    pub fn new(
        storage: Arc<StorageEngine>,
        replication: Option<Arc<crate::sync::log::SyncLog>>,
    ) -> Self {
        Self {
            storage,
            replication,
            transactions: HashMap::new(),
            authenticated_db: None,
            session_subject: String::new(),
            session_permissions: std::collections::HashSet::new(),
            session_roles: Vec::new(),
            session_scoped_databases: None,
            session_api_key_id: None,
            session_validated_at: None,
            peer_ip: String::new(),
        }
    }

    /// Re-check the session's credential against current storage and refresh
    /// its permission snapshot. `Err` means the session must be closed.
    pub(crate) fn revalidate_session(&mut self) -> Result<(), String> {
        let Some(database) = self.authenticated_db.clone() else {
            return Ok(());
        };
        let system_db = self
            .storage
            .get_database("_system")
            .map_err(|e| format!("system database unavailable: {}", e))?;

        let (roles, scoped) = if let Some(key_id) = &self.session_api_key_id {
            let coll = system_db
                .system_collection(crate::server::auth::API_KEYS_COLL)
                .map_err(|_| "API key no longer exists".to_string())?;
            let doc = coll
                .get(key_id)
                .map_err(|_| "API key no longer exists".to_string())?;
            let key: crate::server::auth::ApiKey = serde_json::from_value(doc.to_value())
                .map_err(|_| "API key no longer valid".to_string())?;
            if let Some(ref expires_at) = key.expires_at {
                if let Ok(expiry) = chrono::DateTime::parse_from_rfc3339(expires_at) {
                    if expiry < chrono::Utc::now() {
                        return Err("API key expired".to_string());
                    }
                }
            }
            (
                key.roles,
                key.scoped_databases.filter(|dbs| !dbs.is_empty()),
            )
        } else {
            let admins = system_db
                .system_collection(crate::server::auth::ADMIN_COLL)
                .map_err(|_| "user no longer exists".to_string())?;
            if admins.get(&self.session_subject).is_err() {
                return Err("user no longer exists".to_string());
            }
            (
                crate::server::auth::AuthService::get_user_roles(
                    &self.storage,
                    &self.session_subject,
                )
                .unwrap_or_default(),
                None,
            )
        };

        let permissions = crate::server::AuthorizationService::load_permissions_from_storage(
            &self.storage,
            &roles,
        );
        crate::server::AuthorizationService::check_permission_raw(
            &permissions,
            crate::server::PermissionAction::Read,
            Some(&database),
            scoped.as_deref(),
        )
        .map_err(|e| e.to_string())?;

        self.session_permissions = permissions;
        self.session_roles = roles;
        self.session_scoped_databases = scoped;
        self.session_validated_at = Some(std::time::Instant::now());
        Ok(())
    }

    /// Handle a driver connection.
    ///
    /// Generic over the stream so the protocol runs identically over plain
    /// TCP and over a TLS-terminated connection.
    pub async fn handle_connection<S>(&mut self, mut stream: S, addr: String)
    where
        S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        tracing::info!("Driver connection from {}", addr);
        self.peer_ip = peer_ip_of(&addr);

        // The magic header has already been consumed by the multiplexer
        // Start processing commands immediately

        loop {
            // Read message length (4 bytes, big-endian)
            let mut len_buf = [0u8; 4];
            match stream.read_exact(&mut len_buf).await {
                Ok(_) => {}
                Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
                    tracing::debug!("Driver connection closed: {}", addr);
                    break;
                }
                Err(e) => {
                    tracing::warn!("Driver read error from {}: {}", addr, e);
                    break;
                }
            }

            let msg_len = u32::from_be_bytes(len_buf) as usize;

            // Validate message size. An unauthenticated peer only gets the
            // small pre-auth budget.
            let max_len = if self.authenticated_db.is_some() {
                MAX_MESSAGE_SIZE
            } else {
                MAX_PREAUTH_MESSAGE_SIZE
            };
            if msg_len > max_len {
                let resp = Response::error(DriverError::MessageTooLarge);
                if let Err(e) = self.send_response(&mut stream, &resp).await {
                    tracing::warn!("Failed to send error response: {}", e);
                }
                break;
            }

            // Read message payload. A client that announces a payload and
            // then goes quiet must not hold the connection (and this buffer)
            // open indefinitely; idle time *between* messages is fine.
            let mut payload = vec![0u8; msg_len];
            match tokio::time::timeout(PAYLOAD_READ_TIMEOUT, stream.read_exact(&mut payload)).await
            {
                Ok(Ok(_)) => {}
                Ok(Err(e)) => {
                    tracing::warn!("Driver read payload error from {}: {}", addr, e);
                    break;
                }
                Err(_) => {
                    tracing::warn!(
                        "Driver payload from {} not received within {:?}",
                        addr,
                        PAYLOAD_READ_TIMEOUT
                    );
                    break;
                }
            }

            // Decode command
            let command: Command = match decode_message(&payload) {
                Ok(cmd) => cmd,
                Err(e) => {
                    let resp = Response::error(e);
                    if let Err(e) = self.send_response(&mut stream, &resp).await {
                        tracing::warn!("Failed to send error response: {}", e);
                    }
                    continue;
                }
            };

            // Periodically re-check the credential; a revoked one ends the
            // session (uncommitted transactions are rolled back below).
            if self.authenticated_db.is_some()
                && self
                    .session_validated_at
                    .is_none_or(|t| t.elapsed() >= SESSION_REVALIDATE_INTERVAL)
            {
                if let Err(reason) = self.revalidate_session() {
                    tracing::warn!(
                        target: "audit",
                        user = %self.session_subject,
                        "closing driver session from {}: {}",
                        addr,
                        reason
                    );
                    let resp = Response::error(DriverError::AuthError(format!(
                        "Session no longer valid: {}",
                        reason
                    )));
                    let _ = self.send_response(&mut stream, &resp).await;
                    break;
                }
            }

            // Execute command
            let response = self.execute_command(command).await;

            // Send response
            if let Err(e) = self.send_response(&mut stream, &response).await {
                tracing::warn!("Failed to send response to {}: {}", addr, e);
                break;
            }
        }

        // Cleanup: rollback any uncommitted transactions
        for (tx_id_str, tx_id) in self.transactions.drain() {
            tracing::debug!("Rolling back uncommitted transaction: {}", tx_id_str);
            let _ = self.storage.rollback_transaction(tx_id);
        }
    }

    /// Send a response to the client
    async fn send_response<S>(&self, stream: &mut S, response: &Response) -> Result<(), DriverError>
    where
        S: tokio::io::AsyncWrite + Unpin,
    {
        let data = encode_response(response)?;
        stream
            .write_all(&data)
            .await
            .map_err(|e| DriverError::ConnectionError(e.to_string()))?;
        stream
            .flush()
            .await
            .map_err(|e| DriverError::ConnectionError(e.to_string()))?;
        Ok(())
    }

    /// Execute a command and return a response
    async fn execute_command(&mut self, command: Command) -> Response {
        // Gate every command behind authentication. Only Ping and Auth are
        // allowed before the connection has authenticated. Batch used to be
        // let through too, so one pre-auth frame could carry hundreds of
        // thousands of `Auth` entries — each an Argon2 run, none rate
        // limited (Audit H5).
        if self.authenticated_db.is_none() {
            match &command {
                Command::Ping | Command::Auth { .. } => {}
                _ => {
                    return Response::error(DriverError::AuthError(
                        "Authentication required".to_string(),
                    ));
                }
            }
        }

        // Per-command authorization against the session's resolved
        // permissions, mirroring the HTTP authz middleware. The command's own
        // `database` field is checked — connections are NOT trusted to stay
        // on the database they authenticated against.
        if self.authenticated_db.is_some() {
            if let Some(action) = command.required_action() {
                if let Err(e) = crate::server::AuthorizationService::check_permission_raw(
                    &self.session_permissions,
                    action,
                    command.database(),
                    self.session_scoped_databases.as_deref(),
                ) {
                    tracing::warn!(
                        target: "audit",
                        user = %self.session_subject,
                        "driver command denied: {}",
                        e
                    );
                    return Response::error(DriverError::AuthError(e.to_string()));
                }
            }
        }

        match command {
            // ==================== Auth & Utility ====================
            Command::Ping => Response::pong(),

            Command::Auth {
                database,
                username,
                password,
                api_key,
            } => auth::handle_auth(self, database, username, password, api_key).await,

            // ==================== Database Operations ====================
            Command::ListDatabases => database::handle_list_databases(self),

            Command::CreateDatabase { name } => database::handle_create_database(self, name),

            Command::DeleteDatabase { name } => database::handle_delete_database(self, name),

            // ==================== Collection Operations ====================
            Command::ListCollections { database } => {
                database::handle_list_collections(self, database)
            }

            Command::CreateCollection {
                database,
                name,
                collection_type,
            } => database::handle_create_collection(self, database, name, collection_type),

            Command::DeleteCollection { database, name } => {
                database::handle_delete_collection(self, database, name)
            }

            Command::CollectionStats { database, name } => {
                database::handle_collection_stats(self, database, name)
            }

            // ==================== Document Operations ====================
            Command::Get {
                database,
                collection,
                key,
            } => document::handle_get(self, database, collection, key),

            Command::Insert {
                database,
                collection,
                key,
                document,
            } => document::handle_insert(self, database, collection, key, document),

            Command::Update {
                database,
                collection,
                key,
                document,
                merge,
            } => document::handle_update(self, database, collection, key, document, merge),

            Command::Delete {
                database,
                collection,
                key,
            } => document::handle_delete(self, database, collection, key),

            Command::List {
                database,
                collection,
                limit,
                offset,
            } => document::handle_list(self, database, collection, limit, offset),

            // ==================== Query Operations ====================
            Command::Query {
                database,
                sdbql,
                bind_vars,
                cache,
            } => query::handle_query(self, database, sdbql, bind_vars, cache).await,

            Command::Explain {
                database,
                sdbql,
                bind_vars,
            } => query::handle_explain(self, database, sdbql, bind_vars).await,

            // ==================== Index Operations ====================
            Command::CreateIndex {
                database,
                collection,
                name,
                fields,
                unique,
                sparse: _,
            } => {
                // Off the runtime: a build over a large collection can take
                // minutes, and this task shares its worker with every other
                // connection. No time limit on purpose.
                let coll = self.get_collection(&database, &collection);
                match coll {
                    Ok(coll) => tokio::task::spawn_blocking(move || {
                        index::create_persistent_index(&coll, name, fields, unique)
                    })
                    .await
                    .unwrap_or_else(|e| {
                        Response::error(DriverError::DatabaseError(format!(
                            "Task join error: {}",
                            e
                        )))
                    }),
                    Err(e) => Response::error(e),
                }
            }

            Command::DeleteIndex {
                database,
                collection,
                name,
            } => index::handle_delete_index(self, database, collection, name),

            Command::ListIndexes {
                database,
                collection,
            } => index::handle_list_indexes(self, database, collection),

            // ==================== Transaction Operations ====================
            Command::BeginTransaction {
                database,
                isolation_level,
            } => transaction::handle_begin_transaction(self, database, isolation_level),

            Command::CommitTransaction { tx_id } => {
                transaction::handle_commit_transaction(self, tx_id)
            }

            Command::RollbackTransaction { tx_id } => {
                transaction::handle_rollback_transaction(self, tx_id)
            }

            Command::TransactionCommand { tx_id, command } => {
                transaction::handle_transaction_command(self, tx_id, command).await
            }

            // ==================== Bulk Operations ====================
            Command::Batch { commands } => {
                let mut responses = Vec::with_capacity(commands.len());
                for cmd in commands {
                    // One frame must not carry a stream of password checks:
                    // each `Auth` is an Argon2 run (Audit H5).
                    if matches!(cmd, Command::Auth { .. }) {
                        responses.push(Response::error(DriverError::InvalidCommand(
                            "Auth is not allowed inside a Batch".to_string(),
                        )));
                        continue;
                    }
                    let resp = Box::pin(self.execute_command(cmd)).await;
                    responses.push(resp);
                }
                Response::Batch { responses }
            }

            Command::BulkInsert {
                database,
                collection,
                documents,
            } => document::handle_bulk_insert(self, database, collection, documents),

            // ==================== Script Management ====================
            Command::CreateScript {
                database,
                name,
                path,
                methods,
                code,
                description,
                collection,
            } => {
                scheduler::handle_script_create(
                    self,
                    database,
                    scheduler::ScriptCreateConfig {
                        name,
                        path,
                        methods,
                        code,
                        description,
                        collection,
                    },
                )
                .await
            }

            Command::ListScripts { database } => {
                scheduler::handle_script_list(self, database).await
            }

            Command::GetScript {
                database,
                script_id,
            } => scheduler::handle_script_get(self, database, script_id).await,

            Command::UpdateScript {
                database,
                script_id,
                name,
                path,
                methods,
                code,
                description,
            } => {
                scheduler::handle_script_update(
                    self,
                    database,
                    script_id,
                    scheduler::ScriptUpdateConfig {
                        name,
                        path,
                        methods,
                        code,
                        description,
                    },
                )
                .await
            }

            Command::DeleteScript {
                database,
                script_id,
            } => scheduler::handle_script_delete(self, database, script_id).await,

            Command::GetScriptStats => {
                Response::ok(serde_json::json!({"message": "Script stats available via HTTP API"}))
            }

            // ==================== Trigger Management ====================
            Command::ListTriggers { database } => {
                scheduler::handle_list_triggers(self, database).await
            }

            Command::ListCollectionTriggers {
                database,
                collection,
            } => scheduler::handle_list_collection_triggers(self, database, collection).await,

            Command::CreateTrigger {
                database,
                name,
                collection,
                events,
                script_path,
                filter,
                queue,
                priority,
                max_retries,
                enabled,
            } => {
                scheduler::handle_create_trigger(
                    self,
                    database,
                    scheduler::TriggerCreateConfig {
                        name,
                        collection,
                        events,
                        script_path,
                        filter,
                        queue,
                        priority,
                        max_retries,
                        enabled,
                    },
                )
                .await
            }

            Command::GetTrigger {
                database,
                trigger_id,
            } => scheduler::handle_get_trigger(self, database, trigger_id).await,

            Command::UpdateTrigger {
                database,
                trigger_id,
                name,
                events,
                script_path,
                filter,
                queue,
                priority,
                max_retries,
                enabled,
            } => {
                scheduler::handle_update_trigger(
                    self,
                    database,
                    trigger_id,
                    scheduler::TriggerUpdateConfig {
                        name,
                        events,
                        script_path,
                        filter,
                        queue,
                        priority,
                        max_retries,
                        enabled,
                    },
                )
                .await
            }

            Command::DeleteTrigger {
                database,
                trigger_id,
            } => scheduler::handle_delete_trigger(self, database, trigger_id).await,

            Command::ToggleTrigger {
                database,
                trigger_id,
            } => scheduler::handle_toggle_trigger(self, database, trigger_id).await,

            // ==================== Environment Variables ====================
            Command::ListEnvVars { database } => admin::handle_list_env_vars(self, database).await,

            Command::SetEnvVar {
                database,
                key,
                value,
            } => admin::handle_set_env_var(self, database, key, value).await,

            Command::DeleteEnvVar { database, key } => {
                admin::handle_delete_env_var(self, database, key).await
            }

            // ==================== Role Management ====================
            Command::ListRoles => admin::handle_list_roles(self).await,

            Command::CreateRole { name, permissions } => {
                admin::handle_create_role(self, name, permissions).await
            }

            Command::GetRole { name } => admin::handle_get_role(self, name).await,

            Command::UpdateRole { name, permissions } => {
                admin::handle_update_role(self, name, permissions).await
            }

            Command::DeleteRole { name } => admin::handle_delete_role(self, name).await,

            // ==================== User Management ====================
            Command::ListUsers => admin::handle_list_users(self).await,

            Command::CreateUser {
                username,
                password,
                roles,
            } => admin::handle_create_user(self, username, password, roles).await,

            Command::DeleteUser { username } => admin::handle_delete_user(self, username).await,

            Command::GetUserRoles { username } => {
                admin::handle_get_user_roles(self, username).await
            }

            Command::AssignRole {
                username,
                role,
                database,
            } => admin::handle_assign_role(self, username, role, database).await,

            Command::RevokeRole { username, role } => {
                admin::handle_revoke_role(self, username, role).await
            }

            Command::GetCurrentUser => {
                // This requires knowing the current authenticated user context
                Response::ok(serde_json::json!({"database": self.authenticated_db}))
            }

            Command::GetCurrentUserPermissions => {
                // Would need user context
                Response::ok(serde_json::json!({"message": "Permissions available after auth"}))
            }

            // ==================== API Key Management ====================
            Command::ListApiKeys => admin::handle_list_api_keys(self).await,

            Command::CreateApiKey {
                name,
                permissions,
                expires_at,
            } => admin::handle_create_api_key(self, name, permissions, expires_at).await,

            Command::DeleteApiKey { key_id } => admin::handle_delete_api_key(self, key_id).await,

            // ==================== Cluster Management ====================
            Command::ClusterStatus | Command::ClusterInfo => {
                // These need ClusterManager which driver doesn't have access to
                Response::error(DriverError::DatabaseError(
                    "Cluster operations require HTTP API".to_string(),
                ))
            }

            Command::ClusterRemoveNode { .. }
            | Command::ClusterRebalance
            | Command::ClusterCleanup
            | Command::ClusterReshard { .. } => Response::error(DriverError::DatabaseError(
                "Cluster operations require HTTP API".to_string(),
            )),

            // ==================== Advanced Collection Operations ====================
            Command::TruncateCollection {
                database,
                collection,
            } => database::handle_truncate_collection(self, database, collection),

            Command::CompactCollection {
                database,
                collection,
            } => database::handle_compact_collection(self, database, collection),

            Command::PruneCollection { .. } => Response::error(DriverError::InvalidCommand(
                "Prune not supported".to_string(),
            )),

            Command::RecountCollection {
                database,
                collection,
            } => database::handle_recount_collection(self, database, collection),

            Command::RepairCollection { .. } => Response::error(DriverError::InvalidCommand(
                "Repair not supported".to_string(),
            )),

            Command::GetCollectionSharding { .. } => Response::error(DriverError::InvalidCommand(
                "Sharding not supported".to_string(),
            )),

            Command::ExportCollection {
                database,
                collection,
            } => database::handle_export_collection(self, database, collection).await,

            Command::ImportCollection {
                database,
                collection,
                documents,
            } => database::handle_import_collection(self, database, collection, documents),

            Command::SetCollectionSchema {
                database,
                collection,
                schema,
            } => database::handle_set_collection_schema(self, database, collection, schema),

            Command::GetCollectionSchema {
                database,
                collection,
            } => database::handle_get_collection_schema(self, database, collection),

            Command::DeleteCollectionSchema {
                database,
                collection,
            } => database::handle_delete_collection_schema(self, database, collection),

            // ==================== Advanced Index Operations ====================
            Command::RebuildIndexes {
                database,
                collection,
            } => index::handle_rebuild_indexes(self, database, collection),

            Command::HybridSearch {
                database,
                collection,
                vector,
                text_query,
                vector_index,
                fulltext_field,
                vector_weight,
                text_weight,
                limit,
                fusion,
            } => index::handle_hybrid_search(
                self,
                database,
                collection,
                vector,
                text_query,
                vector_index,
                fulltext_field,
                vector_weight,
                text_weight,
                limit,
                fusion,
            ),

            Command::GraphNeighbors {
                database,
                edge_collection,
                seeds,
                options,
            } => graph::handle_graph_neighbors(self, database, edge_collection, seeds, options),

            Command::GraphRag {
                database,
                seed_collection,
                vector_index,
                edge_collection,
                query_vector,
                options,
            } => graph::handle_graph_rag(
                self,
                database,
                seed_collection,
                vector_index,
                edge_collection,
                query_vector,
                options,
            ),

            Command::CommunitySearch {
                database,
                query_text,
                options,
            } => graph::handle_community_search(self, database, query_text, options),

            // ==================== Geo Index Operations ====================
            Command::CreateGeoIndex {
                database,
                collection,
                name,
                field,
            } => index::handle_create_geo_index(self, database, collection, name, field),

            Command::ListGeoIndexes {
                database,
                collection,
            } => index::handle_list_geo_indexes(self, database, collection),

            Command::DeleteGeoIndex {
                database,
                collection,
                name,
            } => index::handle_delete_geo_index(self, database, collection, name),

            Command::GeoNear {
                database,
                collection,
                field,
                latitude,
                longitude,
                radius,
                limit,
            } => index::handle_geo_near(
                self,
                database,
                index::GeoNearConfig {
                    collection,
                    field,
                    latitude,
                    longitude,
                    radius,
                    limit,
                },
            ),

            Command::GeoWithin { .. } => Response::error(DriverError::InvalidCommand(
                "Geo polygon search not supported".to_string(),
            )),

            // ==================== Vector Index Operations ====================
            Command::CreateVectorIndex {
                database,
                collection,
                name,
                field,
                dimensions,
                metric,
                ef_construction,
                m,
            } => index::handle_create_vector_index(
                self,
                database,
                index::VectorIndexCreateConfig {
                    collection,
                    name,
                    field,
                    dimensions,
                    metric,
                    ef_construction,
                    m,
                },
            ),

            Command::ListVectorIndexes {
                database,
                collection,
            } => index::handle_list_vector_indexes(self, database, collection),

            Command::DeleteVectorIndex {
                database,
                collection,
                name,
            } => index::handle_delete_vector_index(self, database, collection, name),

            Command::VectorSearch {
                database,
                collection,
                index_name,
                vector,
                limit,
                ef_search,
                filter: _,
            } => index::handle_vector_search(
                self, database, collection, index_name, vector, limit, ef_search,
            ),

            Command::QuantizeVectorIndex {
                database,
                collection,
                index_name,
            } => index::handle_quantize_vector_index(self, database, collection, index_name),

            Command::DequantizeVectorIndex {
                database,
                collection,
                index_name,
            } => index::handle_dequantize_vector_index(self, database, collection, index_name),

            // ==================== TTL Index Operations ====================
            Command::CreateTtlIndex {
                database,
                collection,
                name,
                field,
                expire_after_seconds,
            } => index::handle_create_ttl_index(
                self,
                database,
                collection,
                name,
                field,
                expire_after_seconds,
            ),

            Command::ListTtlIndexes {
                database,
                collection,
            } => index::handle_list_ttl_indexes(self, database, collection),

            Command::DeleteTtlIndex {
                database,
                collection,
                name,
            } => index::handle_delete_ttl_index(self, database, collection, name),

            // ==================== Columnar Storage ====================
            Command::CreateColumnar {
                database,
                name,
                columns,
            } => database::handle_create_columnar(self, database, name, columns),

            Command::ListColumnar { database } => database::handle_list_columnar(self, database),

            Command::GetColumnar {
                database,
                collection,
            } => database::handle_get_columnar(self, database, collection),

            Command::DeleteColumnar {
                database,
                collection,
            } => database::handle_delete_columnar(self, database, collection),

            Command::InsertColumnar {
                database,
                collection,
                rows,
            } => database::handle_insert_columnar(self, database, collection, rows),

            Command::AggregateColumnar {
                database,
                collection,
                aggregations,
                group_by,
                filter,
            } => database::handle_aggregate_columnar(
                self,
                database,
                collection,
                aggregations,
                group_by,
                filter,
            ),

            Command::QueryColumnar {
                database,
                collection,
                columns,
                filter,
                order_by,
                limit,
            } => database::handle_query_columnar(
                self, database, collection, columns, filter, order_by, limit,
            ),

            Command::CreateColumnarIndex {
                database,
                collection,
                column,
            } => database::handle_create_columnar_index(self, database, collection, column),

            Command::ListColumnarIndexes {
                database,
                collection,
            } => database::handle_list_columnar_indexes(self, database, collection),

            Command::DeleteColumnarIndex {
                database,
                collection,
                column,
            } => database::handle_delete_columnar_index(self, database, collection, column),
        }
    }

    /// Log a driver mutation to the replication log, mirroring the HTTP
    /// handlers' behavior (physical shard collections are partitioned, not
    /// replicated — same rule as `server/handlers/documents.rs`).
    pub(crate) fn log_replication(
        &self,
        database: &str,
        collection: &str,
        operation: crate::sync::protocol::Operation,
        key: &str,
        data: Option<&serde_json::Value>,
    ) {
        if crate::server::handlers::system::is_physical_shard_collection(collection) {
            return;
        }
        if let Some(ref log) = self.replication {
            log.log_document_op(
                database,
                collection,
                operation,
                key,
                data.and_then(|v| serde_json::to_vec(v).ok()),
            );
        }
    }

    /// Batch variant of `log_replication` (single fsync via `append_batch`).
    pub(crate) fn log_replication_batch(
        &self,
        database: &str,
        collection: &str,
        operation: crate::sync::protocol::Operation,
        docs: &[crate::storage::Document],
    ) {
        if crate::server::handlers::system::is_physical_shard_collection(collection) {
            return;
        }
        if let Some(ref log) = self.replication {
            let entries = docs
                .iter()
                .map(|doc| {
                    crate::sync::log::LogEntry::new_op(
                        database,
                        collection,
                        operation,
                        doc.key.clone(),
                        serde_json::to_vec(&doc.to_value()).ok(),
                    )
                })
                .collect();
            log.append_batch(entries);
        }
    }

    /// A collection the connection's principal is about to write by name.
    ///
    /// The driver used to reach every collection through [`Self::get_collection`],
    /// so the write tiers the HTTP document API enforces did not exist here:
    /// a driver `Insert` into `_scripts` went through.
    pub(crate) fn get_collection_for_write(
        &self,
        database: &str,
        collection: &str,
    ) -> Result<crate::storage::Collection, DriverError> {
        let actor = crate::storage::WriteActor::client(self.query_principal(database).can_admin);
        let db = self
            .storage
            .get_database(database)
            .map_err(|e| DriverError::DatabaseError(e.to_string()))?;
        db.get_collection_for_write(collection, actor)
            .map_err(|e| DriverError::DatabaseError(e.to_string()))
    }

    /// Helper to get a collection
    pub(crate) fn get_collection(
        &self,
        database: &str,
        collection: &str,
    ) -> Result<crate::storage::Collection, DriverError> {
        let db = self
            .storage
            .get_database(database)
            .map_err(|e| DriverError::DatabaseError(e.to_string()))?;
        db.get_collection(collection)
            .map_err(|e| DriverError::DatabaseError(e.to_string()))
    }
}

/// Object-safe stream bound for driver connections.
pub trait DriverConnTrait: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin {}
impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + ?Sized> DriverConnTrait
    for T
{
}

/// A driver connection stream, boxed so the multiplexer can dispatch either
/// a plain TCP connection or a TLS-terminated one to the same handler.
pub type DriverConn = Box<dyn DriverConnTrait>;

/// Spawn a handler for incoming driver connections
pub fn spawn_driver_handler(
    storage: Arc<StorageEngine>,
    replication: Option<Arc<crate::sync::log::SyncLog>>,
) -> tokio::sync::mpsc::Sender<(DriverConn, String)> {
    let (tx, mut rx) = tokio::sync::mpsc::channel::<(DriverConn, String)>(100);

    tokio::spawn(async move {
        while let Some((mut stream, addr)) = rx.recv().await {
            let storage = storage.clone();
            let replication = replication.clone();
            tokio::spawn(async move {
                let mut handler = DriverHandler::new(storage, replication);
                handler.handle_connection(&mut *stream, addr).await;
            });
        }
    });

    tx
}