risingwave_stream/from_proto/

sink.rs

// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::BTreeMap;
use std::sync::Arc;

use anyhow::anyhow;
use futures::FutureExt;
use risingwave_common::bail;
use risingwave_common::catalog::{ColumnCatalog, Schema};
use risingwave_common::secret::LocalSecretManager;
use risingwave_common::types::DataType;
use risingwave_connector::match_sink_name_str;
use risingwave_connector::sink::catalog::{SinkFormat, SinkFormatDesc, SinkId, SinkType};
use risingwave_connector::sink::file_sink::fs::FsSink;
use risingwave_connector::sink::iceberg::ICEBERG_SINK;
use risingwave_connector::sink::{
    CONNECTOR_TYPE_KEY, SINK_TYPE_OPTION, SinkError, SinkMetaClient, SinkParam, SinkWriterParam,
    build_sink,
};
use risingwave_hummock_sdk::HummockReadEpoch;
use risingwave_pb::catalog::Table;
use risingwave_pb::plan_common::PbColumnCatalog;
use risingwave_pb::stream_plan::{PbSinkDesc, SinkLogStoreType, SinkNode};
use risingwave_storage::store::TryWaitEpochOptions;
use url::Url;

use super::*;
use crate::common::log_store_impl::in_mem::BoundedInMemLogStoreFactory;
use crate::common::log_store_impl::kv_log_store::{
    KV_LOG_STORE_V2_INFO, KvLogStoreFactory, KvLogStoreMetrics, KvLogStorePkInfo,
};
use crate::executor::{SinkExecutor, StreamExecutorError};

/// Build [`SinkParam`] from a [`PbSinkDesc`], along with the full [`ColumnCatalog`] list for
/// downstream use. `connector` is the resolved connector name (e.g. `ICEBERG_SINK`), used for
/// legacy `type = '...'` format/encode fallback and connector-specific backward-compatibility
/// conversions.
///
/// The caller is responsible for preparing `properties_with_secret` (i.e. the result of
/// [`LocalSecretManager::fill_secrets`] plus any connector-specific rewrites such as the
/// JDBC-to-native-postgres switch).
pub(super) fn build_sink_param(
    sink_desc: &PbSinkDesc,
    properties_with_secret: BTreeMap<String, String>,
    connector: &str,
) -> Result<(SinkParam, Vec<ColumnCatalog>), StreamExecutorError> {
    let sink_id: SinkId = sink_desc.get_id();
    let sink_name = sink_desc.get_name().to_owned();
    let db_name = sink_desc.get_db_name().into();
    let sink_from_name = sink_desc.get_sink_from_name().into();
    let downstream_pk = if sink_desc.downstream_pk.is_empty() {
        None
    } else {
        Some(
            (sink_desc.downstream_pk.iter())
                .map(|idx| *idx as usize)
                .collect_vec(),
        )
    };
    let columns = sink_desc
        .column_catalogs
        .clone()
        .into_iter()
        .map(ColumnCatalog::from)
        .collect_vec();

    let format_desc = match &sink_desc.format_desc {
        // Case A: new syntax `format ... encode ...`
        Some(f) => Some(
            f.clone()
                .try_into()
                .map_err(|e| StreamExecutorError::from((e, sink_id)))?,
        ),
        None => match properties_with_secret.get(SINK_TYPE_OPTION) {
            // Case B: old syntax `type = '...'`
            Some(t) => SinkFormatDesc::from_legacy_type(connector, t)
                .map_err(|e| StreamExecutorError::from((e, sink_id)))?,
            // Case C: no format + encode required
            None => None,
        },
    };
    let format_desc = SinkParam::fill_secret_for_format_desc(format_desc)
        .map_err(|e| StreamExecutorError::from((e, sink_id)))?;

    // Backward compatibility: DEBEZIUM format should be treated as `Retract` type instead of `Upsert`.
    let sink_type = if let Some(format_desc) = &format_desc
        && format_desc.format == SinkFormat::Debezium
    {
        SinkType::Retract
    } else {
        let sink_type_from_proto = SinkType::from_proto(sink_desc.get_sink_type().unwrap());
        // For backward compatibility: Iceberg sink with Upsert type should be treated as Retract type.
        if connector.eq_ignore_ascii_case(ICEBERG_SINK)
            && matches!(sink_type_from_proto, SinkType::Upsert)
        {
            SinkType::Retract
        } else {
            sink_type_from_proto
        }
    };
    let ignore_delete = sink_desc.ignore_delete();

    let sink_param = SinkParam {
        sink_id,
        sink_name,
        properties: properties_with_secret,
        columns: columns
            .iter()
            .filter(|col| !col.is_hidden)
            .map(|col| col.column_desc.clone())
            .collect(),
        downstream_pk,
        sink_type,
        ignore_delete,
        format_desc,
        db_name,
        sink_from_name,
    };

    Ok((sink_param, columns))
}

pub struct SinkExecutorBuilder;

fn resolve_pk_info(
    input_schema: &Schema,
    log_store_table: &Table,
) -> StreamResult<&'static KvLogStorePkInfo> {
    let predefined_column_len = log_store_table.columns.len() - input_schema.fields.len();

    #[expect(deprecated)]
    let info = match predefined_column_len {
        len if len
            == crate::common::log_store_impl::kv_log_store::KV_LOG_STORE_V1_INFO
                .predefined_column_len() =>
        {
            Ok(&crate::common::log_store_impl::kv_log_store::KV_LOG_STORE_V1_INFO)
        }
        len if len == KV_LOG_STORE_V2_INFO.predefined_column_len() => Ok(&KV_LOG_STORE_V2_INFO),
        other_len => Err(anyhow!(
            "invalid log store predefined len {}. log store table: {:?}, input_schema: {:?}",
            other_len,
            log_store_table,
            input_schema
        )),
    }?;
    validate_payload_schema(
        &log_store_table.columns[predefined_column_len..],
        input_schema,
    )?;
    Ok(info)
}

fn validate_payload_schema(
    log_store_payload_schema: &[PbColumnCatalog],
    input_schema: &Schema,
) -> StreamResult<()> {
    if log_store_payload_schema
        .iter()
        .zip_eq(input_schema.fields.iter())
        .all(|(log_store_col, input_field)| {
            let log_store_col_type = DataType::from(
                log_store_col
                    .column_desc
                    .as_ref()
                    .unwrap()
                    .column_type
                    .as_ref()
                    .unwrap(),
            );
            log_store_col_type.equals_datatype(&input_field.data_type)
        })
    {
        Ok(())
    } else {
        Err(anyhow!(
            "mismatch schema: log store: {:?}, input: {:?}",
            log_store_payload_schema,
            input_schema
        )
        .into())
    }
}

impl_stream_node_body!(Sink(SinkNode) => SinkExecutorBuilder);

impl ExecutorBuilder for SinkExecutorBuilder {
    type Node = SinkNode;

    async fn new_boxed_executor(
        params: ExecutorParams,
        node: &Self::Node,
        state_store: impl StateStore,
    ) -> StreamResult<Executor> {
        let [input_executor]: [_; 1] = params.input.try_into().unwrap();
        let input_data_types = input_executor.info().schema.data_types();
        let chunk_size = params.config.developer.chunk_size;

        let sink_desc = node.sink_desc.as_ref().unwrap();
        let sink_id: SinkId = sink_desc.get_id();
        let sink_name = sink_desc.get_name().to_owned();
        let properties = sink_desc.get_properties().clone();
        let secret_refs = sink_desc.get_secret_refs().clone();

        let mut properties_with_secret =
            LocalSecretManager::global().fill_secrets(properties, secret_refs)?;

        if params.config.developer.switch_jdbc_pg_to_native
            && let Some(connector_type) = properties_with_secret.get(CONNECTOR_TYPE_KEY)
            && connector_type == "jdbc"
            && let Some(url) = properties_with_secret.get("jdbc.url")
            && url.starts_with("jdbc:postgresql:")
        {
            tracing::info!("switching to native postgres connector");
            let jdbc_url = parse_jdbc_url(url)
                .map_err(|e| StreamExecutorError::from((SinkError::Config(e), sink_id)))?;
            properties_with_secret.insert(CONNECTOR_TYPE_KEY.to_owned(), "postgres".to_owned());
            properties_with_secret.insert("host".to_owned(), jdbc_url.host);
            properties_with_secret.insert("port".to_owned(), jdbc_url.port.to_string());
            properties_with_secret.insert("database".to_owned(), jdbc_url.db_name);
            if let Some(username) = jdbc_url.username {
                properties_with_secret.insert("user".to_owned(), username);
            }
            if let Some(password) = jdbc_url.password {
                properties_with_secret.insert("password".to_owned(), password);
            }
            if let Some(table_name) = properties_with_secret.get("table.name") {
                properties_with_secret.insert("table".to_owned(), table_name.clone());
            }
            if let Some(schema_name) = properties_with_secret.get("schema.name") {
                properties_with_secret.insert("schema".to_owned(), schema_name.clone());
            }
            // TODO(kwannoel): Do we need to handle jdbc.query.timeout?
        }

        let connector = {
            let sink_type = properties_with_secret
                .get(CONNECTOR_TYPE_KEY)
                .ok_or_else(|| {
                    StreamExecutorError::from((
                        SinkError::Config(anyhow!("missing config: {}", CONNECTOR_TYPE_KEY)),
                        sink_id,
                    ))
                })?;

            let sink_type_str = sink_type.to_lowercase();
            match_sink_name_str!(
                sink_type_str.as_str(),
                SinkType,
                Ok(SinkType::SINK_NAME),
                |other: &str| {
                    Err(StreamExecutorError::from((
                        SinkError::Config(anyhow!("unsupported sink connector {}", other)),
                        sink_id,
                    )))
                }
            )?
        };

        let (sink_param, columns) = build_sink_param(sink_desc, properties_with_secret, connector)?;

        let sink_write_param = SinkWriterParam {
            executor_id: params.executor_id,
            vnode_bitmap: params.vnode_bitmap.clone(),
            meta_client: params.env.meta_client().map(SinkMetaClient::MetaClient),
            extra_partition_col_idx: sink_desc.extra_partition_col_idx.map(|v| v as usize),

            actor_id: params.actor_context.id,
            sink_id,
            sink_name,
            connector: connector.to_owned(),
            streaming_config: params.config.as_ref().clone(),
            time_zone: params.actor_context.time_zone,
        };

        let log_store_identity = format!(
            "sink[{}]-[{}]-executor[{}]",
            connector, sink_id, params.executor_id
        );

        let sink = build_sink(sink_param.clone())
            .map_err(|e| StreamExecutorError::from((e, sink_param.sink_id)))?;

        let exec = match node.log_store_type() {
            // Default value is the normal in memory log store to be backward compatible with the
            // previously unset value
            SinkLogStoreType::InMemoryLogStore | SinkLogStoreType::Unspecified => {
                let factory = BoundedInMemLogStoreFactory::new(1, {
                    let state_store = state_store;
                    let table_id = node.table.as_ref().map(|table| table.id);
                    move |epoch| {
                        async move {
                            if let Some(table_id) = table_id {
                                state_store
                                    .try_wait_epoch(
                                        HummockReadEpoch::Committed(epoch.prev),
                                        TryWaitEpochOptions { table_id },
                                    )
                                    .await?;
                            }
                            Ok(())
                        }
                        .boxed()
                    }
                });
                SinkExecutor::new(
                    params.actor_context,
                    params.info.clone(),
                    input_executor,
                    sink_write_param,
                    sink,
                    sink_param,
                    columns,
                    factory,
                    chunk_size,
                    input_data_types,
                    node.rate_limit.map(|x| x as _),
                )?
                .boxed()
            }
            SinkLogStoreType::KvLogStore => {
                let metrics = KvLogStoreMetrics::new(
                    &params.executor_stats,
                    params.actor_context.id,
                    &sink_param,
                    connector,
                );

                let table = node.table.as_ref().unwrap().clone();
                let input_schema = input_executor.schema();
                let pk_info = resolve_pk_info(input_schema, &table)?;

                // TODO: support setting max row count in config
                let factory = KvLogStoreFactory::new(
                    state_store,
                    table,
                    params.vnode_bitmap.clone().map(Arc::new),
                    65536,
                    params.config.developer.chunk_size,
                    metrics,
                    log_store_identity,
                    params.env.kv_log_store_historical_read_semaphore(),
                    pk_info,
                );

                SinkExecutor::new(
                    params.actor_context,
                    params.info.clone(),
                    input_executor,
                    sink_write_param,
                    sink,
                    sink_param,
                    columns,
                    factory,
                    chunk_size,
                    input_data_types,
                    node.rate_limit.map(|x| x as _),
                )?
                .boxed()
            }
        };

        Ok((params.info, exec).into())
    }
}

struct JdbcUrl {
    host: String,
    port: u16,
    db_name: String,
    username: Option<String>,
    password: Option<String>,
}

fn parse_jdbc_url(url: &str) -> anyhow::Result<JdbcUrl> {
    if !url.starts_with("jdbc:postgresql") {
        bail!(
            "invalid jdbc url, to switch to postgres rust connector, we need to use the url jdbc:postgresql://..."
        )
    }

    // trim the "jdbc:" prefix to make it a valid url
    let url = url.replace("jdbc:", "");

    // parse the url
    let url = Url::parse(&url).map_err(|e| anyhow!(e).context("failed to parse jdbc url"))?;

    let scheme = url.scheme();
    assert_eq!("postgresql", scheme, "jdbc scheme should be postgresql");
    let host = url
        .host_str()
        .ok_or_else(|| anyhow!("missing host in jdbc url"))?;
    let port = url
        .port()
        .ok_or_else(|| anyhow!("missing port in jdbc url"))?;
    let Some(db_name) = url.path().strip_prefix('/') else {
        bail!("missing db_name in jdbc url");
    };
    let mut username = None;
    let mut password = None;
    for (key, value) in url.query_pairs() {
        if key == "user" {
            username = Some(value.to_string());
        }
        if key == "password" {
            password = Some(value.to_string());
        }
    }

    Ok(JdbcUrl {
        host: host.to_owned(),
        port,
        db_name: db_name.to_owned(),
        username,
        password,
    })
}

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

    #[test]
    fn test_parse_jdbc_url() {
        let url = "jdbc:postgresql://localhost:5432/test?user=postgres&password=postgres";
        let jdbc_url = parse_jdbc_url(url).unwrap();
        assert_eq!(jdbc_url.host, "localhost");
        assert_eq!(jdbc_url.port, 5432);
        assert_eq!(jdbc_url.db_name, "test");
        assert_eq!(jdbc_url.username, Some("postgres".to_owned()));
        assert_eq!(jdbc_url.password, Some("postgres".to_owned()));
    }
}