Repository: risingwavelabs/risingwave
File: src/utils/pgwire/src/pg_protocol.rs (Rust)

Code before:

    while let Some(arg) = args_iter.next() {
        if arg == "-c" {
            if let Some(config_str) = args_iter.next() {
                if let Some((key, value)) = config_str.split_once('=') {
                    let key = key.replace("-", "_");
                    config.push((key, value.to_owned()));
                } else {
                    return Err(PsqlError::StartupError(
                        format!("invalid config format: {}", config_str).into(),
                    ));
                }
            } else {
                return Err(PsqlError::StartupError("missing argument for -c".into()));
            }
        } else if let Some(config_str) = arg.strip_prefix("--") {
            if let Some((key, value)) = config_str.split_once('=') {
                let key = key.replace("-", "_");
                config.push((key, value.to_owned()));
            } else {
                return Err(PsqlError::StartupError(
                    format!("invalid config format: {}", config_str).into(),
                ));
            }
        } else {
            tracing::warn!("ignoring unrecognized option for backward compatibility: {arg}");
        }
    }
    Ok(config)
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;

    #[test]
    fn test_redact_parsable_sql() {
        let keywords = Arc::new(HashSet::from(["v2".into(), "v4".into(), "b".into()]));
        let sql = r"
        create source temp (k bigint, v varchar) with (
            connector = 'datagen',
            v1 = 123,
            v2 = 'with',
            v3 = false,
            v4 = '',
        ) FORMAT plain ENCODE json (a='1',b='2')
        ";
        assert_eq!(
            redact_sql(sql, keywords),

Review comment (bug):
Logging unknown `options` args at `warn` can create noisy logs (and can be abused for log-spam since this is client-controlled during startup). Consider downgrading this to `debug`/`trace`, adding rate-limiting, or only warning when the option looks like a supported flag but has an invalid format.

Code after:

    while let Some(arg) = args_iter.next() {
        if arg == "-c" {
            if let Some(config_str) = args_iter.next() {
                if let Some((key, value)) = config_str.split_once('=') {
                    let key = key.replace("-", "_");
                    config.push((key, value.to_owned()));
                } else {
                    return Err(PsqlError::StartupError(
                        format!("invalid config format: {}", config_str).into(),
                    ));
                }
            } else {
                return Err(PsqlError::StartupError("missing argument for -c".into()));
            }
        } else if let Some(config_str) = arg.strip_prefix("--") {
            if let Some((key, value)) = config_str.split_once('=') {
                let key = key.replace("-", "_");
                config.push((key, value.to_owned()));
            } else {
                return Err(PsqlError::StartupError(
                    format!("invalid config format: {}", config_str).into(),
                ));
            }
        } else {
            tracing::warn!(
                arg,
                "ignoring unrecognized option for backward compatibility"
            );
        }
    }
    Ok(config)
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;

    #[test]
    fn test_redact_parsable_sql() {
        let keywords = Arc::new(HashSet::from(["v2".into(), "v4".into(), "b".into()]));
        let sql = r"
        create source temp (k bigint, v varchar) with (
            connector = 'datagen',
            v1 = 123,
            v2 = 'with',
            v3 = false,
            v4 = '',
        ) FORMAT plain ENCODE json (a='1',b='2')