nodedb 0.4.0

Local-first, real-time, edge-to-cloud hybrid database for multi-modal workloads
Documentation
// SPDX-License-Identifier: BUSL-1.1

//! Protocol-neutral `SELECT * FROM STREAM` handler.
//!
//! Syntax:
//! ```sql
//! SELECT * FROM STREAM <stream> CONSUMER GROUP <group> [PARTITION <p>] [LIMIT <n>]
//! ```
//!
//! Ported from the pgwire `ddl::stream_select::select_from_stream` handler.
//! Reads events from the stream buffer starting after the consumer group's
//! committed offsets and returns a materialized result set — there is no
//! per-connection push stream — so this handler carries no per-connection
//! state. The partition/limit token parsing, the cluster-aware forwarding to
//! the leader node on `ConsumeError::RemotePartition`, the empty-result
//! short-circuit on `ConsumeError::BufferEmpty`, and the column layout are
//! preserved verbatim; only the result construction changed from pgwire
//! `Response` / `PgWireError` to the protocol-neutral [`DdlResult`] /
//! [`DdlError`].

use serde_json::{Map, Value as JsonValue};
use sonic_rs;

use crate::control::security::identity::AuthenticatedIdentity;
use crate::control::server::response_shape::types::ShapedRows;
use crate::control::state::SharedState;
use crate::event::cdc::consume::{ConsumeError, ConsumeParams, consume_stream};

use super::super::result::{DdlError, DdlResult};

fn err(sqlstate: &str, message: impl Into<String>) -> DdlError {
    DdlError {
        sqlstate: sqlstate.to_string(),
        message: message.into(),
    }
}

/// Handle `SELECT * FROM STREAM <stream> CONSUMER GROUP <group> [PARTITION <p>] [LIMIT <n>]`
///
/// Cluster-aware: if the requested partition is on a remote node, forwards
/// the consume request to the leader via the gateway (C-δ.6: `ExecuteRequest`).
pub async fn select_from_stream(
    state: &SharedState,
    identity: &AuthenticatedIdentity,
    parts: &[&str],
) -> Result<Vec<DdlResult>, DdlError> {
    let tenant_id = identity.tenant_id.as_u64();

    // Parse: SELECT * FROM STREAM <stream> CONSUMER GROUP <group> [PARTITION <p>] [LIMIT <n>]
    // parts: [SELECT, *, FROM, STREAM, <stream>, CONSUMER, GROUP, <group>, ...]
    //         0       1  2     3       4         5         6      7
    if parts.len() < 8
        || !parts[3].eq_ignore_ascii_case("STREAM")
        || !parts[5].eq_ignore_ascii_case("CONSUMER")
        || !parts[6].eq_ignore_ascii_case("GROUP")
    {
        return Err(err(
            "42601",
            "expected SELECT * FROM STREAM <stream> CONSUMER GROUP <group> [PARTITION <p>] [LIMIT <n>]",
        ));
    }

    let stream_name = parts[4].to_lowercase();
    let group_name = parts[7].to_lowercase();

    let mut partition: Option<u32> = None;
    let mut limit: usize = 100;
    let mut i = 8;

    while i < parts.len() {
        if parts[i].eq_ignore_ascii_case("PARTITION") && i + 1 < parts.len() {
            partition = Some(
                parts[i + 1]
                    .parse()
                    .map_err(|_| err("42601", format!("invalid partition: '{}'", parts[i + 1])))?,
            );
            i += 2;
        } else if parts[i].eq_ignore_ascii_case("LIMIT") && i + 1 < parts.len() {
            limit = parts[i + 1]
                .parse()
                .map_err(|_| err("42601", format!("invalid limit: '{}'", parts[i + 1])))?;
            i += 2;
        } else {
            i += 1;
        }
    }

    let consume_params = ConsumeParams {
        tenant_id,
        stream_name: &stream_name,
        group_name: &group_name,
        partition,
        limit,
    };

    let result = match consume_stream(state, &consume_params) {
        Ok(r) => r,
        Err(ConsumeError::RemotePartition { leader_node, .. }) => {
            match crate::event::cdc::consume::consume_remote(state, &consume_params, leader_node)
                .await
            {
                Ok(r) => r,
                Err(e) => return Err(err("58000", e.to_string())),
            }
        }
        Err(ConsumeError::BufferEmpty(_)) => {
            // Return empty result set.
            let columns = result_columns();
            let column_types = ShapedRows::text_types(columns.len());
            return Ok(vec![DdlResult::Rows(ShapedRows {
                columns,
                column_types,
                rows: Vec::new(),
                notice: None,
            })]);
        }
        Err(e) => {
            return Err(err("42704", e.to_string()));
        }
    };

    let columns = result_columns();
    let mut rows = Vec::with_capacity(result.events.len());

    for event in &result.events {
        let mut row = Map::new();
        row.insert(
            "sequence".to_string(),
            JsonValue::String(event.sequence.to_string()),
        );
        row.insert(
            "partition".to_string(),
            JsonValue::String(event.partition.to_string()),
        );
        row.insert(
            "collection".to_string(),
            JsonValue::String(event.collection.clone()),
        );
        row.insert(
            "event_type".to_string(),
            JsonValue::String(event.op.clone()),
        );
        row.insert(
            "row_id".to_string(),
            JsonValue::String(event.row_id.clone()),
        );
        row.insert("lsn".to_string(), JsonValue::String(event.lsn.to_string()));
        row.insert(
            "event_time".to_string(),
            JsonValue::String(event.event_time.to_string()),
        );
        let new_val = event
            .new_value
            .as_ref()
            .map(|v| sonic_rs::to_string(v).unwrap_or_default())
            .unwrap_or_default();
        row.insert("new_value".to_string(), JsonValue::String(new_val));
        let old_val = event
            .old_value
            .as_ref()
            .map(|v| sonic_rs::to_string(v).unwrap_or_default())
            .unwrap_or_default();
        row.insert("old_value".to_string(), JsonValue::String(old_val));
        rows.push(row);
    }

    let column_types = ShapedRows::text_types(columns.len());
    Ok(vec![DdlResult::Rows(ShapedRows {
        columns,
        column_types,
        rows,
        notice: None,
    })])
}

/// Column schema for stream SELECT results.
fn result_columns() -> Vec<String> {
    vec![
        "sequence".to_string(),
        "partition".to_string(),
        "collection".to_string(),
        "event_type".to_string(),
        "row_id".to_string(),
        "lsn".to_string(),
        "event_time".to_string(),
        "new_value".to_string(),
        "old_value".to_string(),
    ]
}