crates/fluss/src/config.rs - fluss-rust - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you 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 clap::{Parser, ValueEnum};
 use serde::{Deserialize, Serialize};
 use std::fmt;

 const DEFAULT_BOOTSTRAP_SERVER: &str = "127.0.0.1:9123";
 const DEFAULT_REQUEST_MAX_SIZE: i32 = 10 * 1024 * 1024;
 const DEFAULT_WRITER_BATCH_SIZE: i32 = 2 * 1024 * 1024;
 const DEFAULT_RETRIES: i32 = i32::MAX;
 const DEFAULT_PREFETCH_NUM: usize = 4;
 const DEFAULT_DOWNLOAD_THREADS: usize = 3;
 const DEFAULT_SCANNER_REMOTE_LOG_READ_CONCURRENCY: usize = 4;
 const DEFAULT_MAX_POLL_RECORDS: usize = 500;
 const DEFAULT_WRITER_BATCH_TIMEOUT_MS: i64 = 100;

 const DEFAULT_ACKS: &str = "all";
 const DEFAULT_CONNECT_TIMEOUT_MS: u64 = 120_000;
 const DEFAULT_SECURITY_PROTOCOL: &str = "PLAINTEXT";
 const DEFAULT_SASL_MECHANISM: &str = "PLAIN";

 /// Bucket assigner strategy for tables without bucket keys.
 /// Matches Java `client.writer.bucket.no-key-assigner`.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum NoKeyAssigner {
     /// Sticks to one bucket until the batch is full, then switches.
     Sticky,
     /// Assigns each record to the next bucket in a rotating sequence.
     RoundRobin,
 }

 impl fmt::Display for NoKeyAssigner {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             NoKeyAssigner::Sticky => write!(f, "sticky"),
             NoKeyAssigner::RoundRobin => write!(f, "round_robin"),
         }
     }
 }

 #[derive(Parser, Clone, Deserialize, Serialize)]
 #[command(author, version, about, long_about = None)]
 pub struct Config {
     #[arg(long, default_value_t = String::from(DEFAULT_BOOTSTRAP_SERVER))]
     pub bootstrap_servers: String,

     #[arg(long, default_value_t = DEFAULT_REQUEST_MAX_SIZE)]
     pub writer_request_max_size: i32,

     #[arg(long, default_value_t = String::from(DEFAULT_ACKS))]
     pub writer_acks: String,

     #[arg(long, default_value_t = DEFAULT_RETRIES)]
     pub writer_retries: i32,

     #[arg(long, default_value_t = DEFAULT_WRITER_BATCH_SIZE)]
     pub writer_batch_size: i32,

     #[arg(long, value_enum, default_value_t = NoKeyAssigner::Sticky)]
     pub writer_bucket_no_key_assigner: NoKeyAssigner,

     /// Maximum number of remote log segments to prefetch
     /// Default: 4 (matching Java CLIENT_SCANNER_REMOTE_LOG_PREFETCH_NUM)
     #[arg(long, default_value_t = DEFAULT_PREFETCH_NUM)]
     pub scanner_remote_log_prefetch_num: usize,

     /// Maximum concurrent remote log downloads
     /// Default: 3 (matching Java REMOTE_FILE_DOWNLOAD_THREAD_NUM)
     #[arg(long, default_value_t = DEFAULT_DOWNLOAD_THREADS)]
     pub remote_file_download_thread_num: usize,

     /// Intra-file remote log read concurrency for each remote segment download.
     /// Download path always uses streaming reader.
     #[arg(long, default_value_t = DEFAULT_SCANNER_REMOTE_LOG_READ_CONCURRENCY)]
     pub scanner_remote_log_read_concurrency: usize,

     /// Maximum number of records returned in a single call to poll() for LogScanner.
     /// Default: 500 (matching Java CLIENT_SCANNER_LOG_MAX_POLL_RECORDS)
     #[arg(long, default_value_t = DEFAULT_MAX_POLL_RECORDS)]
     pub scanner_log_max_poll_records: usize,

     /// The maximum time to wait for a batch to be completed in milliseconds.
     /// Default: 100 (matching Java CLIENT_WRITER_BATCH_TIMEOUT)
     #[arg(long, default_value_t = DEFAULT_WRITER_BATCH_TIMEOUT_MS)]
     pub writer_batch_timeout_ms: i64,

     /// Connect timeout in milliseconds for TCP transport connect.
     /// Default: 120000 (120 seconds).
     #[arg(long, default_value_t = DEFAULT_CONNECT_TIMEOUT_MS)]
     pub connect_timeout_ms: u64,

     #[arg(long, default_value_t = String::from(DEFAULT_SECURITY_PROTOCOL))]
     pub security_protocol: String,

     #[arg(long, default_value_t = String::from(DEFAULT_SASL_MECHANISM))]
     pub security_sasl_mechanism: String,

     #[arg(long, default_value_t = String::new())]
     pub security_sasl_username: String,

     #[arg(long, default_value_t = String::new())]
     #[serde(skip_serializing)]
     pub security_sasl_password: String,
 }

 impl std::fmt::Debug for Config {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Config")
             .field("bootstrap_servers", &self.bootstrap_servers)
             .field("writer_request_max_size", &self.writer_request_max_size)
             .field("writer_acks", &self.writer_acks)
             .field("writer_retries", &self.writer_retries)
             .field("writer_batch_size", &self.writer_batch_size)
             .field(
                 "writer_bucket_no_key_assigner",
                 &self.writer_bucket_no_key_assigner,
             )
             .field(
                 "scanner_remote_log_prefetch_num",
                 &self.scanner_remote_log_prefetch_num,
             )
             .field(
                 "remote_file_download_thread_num",
                 &self.remote_file_download_thread_num,
             )
             .field(
                 "scanner_log_max_poll_records",
                 &self.scanner_log_max_poll_records,
             )
             .field("writer_batch_timeout_ms", &self.writer_batch_timeout_ms)
             .field("connect_timeout_ms", &self.connect_timeout_ms)
             .field("security_protocol", &self.security_protocol)
             .field("security_sasl_mechanism", &self.security_sasl_mechanism)
             .field("security_sasl_username", &self.security_sasl_username)
             .field("security_sasl_password", &"[REDACTED]")
             .finish()
     }
 }

 impl Default for Config {
     fn default() -> Self {
         Self {
             bootstrap_servers: String::from(DEFAULT_BOOTSTRAP_SERVER),
             writer_request_max_size: DEFAULT_REQUEST_MAX_SIZE,
             writer_acks: String::from(DEFAULT_ACKS),
             writer_retries: i32::MAX,
             writer_batch_size: DEFAULT_WRITER_BATCH_SIZE,
             writer_bucket_no_key_assigner: NoKeyAssigner::Sticky,
             scanner_remote_log_prefetch_num: DEFAULT_PREFETCH_NUM,
             remote_file_download_thread_num: DEFAULT_DOWNLOAD_THREADS,
             scanner_remote_log_read_concurrency: DEFAULT_SCANNER_REMOTE_LOG_READ_CONCURRENCY,
             scanner_log_max_poll_records: DEFAULT_MAX_POLL_RECORDS,
             writer_batch_timeout_ms: DEFAULT_WRITER_BATCH_TIMEOUT_MS,
             connect_timeout_ms: DEFAULT_CONNECT_TIMEOUT_MS,
             security_protocol: String::from(DEFAULT_SECURITY_PROTOCOL),
             security_sasl_mechanism: String::from(DEFAULT_SASL_MECHANISM),
             security_sasl_username: String::new(),
             security_sasl_password: String::new(),
         }
     }
 }

 impl Config {
     /// Returns true when the security protocol indicates SASL authentication
     /// should be performed. Matches Java's `SaslAuthenticationPlugin` which
     /// registers as `"sasl"` (case-insensitive).
     pub fn is_sasl_enabled(&self) -> bool {
         self.security_protocol.eq_ignore_ascii_case("sasl")
     }

     /// Validates security configuration. Returns `Ok(())` when the config is
     /// consistent, or an error message when SASL is enabled but the config is
     /// incomplete or uses an unsupported mechanism.
     pub fn validate_security(&self) -> Result<(), String> {
         if !self.is_sasl_enabled() {
             return Ok(());
         }
         if !self.security_sasl_mechanism.eq_ignore_ascii_case("PLAIN") {
             return Err(format!(
                 "Unsupported SASL mechanism: '{}'. Only 'PLAIN' is supported.",
                 self.security_sasl_mechanism
             ));
         }
         if self.security_sasl_username.is_empty() {
             return Err(
                 "security_sasl_username must be set when security_protocol is 'sasl'".to_string(),
             );
         }
         if self.security_sasl_password.is_empty() {
             return Err(
                 "security_sasl_password must be set when security_protocol is 'sasl'".to_string(),
             );
         }
         Ok(())
     }
 }

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

     #[test]
     fn test_default_is_not_sasl() {
         let config = Config::default();
         assert!(!config.is_sasl_enabled());
         assert!(config.validate_security().is_ok());
     }

     #[test]
     fn test_sasl_enabled_valid() {
         let config = Config {
             security_protocol: "sasl".to_string(),
             security_sasl_mechanism: "PLAIN".to_string(),
             security_sasl_username: "admin".to_string(),
             security_sasl_password: "secret".to_string(),
             ..Config::default()
         };
         assert!(config.is_sasl_enabled());
         assert!(config.validate_security().is_ok());
     }

     #[test]
     fn test_sasl_enabled_case_insensitive() {
         let config = Config {
             security_protocol: "SASL".to_string(),
             security_sasl_username: "admin".to_string(),
             security_sasl_password: "secret".to_string(),
             ..Config::default()
         };
         assert!(config.is_sasl_enabled());
         assert!(config.validate_security().is_ok());
     }

     #[test]
     fn test_sasl_missing_username() {
         let config = Config {
             security_protocol: "sasl".to_string(),
             security_sasl_password: "secret".to_string(),
             ..Config::default()
         };
         assert!(config.validate_security().is_err());
     }

     #[test]
     fn test_sasl_missing_password() {
         let config = Config {
             security_protocol: "sasl".to_string(),
             security_sasl_username: "admin".to_string(),
             ..Config::default()
         };
         assert!(config.validate_security().is_err());
     }

     #[test]
     fn test_sasl_unsupported_mechanism() {
         let config = Config {
             security_protocol: "sasl".to_string(),
             security_sasl_mechanism: "SCRAM-SHA-256".to_string(),
             security_sasl_username: "admin".to_string(),
             security_sasl_password: "secret".to_string(),
             ..Config::default()
         };
         assert!(config.validate_security().is_err());
     }
 }
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you 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 clap::{Parser, ValueEnum};
	use serde::{Deserialize, Serialize};
	use std::fmt;

	const DEFAULT_BOOTSTRAP_SERVER: &str = "127.0.0.1:9123";
	const DEFAULT_REQUEST_MAX_SIZE: i32 = 10 * 1024 * 1024;
	const DEFAULT_WRITER_BATCH_SIZE: i32 = 2 * 1024 * 1024;
	const DEFAULT_RETRIES: i32 = i32::MAX;
	const DEFAULT_PREFETCH_NUM: usize = 4;
	const DEFAULT_DOWNLOAD_THREADS: usize = 3;
	const DEFAULT_SCANNER_REMOTE_LOG_READ_CONCURRENCY: usize = 4;
	const DEFAULT_MAX_POLL_RECORDS: usize = 500;
	const DEFAULT_WRITER_BATCH_TIMEOUT_MS: i64 = 100;

	const DEFAULT_ACKS: &str = "all";
	const DEFAULT_CONNECT_TIMEOUT_MS: u64 = 120_000;
	const DEFAULT_SECURITY_PROTOCOL: &str = "PLAINTEXT";
	const DEFAULT_SASL_MECHANISM: &str = "PLAIN";

	/// Bucket assigner strategy for tables without bucket keys.
	/// Matches Java `client.writer.bucket.no-key-assigner`.
	#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, Deserialize, Serialize)]
	#[serde(rename_all = "snake_case")]
	pub enum NoKeyAssigner {
	/// Sticks to one bucket until the batch is full, then switches.
	Sticky,
	/// Assigns each record to the next bucket in a rotating sequence.
	RoundRobin,
	}

	impl fmt::Display for NoKeyAssigner {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
	match self {
	NoKeyAssigner::Sticky => write!(f, "sticky"),
	NoKeyAssigner::RoundRobin => write!(f, "round_robin"),
	}
	}
	}

	#[derive(Parser, Clone, Deserialize, Serialize)]
	#[command(author, version, about, long_about = None)]
	pub struct Config {
	#[arg(long, default_value_t = String::from(DEFAULT_BOOTSTRAP_SERVER))]
	pub bootstrap_servers: String,

	#[arg(long, default_value_t = DEFAULT_REQUEST_MAX_SIZE)]
	pub writer_request_max_size: i32,

	#[arg(long, default_value_t = String::from(DEFAULT_ACKS))]
	pub writer_acks: String,

	#[arg(long, default_value_t = DEFAULT_RETRIES)]
	pub writer_retries: i32,

	#[arg(long, default_value_t = DEFAULT_WRITER_BATCH_SIZE)]
	pub writer_batch_size: i32,

	#[arg(long, value_enum, default_value_t = NoKeyAssigner::Sticky)]
	pub writer_bucket_no_key_assigner: NoKeyAssigner,

	/// Maximum number of remote log segments to prefetch
	/// Default: 4 (matching Java CLIENT_SCANNER_REMOTE_LOG_PREFETCH_NUM)
	#[arg(long, default_value_t = DEFAULT_PREFETCH_NUM)]
	pub scanner_remote_log_prefetch_num: usize,

	/// Maximum concurrent remote log downloads
	/// Default: 3 (matching Java REMOTE_FILE_DOWNLOAD_THREAD_NUM)
	#[arg(long, default_value_t = DEFAULT_DOWNLOAD_THREADS)]
	pub remote_file_download_thread_num: usize,

	/// Intra-file remote log read concurrency for each remote segment download.
	/// Download path always uses streaming reader.
	#[arg(long, default_value_t = DEFAULT_SCANNER_REMOTE_LOG_READ_CONCURRENCY)]
	pub scanner_remote_log_read_concurrency: usize,

	/// Maximum number of records returned in a single call to poll() for LogScanner.
	/// Default: 500 (matching Java CLIENT_SCANNER_LOG_MAX_POLL_RECORDS)
	#[arg(long, default_value_t = DEFAULT_MAX_POLL_RECORDS)]
	pub scanner_log_max_poll_records: usize,

	/// The maximum time to wait for a batch to be completed in milliseconds.
	/// Default: 100 (matching Java CLIENT_WRITER_BATCH_TIMEOUT)
	#[arg(long, default_value_t = DEFAULT_WRITER_BATCH_TIMEOUT_MS)]
	pub writer_batch_timeout_ms: i64,

	/// Connect timeout in milliseconds for TCP transport connect.
	/// Default: 120000 (120 seconds).
	#[arg(long, default_value_t = DEFAULT_CONNECT_TIMEOUT_MS)]
	pub connect_timeout_ms: u64,

	#[arg(long, default_value_t = String::from(DEFAULT_SECURITY_PROTOCOL))]
	pub security_protocol: String,

	#[arg(long, default_value_t = String::from(DEFAULT_SASL_MECHANISM))]
	pub security_sasl_mechanism: String,

	#[arg(long, default_value_t = String::new())]
	pub security_sasl_username: String,

	#[arg(long, default_value_t = String::new())]
	#[serde(skip_serializing)]
	pub security_sasl_password: String,
	}

	impl std::fmt::Debug for Config {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
	f.debug_struct("Config")
	.field("bootstrap_servers", &self.bootstrap_servers)
	.field("writer_request_max_size", &self.writer_request_max_size)
	.field("writer_acks", &self.writer_acks)
	.field("writer_retries", &self.writer_retries)
	.field("writer_batch_size", &self.writer_batch_size)
	.field(
	"writer_bucket_no_key_assigner",
	&self.writer_bucket_no_key_assigner,
	)
	.field(
	"scanner_remote_log_prefetch_num",
	&self.scanner_remote_log_prefetch_num,
	)
	.field(
	"remote_file_download_thread_num",
	&self.remote_file_download_thread_num,
	)
	.field(
	"scanner_log_max_poll_records",
	&self.scanner_log_max_poll_records,
	)
	.field("writer_batch_timeout_ms", &self.writer_batch_timeout_ms)
	.field("connect_timeout_ms", &self.connect_timeout_ms)
	.field("security_protocol", &self.security_protocol)
	.field("security_sasl_mechanism", &self.security_sasl_mechanism)
	.field("security_sasl_username", &self.security_sasl_username)
	.field("security_sasl_password", &"[REDACTED]")
	.finish()
	}
	}

	impl Default for Config {
	fn default() -> Self {
	Self {
	bootstrap_servers: String::from(DEFAULT_BOOTSTRAP_SERVER),
	writer_request_max_size: DEFAULT_REQUEST_MAX_SIZE,
	writer_acks: String::from(DEFAULT_ACKS),
	writer_retries: i32::MAX,
	writer_batch_size: DEFAULT_WRITER_BATCH_SIZE,
	writer_bucket_no_key_assigner: NoKeyAssigner::Sticky,
	scanner_remote_log_prefetch_num: DEFAULT_PREFETCH_NUM,
	remote_file_download_thread_num: DEFAULT_DOWNLOAD_THREADS,
	scanner_remote_log_read_concurrency: DEFAULT_SCANNER_REMOTE_LOG_READ_CONCURRENCY,
	scanner_log_max_poll_records: DEFAULT_MAX_POLL_RECORDS,
	writer_batch_timeout_ms: DEFAULT_WRITER_BATCH_TIMEOUT_MS,
	connect_timeout_ms: DEFAULT_CONNECT_TIMEOUT_MS,
	security_protocol: String::from(DEFAULT_SECURITY_PROTOCOL),
	security_sasl_mechanism: String::from(DEFAULT_SASL_MECHANISM),
	security_sasl_username: String::new(),
	security_sasl_password: String::new(),
	}
	}
	}

	impl Config {
	/// Returns true when the security protocol indicates SASL authentication
	/// should be performed. Matches Java's `SaslAuthenticationPlugin` which
	/// registers as `"sasl"` (case-insensitive).
	pub fn is_sasl_enabled(&self) -> bool {
	self.security_protocol.eq_ignore_ascii_case("sasl")
	}

	/// Validates security configuration. Returns `Ok(())` when the config is
	/// consistent, or an error message when SASL is enabled but the config is
	/// incomplete or uses an unsupported mechanism.
	pub fn validate_security(&self) -> Result<(), String> {
	if !self.is_sasl_enabled() {
	return Ok(());
	}
	if !self.security_sasl_mechanism.eq_ignore_ascii_case("PLAIN") {
	return Err(format!(
	"Unsupported SASL mechanism: '{}'. Only 'PLAIN' is supported.",
	self.security_sasl_mechanism
	));
	}
	if self.security_sasl_username.is_empty() {
	return Err(
	"security_sasl_username must be set when security_protocol is 'sasl'".to_string(),
	);
	}
	if self.security_sasl_password.is_empty() {
	return Err(
	"security_sasl_password must be set when security_protocol is 'sasl'".to_string(),
	);
	}
	Ok(())
	}
	}

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

	#[test]
	fn test_default_is_not_sasl() {
	let config = Config::default();
	assert!(!config.is_sasl_enabled());
	assert!(config.validate_security().is_ok());
	}

	#[test]
	fn test_sasl_enabled_valid() {
	let config = Config {
	security_protocol: "sasl".to_string(),
	security_sasl_mechanism: "PLAIN".to_string(),
	security_sasl_username: "admin".to_string(),
	security_sasl_password: "secret".to_string(),
	..Config::default()
	};
	assert!(config.is_sasl_enabled());
	assert!(config.validate_security().is_ok());
	}

	#[test]
	fn test_sasl_enabled_case_insensitive() {
	let config = Config {
	security_protocol: "SASL".to_string(),
	security_sasl_username: "admin".to_string(),
	security_sasl_password: "secret".to_string(),
	..Config::default()
	};
	assert!(config.is_sasl_enabled());
	assert!(config.validate_security().is_ok());
	}

	#[test]
	fn test_sasl_missing_username() {
	let config = Config {
	security_protocol: "sasl".to_string(),
	security_sasl_password: "secret".to_string(),
	..Config::default()
	};
	assert!(config.validate_security().is_err());
	}

	#[test]
	fn test_sasl_missing_password() {
	let config = Config {
	security_protocol: "sasl".to_string(),
	security_sasl_username: "admin".to_string(),
	..Config::default()
	};
	assert!(config.validate_security().is_err());
	}

	#[test]
	fn test_sasl_unsupported_mechanism() {
	let config = Config {
	security_protocol: "sasl".to_string(),
	security_sasl_mechanism: "SCRAM-SHA-256".to_string(),
	security_sasl_username: "admin".to_string(),
	security_sasl_password: "secret".to_string(),
	..Config::default()
	};
	assert!(config.validate_security().is_err());
	}
	}