1use vector_lib::configurable::configurable_component;
4use vector_lib::sensitive_string::SensitiveString;
5
6use crate::config::{
7 AcknowledgementsConfig, DynValidatedSink, GenerateConfig, Input, SinkConfig, SinkContext,
8 ValidatedSink,
9};
10use crate::sinks::{
11 prelude::*,
12 util::{
13 BatchConfig, HttpEndpoint, RealtimeSizeBasedDefaultBatchSettings, TowerRequestSettings,
14 },
15};
16
17use super::{error::ZerobusSinkError, service::ZerobusService, sink::ZerobusSink};
18
19#[configurable_component]
21#[derive(Clone, Debug)]
22#[serde(tag = "strategy", rename_all = "snake_case")]
23#[configurable(metadata(
24 docs::enum_tag_description = "The authentication strategy to use for Databricks."
25))]
26pub enum DatabricksAuthentication {
27 #[serde(rename = "oauth")]
29 OAuth {
30 #[configurable(metadata(docs::examples = "${DATABRICKS_CLIENT_ID}"))]
32 #[configurable(metadata(docs::examples = "abc123..."))]
33 client_id: SensitiveString,
34
35 #[configurable(metadata(docs::examples = "${DATABRICKS_CLIENT_SECRET}"))]
37 #[configurable(metadata(docs::examples = "secret123..."))]
38 client_secret: SensitiveString,
39 },
40}
41
42impl DatabricksAuthentication {
43 pub fn credentials(&self) -> (&str, &str) {
45 match self {
46 DatabricksAuthentication::OAuth {
47 client_id,
48 client_secret,
49 } => (client_id.inner(), client_secret.inner()),
50 }
51 }
52}
53
54#[configurable_component]
56#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
57#[serde(rename_all = "snake_case")]
58pub enum Compression {
59 #[default]
61 None,
62 Lz4Frame,
64 Zstd,
66}
67
68impl From<Compression> for Option<arrow::ipc::CompressionType> {
69 fn from(value: Compression) -> Self {
70 match value {
71 Compression::None => None,
72 Compression::Lz4Frame => Some(arrow::ipc::CompressionType::LZ4_FRAME),
73 Compression::Zstd => Some(arrow::ipc::CompressionType::ZSTD),
74 }
75 }
76}
77
78#[configurable_component]
83#[derive(Clone, Debug)]
84#[serde(deny_unknown_fields)]
85pub struct ZerobusStreamOptions {
86 #[serde(default = "default_flush_timeout_ms")]
88 #[configurable(metadata(docs::examples = 30000))]
89 pub flush_timeout_ms: u64,
90
91 #[serde(default = "default_server_ack_timeout_ms")]
93 #[configurable(metadata(docs::examples = 60000))]
94 pub server_lack_of_ack_timeout_ms: u64,
95
96 #[configurable(derived)]
98 #[serde(default, skip_serializing_if = "crate::serde::is_default")]
99 pub compression: Compression,
100}
101
102impl Default for ZerobusStreamOptions {
103 fn default() -> Self {
104 Self {
105 flush_timeout_ms: default_flush_timeout_ms(),
106 server_lack_of_ack_timeout_ms: default_server_ack_timeout_ms(),
107 compression: Compression::None,
108 }
109 }
110}
111
112#[configurable_component(sink(
114 "databricks_zerobus",
115 "Stream observability data to Databricks Unity Catalog via Zerobus."
116))]
117#[derive(Clone, Debug)]
118#[serde(deny_unknown_fields)]
119pub struct ZerobusSinkConfig {
120 #[configurable(metadata(
129 docs::examples = "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com"
130 ))]
131 #[configurable(metadata(
132 docs::examples = "https://6543210987654321.zerobus.us-east-1.cloud.databricks.com"
133 ))]
134 pub ingestion_endpoint: HttpEndpoint,
135
136 #[configurable(metadata(docs::examples = "main.default.logs"))]
145 #[configurable(metadata(docs::examples = "main.default.vector_logs"))]
146 pub table_name: String,
147
148 #[configurable(metadata(docs::examples = "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com"))]
157 #[configurable(metadata(docs::examples = "https://dbc-f6e5d4c3-b2a1.cloud.databricks.com"))]
158 pub unity_catalog_endpoint: HttpEndpoint,
159
160 #[configurable(derived)]
167 pub auth: DatabricksAuthentication,
168
169 #[serde(default)]
174 #[configurable(metadata(docs::examples = "my-service/1.2"))]
175 pub user_agent: Option<String>,
176
177 #[configurable(derived)]
178 #[serde(default)]
179 pub stream_options: ZerobusStreamOptions,
180
181 #[configurable(derived)]
182 #[serde(default)]
183 pub batch: BatchConfig<RealtimeSizeBasedDefaultBatchSettings>,
184
185 #[configurable(derived)]
186 #[serde(default)]
187 pub request: TowerRequestConfig,
188
189 #[configurable(derived)]
190 #[serde(
191 default,
192 deserialize_with = "crate::serde::bool_or_struct",
193 skip_serializing_if = "crate::serde::is_default"
194 )]
195 pub acknowledgements: AcknowledgementsConfig,
196}
197
198impl GenerateConfig for ZerobusSinkConfig {
199 fn generate_config() -> serde_json::Value {
200 serde_json::to_value(Self {
201 ingestion_endpoint: HttpEndpoint::parse(
202 "https://1234567890123456.zerobus.us-west-2.cloud.databricks.com",
203 )
204 .expect("valid example ingestion endpoint"),
205 table_name: "main.default.logs".to_string(),
206 unity_catalog_endpoint: HttpEndpoint::parse(
207 "https://dbc-a1b2c3d4-e5f6.cloud.databricks.com",
208 )
209 .expect("valid example unity catalog endpoint"),
210 auth: DatabricksAuthentication::OAuth {
211 client_id: SensitiveString::from("${DATABRICKS_CLIENT_ID}".to_string()),
212 client_secret: SensitiveString::from("${DATABRICKS_CLIENT_SECRET}".to_string()),
213 },
214 user_agent: None,
215 stream_options: ZerobusStreamOptions::default(),
216 batch: BatchConfig::default(),
217 request: TowerRequestConfig::default(),
218 acknowledgements: AcknowledgementsConfig::default(),
219 })
220 .unwrap()
221 }
222}
223
224#[async_trait::async_trait]
225#[typetag::serde(name = "databricks_zerobus")]
226impl SinkConfig for ZerobusSinkConfig {
227 fn input(&self) -> Input {
228 Input::log()
229 }
230
231 fn acknowledgements(&self) -> &AcknowledgementsConfig {
232 &self.acknowledgements
233 }
234
235 fn as_dyn_validated(&self) -> Option<&dyn DynValidatedSink> {
236 Some(self)
237 }
238}
239
240#[derive(Clone, Debug)]
241pub struct ValidatedZerobus {
242 batch_settings: BatcherSettings,
243 request_limits: TowerRequestSettings,
244}
245
246#[async_trait::async_trait]
247impl ValidatedSink for ZerobusSinkConfig {
248 type Validated = ValidatedZerobus;
249
250 fn validate(&self) -> crate::Result<ValidatedZerobus> {
251 self.validate()?;
253
254 let batch_settings = self.batch.into_batcher_settings()?;
255
256 let request_limits = self.request.into_settings();
257
258 Ok(ValidatedZerobus {
259 batch_settings,
260 request_limits,
261 })
262 }
263
264 async fn build(
265 &self,
266 validated: &ValidatedZerobus,
267 cx: SinkContext,
268 ) -> crate::Result<(VectorSink, Healthcheck)> {
269 let service = ZerobusService::new(self.clone(), cx.proxy()).await?;
270 let healthcheck_service = service.clone();
271
272 let sink = ZerobusSink::new(
273 service,
274 validated.request_limits.clone(),
275 validated.batch_settings,
276 );
277
278 let healthcheck = async move {
279 healthcheck_service
280 .ensure_stream()
281 .await
282 .map_err(|e| e.into())
283 };
284
285 Ok((
286 VectorSink::from_event_streamsink(sink),
287 Box::pin(healthcheck),
288 ))
289 }
290}
291
292impl ZerobusSinkConfig {
293 pub fn validate(&self) -> Result<(), ZerobusSinkError> {
294 if self.table_name.is_empty() {
299 return Err(ZerobusSinkError::ConfigError {
300 message: "table_name cannot be empty".to_string(),
301 });
302 }
303
304 let parts: Vec<&str> = self.table_name.split('.').collect();
305 if parts.len() != 3 || parts.iter().any(|p| p.is_empty()) {
306 return Err(ZerobusSinkError::ConfigError {
307 message: "table_name must be in format 'catalog.schema.table' (exactly 3 non-empty parts)"
308 .to_string(),
309 });
310 }
311
312 match &self.auth {
320 DatabricksAuthentication::OAuth {
321 client_id,
322 client_secret,
323 } => {
324 if client_id.inner().is_empty() {
325 return Err(ZerobusSinkError::ConfigError {
326 message: "OAuth client_id cannot be empty".to_string(),
327 });
328 }
329 if client_secret.inner().is_empty() {
330 return Err(ZerobusSinkError::ConfigError {
331 message: "OAuth client_secret cannot be empty".to_string(),
332 });
333 }
334 }
335 }
336
337 if let Some(max_bytes) = self.batch.max_bytes {
338 if max_bytes > 10_000_000 {
346 return Err(ZerobusSinkError::ConfigError {
347 message: "max_bytes must be less than or equal to 10MB".to_string(),
348 });
349 }
350 }
351
352 Ok(())
353 }
354
355 pub fn user_agent_suffix(&self) -> String {
359 let vector = format!("Vector/{}", crate::vector_version());
360 match self.user_agent.as_deref().filter(|s| !s.is_empty()) {
361 Some(ua) => format!("{vector} {ua}"),
362 None => vector,
363 }
364 }
365}
366
367const fn default_flush_timeout_ms() -> u64 {
369 30000
370}
371
372const fn default_server_ack_timeout_ms() -> u64 {
373 60000
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379 use vector_lib::sensitive_string::SensitiveString;
380
381 fn create_test_config() -> ZerobusSinkConfig {
382 ZerobusSinkConfig {
383 ingestion_endpoint: HttpEndpoint::parse("https://test.databricks.com").unwrap(),
384 table_name: "test.default.logs".to_string(),
385 unity_catalog_endpoint: HttpEndpoint::parse("https://test-workspace.databricks.com")
386 .unwrap(),
387 auth: DatabricksAuthentication::OAuth {
388 client_id: SensitiveString::from("test-client-id".to_string()),
389 client_secret: SensitiveString::from("test-client-secret".to_string()),
390 },
391 user_agent: None,
392 stream_options: ZerobusStreamOptions::default(),
393 batch: Default::default(),
394 request: Default::default(),
395 acknowledgements: Default::default(),
396 }
397 }
398
399 #[test]
400 fn validate_produces_request_limits_and_batch() {
401 use crate::config::ValidatedSink;
402 let config = create_test_config();
403 let validated = <ZerobusSinkConfig as ValidatedSink>::validate(&config)
406 .expect("validation should succeed");
407 assert_eq!(validated.request_limits.concurrency, None);
409 assert!(validated.request_limits.timeout.as_secs() > 0);
410 assert_eq!(validated.batch_settings.size_limit, 10_000_000);
412 }
413
414 #[test]
415 fn validate_rejects_invalid_batch_settings() {
416 use crate::config::ValidatedSink;
417 let mut config = create_test_config();
420 config.batch.max_events = Some(0);
421 assert!(
422 <ZerobusSinkConfig as ValidatedSink>::validate(&config).is_err(),
423 "max_events = 0 should fail validation"
424 );
425
426 let mut config = create_test_config();
428 config.batch.timeout_secs = Some(0.0);
429 assert!(
430 <ZerobusSinkConfig as ValidatedSink>::validate(&config).is_err(),
431 "timeout_secs <= 0 should fail validation"
432 );
433 }
434
435 #[test]
436 fn test_config_validation_success() {
437 let config = create_test_config();
438 assert!(config.validate().is_ok());
439 }
440
441 #[test]
442 fn test_config_validation_rejects_non_http_ingestion_endpoint() {
443 let config = r#"
446ingestion_endpoint: "ftp://test.databricks.com"
447table_name: "test.default.logs"
448unity_catalog_endpoint: "https://test-workspace.databricks.com"
449auth:
450 strategy: oauth
451 client_id: "test-client-id"
452 client_secret: "test-client-secret"
453"#;
454 let result: Result<ZerobusSinkConfig, _> = serde_yaml::from_str(config);
455 assert!(
456 result.is_err(),
457 "non-http ingestion_endpoint should fail deserialization"
458 );
459
460 let config = r#"
462ingestion_endpoint: "/some/path"
463table_name: "test.default.logs"
464unity_catalog_endpoint: "https://test-workspace.databricks.com"
465auth:
466 strategy: oauth
467 client_id: "test-client-id"
468 client_secret: "test-client-secret"
469"#;
470 let result: Result<ZerobusSinkConfig, _> = serde_yaml::from_str(config);
471 assert!(
472 result.is_err(),
473 "host-less ingestion_endpoint should fail deserialization"
474 );
475 }
476
477 #[test]
478 fn test_config_validation_empty_table_name() {
479 let mut config = create_test_config();
480 config.table_name = "".to_string();
481
482 let result = config.validate();
483 assert!(result.is_err());
484
485 if let Err(crate::sinks::databricks_zerobus::error::ZerobusSinkError::ConfigError {
486 message,
487 }) = result
488 {
489 assert!(message.contains("table_name cannot be empty"));
490 } else {
491 panic!("Expected ConfigError for empty table_name");
492 }
493 }
494
495 #[test]
496 fn test_config_validation_invalid_table_name() {
497 let mut config = create_test_config();
498 config.table_name = "invalid_table".to_string(); let result = config.validate();
501 assert!(result.is_err());
502
503 if let Err(crate::sinks::databricks_zerobus::error::ZerobusSinkError::ConfigError {
504 message,
505 }) = result
506 {
507 assert!(message.contains("catalog.schema.table"));
508 } else {
509 panic!("Expected ConfigError for invalid table_name format");
510 }
511 }
512
513 #[test]
514 fn test_config_validation_table_name_empty_segments() {
515 for bad in [
516 "catalog..table",
517 ".schema.table",
518 "catalog.schema.",
519 "..",
520 "catalog.schema.table.extra",
521 ] {
522 let mut config = create_test_config();
523 config.table_name = bad.to_string();
524 let result = config.validate();
525 assert!(result.is_err(), "expected error for table_name={bad:?}");
526 if let Err(crate::sinks::databricks_zerobus::error::ZerobusSinkError::ConfigError {
527 message,
528 }) = result
529 {
530 assert!(message.contains("catalog.schema.table"));
531 } else {
532 panic!("Expected ConfigError for table_name={bad:?}");
533 }
534 }
535 }
536
537 #[test]
538 fn test_config_rejects_invalid_unity_catalog_endpoint() {
539 for bad in [
545 "",
546 "not a uri",
547 "//workspace.databricks.com",
548 "ftp://workspace.databricks.com",
549 "http://:8080",
550 "http://localhost:notaport",
551 ] {
552 let config = format!(
553 r#"
554ingestion_endpoint: "https://test.databricks.com"
555table_name: "test.default.logs"
556unity_catalog_endpoint: {bad:?}
557auth:
558 strategy: oauth
559 client_id: "test-client-id"
560 client_secret: "test-client-secret"
561"#
562 );
563 let result: Result<ZerobusSinkConfig, _> = serde_yaml::from_str(&config);
564 assert!(
565 result.is_err(),
566 "expected deserialization error for endpoint={bad:?}"
567 );
568 }
569 }
570
571 #[test]
572 fn test_config_validation_empty_oauth_credentials() {
573 let mut config = create_test_config();
574 config.auth = DatabricksAuthentication::OAuth {
575 client_id: SensitiveString::from("".to_string()),
576 client_secret: SensitiveString::from("test-secret".to_string()),
577 };
578
579 let result = config.validate();
580 assert!(result.is_err());
581
582 if let Err(crate::sinks::databricks_zerobus::error::ZerobusSinkError::ConfigError {
583 message,
584 }) = result
585 {
586 assert!(message.contains("OAuth client_id cannot be empty"));
587 } else {
588 panic!("Expected ConfigError for empty OAuth client_id");
589 }
590 }
591
592 #[test]
593 fn test_stream_options_compression_deserializes() {
594 let opts: ZerobusStreamOptions =
595 serde_json::from_str(r#"{"compression":"zstd"}"#).expect("should parse zstd");
596 assert_eq!(opts.compression, Compression::Zstd);
597
598 let opts: ZerobusStreamOptions =
599 serde_json::from_str(r#"{"compression":"lz4_frame"}"#).expect("should parse lz4_frame");
600 assert_eq!(opts.compression, Compression::Lz4Frame);
601
602 let opts: ZerobusStreamOptions =
603 serde_json::from_str(r#"{"compression":"none"}"#).expect("should parse none");
604 assert_eq!(opts.compression, Compression::None);
605
606 let opts: ZerobusStreamOptions = serde_json::from_str("{}").expect("should parse empty");
608 assert_eq!(opts.compression, Compression::None);
609 }
610
611 #[test]
612 fn test_compression_maps_to_arrow_ipc() {
613 assert_eq!(
614 Option::<arrow::ipc::CompressionType>::from(Compression::None),
615 None,
616 );
617 assert_eq!(
618 Option::<arrow::ipc::CompressionType>::from(Compression::Lz4Frame),
619 Some(arrow::ipc::CompressionType::LZ4_FRAME),
620 );
621 assert_eq!(
622 Option::<arrow::ipc::CompressionType>::from(Compression::Zstd),
623 Some(arrow::ipc::CompressionType::ZSTD),
624 );
625 }
626
627 #[test]
631 fn test_arrow_ipc_compression_codecs_are_enabled() {
632 use std::sync::Arc;
633
634 use arrow::array::Int32Array;
635 use arrow::datatypes::{DataType, Field, Schema};
636 use arrow::ipc::writer::{IpcWriteOptions, StreamWriter};
637 use arrow::record_batch::RecordBatch;
638
639 let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int32, false)]));
640 let batch = RecordBatch::try_new(
641 Arc::clone(&schema),
642 vec![Arc::new(Int32Array::from((0..1024).collect::<Vec<_>>()))],
643 )
644 .expect("batch should build");
645
646 for codec in [Compression::Lz4Frame, Compression::Zstd] {
647 let compression: Option<arrow::ipc::CompressionType> = codec.into();
648 let options = IpcWriteOptions::default()
649 .try_with_compression(compression)
650 .unwrap_or_else(|e| panic!("{codec:?} not enabled in arrow build: {e}"));
651
652 let mut buf = Vec::new();
653 let mut writer = StreamWriter::try_new_with_options(&mut buf, &schema, options)
654 .unwrap_or_else(|e| panic!("writer for {codec:?} should build: {e}"));
655 writer
656 .write(&batch)
657 .unwrap_or_else(|e| panic!("writing compressed batch for {codec:?} failed: {e}"));
658 writer
659 .finish()
660 .unwrap_or_else(|e| panic!("finishing stream for {codec:?} failed: {e}"));
661
662 assert!(!buf.is_empty(), "{codec:?} produced no output");
663 }
664 }
665
666 #[test]
672 fn test_batch_max_bytes_none_defaults_to_10mb() {
673 let mut config = create_test_config();
674 config.batch.max_bytes = None;
675
676 let settings = config
677 .batch
678 .into_batcher_settings()
679 .expect("batch settings should build");
680
681 assert_eq!(settings.size_limit, 10_000_000);
682 }
683
684 #[test]
685 fn test_user_agent_suffix_without_user_value() {
686 let config = create_test_config();
687 let suffix = config.user_agent_suffix();
688 assert!(
689 suffix.starts_with("Vector/"),
690 "expected Vector/<version> prefix, got {suffix:?}"
691 );
692 assert!(
694 !suffix.contains(' '),
695 "unexpected appended value in {suffix:?}"
696 );
697 }
698
699 #[test]
700 fn test_user_agent_suffix_with_user_value() {
701 let mut config = create_test_config();
702 config.user_agent = Some("my-service/1.2".to_string());
703 let suffix = config.user_agent_suffix();
704 assert!(
705 suffix.starts_with("Vector/"),
706 "expected Vector/<version> prefix, got {suffix:?}"
707 );
708 assert!(
709 suffix.ends_with(" my-service/1.2"),
710 "expected user value appended, got {suffix:?}"
711 );
712 }
713
714 #[test]
715 fn test_user_agent_suffix_empty_user_value_ignored() {
716 let mut config = create_test_config();
717 config.user_agent = Some(String::new());
718 let suffix = config.user_agent_suffix();
719 assert!(
721 !suffix.contains(' '),
722 "empty user_agent should be ignored, got {suffix:?}"
723 );
724 }
725}