netget/server/snowflake/

actions.rs

//! Snowflake protocol actions implementation
//!
//! Snowflake's client/driver protocol is HTTPS + JSON: the driver POSTs to a small
//! set of REST endpoints (`/session/v1/login-request`, `/queries/v1/query-request`,
//! `/session/logout-request`, `/session/token-request`) and reads a JSON envelope
//! back. NetGet serves those endpoints and lets the LLM play the warehouse: it
//! decides the session token minted at login and the rowset returned for a query.
//!
//! **There is no storage.** No tables, no rows, no session table in Rust. The model
//! answers every request. Because there is no session store, the query endpoint
//! cannot *validate* a token against issued ones — it surfaces whether an
//! `Authorization` token was presented and leaves the decision to the model.

use crate::llm::actions::{
    protocol_trait::{ActionResult, Protocol, Server},
    ActionDefinition, Parameter,
};
use crate::protocol::log_template::LogTemplate;
use crate::protocol::EventType;
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use anyhow::{Context, Result};
use serde_json::json;
use std::sync::{Arc, LazyLock};
use tokio::sync::mpsc;
use tracing::debug;

/// Snowflake protocol action handler.
///
/// Like the other database protocols this is a thin, stateless description: it
/// carries no session table (there is no storage), only the channel used for dual
/// logging when an action is executed.
pub struct SnowflakeProtocol {
    _connection_id: ConnectionId,
    _app_state: Arc<AppState>,
    status_tx: mpsc::UnboundedSender<String>,
}

impl SnowflakeProtocol {
    pub fn new(
        connection_id: ConnectionId,
        app_state: Arc<AppState>,
        status_tx: mpsc::UnboundedSender<String>,
    ) -> Self {
        Self {
            _connection_id: connection_id,
            _app_state: app_state,
            status_tx,
        }
    }
}

impl Protocol for SnowflakeProtocol {
    fn get_startup_parameters(&self) -> Vec<crate::llm::actions::ParameterDefinition> {
        // No startup parameters: every request is answered by the model, and there
        // is nothing to configure at spawn time. Declaring a parameter that is never
        // read is a documented trap (the model tries to use it), so declare none.
        vec![]
    }

    fn get_async_actions(&self, _state: &AppState) -> Vec<ActionDefinition> {
        // Snowflake is purely reactive: the model answers login/query/session events.
        vec![]
    }

    fn get_sync_actions(&self) -> Vec<ActionDefinition> {
        vec![
            snowflake_login_success_action(),
            snowflake_query_response_action(),
            snowflake_session_response_action(),
            snowflake_error_action(),
        ]
    }

    fn protocol_name(&self) -> &'static str {
        "Snowflake"
    }

    fn get_event_types(&self) -> Vec<EventType> {
        get_snowflake_event_types()
    }

    fn stack_name(&self) -> &'static str {
        "ETH>IP>TCP>HTTP>Snowflake"
    }

    fn keywords(&self) -> Vec<&'static str> {
        // Distinctive, multi-word keywords — no bare "db"/"sql"/"database".
        vec!["snowflake", "snowflake warehouse", "snowflake data cloud"]
    }

    fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
        use crate::protocol::metadata::{DevelopmentState, ProtocolMetadataV2};

        ProtocolMetadataV2::builder()
            .state(DevelopmentState::Experimental)
            .implementation("Hand-written Snowflake REST/JSON endpoints over hyper (login, query, session)")
            .llm_control("Session tokens at login, rowset/rowtype for queries, error envelopes")
            .e2e_testing("tests/server/snowflake/real_client_test.rs drives the server with snowflake-connector-python 4.8.0, Snowflake's official Python driver (gzipped request bodies): a login the model grants and one it refuses (the driver raises 390100), a six-column rowset typed by the driver's own JSON converters (FIXED with a scale, BOOLEAN, DATE, TIMESTAMP_NTZ, NULL), a query error raised with the model's errno and SQLSTATE, and the driver's POST /session?delete=true logout. e2e_test.rs and llm_failure_test.rs drive the envelopes with reqwest")
            .notes(
                "HTTP/JSON only; TLS termination is out of scope (real drivers use HTTPS — front with a \
                 TLS proxy if needed). One real driver so far (snowflake-connector-python over plain \
                 HTTP); JDBC, ODBC and gosnowflake are untested, and results are JSON only (no Arrow, \
                 no chunks). No session store, so tokens are not validated against issued ones. \
                 Fail-closed: an LLM outage on login is a refusal (no token issued), never a \
                 success-shaped empty result.",
            )
            .max_inbound_bytes(crate::server::snowflake::MAX_REQUEST_BYTES)
            .request_only("The Snowflake API is HTTP request/response; a server cannot send a peer anything unprompted")
            .build()
    }

    fn description(&self) -> &'static str {
        "Snowflake data-warehouse server (REST/JSON login + query protocol)"
    }

    fn example_prompt(&self) -> &'static str {
        "Start a Snowflake server on port 8085 that logs clients in and answers SELECT queries"
    }

    fn group_name(&self) -> &'static str {
        "Database"
    }

    fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
        use crate::llm::actions::StartupExamples;

        // Deterministic: answer every SQL query with a single-column "OK"
        // rowset, no LLM call.
        let script = r#"import json, sys
data = json.load(sys.stdin)
if data["event_type_id"] == "snowflake_query":
    actions = [{"type": "snowflake_query_response",
                "rowtype": [{"name": "STATUS", "type": "text"}],
                "rowset": [["OK"]]}]
else:
    actions = []
print(json.dumps({"actions": actions}))"#;

        StartupExamples::new(
            // LLM mode
            json!({
                "type": "open_server",
                "port": 8085,
                "base_stack": "snowflake",
                "instruction": "Snowflake warehouse: accept logins and answer SQL queries with plausible rowsets"
            }),
            // Script mode
            json!({
                "type": "open_server",
                "port": 8085,
                "base_stack": "snowflake",
                "event_handlers": [{
                    "event_pattern": "snowflake_query",
                    "handler": {
                        "type": "script",
                        "language": "python",
                        "code": script
                    }
                }]
            }),
            // Static mode
            json!({
                "type": "open_server",
                "port": 8085,
                "base_stack": "snowflake",
                "event_handlers": [{
                    "event_pattern": "snowflake_login",
                    "handler": {
                        "type": "static",
                        "actions": [{
                            "type": "snowflake_login_success",
                            "token": "SESSION_TOKEN_ABC",
                            "master_token": "MASTER_TOKEN_ABC",
                            "session_id": 100
                        }]
                    }
                }]
            }),
        )
    }
}

impl Server for SnowflakeProtocol {
    fn spawn(
        &self,
        ctx: crate::protocol::SpawnContext,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = anyhow::Result<std::net::SocketAddr>> + Send>,
    > {
        Box::pin(async move {
            use crate::server::snowflake::SnowflakeServer;
            SnowflakeServer::spawn_with_llm_actions(
                ctx.legacy_listen_addr(),
                ctx.llm_client,
                ctx.state,
                ctx.status_tx,
                ctx.server_id,
            )
            .await
        })
    }

    fn execute_action(&self, action: serde_json::Value) -> Result<ActionResult> {
        let action_type = action
            .get("type")
            .and_then(|v| v.as_str())
            .context("Missing 'type' field in action")?;

        match action_type {
            "snowflake_login_success" => self.execute_login_success(action),
            "snowflake_query_response" => self.execute_query_response(action),
            "snowflake_session_response" => self.execute_session_response(action),
            "snowflake_error" => self.execute_error(action),
            _ => Err(anyhow::anyhow!("Unknown Snowflake action: {}", action_type)),
        }
    }
}

impl SnowflakeProtocol {
    fn execute_login_success(&self, action: serde_json::Value) -> Result<ActionResult> {
        let token = action
            .get("token")
            .and_then(|v| v.as_str())
            .context("Missing 'token' in snowflake_login_success")?
            .to_string();
        let master_token = action
            .get("master_token")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());
        let session_id = action.get("session_id").cloned();
        let validity_seconds = action.get("validity_seconds").and_then(|v| v.as_u64());

        debug!("Snowflake login success: session token minted");
        let _ = self
            .status_tx
            .send("[DEBUG] Snowflake → login success (token issued)".to_string());

        Ok(ActionResult::Custom {
            name: "snowflake_login_success".to_string(),
            data: json!({
                "token": token,
                "master_token": master_token,
                "session_id": session_id,
                "validity_seconds": validity_seconds,
            }),
        })
    }

    fn execute_query_response(&self, action: serde_json::Value) -> Result<ActionResult> {
        let rowtype = action
            .get("rowtype")
            .and_then(|v| v.as_array())
            .context("Missing 'rowtype' array in snowflake_query_response")?;
        let rowset = action
            .get("rowset")
            .and_then(|v| v.as_array())
            .context("Missing 'rowset' array in snowflake_query_response")?;
        let query_id = action
            .get("query_id")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        debug!(
            "Snowflake query response: {} columns, {} rows",
            rowtype.len(),
            rowset.len()
        );
        let _ = self.status_tx.send(format!(
            "[DEBUG] Snowflake → rowset: {} columns, {} rows",
            rowtype.len(),
            rowset.len()
        ));

        Ok(ActionResult::Custom {
            name: "snowflake_query_response".to_string(),
            data: json!({
                "rowtype": rowtype,
                "rowset": rowset,
                "query_id": query_id,
            }),
        })
    }

    fn execute_session_response(&self, action: serde_json::Value) -> Result<ActionResult> {
        let data = action.get("data").cloned().unwrap_or(json!({}));
        debug!("Snowflake session response");
        let _ = self
            .status_tx
            .send("[DEBUG] Snowflake → session response".to_string());
        Ok(ActionResult::Custom {
            name: "snowflake_session_response".to_string(),
            data: json!({ "data": data }),
        })
    }

    fn execute_error(&self, action: serde_json::Value) -> Result<ActionResult> {
        let code = action
            .get("code")
            .and_then(|v| v.as_str())
            .unwrap_or("000603")
            .to_string();
        let message = action
            .get("message")
            .and_then(|v| v.as_str())
            .unwrap_or("Unknown error")
            .to_string();
        let sql_state = action
            .get("sql_state")
            .and_then(|v| v.as_str())
            .filter(|s| !s.is_empty());
        debug!("Snowflake error response: {} - {}", code, message);
        let _ = self
            .status_tx
            .send(format!("[DEBUG] Snowflake ✗ error {}: {}", code, message));
        let mut data = json!({ "code": code, "message": message });
        if let Some(state) = sql_state {
            data["sql_state"] = json!(state);
        }
        Ok(ActionResult::Custom {
            name: "snowflake_error".to_string(),
            data,
        })
    }
}

// ============================================================================
// Action definitions
// ============================================================================

/// `snowflake_login_success` — issue a session token in response to a login.
pub fn snowflake_login_success_action() -> ActionDefinition {
    ActionDefinition {
        name: "snowflake_login_success".to_string(),
        description: "Authenticate the client and issue a Snowflake session token. Only emit this \
            when the login should succeed; to refuse a login use snowflake_error instead (never \
            leave the event unanswered — an unanswered login is treated as a refusal)."
            .to_string(),
        parameters: vec![
            Parameter {
                name: "token".to_string(),
                type_hint: "string".to_string(),
                description: "The session token the driver will send back as \
                    Authorization: Snowflake Token=\"...\" on subsequent query requests."
                    .to_string(),
                required: true,
            },
            Parameter {
                name: "master_token".to_string(),
                type_hint: "string".to_string(),
                description: "Master token used by the driver to renew the session (optional)."
                    .to_string(),
                required: false,
            },
            Parameter {
                name: "session_id".to_string(),
                type_hint: "number".to_string(),
                description: "Numeric session id (optional)".to_string(),
                required: false,
            },
            Parameter {
                name: "validity_seconds".to_string(),
                type_hint: "number".to_string(),
                description: "Session token validity in seconds (optional, default 3600)"
                    .to_string(),
                required: false,
            },
        ],
        example: json!({
            "type": "snowflake_login_success",
            "token": "ver:1-hint:abc-ETMsDgAAA...",
            "master_token": "ver:1-hint:def-ETMsDgAAA...",
            "session_id": 123456789,
            "validity_seconds": 3600
        }),
        log_template: Some(
            LogTemplate::new()
                .with_info("-> Snowflake login OK (session {session_id})")
                .with_debug("Snowflake snowflake_login_success: session_id={session_id}"),
        ),
    }
}

/// `snowflake_query_response` — return a result set for a SQL query.
pub fn snowflake_query_response_action() -> ActionDefinition {
    ActionDefinition {
        name: "snowflake_query_response".to_string(),
        description: "Return a result set for a SQL query. `rowtype` describes the columns and \
            `rowset` is the rows. Snowflake's JSON result format transmits every value as a \
            string, so all cell values are stringified on the wire."
            .to_string(),
        parameters: vec![
            Parameter {
                name: "rowtype".to_string(),
                type_hint: "array".to_string(),
                description: "Array of column descriptors. Each is an object with at least \
                    'name' and 'type' (Snowflake logical types: text, fixed, real, boolean, date, \
                    time, timestamp_ntz, variant, ...). Optional fields (nullable, length, \
                    precision, scale) are passed through; sensible defaults are filled in."
                    .to_string(),
                required: true,
            },
            Parameter {
                name: "rowset".to_string(),
                type_hint: "array".to_string(),
                description: "Array of rows; each row is an array of cell values matching the \
                    column order. Values are sent as strings (Snowflake JSON result format)."
                    .to_string(),
                required: true,
            },
            Parameter {
                name: "query_id".to_string(),
                type_hint: "string".to_string(),
                description: "Optional query id (UUID) echoed back to the driver".to_string(),
                required: false,
            },
        ],
        example: json!({
            "type": "snowflake_query_response",
            "rowtype": [
                {"name": "ID", "type": "fixed"},
                {"name": "NAME", "type": "text"}
            ],
            "rowset": [["1", "Alice"], ["2", "Bob"]],
            "query_id": "01b2c3d4-0000-0000-0000-000000000001"
        }),
        log_template: Some(
            LogTemplate::new()
                .with_info("-> Snowflake {rowtype_len} cols, {rowset_len} rows")
                .with_debug(
                    "Snowflake snowflake_query_response: {rowtype_len} cols, {rowset_len} rows",
                ),
        ),
    }
}

/// `snowflake_session_response` — answer a logout / token-renew request.
pub fn snowflake_session_response_action() -> ActionDefinition {
    ActionDefinition {
        name: "snowflake_session_response".to_string(),
        description: "Acknowledge a session-management request (logout or token renewal) with a \
            success envelope. For logout, omit `data`. For token renewal, put the renewed \
            {sessionToken, masterToken, validityInSeconds} in `data`."
            .to_string(),
        parameters: vec![Parameter {
            name: "data".to_string(),
            type_hint: "object".to_string(),
            description: "Optional data object placed in the response 'data' field (e.g. renewed \
                tokens for a token-request). Omit for a plain logout acknowledgement."
                .to_string(),
            required: false,
        }],
        example: json!({
            "type": "snowflake_session_response",
            "data": {"sessionToken": "ver:1-...", "validityInSeconds": 3600}
        }),
        log_template: Some(
            LogTemplate::new()
                .with_info("-> Snowflake session OK")
                .with_debug("Snowflake snowflake_session_response"),
        ),
    }
}

/// `snowflake_error` — refuse a login/query/session with a Snowflake error envelope.
pub fn snowflake_error_action() -> ActionDefinition {
    ActionDefinition {
        name: "snowflake_error".to_string(),
        description: "Refuse the current request with a Snowflake error envelope (success:false). \
            This is the only way to deny a login or reject a query. The client receives HTTP 200 \
            with success:false and the code/message you provide."
            .to_string(),
        parameters: vec![
            Parameter {
                name: "code".to_string(),
                type_hint: "string".to_string(),
                description: "Snowflake error code string, e.g. \"390100\" (incorrect username or \
                    password), \"390114\" (authentication token expired), \"002003\" (object does \
                    not exist), \"001003\" (SQL compilation error)."
                    .to_string(),
                required: true,
            },
            Parameter {
                name: "message".to_string(),
                type_hint: "string".to_string(),
                description: "Human-readable error message shown to the client".to_string(),
                required: true,
            },
            Parameter {
                name: "sql_state".to_string(),
                type_hint: "string".to_string(),
                description: "Five-character SQLSTATE for a refused query, e.g. \"42S02\" \
                    (table not found) or \"42000\" (syntax error); drivers raise it with the \
                    error. Ignored for a refused login."
                    .to_string(),
                required: false,
            },
        ],
        example: json!({
            "type": "snowflake_error",
            "code": "390100",
            "message": "Incorrect username or password was specified."
        }),
        log_template: Some(
            LogTemplate::new()
                .with_info("-> Snowflake ERR {code}: {message}")
                .with_debug("Snowflake snowflake_error: code={code}, message={message}"),
        ),
    }
}

// ============================================================================
// Action constants (attached to event types)
// ============================================================================

pub static SNOWFLAKE_LOGIN_SUCCESS_ACTION: LazyLock<ActionDefinition> =
    LazyLock::new(snowflake_login_success_action);
pub static SNOWFLAKE_QUERY_RESPONSE_ACTION: LazyLock<ActionDefinition> =
    LazyLock::new(snowflake_query_response_action);
pub static SNOWFLAKE_SESSION_RESPONSE_ACTION: LazyLock<ActionDefinition> =
    LazyLock::new(snowflake_session_response_action);
pub static SNOWFLAKE_ERROR_ACTION: LazyLock<ActionDefinition> =
    LazyLock::new(snowflake_error_action);

// ============================================================================
// Event types
// ============================================================================

/// Emitted on `POST /session/v1/login-request` — the client is authenticating.
pub static SNOWFLAKE_LOGIN_EVENT: LazyLock<EventType> = LazyLock::new(|| {
    EventType::new(
        "snowflake_login",
        "Snowflake client login request",
        json!({
            "type": "snowflake_login_success",
            "token": "SESSION_TOKEN",
            "session_id": 100
        }),
    )
    .with_parameters(vec![
        Parameter {
            name: "login_name".to_string(),
            type_hint: "string".to_string(),
            description: "The username the client is logging in as".to_string(),
            required: false,
        },
        Parameter {
            name: "account".to_string(),
            type_hint: "string".to_string(),
            description: "The Snowflake account name the client is targeting".to_string(),
            required: false,
        },
        Parameter {
            name: "client_app_id".to_string(),
            type_hint: "string".to_string(),
            description: "The driver/app id (e.g. PythonConnector, JDBC)".to_string(),
            required: false,
        },
        Parameter {
            name: "client_app_version".to_string(),
            type_hint: "string".to_string(),
            description: "The driver/app version".to_string(),
            required: false,
        },
        Parameter {
            name: "has_password".to_string(),
            type_hint: "boolean".to_string(),
            description: "True if a password was supplied in the login request".to_string(),
            required: false,
        },
    ])
    .with_actions(vec![
        SNOWFLAKE_LOGIN_SUCCESS_ACTION.clone(),
        SNOWFLAKE_ERROR_ACTION.clone(),
    ])
    .with_log_template(
        LogTemplate::new()
            .with_info("Snowflake login: {login_name}@{account}")
            .with_debug("Snowflake login request: {json_pretty(.)}"),
    )
});

/// Emitted on `POST /queries/v1/query-request` — the client runs SQL.
pub static SNOWFLAKE_QUERY_EVENT: LazyLock<EventType> = LazyLock::new(|| {
    EventType::new(
        "snowflake_query",
        "Snowflake SQL query request",
        json!({
            "type": "snowflake_query_response",
            "rowtype": [{"name": "COL", "type": "text"}],
            "rowset": [["value"]]
        }),
    )
    .with_parameters(vec![
        Parameter {
            name: "sql_text".to_string(),
            type_hint: "string".to_string(),
            description: "The SQL statement the client wants to run".to_string(),
            required: true,
        },
        Parameter {
            name: "has_auth_token".to_string(),
            type_hint: "boolean".to_string(),
            description: "True if the client presented an Authorization session token. There is \
                no session store, so the token is not validated — decide based on this signal."
                .to_string(),
            required: false,
        },
        Parameter {
            name: "request_id".to_string(),
            type_hint: "string".to_string(),
            description: "The requestId query parameter, if present".to_string(),
            required: false,
        },
    ])
    .with_actions(vec![
        SNOWFLAKE_QUERY_RESPONSE_ACTION.clone(),
        SNOWFLAKE_ERROR_ACTION.clone(),
    ])
    .with_log_template(
        LogTemplate::new()
            .with_info("Snowflake query: {preview(sql_text,80)}")
            .with_debug("Snowflake query: {sql_text}"),
    )
});

/// Emitted on `POST /session/logout-request` and `POST /session/token-request`.
pub static SNOWFLAKE_SESSION_EVENT: LazyLock<EventType> = LazyLock::new(|| {
    EventType::new(
        "snowflake_session",
        "Snowflake session management request (logout or token renewal)",
        json!({
            "type": "snowflake_session_response"
        }),
    )
    .with_parameters(vec![
        Parameter {
            name: "operation".to_string(),
            type_hint: "string".to_string(),
            description: "Either \"logout\" or \"token_renew\"".to_string(),
            required: true,
        },
        Parameter {
            name: "has_auth_token".to_string(),
            type_hint: "boolean".to_string(),
            description: "True if the client presented an Authorization session token".to_string(),
            required: false,
        },
    ])
    .with_actions(vec![
        SNOWFLAKE_SESSION_RESPONSE_ACTION.clone(),
        SNOWFLAKE_ERROR_ACTION.clone(),
    ])
    .with_log_template(
        LogTemplate::new()
            .with_info("Snowflake session: {operation}")
            .with_debug("Snowflake session request: {operation}"),
    )
});

/// All Snowflake event types.
pub fn get_snowflake_event_types() -> Vec<EventType> {
    vec![
        SNOWFLAKE_LOGIN_EVENT.clone(),
        SNOWFLAKE_QUERY_EVENT.clone(),
        SNOWFLAKE_SESSION_EVENT.clone(),
    ]
}