From f6c204f992db88b1842548241a2e222f2a0a1ef5 Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:48:38 -0700 Subject: [PATCH 1/9] Start building new adapter stages. --- sentry_streams/Cargo.lock | 58 ++++- sentry_streams/Cargo.toml | 3 +- sentry_streams/src/lib.rs | 2 + sentry_streams/src/pull/mod.rs | 235 ++++++++++++++++++ sentry_streams/src/pull/stages/batch.rs | 96 +++++++ .../src/pull/stages/header_filter.rs | 63 +++++ sentry_streams/src/pull/stages/mod.rs | 2 + 7 files changed, 455 insertions(+), 4 deletions(-) create mode 100644 sentry_streams/src/pull/mod.rs create mode 100644 sentry_streams/src/pull/stages/batch.rs create mode 100644 sentry_streams/src/pull/stages/header_filter.rs create mode 100644 sentry_streams/src/pull/stages/mod.rs diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index e21a054a..af588663 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -669,6 +669,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.31" @@ -685,12 +700,34 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.31" @@ -709,8 +746,10 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ + "futures-channel", "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -2139,6 +2178,7 @@ dependencies = [ "chrono", "clap", "ctrlc", + "futures", "gcp_auth", "log", "metrics", @@ -2430,12 +2470,11 @@ dependencies = [ [[package]] name = "sentry_arroyo" -version = "2.40.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ead64064bc67db796b37b2ad465c6798a80343a11fc23186cb78f06c116added" +version = "2.42.0" dependencies = [ "chrono", "coarsetime", + "futures", "once_cell", "parking_lot", "rand 0.8.7", @@ -2445,6 +2484,8 @@ dependencies = [ "serde_json", "thiserror 1.0.69", "tokio", + "tokio-stream", + "tokio-util", "tracing", "uuid", ] @@ -2797,6 +2838,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.15" diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index 07847758..3c2d6d05 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -6,7 +6,8 @@ edition = "2021" [dependencies] pyo3 = { version = "0.29.0" } serde = { version = "1.0", features = ["derive"] } -sentry_arroyo = { version = "2.40.0", features = ["ssl"] } +sentry_arroyo = { path = "../../arroyo", features = ["ssl"] } +futures = "0.3" chrono = "0.4.40" tracing = "0.1.40" tracing-subscriber = "0.3.20" diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index b8819a2b..bb2e3b59 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -26,6 +26,8 @@ mod transformer; mod utils; mod watermark; +pub mod pull; + #[doc(hidden)] pub mod ffi; pub use ffi::Message; diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs new file mode 100644 index 00000000..290f1330 --- /dev/null +++ b/sentry_streams/src/pull/mod.rs @@ -0,0 +1,235 @@ +pub mod stages; + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use futures::stream; + use sentry_arroyo::backends::kafka::types::{Headers, KafkaPayload}; + use sentry_arroyo::processing::strategies::offset_tracker::{OffsetCommitter, OffsetTracker}; + use sentry_arroyo::processing::stream::{ + LogHandler, MessageMetadata, PipelineEnvelope, PipelineExt, Stage, StageResult, + }; + use sentry_arroyo::types::{Partition, Topic}; + + use super::stages::batch::BatchAccumulatorStage; + use super::stages::header_filter::HeaderFilterStage; + + /// Mock committer that records what was committed. + struct MockCommitter { + committed: Mutex>>, + } + + impl MockCommitter { + fn new() -> Self { + Self { + committed: Mutex::new(Vec::new()), + } + } + + fn committed(&self) -> Vec> { + self.committed.lock().unwrap().clone() + } + } + + impl OffsetCommitter for MockCommitter { + fn commit_offsets( + &self, + positions: &HashMap, + ) -> Result<(), Box> { + self.committed.lock().unwrap().push(positions.clone()); + Ok(()) + } + } + + /// Helper to create a test envelope with no headers. + fn make_envelope(payload: &[u8], offset: u64) -> StageResult { + let kafka_payload = KafkaPayload::new(None, None, Some(payload.to_vec())); + let metadata = MessageMetadata { + partition: Partition::new(Topic::new("test"), 0), + offset, + timestamp: chrono::Utc::now(), + }; + StageResult::Emit(PipelineEnvelope::new( + kafka_payload.clone(), + metadata, + Arc::new(kafka_payload), + )) + } + + /// Helper to create a test envelope with a header. + fn make_envelope_with_header( + payload: &[u8], + offset: u64, + header_name: &str, + header_value: i64, + ) -> StageResult { + let headers = Headers::new().insert( + header_name, + Some(header_value.to_string().into_bytes()), + ); + let kafka_payload = KafkaPayload::new(None, Some(headers), Some(payload.to_vec())); + let metadata = MessageMetadata { + partition: Partition::new(Topic::new("test"), 0), + offset, + timestamp: chrono::Utc::now(), + }; + StageResult::Emit(PipelineEnvelope::new( + kafka_payload.clone(), + metadata, + Arc::new(kafka_payload), + )) + } + + /// Helper to create a test envelope with an invalid (non-integer) header. + fn make_envelope_with_bad_header( + payload: &[u8], + offset: u64, + header_name: &str, + header_value: &[u8], + ) -> StageResult { + let headers = Headers::new().insert( + header_name, + Some(header_value.to_vec()), + ); + let kafka_payload = KafkaPayload::new(None, Some(headers), Some(payload.to_vec())); + let metadata = MessageMetadata { + partition: Partition::new(Topic::new("test"), 0), + offset, + timestamp: chrono::Utc::now(), + }; + StageResult::Emit(PipelineEnvelope::new( + kafka_payload.clone(), + metadata, + Arc::new(kafka_payload), + )) + } + + #[tokio::test] + async fn test_header_filter_passes_matching() { + let committer = MockCommitter::new(); + let mut tracker = OffsetTracker::new(Duration::from_millis(1), &committer); + let filter = HeaderFilterStage::new("item_type", 1); + let error_handler = LogHandler; + + // 3 messages with matching header, 2 without + let messages = vec![ + make_envelope_with_header(b"span-1", 0, "item_type", 1), // match + make_envelope(b"no-header", 1), // no header → drop + make_envelope_with_header(b"span-2", 2, "item_type", 1), // match + make_envelope_with_header(b"log-1", 3, "item_type", 2), // wrong value → drop + make_envelope_with_header(b"span-3", 4, "item_type", 1), // match + ]; + + let result = stream::iter(messages) + .apply(&filter) + .on_reject(&error_handler) + .commit(&mut tracker) + .await; + + assert!(result.is_ok()); + + // All 5 offsets should be tracked (3 Emit + 2 Drop) + let committed = committer.committed(); + assert!(!committed.is_empty(), "Expected at least one commit"); + let last = committed.last().unwrap(); + let partition = Partition::new(Topic::new("test"), 0); + assert_eq!(last.get(&partition), Some(&5)); + } + + #[tokio::test] + async fn test_pipeline_filter_then_batch() { + let committer = MockCommitter::new(); + let mut tracker = OffsetTracker::new(Duration::from_millis(1), &committer); + let filter = HeaderFilterStage::new("item_type", 1); + let batch = BatchAccumulatorStage::new(3); + let error_handler = LogHandler; + + // 5 messages, all with matching header + let messages: Vec> = (0..5) + .map(|i| make_envelope_with_header( + format!("msg-{i}").as_bytes(), i, "item_type", 1, + )) + .collect(); + + // Track how many batches we receive and their sizes + let batch_sizes: Arc>> = Arc::new(Mutex::new(Vec::new())); + let batch_sizes_clone = batch_sizes.clone(); + + // Use a counting stage after the batch to record batch sizes + struct CountBatchStage { + sizes: Arc>>, + } + impl Stage for CountBatchStage { + type In = Vec; + type Out = Vec; + async fn process( + &self, + envelope: PipelineEnvelope>, + ) -> StageResult> { + self.sizes.lock().unwrap().push(envelope.payload.len()); + StageResult::Emit(envelope) + } + fn name(&self) -> &'static str { "count_batch" } + } + + let counter = CountBatchStage { sizes: batch_sizes_clone }; + + let result = stream::iter(messages) + .apply(&filter) + .apply(&batch) + .apply(&counter) + .on_reject(&error_handler) + .commit(&mut tracker) + .await; + + assert!(result.is_ok()); + + // With batch size 3 and 5 messages: first batch has 3, remaining 2 are + // still in the accumulator (not flushed since stream ended without + // reaching batch size again). + let sizes = batch_sizes.lock().unwrap(); + assert_eq!(*sizes, vec![3], "Expected one batch of 3 (remaining 2 not flushed)"); + + // Offsets: batch of 3 emitted with last message's offset (2), so + // tracker sees offset 3. The 2 remaining messages returned Skip, + // so their offsets are not tracked. + let committed = committer.committed(); + assert!(!committed.is_empty()); + let last = committed.last().unwrap(); + let partition = Partition::new(Topic::new("test"), 0); + assert_eq!(last.get(&partition), Some(&3), + "Expected offset 3 (batch last offset 2 + 1)"); + } + + #[tokio::test] + async fn test_header_filter_rejects_invalid_header() { + let committer = MockCommitter::new(); + let mut tracker = OffsetTracker::new(Duration::from_millis(1), &committer); + let filter = HeaderFilterStage::new("item_type", 1); + let error_handler = LogHandler; + + let messages = vec![ + make_envelope_with_header(b"good", 0, "item_type", 1), // match → Emit + make_envelope_with_bad_header(b"bad", 1, "item_type", b"not-an-int"), // invalid → Reject + make_envelope_with_header(b"also-good", 2, "item_type", 1), // match → Emit + ]; + + let result = stream::iter(messages) + .apply(&filter) + .on_reject(&error_handler) + .commit(&mut tracker) + .await; + + assert!(result.is_ok()); + + // All 3 offsets should be tracked (2 Emit + 1 Reject) + let committed = committer.committed(); + assert!(!committed.is_empty()); + let last = committed.last().unwrap(); + let partition = Partition::new(Topic::new("test"), 0); + assert_eq!(last.get(&partition), Some(&3)); + } +} diff --git a/sentry_streams/src/pull/stages/batch.rs b/sentry_streams/src/pull/stages/batch.rs new file mode 100644 index 00000000..e75281dc --- /dev/null +++ b/sentry_streams/src/pull/stages/batch.rs @@ -0,0 +1,96 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use sentry_arroyo::backends::kafka::types::KafkaPayload; +use sentry_arroyo::processing::stream::{MessageMetadata, PipelineEnvelope, Stage, StageResult}; +use sentry_arroyo::types::Partition; + +/// Accumulates envelopes into batches and emits when the batch reaches +/// max_batch_size. Returns Skip while accumulating, Emit when flushing. +/// +/// The emitted envelope contains a Vec of the accumulated payloads. +/// Offsets are merged across the batch — highest offset per partition. +pub struct BatchAccumulatorStage { + max_batch_size: usize, + state: Mutex, +} + +struct BatchState { + payloads: Vec, + offsets: HashMap, + last_metadata: Option, + last_raw: Option>, +} + +impl BatchState { + fn new() -> Self { + Self { + payloads: Vec::new(), + offsets: HashMap::new(), + last_metadata: None, + last_raw: None, + } + } + + /// Add a message to the batch, merging its offset. + fn accumulate(&mut self, envelope: PipelineEnvelope) { + self.offsets + .entry(envelope.metadata.partition) + .and_modify(|o| *o = (*o).max(envelope.metadata.offset)) + .or_insert(envelope.metadata.offset); + self.payloads.push(envelope.payload); + self.last_metadata = Some(envelope.metadata); + self.last_raw = Some(envelope.raw); + } + + /// Check if the batch has reached the size threshold. + fn is_full(&self, max_batch_size: usize) -> bool { + self.payloads.len() >= max_batch_size + } + + /// Drain the batch into an Envelope, clearing internal state. + fn flush(&mut self) -> PipelineEnvelope> { + let payloads = std::mem::take(&mut self.payloads); + let mut metadata = self.last_metadata.take().unwrap(); + let raw = self.last_raw.take().unwrap(); + + if let Some(&max_offset) = self.offsets.get(&metadata.partition) { + metadata.offset = max_offset; + } + self.offsets.clear(); + + PipelineEnvelope::new(payloads, metadata, raw) + } +} + +impl BatchAccumulatorStage { + pub fn new(max_batch_size: usize) -> Self { + Self { + max_batch_size, + state: Mutex::new(BatchState::new()), + } + } +} + +impl Stage for BatchAccumulatorStage { + type In = KafkaPayload; + type Out = Vec; + + async fn process( + &self, + envelope: PipelineEnvelope, + ) -> StageResult> { + let mut state = self.state.lock().unwrap(); + state.accumulate(envelope); + + if state.is_full(self.max_batch_size) { + StageResult::Emit(state.flush()) + } else { + StageResult::Skip + } + } + + fn name(&self) -> &'static str { + "batch_accumulator" + } +} diff --git a/sentry_streams/src/pull/stages/header_filter.rs b/sentry_streams/src/pull/stages/header_filter.rs new file mode 100644 index 00000000..702a7cb7 --- /dev/null +++ b/sentry_streams/src/pull/stages/header_filter.rs @@ -0,0 +1,63 @@ +use sentry_arroyo::backends::kafka::types::KafkaPayload; +use sentry_arroyo::processing::stream::{PipelineEnvelope, RejectionReason, Stage, StageResult}; + +/// Filters messages by checking a Kafka header for an expected integer value. +/// Messages with a matching header pass through (Emit). +/// Messages without the header or with a non-matching value are dropped (Drop). +/// Messages with an unparseable header value are rejected (Reject → DLQ). +/// +/// Header values are treated as UTF-8 ASCII decimal integers (matching the +/// existing push-based HeaderIntEqualityFilter in streams). +pub struct HeaderFilterStage { + header_name: String, + expected_value: i64, +} + +impl HeaderFilterStage { + pub fn new(header_name: impl Into, expected_value: i64) -> Self { + Self { + header_name: header_name.into(), + expected_value, + } + } + + /// `Ok(true)` — header matches expected value. + /// `Ok(false)` — header missing, empty, or different value. + /// `Err(())` — header present but not a valid decimal integer. + fn check_header(&self, payload: &KafkaPayload) -> Result { + let Some(headers) = payload.headers() else { + return Ok(false); + }; + let Some(bytes) = headers.get(&self.header_name) else { + return Ok(false); + }; + if bytes.is_empty() { + return Ok(false); + } + let parsed = std::str::from_utf8(bytes) + .map_err(|_| ())? + .parse::() + .map_err(|_| ())?; + Ok(parsed == self.expected_value) + } +} + +impl Stage for HeaderFilterStage { + type In = KafkaPayload; + type Out = KafkaPayload; + + async fn process( + &self, + envelope: PipelineEnvelope, + ) -> StageResult { + match self.check_header(&envelope.payload) { + Ok(true) => StageResult::Emit(envelope), + Ok(false) => StageResult::drop(envelope), + Err(()) => StageResult::reject(envelope, RejectionReason::Invalid), + } + } + + fn name(&self) -> &'static str { + "header_filter" + } +} diff --git a/sentry_streams/src/pull/stages/mod.rs b/sentry_streams/src/pull/stages/mod.rs new file mode 100644 index 00000000..fba98482 --- /dev/null +++ b/sentry_streams/src/pull/stages/mod.rs @@ -0,0 +1,2 @@ +pub mod batch; +pub mod header_filter; From 218ae9c8d7e23d015e17baf4f2bc669ea1f34792 Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:28:14 -0700 Subject: [PATCH 2/9] Define rust / python boundary. Add requisite canned stages. --- sentry_streams/Cargo.lock | 39 +++++ sentry_streams/Cargo.toml | 1 + sentry_streams/src/pull/gcs_client.rs | 123 ++++++++++++++ sentry_streams/src/pull/gcs_sink_handler.rs | 70 ++++++++ sentry_streams/src/pull/mod.rs | 153 +++++++++++++----- sentry_streams/src/pull/pipeline_value.rs | 137 ++++++++++++++++ sentry_streams/src/pull/stages/batch.rs | 47 +++--- .../src/pull/stages/header_filter.rs | 66 ++++---- sentry_streams/src/pull/stages/mod.rs | 1 + sentry_streams/src/pull/stages/py_callable.rs | 104 ++++++++++++ 10 files changed, 650 insertions(+), 91 deletions(-) create mode 100644 sentry_streams/src/pull/gcs_client.rs create mode 100644 sentry_streams/src/pull/gcs_sink_handler.rs create mode 100644 sentry_streams/src/pull/pipeline_value.rs create mode 100644 sentry_streams/src/pull/stages/py_callable.rs diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index af588663..0664da83 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -578,6 +578,12 @@ dependencies = [ "syn", ] +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + [[package]] name = "encoding_rs" version = "0.8.35" @@ -1206,6 +1212,15 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" @@ -1822,6 +1837,29 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pyo3" version = "0.29.0" @@ -2184,6 +2222,7 @@ dependencies = [ "metrics", "metrics-exporter-dogstatsd", "parking_lot", + "prost", "pyo3", "rand 0.9.4", "rdkafka", diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index 3c2d6d05..64cdcd85 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -8,6 +8,7 @@ pyo3 = { version = "0.29.0" } serde = { version = "1.0", features = ["derive"] } sentry_arroyo = { path = "../../arroyo", features = ["ssl"] } futures = "0.3" +prost = "0.13" chrono = "0.4.40" tracing = "0.1.40" tracing-subscriber = "0.3.20" diff --git a/sentry_streams/src/pull/gcs_client.rs b/sentry_streams/src/pull/gcs_client.rs new file mode 100644 index 00000000..57764704 --- /dev/null +++ b/sentry_streams/src/pull/gcs_client.rs @@ -0,0 +1,123 @@ +use std::sync::Arc; + +use gcp_auth::{provider, TokenProvider}; +use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; +use reqwest::{Client, ClientBuilder}; +use tokio::sync::OnceCell; + +/// Client for uploading bytes to Google Cloud Storage via the JSON API. +/// +/// Handles authentication (lazily initialized via gcp_auth), token refresh, +/// and HTTP upload. Inject this into handlers/stages that need GCS access. +pub struct GcsClient { + client: Client, + bucket: String, + auth_provider: Arc>>, +} + +impl GcsClient { + pub fn new(client: Client, bucket: impl Into) -> Self { + Self { + client, + bucket: bucket.into(), + auth_provider: Arc::new(OnceCell::new()), + } + } + + /// Create a client with default reqwest settings. + pub fn with_defaults(bucket: impl Into) -> Self { + let mut headers = HeaderMap::with_capacity(1); + headers.insert( + CONTENT_TYPE, + HeaderValue::from_static("application/octet-stream"), + ); + let client = ClientBuilder::new() + .default_headers(headers) + .build() + .expect("Failed to build reqwest client"); + + Self::new(client, bucket) + } + + /// Upload bytes to GCS as the given object name. + pub async fn upload( + &self, + object_name: &str, + bytes: &[u8], + ) -> Result<(), GcsError> { + let auth_provider = self + .auth_provider + .get_or_init(|| async { + provider().await.expect("Failed to get gcp_auth provider") + }) + .await; + + let scopes = &["https://www.googleapis.com/auth/devstorage.read_write"]; + let token = auth_provider + .token(scopes) + .await + .map_err(|e| GcsError::Auth(format!("{e}")))?; + + let url = format!( + "https://storage.googleapis.com/upload/storage/v1/b/{}/o?uploadType=media&name={}", + self.bucket, object_name, + ); + + let response = self + .client + .post(&url) + .header( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {}", token.as_str())).unwrap(), + ) + .body(bytes.to_vec()) + .send() + .await + .map_err(|e| GcsError::Network(format!("{e}")))?; + + let status = response.status(); + if status.is_success() { + tracing::info!( + "GCS upload complete: bucket={}, object={}, bytes={}", + self.bucket, + object_name, + bytes.len(), + ); + metrics::histogram!("streams.pipeline.sink.gcs_writer.bytes") + .record(bytes.len() as f64); + Ok(()) + } else if status.is_client_error() { + let body = response.text().await.unwrap_or_default(); + Err(GcsError::ClientError(format!( + "status={status}, body={body}" + ))) + } else { + Err(GcsError::ServerError(format!("status={status}"))) + } + } +} + +#[derive(Debug)] +pub enum GcsError { + /// Failed to obtain auth token. + Auth(String), + /// Network/transport error (retryable). + Network(String), + /// 4xx — fatal, bad request. + ClientError(String), + /// 5xx — retryable server error. + ServerError(String), +} + +impl std::fmt::Display for GcsError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + GcsError::Auth(msg) => write!(f, "GCS auth error: {msg}"), + GcsError::Network(msg) => write!(f, "GCS network error: {msg}"), + GcsError::ClientError(msg) => write!(f, "GCS client error: {msg}"), + GcsError::ServerError(msg) => write!(f, "GCS server error: {msg}"), + } + } +} + +impl std::error::Error for GcsError {} diff --git a/sentry_streams/src/pull/gcs_sink_handler.rs b/sentry_streams/src/pull/gcs_sink_handler.rs new file mode 100644 index 00000000..163681d5 --- /dev/null +++ b/sentry_streams/src/pull/gcs_sink_handler.rs @@ -0,0 +1,70 @@ +use pyo3::prelude::*; +use sentry_arroyo::processing::stream::handlers::next::NextHandler; +use sentry_arroyo::processing::stream::PipelineEnvelope; + +use super::gcs_client::GcsClient; +use super::pipeline_value::PipelineValue; + +/// Sink handler that uploads pipeline output to GCS. +/// +/// Used with `.on_next(&gcs_sink)` in the pipeline. Extracts bytes +/// from the envelope payload, generates an object name via a Python +/// callable, and delegates to `GcsClient` for the actual upload. +pub struct GcsSinkHandler { + client: GcsClient, + object_generator: Py, +} + +impl GcsSinkHandler { + pub fn new(client: GcsClient, object_generator: Py) -> Self { + Self { + client, + object_generator, + } + } + + /// Extract bytes from the pipeline value. + fn extract_bytes(value: &PipelineValue) -> Result, Box> { + match value { + PipelineValue::Python(obj) => Python::attach(|py| { + obj.extract::>(py) + .map_err(|e| Box::new(e) as Box) + }), + PipelineValue::Rust(boxed) => { + boxed.downcast_ref::>().cloned().ok_or_else(|| { + Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "GcsSinkHandler expected Rust Vec", + )) as Box + }) + } + PipelineValue::Raw(kp) => Ok(kp.payload().map(|v| v.to_vec()).unwrap_or_default()), + } + } + + fn generate_object_name(&self) -> Result> { + Python::attach(|py| { + let result = self + .object_generator + .call0(py) + .map_err(|e| Box::new(e) as Box)?; + result + .extract::(py) + .map_err(|e| Box::new(e) as Box) + }) + } +} + +impl NextHandler for GcsSinkHandler { + async fn handle( + &self, + envelope: &PipelineEnvelope, + ) -> Result<(), Box> { + let bytes = Self::extract_bytes(&envelope.payload)?; + let object_name = self.generate_object_name()?; + self.client + .upload(&object_name, &bytes) + .await + .map_err(|e| Box::new(e) as _) + } +} diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index 290f1330..95292ca4 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -1,3 +1,6 @@ +pub mod gcs_client; +pub mod gcs_sink_handler; +pub mod pipeline_value; pub mod stages; #[cfg(test)] @@ -14,6 +17,7 @@ mod tests { }; use sentry_arroyo::types::{Partition, Topic}; + use super::pipeline_value::{PipelineValue, PipelineValueCaster}; use super::stages::batch::BatchAccumulatorStage; use super::stages::header_filter::HeaderFilterStage; @@ -44,8 +48,8 @@ mod tests { } } - /// Helper to create a test envelope with no headers. - fn make_envelope(payload: &[u8], offset: u64) -> StageResult { + /// Helper to create a PipelineValue::Raw envelope with no headers. + fn make_envelope(payload: &[u8], offset: u64) -> StageResult { let kafka_payload = KafkaPayload::new(None, None, Some(payload.to_vec())); let metadata = MessageMetadata { partition: Partition::new(Topic::new("test"), 0), @@ -53,19 +57,19 @@ mod tests { timestamp: chrono::Utc::now(), }; StageResult::Emit(PipelineEnvelope::new( - kafka_payload.clone(), + PipelineValue::Raw(kafka_payload.clone()), metadata, Arc::new(kafka_payload), )) } - /// Helper to create a test envelope with a header. + /// Helper to create a PipelineValue::Raw envelope with a header. fn make_envelope_with_header( payload: &[u8], offset: u64, header_name: &str, header_value: i64, - ) -> StageResult { + ) -> StageResult { let headers = Headers::new().insert( header_name, Some(header_value.to_string().into_bytes()), @@ -77,19 +81,19 @@ mod tests { timestamp: chrono::Utc::now(), }; StageResult::Emit(PipelineEnvelope::new( - kafka_payload.clone(), + PipelineValue::Raw(kafka_payload.clone()), metadata, Arc::new(kafka_payload), )) } - /// Helper to create a test envelope with an invalid (non-integer) header. + /// Helper to create a PipelineValue::Raw envelope with an invalid header. fn make_envelope_with_bad_header( payload: &[u8], offset: u64, header_name: &str, header_value: &[u8], - ) -> StageResult { + ) -> StageResult { let headers = Headers::new().insert( header_name, Some(header_value.to_vec()), @@ -101,12 +105,14 @@ mod tests { timestamp: chrono::Utc::now(), }; StageResult::Emit(PipelineEnvelope::new( - kafka_payload.clone(), + PipelineValue::Raw(kafka_payload.clone()), metadata, Arc::new(kafka_payload), )) } + // ── HeaderFilterStage tests ───────────────────────────────────── + #[tokio::test] async fn test_header_filter_passes_matching() { let committer = MockCommitter::new(); @@ -114,7 +120,6 @@ mod tests { let filter = HeaderFilterStage::new("item_type", 1); let error_handler = LogHandler; - // 3 messages with matching header, 2 without let messages = vec![ make_envelope_with_header(b"span-1", 0, "item_type", 1), // match make_envelope(b"no-header", 1), // no header → drop @@ -131,7 +136,6 @@ mod tests { assert!(result.is_ok()); - // All 5 offsets should be tracked (3 Emit + 2 Drop) let committed = committer.committed(); assert!(!committed.is_empty(), "Expected at least one commit"); let last = committed.last().unwrap(); @@ -139,6 +143,36 @@ mod tests { assert_eq!(last.get(&partition), Some(&5)); } + #[tokio::test] + async fn test_header_filter_rejects_invalid_header() { + let committer = MockCommitter::new(); + let mut tracker = OffsetTracker::new(Duration::from_millis(1), &committer); + let filter = HeaderFilterStage::new("item_type", 1); + let error_handler = LogHandler; + + let messages = vec![ + make_envelope_with_header(b"good", 0, "item_type", 1), + make_envelope_with_bad_header(b"bad", 1, "item_type", b"not-an-int"), + make_envelope_with_header(b"also-good", 2, "item_type", 1), + ]; + + let result = stream::iter(messages) + .apply(&filter) + .on_reject(&error_handler) + .commit(&mut tracker) + .await; + + assert!(result.is_ok()); + + let committed = committer.committed(); + assert!(!committed.is_empty()); + let last = committed.last().unwrap(); + let partition = Partition::new(Topic::new("test"), 0); + assert_eq!(last.get(&partition), Some(&3)); + } + + // ── BatchAccumulatorStage tests ───────────────────────────────── + #[tokio::test] async fn test_pipeline_filter_then_batch() { let committer = MockCommitter::new(); @@ -147,35 +181,41 @@ mod tests { let batch = BatchAccumulatorStage::new(3); let error_handler = LogHandler; - // 5 messages, all with matching header - let messages: Vec> = (0..5) + let messages: Vec> = (0..5) .map(|i| make_envelope_with_header( format!("msg-{i}").as_bytes(), i, "item_type", 1, )) .collect(); - // Track how many batches we receive and their sizes + // Count batch sizes via a simple stage that downcasts the Rust batch let batch_sizes: Arc>> = Arc::new(Mutex::new(Vec::new())); - let batch_sizes_clone = batch_sizes.clone(); + let bs = batch_sizes.clone(); - // Use a counting stage after the batch to record batch sizes struct CountBatchStage { sizes: Arc>>, } impl Stage for CountBatchStage { - type In = Vec; - type Out = Vec; + type In = PipelineValue; + type Out = PipelineValue; async fn process( &self, - envelope: PipelineEnvelope>, - ) -> StageResult> { - self.sizes.lock().unwrap().push(envelope.payload.len()); - StageResult::Emit(envelope) + envelope: PipelineEnvelope, + ) -> StageResult { + let typed = match envelope.downcast_rust::>() { + Ok(t) => t, + Err(fail) => return fail, + }; + self.sizes.lock().unwrap().push(typed.payload.len()); + StageResult::Emit(PipelineEnvelope::new( + PipelineValue::Rust(Box::new(typed.payload)), + typed.metadata, + typed.raw, + )) } fn name(&self) -> &'static str { "count_batch" } } - let counter = CountBatchStage { sizes: batch_sizes_clone }; + let counter = CountBatchStage { sizes: bs }; let result = stream::iter(messages) .apply(&filter) @@ -187,15 +227,9 @@ mod tests { assert!(result.is_ok()); - // With batch size 3 and 5 messages: first batch has 3, remaining 2 are - // still in the accumulator (not flushed since stream ended without - // reaching batch size again). let sizes = batch_sizes.lock().unwrap(); assert_eq!(*sizes, vec![3], "Expected one batch of 3 (remaining 2 not flushed)"); - // Offsets: batch of 3 emitted with last message's offset (2), so - // tracker sees offset 3. The 2 remaining messages returned Skip, - // so their offsets are not tracked. let committed = committer.committed(); assert!(!committed.is_empty()); let last = committed.last().unwrap(); @@ -204,32 +238,69 @@ mod tests { "Expected offset 3 (batch last offset 2 + 1)"); } + // ── Full pipeline integration test ────────────────────────────── + #[tokio::test] - async fn test_header_filter_rejects_invalid_header() { + async fn test_full_pipeline_filter_batch() { let committer = MockCommitter::new(); let mut tracker = OffsetTracker::new(Duration::from_millis(1), &committer); let filter = HeaderFilterStage::new("item_type", 1); + let batch = BatchAccumulatorStage::new(2); let error_handler = LogHandler; - let messages = vec![ - make_envelope_with_header(b"good", 0, "item_type", 1), // match → Emit - make_envelope_with_bad_header(b"bad", 1, "item_type", b"not-an-int"), // invalid → Reject - make_envelope_with_header(b"also-good", 2, "item_type", 1), // match → Emit - ]; + // Track batches and their content + let batches: Arc>>>> = Arc::new(Mutex::new(Vec::new())); + let batches_clone = batches.clone(); + + struct CollectBatchStage { + batches: Arc>>>>, + } + impl Stage for CollectBatchStage { + type In = PipelineValue; + type Out = PipelineValue; + async fn process( + &self, + envelope: PipelineEnvelope, + ) -> StageResult { + let typed = match envelope.downcast_rust::>() { + Ok(t) => t, + Err(fail) => return fail, + }; + let contents: Vec> = typed.payload.iter() + .map(|kp| kp.payload().map(|v| v.to_vec()).unwrap_or_default()) + .collect(); + self.batches.lock().unwrap().push(contents); + StageResult::Emit(PipelineEnvelope::new( + PipelineValue::Rust(Box::new(typed.payload)), + typed.metadata, + typed.raw, + )) + } + fn name(&self) -> &'static str { "collect_batch" } + } + + let collector = CollectBatchStage { batches: batches_clone }; + + // 4 messages with matching header → 2 batches of 2 + let messages: Vec> = (0..4) + .map(|i| make_envelope_with_header( + format!("span-{i}").as_bytes(), i, "item_type", 1, + )) + .collect(); let result = stream::iter(messages) .apply(&filter) + .apply(&batch) + .apply(&collector) .on_reject(&error_handler) .commit(&mut tracker) .await; assert!(result.is_ok()); - // All 3 offsets should be tracked (2 Emit + 1 Reject) - let committed = committer.committed(); - assert!(!committed.is_empty()); - let last = committed.last().unwrap(); - let partition = Partition::new(Topic::new("test"), 0); - assert_eq!(last.get(&partition), Some(&3)); + let collected = batches.lock().unwrap(); + assert_eq!(collected.len(), 2, "Expected 2 batches of 2"); + assert_eq!(collected[0], vec![b"span-0".to_vec(), b"span-1".to_vec()]); + assert_eq!(collected[1], vec![b"span-2".to_vec(), b"span-3".to_vec()]); } } diff --git a/sentry_streams/src/pull/pipeline_value.rs b/sentry_streams/src/pull/pipeline_value.rs new file mode 100644 index 00000000..af28fe5a --- /dev/null +++ b/sentry_streams/src/pull/pipeline_value.rs @@ -0,0 +1,137 @@ +use std::any::Any; +use std::fmt; + +use pyo3::prelude::*; +use sentry_arroyo::backends::kafka::types::KafkaPayload; +use sentry_arroyo::processing::stream::{PipelineEnvelope, StageResult}; + +/// The universal payload type for pull-based pipeline stages in streams. +/// +/// Stages are `Stage`, allowing +/// dynamic pipeline construction from the Python DSL without compile-time +/// generics for every pipeline shape. +/// +/// Each stage knows which variant it expects and uses the downcast helpers +/// on `PipelineEnvelope` to unwrap. A variant mismatch is +/// a pipeline construction bug and results in `StageResult::Fail`. +pub enum PipelineValue { + /// Raw Kafka payload — from source, pre-parse. + Raw(KafkaPayload), + + /// Typed Rust data. + /// Downcast via `PipelineEnvelope::downcast_rust::()`. + Rust(Box), + + /// Python heap object. + Python(Py), +} + +impl fmt::Debug for PipelineValue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + PipelineValue::Raw(_) => write!(f, "PipelineValue::Raw(...)"), + PipelineValue::Rust(_) => write!(f, "PipelineValue::Rust(...)"), + PipelineValue::Python(_) => write!(f, "PipelineValue::Python(...)"), + } + } +} + +// PipelineValue is Send because all variants are Send: +// - KafkaPayload is Send +// - Box is Send +// - Py is Send (PyO3 guarantees this) +unsafe impl Send for PipelineValue {} + +/// Error returned when a stage receives an unexpected PipelineValue variant. +#[derive(Debug)] +pub struct DowncastError { + expected: &'static str, + actual: &'static str, +} + +impl fmt::Display for DowncastError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "Pipeline misconfigured: expected PipelineValue::{}, got PipelineValue::{}", + self.expected, self.actual + ) + } +} + +impl std::error::Error for DowncastError {} + +impl PipelineValue { + fn variant_name(&self) -> &'static str { + match self { + PipelineValue::Raw(_) => "Raw", + PipelineValue::Rust(_) => "Rust", + PipelineValue::Python(_) => "Python", + } + } +} + +/// Extension trait for downcasting PipelineEnvelope. +pub trait PipelineValueCaster { + /// Downcast to a Raw kafka payload. + fn downcast_raw(self) -> Result, StageResult>; + + /// Downcast to a typed Rust value. + fn downcast_rust( + self, + ) -> Result, StageResult>; + + /// Downcast to a Python object. + fn downcast_python(self) -> Result>, StageResult>; +} + +impl PipelineValueCaster for PipelineEnvelope { + fn downcast_raw(self) -> Result, StageResult> { + match self.payload { + PipelineValue::Raw(kp) => Ok(PipelineEnvelope::new(kp, self.metadata, self.raw)), + other => Err(StageResult::Fail(Box::new(DowncastError { + expected: "Raw", + actual: other.variant_name(), + }))), + } + } + + fn downcast_rust( + self, + ) -> Result, StageResult> { + match self.payload { + PipelineValue::Rust(boxed) => match boxed.downcast::() { + Ok(val) => Ok(PipelineEnvelope::new(*val, self.metadata, self.raw)), + Err(_) => Err(StageResult::Fail(Box::new(DowncastError { + expected: std::any::type_name::(), + actual: "Rust(wrong type)", + }))), + }, + other => Err(StageResult::Fail(Box::new(DowncastError { + expected: "Rust", + actual: other.variant_name(), + }))), + } + } + + fn downcast_python(self) -> Result>, StageResult> { + match self.payload { + PipelineValue::Python(obj) => Ok(PipelineEnvelope::new(obj, self.metadata, self.raw)), + other => Err(StageResult::Fail(Box::new(DowncastError { + expected: "Python", + actual: other.variant_name(), + }))), + } + } +} + +/// Helper to wrap a typed envelope back into a PipelineValue envelope. +pub trait IntoPipelineValue { + fn into_pipeline_value(self) -> PipelineEnvelope; +} + +impl IntoPipelineValue for PipelineEnvelope { + fn into_pipeline_value(self) -> PipelineEnvelope { + self.map_payload(PipelineValue::Raw) + } +} diff --git a/sentry_streams/src/pull/stages/batch.rs b/sentry_streams/src/pull/stages/batch.rs index e75281dc..085e4d31 100644 --- a/sentry_streams/src/pull/stages/batch.rs +++ b/sentry_streams/src/pull/stages/batch.rs @@ -5,11 +5,13 @@ use sentry_arroyo::backends::kafka::types::KafkaPayload; use sentry_arroyo::processing::stream::{MessageMetadata, PipelineEnvelope, Stage, StageResult}; use sentry_arroyo::types::Partition; +use crate::pull::pipeline_value::{PipelineValue, PipelineValueCaster}; + /// Accumulates envelopes into batches and emits when the batch reaches /// max_batch_size. Returns Skip while accumulating, Emit when flushing. /// -/// The emitted envelope contains a Vec of the accumulated payloads. -/// Offsets are merged across the batch — highest offset per partition. +/// Expects PipelineValue::Raw on input. Emits PipelineValue::Rust(Vec) +/// on flush. Offsets are merged across the batch — highest offset per partition. pub struct BatchAccumulatorStage { max_batch_size: usize, state: Mutex, @@ -32,24 +34,26 @@ impl BatchState { } } - /// Add a message to the batch, merging its offset. - fn accumulate(&mut self, envelope: PipelineEnvelope) { + fn accumulate( + &mut self, + payload: KafkaPayload, + metadata: MessageMetadata, + raw: Arc, + ) { self.offsets - .entry(envelope.metadata.partition) - .and_modify(|o| *o = (*o).max(envelope.metadata.offset)) - .or_insert(envelope.metadata.offset); - self.payloads.push(envelope.payload); - self.last_metadata = Some(envelope.metadata); - self.last_raw = Some(envelope.raw); + .entry(metadata.partition) + .and_modify(|o| *o = (*o).max(metadata.offset)) + .or_insert(metadata.offset); + self.payloads.push(payload); + self.last_metadata = Some(metadata); + self.last_raw = Some(raw); } - /// Check if the batch has reached the size threshold. fn is_full(&self, max_batch_size: usize) -> bool { self.payloads.len() >= max_batch_size } - /// Drain the batch into an Envelope, clearing internal state. - fn flush(&mut self) -> PipelineEnvelope> { + fn flush(&mut self) -> PipelineEnvelope { let payloads = std::mem::take(&mut self.payloads); let mut metadata = self.last_metadata.take().unwrap(); let raw = self.last_raw.take().unwrap(); @@ -59,7 +63,7 @@ impl BatchState { } self.offsets.clear(); - PipelineEnvelope::new(payloads, metadata, raw) + PipelineEnvelope::new(PipelineValue::Rust(Box::new(payloads)), metadata, raw) } } @@ -73,15 +77,20 @@ impl BatchAccumulatorStage { } impl Stage for BatchAccumulatorStage { - type In = KafkaPayload; - type Out = Vec; + type In = PipelineValue; + type Out = PipelineValue; async fn process( &self, - envelope: PipelineEnvelope, - ) -> StageResult> { + envelope: PipelineEnvelope, + ) -> StageResult { + let envelope = match envelope.downcast_raw() { + Ok(e) => e, + Err(fail) => return fail, + }; + let mut state = self.state.lock().unwrap(); - state.accumulate(envelope); + state.accumulate(envelope.payload, envelope.metadata, envelope.raw); if state.is_full(self.max_batch_size) { StageResult::Emit(state.flush()) diff --git a/sentry_streams/src/pull/stages/header_filter.rs b/sentry_streams/src/pull/stages/header_filter.rs index 702a7cb7..ef708629 100644 --- a/sentry_streams/src/pull/stages/header_filter.rs +++ b/sentry_streams/src/pull/stages/header_filter.rs @@ -1,13 +1,13 @@ -use sentry_arroyo::backends::kafka::types::KafkaPayload; use sentry_arroyo::processing::stream::{PipelineEnvelope, RejectionReason, Stage, StageResult}; +use crate::pull::pipeline_value::{IntoPipelineValue, PipelineValue, PipelineValueCaster}; + /// Filters messages by checking a Kafka header for an expected integer value. /// Messages with a matching header pass through (Emit). /// Messages without the header or with a non-matching value are dropped (Drop). /// Messages with an unparseable header value are rejected (Reject → DLQ). /// -/// Header values are treated as UTF-8 ASCII decimal integers (matching the -/// existing push-based HeaderIntEqualityFilter in streams). +/// Expects PipelineValue::Raw on input, produces PipelineValue::Raw on output. pub struct HeaderFilterStage { header_name: String, expected_value: i64, @@ -20,40 +20,44 @@ impl HeaderFilterStage { expected_value, } } - - /// `Ok(true)` — header matches expected value. - /// `Ok(false)` — header missing, empty, or different value. - /// `Err(())` — header present but not a valid decimal integer. - fn check_header(&self, payload: &KafkaPayload) -> Result { - let Some(headers) = payload.headers() else { - return Ok(false); - }; - let Some(bytes) = headers.get(&self.header_name) else { - return Ok(false); - }; - if bytes.is_empty() { - return Ok(false); - } - let parsed = std::str::from_utf8(bytes) - .map_err(|_| ())? - .parse::() - .map_err(|_| ())?; - Ok(parsed == self.expected_value) - } } impl Stage for HeaderFilterStage { - type In = KafkaPayload; - type Out = KafkaPayload; + type In = PipelineValue; + type Out = PipelineValue; async fn process( &self, - envelope: PipelineEnvelope, - ) -> StageResult { - match self.check_header(&envelope.payload) { - Ok(true) => StageResult::Emit(envelope), - Ok(false) => StageResult::drop(envelope), - Err(()) => StageResult::reject(envelope, RejectionReason::Invalid), + envelope: PipelineEnvelope, + ) -> StageResult { + let envelope = match envelope.downcast_raw() { + Ok(e) => e, + Err(fail) => return fail, + }; + + let payload = &envelope.payload; + let check = match payload.headers() { + None => Ok(false), + Some(headers) => match headers.get(&self.header_name) { + None => Ok(false), + Some(bytes) if bytes.is_empty() => Ok(false), + Some(bytes) => std::str::from_utf8(bytes) + .map_err(|_| ()) + .and_then(|s| s.parse::().map_err(|_| ())) + .map(|v| v == self.expected_value), + }, + }; + + match check { + Ok(true) => StageResult::Emit(envelope.into_pipeline_value()), + Ok(false) => StageResult::Drop { + metadata: envelope.metadata, + }, + Err(()) => StageResult::Reject { + metadata: envelope.metadata, + raw: envelope.raw, + reason: RejectionReason::Invalid, + }, } } diff --git a/sentry_streams/src/pull/stages/mod.rs b/sentry_streams/src/pull/stages/mod.rs index fba98482..d6579791 100644 --- a/sentry_streams/src/pull/stages/mod.rs +++ b/sentry_streams/src/pull/stages/mod.rs @@ -1,2 +1,3 @@ pub mod batch; pub mod header_filter; +pub mod py_callable; diff --git a/sentry_streams/src/pull/stages/py_callable.rs b/sentry_streams/src/pull/stages/py_callable.rs new file mode 100644 index 00000000..7a84d319 --- /dev/null +++ b/sentry_streams/src/pull/stages/py_callable.rs @@ -0,0 +1,104 @@ +use pyo3::prelude::*; +use pyo3::types::PyList; +use sentry_arroyo::backends::kafka::types::KafkaPayload; +use sentry_arroyo::processing::stream::{PipelineEnvelope, Stage, StageResult}; + +use crate::pull::pipeline_value::PipelineValue; + +/// Calls a Python callable as a pipeline stage. +/// +/// Handles three input scenarios: +/// - PipelineValue::Rust(Vec) — converts to Python list of bytes, +/// calls the callable. Used for batch_parser. +/// - PipelineValue::Python — passes the Python object directly to the callable. +/// Used for processor, serializer. +/// - PipelineValue::Raw — converts single KafkaPayload to Python bytes, +/// calls the callable. Used for single-message transforms. +/// +/// Output is always PipelineValue::Python (the callable's return value). +/// +/// Python exceptions become StageResult::Fail (no DLQ for now). +pub struct PyCallableStage { + callable: Py, + stage_name: String, +} + +impl PyCallableStage { + pub fn new(callable: Py, name: impl Into) -> Self { + Self { + callable, + stage_name: name.into(), + } + } + + /// Convert a Vec to a Python list of bytes objects. + fn batch_to_python<'py>( + py: Python<'py>, + payloads: Vec, + ) -> PyResult> { + let items: Vec> = payloads + .iter() + .map(|kp| { + let bytes = kp.payload().map(|v| v.as_slice()).unwrap_or(&[]); + pyo3::types::PyBytes::new(py, bytes).into_any() + }) + .collect(); + Ok(PyList::new(py, &items)?.into_any()) + } + + /// Convert a single KafkaPayload to Python bytes. + fn raw_to_python<'py>( + py: Python<'py>, + payload: &KafkaPayload, + ) -> PyResult> { + let bytes = payload.payload().map(|v| v.as_slice()).unwrap_or(&[]); + Ok(pyo3::types::PyBytes::new(py, bytes).into_any()) + } +} + +impl Stage for PyCallableStage { + type In = PipelineValue; + type Out = PipelineValue; + + async fn process( + &self, + envelope: PipelineEnvelope, + ) -> StageResult { + let result = Python::attach(|py| -> PyResult> { + let input: Bound<'_, PyAny> = match envelope.payload { + PipelineValue::Rust(ref boxed) => { + // Try to downcast as Vec (batch) + if let Some(payloads) = boxed.downcast_ref::>() { + Self::batch_to_python(py, payloads.clone())? + } else { + return Err(pyo3::exceptions::PyTypeError::new_err( + "PyCallableStage received unsupported Rust type", + )); + } + } + PipelineValue::Python(ref obj) => obj.bind(py).clone().into_any(), + PipelineValue::Raw(ref kp) => Self::raw_to_python(py, kp)?, + }; + + self.callable.call1(py, (input,)) + }); + + match result { + Ok(output) => StageResult::Emit(PipelineEnvelope::new( + PipelineValue::Python(output), + envelope.metadata, + envelope.raw, + )), + Err(py_err) => { + // For now, all Python errors are fatal (no DLQ) + StageResult::Fail(Box::new(py_err)) + } + } + } + + fn name(&self) -> &'static str { + // Leak the string to get a &'static str. + // This is fine — stages are long-lived, created once at pipeline build time. + Box::leak(self.stage_name.clone().into_boxed_str()) + } +} From 51f8c27bc98e87a043886f93d08072df2317391d Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Thu, 6 Aug 2026 11:06:51 -0700 Subject: [PATCH 3/9] Adapter draft. --- .../adapters/arroyo/pull_adapter.py | 214 ++++++++++++++++++ .../sentry_streams/adapters/loader.py | 9 + sentry_streams/src/lib.rs | 2 + sentry_streams/src/pull/mod.rs | 3 + 4 files changed, 228 insertions(+) create mode 100644 sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py diff --git a/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py b/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py new file mode 100644 index 00000000..ace4cc29 --- /dev/null +++ b/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py @@ -0,0 +1,214 @@ +""" +Pull-based adapter for the streams pipeline DSL. + +Maps pipeline steps to PullOperator variants, which the Rust PullConsumer +converts to concrete pull-based stages. This is a drop-in replacement for +RustArroyoAdapter for pipelines that can run on the pull-based runtime. + +Currently supports the items_span steel thread pipeline: + KafkaSource → HeaderFilter → Batch → PyCallable(parse) → PyCallable(process) + → PyCallable(serialize) → GcsSink +""" + +from __future__ import annotations + +import logging +from typing import Any, Callable, Mapping, Self, Type + +from sentry_streams.pipeline.function_template import InputType, OutputType +from sentry_streams.pipeline.pipeline import ( + Batch, + Broadcast, + ComplexStep, + Filter, + FlatMap, + GCSSink, + HeadersFilter, + Map, + Reduce, + Router, + RoutingFuncReturnType, + Sink, + Source, + StreamSource, +) +from sentry_streams.pipeline.window import MeasurementUnit + +from sentry_streams.adapters.stream_adapter import PipelineConfig, StreamAdapter + +from sentry_streams.rust_streams import ( + PullConsumer, + PullOperator, + PyKafkaConsumerConfig, +) + +logger = logging.getLogger(__name__) + + +def build_kafka_consumer_config( + source_name: str, + source_config: Mapping[str, Any], + consumer_group: str | None, +) -> PyKafkaConsumerConfig: + """Build a PyKafkaConsumerConfig from pipeline step config.""" + bootstrap_servers = source_config.get( + "bootstrap_servers", source_config.get("broker_config", {}).get("bootstrap.servers", "") + ) + if isinstance(bootstrap_servers, str): + bootstrap_servers = [bootstrap_servers] + + group_id = consumer_group or source_config.get("consumer_group", f"{source_name}-consumer") + auto_offset_reset = source_config.get("auto_offset_reset", "earliest") + + # Map string offset reset to the Rust enum + from sentry_streams.rust_streams import InitialOffset + offset_map = { + "earliest": InitialOffset.Earliest, + "latest": InitialOffset.Latest, + "error": InitialOffset.Error, + } + initial_offset = offset_map.get(auto_offset_reset, InitialOffset.Earliest) + + override_params = source_config.get("override_params", None) + + return PyKafkaConsumerConfig( + bootstrap_servers=bootstrap_servers, + group_id=group_id, + auto_offset_reset=initial_offset, + strict_offset_reset=source_config.get("strict_offset_reset", False), + max_poll_interval_ms=source_config.get("max_poll_interval_ms", 300000), + override_params=override_params, + ) + + +class PullBasedAdapter(StreamAdapter[str, str]): + """ + Pull-based adapter that translates pipeline DSL steps into + PullOperator variants for the Rust pull-based runtime. + + StreamT = str (just a source name identifier, like Route in push model) + StreamSinkT = str + """ + + def __init__( + self, + steps_config: Mapping[str, Any], + ) -> None: + self._steps_config = steps_config + self._consumer: PullConsumer | None = None + self._steps: list[PullOperator] = [] + self._sink: PullOperator | None = None + + @classmethod + def build(cls, config: PipelineConfig) -> Self: # type: ignore[override] + steps_config = config.get("steps_config", {}) + return cls(steps_config) + + def complex_step_override( + self, + ) -> dict[Type[ComplexStep[Any, Any]], Callable[[ComplexStep[Any, Any]], str]]: + return {} + + def source(self, step: Source[Any]) -> str: + assert isinstance(step, StreamSource) + source_name = step.name + source_config = self._steps_config.get(source_name) + assert source_config is not None, f"Config not provided for source {source_name}" + + step_config: Mapping[str, Any] = self._steps_config.get(source_name, {}) + step.override_config(step_config) + step.validate() + + kafka_config = build_kafka_consumer_config( + source_name, source_config, step.consumer_group + ) + self._consumer = PullConsumer( + consumer_config=kafka_config, + topic=step.stream_name, + ) + + return source_name + + def sink(self, step: Sink[Any], stream: str) -> str: + if isinstance(step, GCSSink): + self._sink = PullOperator.GcsSink( + bucket=step.bucket, + object_generator=step.object_generator, + ) + else: + raise NotImplementedError( + f"PullBasedAdapter does not support sink type: {type(step).__name__}" + ) + return stream + + def map(self, step: Map[Any, Any], stream: str) -> str: + self._steps.append( + PullOperator.PyCallable( + callable=step.function, + name=step.name, + ) + ) + return stream + + def flat_map(self, step: FlatMap[Any, Any], stream: str) -> str: + raise NotImplementedError("PullBasedAdapter does not support flat_map") + + def filter(self, step: Filter[Any], stream: str) -> str: + if isinstance(step, HeadersFilter): + self._steps.append( + PullOperator.HeaderFilter( + header_name=step.header_name, + expected_value=step.value, + ) + ) + else: + # PredicateFilter — wrap as PyCallable + self._steps.append( + PullOperator.PyCallable( + callable=step.function, + name=step.name, + ) + ) + return stream + + def reduce( + self, + step: Reduce[MeasurementUnit, InputType, OutputType], + stream: str, + ) -> str: + if isinstance(step, Batch): + self._steps.append( + PullOperator.Batch(max_batch_size=step.batch_size) + ) + else: + raise NotImplementedError( + f"PullBasedAdapter does not support reduce type: {type(step).__name__}" + ) + return stream + + def router( + self, + step: Router[RoutingFuncReturnType, Any], + stream: str, + ) -> Mapping[str, str]: + raise NotImplementedError("PullBasedAdapter does not support router") + + def broadcast( + self, + step: Broadcast[Any], + stream: str, + ) -> Mapping[str, str]: + raise NotImplementedError("PullBasedAdapter does not support broadcast") + + def run(self) -> None: + assert self._consumer is not None, "No source configured" + assert self._sink is not None, "No sink configured" + + logger.info( + "Starting pull-based pipeline with %d steps", len(self._steps) + ) + self._consumer.run(steps=self._steps, sink=self._sink) + + def shutdown(self) -> None: + # TODO: signal the pipeline to stop gracefully + pass diff --git a/sentry_streams/sentry_streams/adapters/loader.py b/sentry_streams/sentry_streams/adapters/loader.py index 0bc04d53..e66e0f3e 100644 --- a/sentry_streams/sentry_streams/adapters/loader.py +++ b/sentry_streams/sentry_streams/adapters/loader.py @@ -65,6 +65,15 @@ def load_adapter( StreamAdapter[Stream, Sink], RustArroyoAdapter.build(config, metrics_config), ) + + if adapter_type == "pull": + from sentry_streams.adapters.arroyo.pull_adapter import PullBasedAdapter + + return cast( + StreamAdapter[Stream, Sink], + PullBasedAdapter.build(config), + ) + else: mod, cls = adapter_type.rsplit(".", 1) diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index bb2e3b59..b7072422 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -52,5 +52,7 @@ fn rust_streams(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index 95292ca4..9bf52768 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -1,6 +1,9 @@ pub mod gcs_client; pub mod gcs_sink_handler; +pub mod pipeline_stage; pub mod pipeline_value; +pub mod pull_consumer; +pub mod pull_operator; pub mod stages; #[cfg(test)] From fc6d776da7f599c6821df4381bf930867a769225 Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:15:37 -0700 Subject: [PATCH 4/9] Formalize up SBC stages with tests. --- .../adapters/arroyo/pull_adapter.py | 59 +-- .../sentry_streams/rust_streams.pyi | 31 ++ sentry_streams/src/lib.rs | 2 + sentry_streams/src/pull/gcs_client.rs | 10 +- sentry_streams/src/pull/gcs_sink_handler.rs | 22 +- sentry_streams/src/pull/message_wrapper.rs | 53 +++ sentry_streams/src/pull/mod.rs | 217 ++++++++-- sentry_streams/src/pull/pipeline_sink.rs | 56 +++ sentry_streams/src/pull/pipeline_stage.rs | 41 ++ .../src/pull/pipeline_value_converter.rs | 73 ++++ sentry_streams/src/pull/pull_consumer.rs | 124 ++++++ sentry_streams/src/pull/pull_operator.rs | 81 ++++ sentry_streams/src/pull/pull_source.rs | 144 +++++++ sentry_streams/src/pull/stages/py_callable.rs | 83 +--- .../adapters/arroyo/test_pull_adapter.py | 375 ++++++++++++++++++ 15 files changed, 1226 insertions(+), 145 deletions(-) create mode 100644 sentry_streams/src/pull/message_wrapper.rs create mode 100644 sentry_streams/src/pull/pipeline_sink.rs create mode 100644 sentry_streams/src/pull/pipeline_stage.rs create mode 100644 sentry_streams/src/pull/pipeline_value_converter.rs create mode 100644 sentry_streams/src/pull/pull_consumer.rs create mode 100644 sentry_streams/src/pull/pull_operator.rs create mode 100644 sentry_streams/src/pull/pull_source.rs create mode 100644 sentry_streams/tests/adapters/arroyo/test_pull_adapter.py diff --git a/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py b/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py index ace4cc29..b634e5e2 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py +++ b/sentry_streams/sentry_streams/adapters/arroyo/pull_adapter.py @@ -15,6 +15,7 @@ import logging from typing import Any, Callable, Mapping, Self, Type +from sentry_streams.adapters.stream_adapter import PipelineConfig, StreamAdapter from sentry_streams.pipeline.function_template import InputType, OutputType from sentry_streams.pipeline.pipeline import ( Batch, @@ -33,12 +34,10 @@ StreamSource, ) from sentry_streams.pipeline.window import MeasurementUnit - -from sentry_streams.adapters.stream_adapter import PipelineConfig, StreamAdapter - from sentry_streams.rust_streams import ( PullConsumer, PullOperator, + PullSourceConfig, PyKafkaConsumerConfig, ) @@ -62,12 +61,13 @@ def build_kafka_consumer_config( # Map string offset reset to the Rust enum from sentry_streams.rust_streams import InitialOffset + offset_map = { - "earliest": InitialOffset.Earliest, - "latest": InitialOffset.Latest, - "error": InitialOffset.Error, + "earliest": InitialOffset.earliest, + "latest": InitialOffset.latest, + "error": InitialOffset.error, } - initial_offset = offset_map.get(auto_offset_reset, InitialOffset.Earliest) + initial_offset = offset_map.get(auto_offset_reset, InitialOffset.earliest) override_params = source_config.get("override_params", None) @@ -95,7 +95,9 @@ def __init__( steps_config: Mapping[str, Any], ) -> None: self._steps_config = steps_config - self._consumer: PullConsumer | None = None + self._kafka_config: PyKafkaConsumerConfig | None = None + self._topic: str = "" + self._schema: str | None = None self._steps: list[PullOperator] = [] self._sink: PullOperator | None = None @@ -119,13 +121,11 @@ def source(self, step: Source[Any]) -> str: step.override_config(step_config) step.validate() - kafka_config = build_kafka_consumer_config( + self._kafka_config = build_kafka_consumer_config( source_name, source_config, step.consumer_group ) - self._consumer = PullConsumer( - consumer_config=kafka_config, - topic=step.stream_name, - ) + self._topic = step.stream_name + self._schema = step.stream_name # schema name matches topic for codec lookup return source_name @@ -142,10 +142,12 @@ def sink(self, step: Sink[Any], stream: str) -> str: return stream def map(self, step: Map[Any, Any], stream: str) -> str: + assert self._schema is not None, "source() must be called before map()" self._steps.append( PullOperator.PyCallable( - callable=step.function, + callable=step.resolved_function, name=step.name, + schema=self._schema, ) ) return stream @@ -162,12 +164,8 @@ def filter(self, step: Filter[Any], stream: str) -> str: ) ) else: - # PredicateFilter — wrap as PyCallable - self._steps.append( - PullOperator.PyCallable( - callable=step.function, - name=step.name, - ) + raise NotImplementedError( + f"PullBasedAdapter does not support filter type: {type(step).__name__}" ) return stream @@ -177,9 +175,8 @@ def reduce( stream: str, ) -> str: if isinstance(step, Batch): - self._steps.append( - PullOperator.Batch(max_batch_size=step.batch_size) - ) + assert step.batch_size is not None, "Batch requires batch_size" + self._steps.append(PullOperator.Batch(max_batch_size=step.batch_size)) else: raise NotImplementedError( f"PullBasedAdapter does not support reduce type: {type(step).__name__}" @@ -201,13 +198,19 @@ def broadcast( raise NotImplementedError("PullBasedAdapter does not support broadcast") def run(self) -> None: - assert self._consumer is not None, "No source configured" - assert self._sink is not None, "No sink configured" + assert self._kafka_config is not None, "No source configured" - logger.info( - "Starting pull-based pipeline with %d steps", len(self._steps) + logger.info("Starting pull-based pipeline with %d steps", len(self._steps)) + source = PullSourceConfig.Kafka( + config=self._kafka_config, + topic=self._topic, + ) + consumer = PullConsumer( + source=source, + steps=self._steps, + sink=self._sink, ) - self._consumer.run(steps=self._steps, sink=self._sink) + consumer.run() def shutdown(self) -> None: # TODO: signal the pipeline to stop gracefully diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index e4c359cf..c9cbdfb4 100644 --- a/sentry_streams/sentry_streams/rust_streams.pyi +++ b/sentry_streams/sentry_streams/rust_streams.pyi @@ -211,3 +211,34 @@ class PyWatermark: def timestamp(self) -> int: ... @property def last_message_time(self) -> float | None: ... + +class PullOperator: + @classmethod + def HeaderFilter(cls, header_name: str, expected_value: int) -> Self: ... + @classmethod + def Batch(cls, max_batch_size: int) -> Self: ... + @classmethod + def PyCallable(cls, callable: Callable[..., Any], name: str, schema: str) -> Self: ... + @classmethod + def GcsSink(cls, bucket: str, object_generator: Callable[[], str]) -> Self: ... + @classmethod + def MockSink(cls) -> Self: ... + +class PyTestMessage: + def __init__(self, payload: bytes, headers: Mapping[str, bytes] | None = None) -> None: ... + +class PullSourceConfig: + @classmethod + def Kafka(cls, config: PyKafkaConsumerConfig, topic: str) -> Self: ... + @classmethod + def Test(cls, messages: Sequence[PyTestMessage]) -> Self: ... + +class PullConsumer: + def __init__( + self, + source: PullSourceConfig, + steps: Sequence[PullOperator], + sink: PullOperator | None = None, + ) -> None: ... + def run(self) -> None: ... + def get_mock_sink_results(self) -> list[bytes]: ... diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index b7072422..4b1f092e 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -54,5 +54,7 @@ fn rust_streams(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/sentry_streams/src/pull/gcs_client.rs b/sentry_streams/src/pull/gcs_client.rs index 57764704..8d7417c2 100644 --- a/sentry_streams/src/pull/gcs_client.rs +++ b/sentry_streams/src/pull/gcs_client.rs @@ -40,16 +40,10 @@ impl GcsClient { } /// Upload bytes to GCS as the given object name. - pub async fn upload( - &self, - object_name: &str, - bytes: &[u8], - ) -> Result<(), GcsError> { + pub async fn upload(&self, object_name: &str, bytes: &[u8]) -> Result<(), GcsError> { let auth_provider = self .auth_provider - .get_or_init(|| async { - provider().await.expect("Failed to get gcp_auth provider") - }) + .get_or_init(|| async { provider().await.expect("Failed to get gcp_auth provider") }) .await; let scopes = &["https://www.googleapis.com/auth/devstorage.read_write"]; diff --git a/sentry_streams/src/pull/gcs_sink_handler.rs b/sentry_streams/src/pull/gcs_sink_handler.rs index 163681d5..d7343fd6 100644 --- a/sentry_streams/src/pull/gcs_sink_handler.rs +++ b/sentry_streams/src/pull/gcs_sink_handler.rs @@ -4,6 +4,7 @@ use sentry_arroyo::processing::stream::PipelineEnvelope; use super::gcs_client::GcsClient; use super::pipeline_value::PipelineValue; +use super::pipeline_value_converter::PipelineValueConverter; /// Sink handler that uploads pipeline output to GCS. /// @@ -23,25 +24,6 @@ impl GcsSinkHandler { } } - /// Extract bytes from the pipeline value. - fn extract_bytes(value: &PipelineValue) -> Result, Box> { - match value { - PipelineValue::Python(obj) => Python::attach(|py| { - obj.extract::>(py) - .map_err(|e| Box::new(e) as Box) - }), - PipelineValue::Rust(boxed) => { - boxed.downcast_ref::>().cloned().ok_or_else(|| { - Box::new(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "GcsSinkHandler expected Rust Vec", - )) as Box - }) - } - PipelineValue::Raw(kp) => Ok(kp.payload().map(|v| v.to_vec()).unwrap_or_default()), - } - } - fn generate_object_name(&self) -> Result> { Python::attach(|py| { let result = self @@ -60,7 +42,7 @@ impl NextHandler for GcsSinkHandler { &self, envelope: &PipelineEnvelope, ) -> Result<(), Box> { - let bytes = Self::extract_bytes(&envelope.payload)?; + let bytes = PipelineValueConverter::extract_bytes(&envelope.payload)?; let object_name = self.generate_object_name()?; self.client .upload(&object_name, &bytes) diff --git a/sentry_streams/src/pull/message_wrapper.rs b/sentry_streams/src/pull/message_wrapper.rs new file mode 100644 index 00000000..49f0edd3 --- /dev/null +++ b/sentry_streams/src/pull/message_wrapper.rs @@ -0,0 +1,53 @@ +use pyo3::prelude::*; +use pyo3::types::PyList; + +/// Handles wrapping/unwrapping of streams `PyMessage` objects. +/// +/// The streams pipeline convention is that all Python callables receive +/// a `Message[T]` wrapper with `.payload`, `.headers`, `.timestamp`, +/// `.schema` attributes. This struct encapsulates that protocol. +pub struct MessageWrapper; + +impl MessageWrapper { + /// Wrap a Python value in a PyMessage if not already wrapped. + /// Returns the input unchanged if it's already a PyMessage instance. + pub fn ensure<'py>( + py: Python<'py>, + input: Bound<'py, PyAny>, + timestamp: f64, + schema: &str, + ) -> PyResult> { + let message_cls = Self::get_message_class(py)?; + + if input.is_instance(&message_cls)? { + Ok(input) + } else { + let headers = PyList::empty(py); + let schema_py = schema.into_pyobject(py)?.into_any(); + message_cls.call1((input, headers, timestamp, schema_py)) + } + } + + /// Re-wrap a callable's return value in a new PyMessage, preserving + /// headers, timestamp, and schema from the original message. + pub fn rewrap<'py>( + py: Python<'py>, + result: Bound<'py, PyAny>, + original: &Bound<'py, PyAny>, + ) -> PyResult> { + let message_cls = Self::get_message_class(py)?; + + let headers = original.getattr("headers")?; + let timestamp = original.getattr("timestamp")?; + let schema = original.getattr("schema")?; + + message_cls + .call1((result, headers, timestamp, schema)) + .map(|r| r.unbind()) + } + + fn get_message_class<'py>(py: Python<'py>) -> PyResult> { + py.import("sentry_streams.pipeline.message")? + .getattr("PyMessage") + } +} diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index 9bf52768..101f9f71 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -1,9 +1,13 @@ pub mod gcs_client; pub mod gcs_sink_handler; +pub mod message_wrapper; +pub mod pipeline_sink; pub mod pipeline_stage; pub mod pipeline_value; +pub mod pipeline_value_converter; pub mod pull_consumer; pub mod pull_operator; +pub mod pull_source; pub mod stages; #[cfg(test)] @@ -73,10 +77,8 @@ mod tests { header_name: &str, header_value: i64, ) -> StageResult { - let headers = Headers::new().insert( - header_name, - Some(header_value.to_string().into_bytes()), - ); + let headers = + Headers::new().insert(header_name, Some(header_value.to_string().into_bytes())); let kafka_payload = KafkaPayload::new(None, Some(headers), Some(payload.to_vec())); let metadata = MessageMetadata { partition: Partition::new(Topic::new("test"), 0), @@ -97,10 +99,7 @@ mod tests { header_name: &str, header_value: &[u8], ) -> StageResult { - let headers = Headers::new().insert( - header_name, - Some(header_value.to_vec()), - ); + let headers = Headers::new().insert(header_name, Some(header_value.to_vec())); let kafka_payload = KafkaPayload::new(None, Some(headers), Some(payload.to_vec())); let metadata = MessageMetadata { partition: Partition::new(Topic::new("test"), 0), @@ -124,11 +123,11 @@ mod tests { let error_handler = LogHandler; let messages = vec![ - make_envelope_with_header(b"span-1", 0, "item_type", 1), // match - make_envelope(b"no-header", 1), // no header → drop - make_envelope_with_header(b"span-2", 2, "item_type", 1), // match - make_envelope_with_header(b"log-1", 3, "item_type", 2), // wrong value → drop - make_envelope_with_header(b"span-3", 4, "item_type", 1), // match + make_envelope_with_header(b"span-1", 0, "item_type", 1), // match + make_envelope(b"no-header", 1), // no header → drop + make_envelope_with_header(b"span-2", 2, "item_type", 1), // match + make_envelope_with_header(b"log-1", 3, "item_type", 2), // wrong value → drop + make_envelope_with_header(b"span-3", 4, "item_type", 1), // match ]; let result = stream::iter(messages) @@ -185,9 +184,7 @@ mod tests { let error_handler = LogHandler; let messages: Vec> = (0..5) - .map(|i| make_envelope_with_header( - format!("msg-{i}").as_bytes(), i, "item_type", 1, - )) + .map(|i| make_envelope_with_header(format!("msg-{i}").as_bytes(), i, "item_type", 1)) .collect(); // Count batch sizes via a simple stage that downcasts the Rust batch @@ -215,7 +212,9 @@ mod tests { typed.raw, )) } - fn name(&self) -> &'static str { "count_batch" } + fn name(&self) -> &'static str { + "count_batch" + } } let counter = CountBatchStage { sizes: bs }; @@ -231,14 +230,21 @@ mod tests { assert!(result.is_ok()); let sizes = batch_sizes.lock().unwrap(); - assert_eq!(*sizes, vec![3], "Expected one batch of 3 (remaining 2 not flushed)"); + assert_eq!( + *sizes, + vec![3], + "Expected one batch of 3 (remaining 2 not flushed)" + ); let committed = committer.committed(); assert!(!committed.is_empty()); let last = committed.last().unwrap(); let partition = Partition::new(Topic::new("test"), 0); - assert_eq!(last.get(&partition), Some(&3), - "Expected offset 3 (batch last offset 2 + 1)"); + assert_eq!( + last.get(&partition), + Some(&3), + "Expected offset 3 (batch last offset 2 + 1)" + ); } // ── Full pipeline integration test ────────────────────────────── @@ -269,7 +275,9 @@ mod tests { Ok(t) => t, Err(fail) => return fail, }; - let contents: Vec> = typed.payload.iter() + let contents: Vec> = typed + .payload + .iter() .map(|kp| kp.payload().map(|v| v.to_vec()).unwrap_or_default()) .collect(); self.batches.lock().unwrap().push(contents); @@ -279,16 +287,18 @@ mod tests { typed.raw, )) } - fn name(&self) -> &'static str { "collect_batch" } + fn name(&self) -> &'static str { + "collect_batch" + } } - let collector = CollectBatchStage { batches: batches_clone }; + let collector = CollectBatchStage { + batches: batches_clone, + }; // 4 messages with matching header → 2 batches of 2 let messages: Vec> = (0..4) - .map(|i| make_envelope_with_header( - format!("span-{i}").as_bytes(), i, "item_type", 1, - )) + .map(|i| make_envelope_with_header(format!("span-{i}").as_bytes(), i, "item_type", 1)) .collect(); let result = stream::iter(messages) @@ -306,4 +316,159 @@ mod tests { assert_eq!(collected[0], vec![b"span-0".to_vec(), b"span-1".to_vec()]); assert_eq!(collected[1], vec![b"span-2".to_vec(), b"span-3".to_vec()]); } + + // ── PullConsumer e2e test ─────────────────────────────────────── + + use super::pull_consumer::PullConsumer; + use super::pull_operator::PullOperator; + use super::pull_source::PullSource; + use futures::stream::Stream; + use std::pin::Pin; + + /// Test source that drains its messages on first stream() call. + /// Wraps committer in Arc so it can be inspected after run. + struct TestSource { + messages: Mutex>>, + committer: Arc, + } + + impl TestSource { + fn new(messages: Vec>) -> (Self, Arc) { + let committer = Arc::new(MockCommitter::new()); + let source = Self { + messages: Mutex::new(messages), + committer: committer.clone(), + }; + (source, committer) + } + } + + impl PullSource for TestSource { + fn stream(&self) -> Pin> + '_>> { + let messages: Vec<_> = self.messages.lock().unwrap().drain(..).collect(); + Box::pin(futures::stream::iter(messages)) + } + + fn committer(&self) -> &dyn OffsetCommitter { + self.committer.as_ref() + } + } + + /// Helper to create a raw StageResult (not wrapped in PipelineValue). + fn make_raw_envelope(payload: &[u8], offset: u64) -> StageResult { + let kp = KafkaPayload::new(None, None, Some(payload.to_vec())); + let md = MessageMetadata { + partition: Partition::new(Topic::new("test"), 0), + offset, + timestamp: chrono::Utc::now(), + }; + StageResult::Emit(PipelineEnvelope::new(kp.clone(), md, Arc::new(kp))) + } + + fn make_raw_envelope_with_header( + payload: &[u8], + offset: u64, + header_name: &str, + header_value: i64, + ) -> StageResult { + let headers = + Headers::new().insert(header_name, Some(header_value.to_string().into_bytes())); + let kp = KafkaPayload::new(None, Some(headers), Some(payload.to_vec())); + let md = MessageMetadata { + partition: Partition::new(Topic::new("test"), 0), + offset, + timestamp: chrono::Utc::now(), + }; + StageResult::Emit(PipelineEnvelope::new(kp.clone(), md, Arc::new(kp))) + } + + #[tokio::test] + async fn test_pull_consumer_e2e_filter_and_batch() { + // 6 messages: 4 with matching header, 2 without + let messages = vec![ + make_raw_envelope_with_header(b"span-0", 0, "item_type", 1), + make_raw_envelope(b"no-header", 1), + make_raw_envelope_with_header(b"span-1", 2, "item_type", 1), + make_raw_envelope_with_header(b"span-2", 3, "item_type", 1), + make_raw_envelope_with_header(b"span-3", 4, "item_type", 2), // wrong value + make_raw_envelope_with_header(b"span-4", 5, "item_type", 1), + ]; + + let (source, committer) = TestSource::new(messages); + + let stages = pyo3::Python::attach(|py| { + vec![ + PullOperator::HeaderFilter { + header_name: "item_type".into(), + expected_value: 1, + }, + PullOperator::Batch { max_batch_size: 2 }, + ] + .iter() + .map(|op| op.build_stage(py)) + .collect() + }); + + let consumer = PullConsumer::with_source(source, stages, None); + + let result = consumer.run_pipeline().await; + assert!(result.is_ok()); + + let committed = committer.committed(); + assert!(!committed.is_empty(), "Expected at least one commit"); + let last = committed.last().unwrap(); + let partition = Partition::new(Topic::new("test"), 0); + assert_eq!(last.get(&partition), Some(&6)); + } + + // ── Test #3: e2e with sink ────────────────────────────────────── + + use super::pipeline_sink::{MockSinkHandler, PipelineSink}; + + #[tokio::test] + async fn test_pull_consumer_e2e_with_sink() { + // 4 messages with matching header → batch of 2 → 2 batches emitted → 2 sink calls + let messages = vec![ + make_raw_envelope_with_header(b"span-0", 0, "item_type", 1), + make_raw_envelope_with_header(b"span-1", 1, "item_type", 1), + make_raw_envelope_with_header(b"span-2", 2, "item_type", 1), + make_raw_envelope_with_header(b"span-3", 3, "item_type", 1), + ]; + + let (source, committer) = TestSource::new(messages); + + let stages = pyo3::Python::attach(|py| { + vec![ + PullOperator::HeaderFilter { + header_name: "item_type".into(), + expected_value: 1, + }, + PullOperator::Batch { max_batch_size: 2 }, + ] + .iter() + .map(|op| op.build_stage(py)) + .collect() + }); + + let mock_sink = MockSinkHandler::new(); + let consumer = + PullConsumer::with_source(source, stages, Some(PipelineSink::Mock(mock_sink))); + + let result = consumer.run_pipeline().await; + assert!(result.is_ok()); + + // Verify sink was called twice (2 batches of 2) + let results = match &consumer.sink { + Some(PipelineSink::Mock(h)) => h.get_results(), + _ => panic!("Expected MockSink"), + }; + assert_eq!(results.len(), 2, "Expected 2 sink calls (2 batches of 2)"); + + // Verify offsets committed + let committed = committer.committed(); + assert!(!committed.is_empty()); + let last = committed.last().unwrap(); + let partition = Partition::new(Topic::new("test"), 0); + assert_eq!(last.get(&partition), Some(&4)); + } } diff --git a/sentry_streams/src/pull/pipeline_sink.rs b/sentry_streams/src/pull/pipeline_sink.rs new file mode 100644 index 00000000..30b97cd7 --- /dev/null +++ b/sentry_streams/src/pull/pipeline_sink.rs @@ -0,0 +1,56 @@ +use std::sync::{Arc, Mutex}; + +use sentry_arroyo::processing::stream::handlers::next::NextHandler; +use sentry_arroyo::processing::stream::PipelineEnvelope; + +use super::gcs_sink_handler::GcsSinkHandler; +use super::pipeline_value::PipelineValue; +use super::pipeline_value_converter::PipelineValueConverter; + +/// Enum dispatch for pipeline sinks. Mirrors PipelineStage pattern. +pub enum PipelineSink { + Gcs(GcsSinkHandler), + /// Test sink that captures received payloads. + Mock(MockSinkHandler), +} + +impl NextHandler for PipelineSink { + async fn handle( + &self, + envelope: &PipelineEnvelope, + ) -> Result<(), Box> { + match self { + PipelineSink::Gcs(h) => h.handle(envelope).await, + PipelineSink::Mock(h) => h.handle(envelope).await, + } + } +} + +/// Mock sink handler that records what it receives as extracted bytes. +pub struct MockSinkHandler { + results: Arc>>>, +} + +impl MockSinkHandler { + pub fn new() -> Self { + Self { + results: Arc::new(Mutex::new(Vec::new())), + } + } + + pub fn get_results(&self) -> Vec> { + self.results.lock().unwrap().clone() + } +} + +impl NextHandler for MockSinkHandler { + async fn handle( + &self, + envelope: &PipelineEnvelope, + ) -> Result<(), Box> { + let bytes = PipelineValueConverter::extract_bytes(&envelope.payload) + .unwrap_or_else(|_| b"".to_vec()); + self.results.lock().unwrap().push(bytes); + Ok(()) + } +} diff --git a/sentry_streams/src/pull/pipeline_stage.rs b/sentry_streams/src/pull/pipeline_stage.rs new file mode 100644 index 00000000..50d81882 --- /dev/null +++ b/sentry_streams/src/pull/pipeline_stage.rs @@ -0,0 +1,41 @@ +use sentry_arroyo::processing::stream::{PipelineEnvelope, Stage, StageResult}; + +use super::pipeline_value::PipelineValue; +use super::stages::batch::BatchAccumulatorStage; +use super::stages::header_filter::HeaderFilterStage; +use super::stages::py_callable::PyCallableStage; + +/// Enum dispatch for pipeline stages, avoiding trait object limitations. +/// +/// `Stage` is not object-safe (returns `impl Future`), so we can't use +/// `Box`. Instead, this enum wraps all concrete stage types +/// and delegates `process()` via match. Zero-cost dispatch. +pub enum PipelineStage { + HeaderFilter(HeaderFilterStage), + Batch(BatchAccumulatorStage), + PyCallable(PyCallableStage), +} + +impl Stage for PipelineStage { + type In = PipelineValue; + type Out = PipelineValue; + + async fn process( + &self, + envelope: PipelineEnvelope, + ) -> StageResult { + match self { + PipelineStage::HeaderFilter(s) => s.process(envelope).await, + PipelineStage::Batch(s) => s.process(envelope).await, + PipelineStage::PyCallable(s) => s.process(envelope).await, + } + } + + fn name(&self) -> &'static str { + match self { + PipelineStage::HeaderFilter(s) => s.name(), + PipelineStage::Batch(s) => s.name(), + PipelineStage::PyCallable(s) => s.name(), + } + } +} diff --git a/sentry_streams/src/pull/pipeline_value_converter.rs b/sentry_streams/src/pull/pipeline_value_converter.rs new file mode 100644 index 00000000..a06b07a3 --- /dev/null +++ b/sentry_streams/src/pull/pipeline_value_converter.rs @@ -0,0 +1,73 @@ +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyList}; +use sentry_arroyo::backends::kafka::types::KafkaPayload; + +use super::pipeline_value::PipelineValue; + +/// Converts PipelineValue to/from Python objects and byte arrays. +/// Keeps conversion logic out of the PipelineValue enum and the stages. +pub struct PipelineValueConverter; + +impl PipelineValueConverter { + /// Convert a PipelineValue to a Python object. + pub fn to_python<'py>(value: &PipelineValue, py: Python<'py>) -> PyResult> { + match value { + PipelineValue::Raw(kp) => { + let bytes = kp.payload().map(|v| v.as_slice()).unwrap_or(&[]); + Ok(PyBytes::new(py, bytes).into_any()) + } + PipelineValue::Rust(boxed) => { + if let Some(payloads) = boxed.downcast_ref::>() { + let items: PyResult>> = payloads + .iter() + .map(|kp| { + let bytes = kp.payload().map(|v| v.as_slice()).unwrap_or(&[]); + Ok(PyBytes::new(py, bytes).into_any()) + }) + .collect(); + Ok(PyList::new(py, &items?)?.into_any()) + } else { + Err(pyo3::exceptions::PyTypeError::new_err( + "PipelineValueConverter: unsupported Rust type for Python conversion", + )) + } + } + PipelineValue::Python(obj) => Ok(obj.bind(py).clone().into_any()), + } + } + + /// Extract raw bytes from a PipelineValue. + /// Unwraps Message wrappers (objects with .payload attribute) automatically. + pub fn extract_bytes( + value: &PipelineValue, + ) -> Result, Box> { + match value { + PipelineValue::Raw(kp) => Ok(kp.payload().map(|v| v.to_vec()).unwrap_or_default()), + PipelineValue::Rust(boxed) => { + boxed.downcast_ref::>().cloned().ok_or_else(|| { + Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "PipelineValueConverter: expected Rust Vec", + )) as Box + }) + } + PipelineValue::Python(obj) => Python::attach(|py| { + let bound = obj.bind(py); + + // Unwrap Message wrapper if present (.payload attribute) + let inner = if let Ok(payload) = bound.getattr("payload") { + payload + } else { + bound.clone() + }; + + inner.extract::>().map_err(|e| { + Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("PipelineValueConverter: failed to extract bytes: {e}"), + )) as Box + }) + }), + } + } +} diff --git a/sentry_streams/src/pull/pull_consumer.rs b/sentry_streams/src/pull/pull_consumer.rs new file mode 100644 index 00000000..c7cdf23a --- /dev/null +++ b/sentry_streams/src/pull/pull_consumer.rs @@ -0,0 +1,124 @@ +use std::pin::Pin; +use std::time::Duration; + +use futures::stream::Stream; +use futures::StreamExt; +use pyo3::prelude::*; +use sentry_arroyo::processing::strategies::offset_tracker::OffsetTracker; +use sentry_arroyo::processing::stream::{LogHandler, PipelineExt, StageResult}; + +use super::pipeline_sink::PipelineSink; +use super::pipeline_stage::PipelineStage; +use super::pipeline_value::PipelineValue; +use super::pull_operator::PullOperator; +use super::pull_source::{PullSource, PullSourceConfig}; + +/// Pull-based pipeline consumer. Fully configured at construction time. +/// `run()` is parameterless. +#[pyclass] +pub struct PullConsumer { + source: Box, + stages: Vec, + pub sink: Option, +} + +#[pymethods] +impl PullConsumer { + #[new] + fn new( + py: Python<'_>, + source: Py, + steps: Vec>, + sink: Option>, + ) -> Self { + let source = source.get().build(py); + let stages = steps.iter().map(|op| op.get().build_stage(py)).collect(); + let sink = sink.map(|s| s.get().build_sink(py)); + + Self { + source, + stages, + sink, + } + } + + /// Run the pipeline. Blocks until completion or fatal error. + fn run(&self) -> PyResult<()> { + let rt = tokio::runtime::Runtime::new().map_err(|e| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "Failed to create tokio runtime: {e}" + )) + })?; + + rt.block_on(self.run_pipeline()) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}"))) + } + + /// Get results captured by a MockSink. Returns a list of byte arrays. + /// Raises RuntimeError if the sink is not a MockSink. + fn get_mock_sink_results(&self) -> PyResult>> { + match &self.sink { + Some(PipelineSink::Mock(handler)) => Ok(handler.get_results()), + _ => Err(pyo3::exceptions::PyRuntimeError::new_err( + "Sink is not a MockSink", + )), + } + } +} + +impl PullConsumer { + /// Construct from Rust with an injected source (for testing). + pub fn with_source( + source: impl PullSource + 'static, + stages: Vec, + sink: Option, + ) -> Self { + Self { + source: Box::new(source), + stages, + sink, + } + } + + pub async fn run_pipeline(&self) -> Result<(), Box> { + let committer = self.source.committer(); + let error_handler = LogHandler; + let commit_interval = Duration::from_secs(5); + let mut tracker = OffsetTracker::new(commit_interval, committer); + + let source_stream = self.source.stream().map(|result| match result { + StageResult::Emit(envelope) => { + StageResult::Emit(envelope.map_payload(PipelineValue::Raw)) + } + StageResult::Drop { metadata } => StageResult::Drop { metadata }, + StageResult::Skip => StageResult::Skip, + StageResult::Reject { + metadata, + raw, + reason, + } => StageResult::Reject { + metadata, + raw, + reason, + }, + StageResult::Fail(e) => StageResult::Fail(e), + }); + + let mut stream: Pin> + '_>> = + Box::pin(source_stream); + + for stage in &self.stages { + stream = Box::pin(stream.apply(stage)); + } + + if let Some(sink) = &self.sink { + stream + .on_next(sink) + .on_reject(&error_handler) + .commit(&mut tracker) + .await + } else { + stream.on_reject(&error_handler).commit(&mut tracker).await + } + } +} diff --git a/sentry_streams/src/pull/pull_operator.rs b/sentry_streams/src/pull/pull_operator.rs new file mode 100644 index 00000000..e605b6dd --- /dev/null +++ b/sentry_streams/src/pull/pull_operator.rs @@ -0,0 +1,81 @@ +use pyo3::prelude::*; + +use super::gcs_client::GcsClient; +use super::gcs_sink_handler::GcsSinkHandler; +use super::pipeline_sink::{MockSinkHandler, PipelineSink}; +use super::pipeline_stage::PipelineStage; +use super::stages::batch::BatchAccumulatorStage; +use super::stages::header_filter::HeaderFilterStage; +use super::stages::py_callable::PyCallableStage; + +/// Operator enum passed from Python to describe a pipeline step. +#[pyclass] +pub enum PullOperator { + #[pyo3(constructor = (header_name, expected_value))] + HeaderFilter { + header_name: String, + expected_value: i64, + }, + + #[pyo3(constructor = (max_batch_size))] + Batch { max_batch_size: usize }, + + #[pyo3(constructor = (callable, name, schema))] + PyCallable { + callable: Py, + name: String, + schema: String, + }, + + #[pyo3(constructor = (bucket, object_generator))] + GcsSink { + bucket: String, + object_generator: Py, + }, + + #[pyo3(constructor = ())] + MockSink {}, +} + +impl PullOperator { + pub fn build_stage(&self, py: Python<'_>) -> PipelineStage { + match self { + PullOperator::HeaderFilter { + header_name, + expected_value, + } => PipelineStage::HeaderFilter(HeaderFilterStage::new( + header_name.clone(), + *expected_value, + )), + PullOperator::Batch { max_batch_size } => { + PipelineStage::Batch(BatchAccumulatorStage::new(*max_batch_size)) + } + PullOperator::PyCallable { + callable, + name, + schema, + } => PipelineStage::PyCallable(PyCallableStage::new( + callable.clone_ref(py), + name.clone(), + schema.clone(), + )), + PullOperator::GcsSink { .. } | PullOperator::MockSink { .. } => { + panic!("Sink operators are not stages — use build_sink()") + } + } + } + + pub fn build_sink(&self, py: Python<'_>) -> PipelineSink { + match self { + PullOperator::GcsSink { + bucket, + object_generator, + } => { + let client = GcsClient::with_defaults(bucket.clone()); + PipelineSink::Gcs(GcsSinkHandler::new(client, object_generator.clone_ref(py))) + } + PullOperator::MockSink {} => PipelineSink::Mock(MockSinkHandler::new()), + _ => panic!("build_sink() called on non-sink operator"), + } + } +} diff --git a/sentry_streams/src/pull/pull_source.rs b/sentry_streams/src/pull/pull_source.rs new file mode 100644 index 00000000..f6bb0f69 --- /dev/null +++ b/sentry_streams/src/pull/pull_source.rs @@ -0,0 +1,144 @@ +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; + +use futures::Stream; +use pyo3::prelude::*; +use sentry_arroyo::backends::kafka::types::{Headers, KafkaPayload}; +use sentry_arroyo::processing::strategies::offset_tracker::OffsetCommitter; +use sentry_arroyo::processing::stream::source::KafkaSource; +use sentry_arroyo::processing::stream::{MessageMetadata, PipelineEnvelope, StageResult}; +use sentry_arroyo::types::{Partition, Topic}; + +/// Trait for pipeline sources. Provides a stream of raw Kafka payloads +/// and an offset committer. +pub trait PullSource: Send + Sync { + fn stream(&self) -> Pin> + '_>>; + fn committer(&self) -> &dyn OffsetCommitter; +} + +impl PullSource for KafkaSource { + fn stream(&self) -> Pin> + '_>> { + Box::pin(KafkaSource::stream(self)) + } + + fn committer(&self) -> &dyn OffsetCommitter { + self + } +} + +// ── Test infrastructure ───────────────────────────────────────────── + +/// A test message with payload bytes and optional headers. +#[pyclass] +#[derive(Clone)] +pub struct PyTestMessage { + #[pyo3(get)] + pub payload: Vec, + #[pyo3(get)] + pub headers: HashMap>, +} + +#[pymethods] +impl PyTestMessage { + #[new] + fn new(payload: Vec, headers: Option>>) -> Self { + Self { + payload, + headers: headers.unwrap_or_default(), + } + } +} + +/// Source config enum — Python constructs this, Rust builds the source. +#[pyclass] +pub enum PullSourceConfig { + #[pyo3(constructor = (config, topic))] + Kafka { + config: crate::kafka_config::PyKafkaConsumerConfig, + topic: String, + }, + + #[pyo3(constructor = (messages))] + Test { messages: Vec }, +} + +impl PullSourceConfig { + pub fn build(&self, py: Python<'_>) -> Box { + match self { + PullSourceConfig::Kafka { config, topic } => { + let kafka_config = config.clone().into(); + Box::new(KafkaSource::new(kafka_config, &[Topic::new(topic)])) + } + PullSourceConfig::Test { messages } => { + Box::new(VecSource::from_test_messages(messages.clone())) + } + } + } +} + +/// In-memory source for testing. +pub struct VecSource { + messages: Mutex>>, + committer: NoOpCommitter, +} + +impl VecSource { + pub fn new(messages: Vec>) -> Self { + Self { + messages: Mutex::new(messages), + committer: NoOpCommitter, + } + } + + pub fn from_test_messages(test_messages: Vec) -> Self { + let messages = test_messages + .into_iter() + .enumerate() + .map(|(i, msg)| { + let headers = if msg.headers.is_empty() { + None + } else { + let mut h = Headers::new(); + for (key, value) in &msg.headers { + h = h.insert(key, Some(value.clone())); + } + Some(h) + }; + + let kp = KafkaPayload::new(None, headers, Some(msg.payload)); + let md = MessageMetadata { + partition: Partition::new(Topic::new("test"), 0), + offset: i as u64, + timestamp: chrono::Utc::now(), + }; + StageResult::Emit(PipelineEnvelope::new(kp.clone(), md, Arc::new(kp))) + }) + .collect(); + + Self::new(messages) + } +} + +impl PullSource for VecSource { + fn stream(&self) -> Pin> + '_>> { + let messages: Vec<_> = self.messages.lock().unwrap().drain(..).collect(); + Box::pin(futures::stream::iter(messages)) + } + + fn committer(&self) -> &dyn OffsetCommitter { + &self.committer + } +} + +/// No-op committer for testing. +struct NoOpCommitter; + +impl OffsetCommitter for NoOpCommitter { + fn commit_offsets( + &self, + _positions: &HashMap, + ) -> Result<(), Box> { + Ok(()) + } +} diff --git a/sentry_streams/src/pull/stages/py_callable.rs b/sentry_streams/src/pull/stages/py_callable.rs index 7a84d319..3d8c7820 100644 --- a/sentry_streams/src/pull/stages/py_callable.rs +++ b/sentry_streams/src/pull/stages/py_callable.rs @@ -1,59 +1,30 @@ use pyo3::prelude::*; -use pyo3::types::PyList; -use sentry_arroyo::backends::kafka::types::KafkaPayload; use sentry_arroyo::processing::stream::{PipelineEnvelope, Stage, StageResult}; +use crate::pull::message_wrapper::MessageWrapper; use crate::pull::pipeline_value::PipelineValue; +use crate::pull::pipeline_value_converter::PipelineValueConverter; /// Calls a Python callable as a pipeline stage. /// -/// Handles three input scenarios: -/// - PipelineValue::Rust(Vec) — converts to Python list of bytes, -/// calls the callable. Used for batch_parser. -/// - PipelineValue::Python — passes the Python object directly to the callable. -/// Used for processor, serializer. -/// - PipelineValue::Raw — converts single KafkaPayload to Python bytes, -/// calls the callable. Used for single-message transforms. -/// -/// Output is always PipelineValue::Python (the callable's return value). -/// -/// Python exceptions become StageResult::Fail (no DLQ for now). +/// Wraps the input in a streams `Message` object before calling, +/// and re-wraps the output. All streams pipeline callables expect +/// `Message[T]` with `.payload`, `.headers`, `.timestamp`, `.schema`. pub struct PyCallableStage { callable: Py, - stage_name: String, + stage_name: &'static str, + schema: String, } impl PyCallableStage { - pub fn new(callable: Py, name: impl Into) -> Self { + pub fn new(callable: Py, name: impl Into, schema: impl Into) -> Self { + let leaked: &'static str = Box::leak(name.into().into_boxed_str()); Self { callable, - stage_name: name.into(), + stage_name: leaked, + schema: schema.into(), } } - - /// Convert a Vec to a Python list of bytes objects. - fn batch_to_python<'py>( - py: Python<'py>, - payloads: Vec, - ) -> PyResult> { - let items: Vec> = payloads - .iter() - .map(|kp| { - let bytes = kp.payload().map(|v| v.as_slice()).unwrap_or(&[]); - pyo3::types::PyBytes::new(py, bytes).into_any() - }) - .collect(); - Ok(PyList::new(py, &items)?.into_any()) - } - - /// Convert a single KafkaPayload to Python bytes. - fn raw_to_python<'py>( - py: Python<'py>, - payload: &KafkaPayload, - ) -> PyResult> { - let bytes = payload.payload().map(|v| v.as_slice()).unwrap_or(&[]); - Ok(pyo3::types::PyBytes::new(py, bytes).into_any()) - } } impl Stage for PyCallableStage { @@ -64,23 +35,14 @@ impl Stage for PyCallableStage { &self, envelope: PipelineEnvelope, ) -> StageResult { - let result = Python::attach(|py| -> PyResult> { - let input: Bound<'_, PyAny> = match envelope.payload { - PipelineValue::Rust(ref boxed) => { - // Try to downcast as Vec (batch) - if let Some(payloads) = boxed.downcast_ref::>() { - Self::batch_to_python(py, payloads.clone())? - } else { - return Err(pyo3::exceptions::PyTypeError::new_err( - "PyCallableStage received unsupported Rust type", - )); - } - } - PipelineValue::Python(ref obj) => obj.bind(py).clone().into_any(), - PipelineValue::Raw(ref kp) => Self::raw_to_python(py, kp)?, - }; + let timestamp = envelope.metadata.timestamp.timestamp_millis() as f64 / 1000.0; + let schema = &self.schema; - self.callable.call1(py, (input,)) + let result = Python::attach(|py| -> PyResult> { + let input = PipelineValueConverter::to_python(&envelope.payload, py)?; + let message = MessageWrapper::ensure(py, input, timestamp, schema)?; + let result = self.callable.call1(py, (&message,))?; + MessageWrapper::rewrap(py, result.bind(py).clone(), &message) }); match result { @@ -89,16 +51,11 @@ impl Stage for PyCallableStage { envelope.metadata, envelope.raw, )), - Err(py_err) => { - // For now, all Python errors are fatal (no DLQ) - StageResult::Fail(Box::new(py_err)) - } + Err(py_err) => StageResult::Fail(Box::new(py_err)), } } fn name(&self) -> &'static str { - // Leak the string to get a &'static str. - // This is fine — stages are long-lived, created once at pipeline build time. - Box::leak(self.stage_name.clone().into_boxed_str()) + self.stage_name } } diff --git a/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py b/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py new file mode 100644 index 00000000..6c03657e --- /dev/null +++ b/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py @@ -0,0 +1,375 @@ +"""Tests for the PullBasedAdapter — verifies DSL steps map to the correct +PullOperator variants, and e2e Python → Rust pipeline execution.""" + +from sentry_streams.adapters.arroyo.pull_adapter import PullBasedAdapter +from sentry_streams.adapters.stream_adapter import RuntimeTranslator +from sentry_streams.pipeline.pipeline import ( + Batch, + BatchParser, + GCSSink, + HeadersFilter, + Map, + Pipeline, + streaming_source, +) +from sentry_streams.runner import iterate_edges +from sentry_streams.rust_streams import ( + PullConsumer, + PullOperator, + PullSourceConfig, + PyTestMessage, +) + + +def _dummy_processor(msg): + """Mock processor function for Map steps.""" + return msg + + +def _dummy_object_generator(): + """Mock GCS object name generator.""" + return "test-object.parquet" + + +def test_pull_adapter_items_span_pipeline(): + """Verify the adapter produces the correct PullOperator list + for a pipeline shaped like items_span.""" + + ITEM_TYPE_SPAN = 1 + + pipeline: Pipeline[bytes] = ( + streaming_source(name="kafka", stream_name="snuba-items") + .apply( + HeadersFilter( + name="span_filter", + header_name="item_type", + value=ITEM_TYPE_SPAN, + ) + ) + .apply(Batch(name="batcher", batch_size=50000)) + .apply(Map(name="processor", function=_dummy_processor)) + .sink( + GCSSink( + name="gcs_sink", + bucket="test-bucket", + object_generator=_dummy_object_generator, + ) + ) + ) + + adapter = PullBasedAdapter.build( + { + "steps_config": { + "kafka": { + "bootstrap_servers": ["localhost:9092"], + "auto_offset_reset": "earliest", + "consumer_group": "test-group", + "override_params": {}, + }, + }, + } + ) + + iterate_edges(pipeline, RuntimeTranslator(adapter)) + + # Verify source was configured + assert adapter._kafka_config is not None + assert adapter._topic == "snuba-items" + + # Verify steps + assert len(adapter._steps) == 3 + + # Step 0: HeaderFilter + step0 = adapter._steps[0] + assert isinstance(step0, PullOperator) + # PullOperator is a PyO3 enum — check the variant by accessing fields + # PyO3 complex enums expose variant names differently, so we check + # by verifying the attributes exist + assert step0.header_name == "item_type" + assert step0.expected_value == ITEM_TYPE_SPAN + + # Step 1: Batch + step1 = adapter._steps[1] + assert isinstance(step1, PullOperator) + assert step1.max_batch_size == 50000 + + # Step 2: PyCallable (Map) + step2 = adapter._steps[2] + assert isinstance(step2, PullOperator) + assert step2.name == "processor" + assert step2.callable is _dummy_processor + + # Verify sink + assert adapter._sink is not None + assert adapter._sink.bucket == "test-bucket" + assert adapter._sink.object_generator is _dummy_object_generator + + +def test_pull_adapter_header_filter_only(): + """Verify a minimal pipeline with just a filter.""" + + pipeline: Pipeline[bytes] = ( + streaming_source(name="src", stream_name="test-topic") + .apply( + HeadersFilter( + name="my_filter", + header_name="type", + value=42, + ) + ) + .sink( + GCSSink( + name="sink", + bucket="bucket", + object_generator=_dummy_object_generator, + ) + ) + ) + + adapter = PullBasedAdapter.build( + { + "steps_config": { + "src": { + "bootstrap_servers": ["localhost:9092"], + "auto_offset_reset": "earliest", + "consumer_group": "test-group", + }, + }, + } + ) + + iterate_edges(pipeline, RuntimeTranslator(adapter)) + + assert len(adapter._steps) == 1 + assert adapter._steps[0].header_name == "type" + assert adapter._steps[0].expected_value == 42 + assert adapter._sink is not None + + +def test_pull_adapter_batch_parser_converts_to_map(): + """BatchParser is a ComplexStep that converts to a Map via convert(). + Verify the adapter receives it as a PyCallable.""" + + pipeline: Pipeline[bytes] = ( + streaming_source(name="src", stream_name="test-topic") + .apply(Batch(name="batcher", batch_size=100)) + .apply(BatchParser[bytes]("parser")) + .sink( + GCSSink( + name="sink", + bucket="bucket", + object_generator=_dummy_object_generator, + ) + ) + ) + + adapter = PullBasedAdapter.build( + { + "steps_config": { + "src": { + "bootstrap_servers": ["localhost:9092"], + "auto_offset_reset": "earliest", + "consumer_group": "test-group", + }, + }, + } + ) + + iterate_edges(pipeline, RuntimeTranslator(adapter)) + + # Batch + BatchParser(converted to Map) = 2 steps + assert len(adapter._steps) == 2 + + # Step 0: Batch + assert adapter._steps[0].max_batch_size == 100 + + # Step 1: PyCallable (from BatchParser.convert() → Map) + step1 = adapter._steps[1] + assert step1.name == "parser" + # The callable should be the batch_msg_parser function + assert callable(step1.callable) + + +# ── E2E: Python → Rust pipeline execution ─────────────────────── + + +def _message_aware_transform(msg): + """Transform that reads msg.payload (like real pipeline callables do). + Appends '-processed' to each bytes item in the batch.""" + items = msg.payload # Message[Sequence[bytes]] -> Sequence[bytes] + return [item + b"-processed" if isinstance(item, bytes) else item for item in items] + + +def test_pull_consumer_e2e_python_to_rust(): + """Full e2e: Python DSL → PullBasedAdapter → PullConsumer → stages → sink. + + Pipeline: HeaderFilter → Batch(2) → PyCallable(transform) → MockSink + + The transform function receives a Message wrapper (because the adapter + sets schema on PyCallable operators), accesses .payload, and returns + the transformed data. PyCallableStage re-wraps the result in a new Message. + """ + + pipeline: Pipeline[bytes] = ( + streaming_source(name="kafka", stream_name="test-topic") + .apply( + HeadersFilter( + name="filter", + header_name="item_type", + value=1, + ) + ) + .apply(Batch(name="batcher", batch_size=2)) + .apply(Map(name="transformer", function=_message_aware_transform)) + .sink( + GCSSink( + name="gcs_sink", + bucket="test-bucket", + object_generator=_dummy_object_generator, + ) + ) + ) + + # Translate DSL to PullOperator list via the adapter + adapter = PullBasedAdapter.build( + { + "steps_config": { + "kafka": { + "bootstrap_servers": ["localhost:9092"], + "auto_offset_reset": "earliest", + "consumer_group": "test-group", + "override_params": {}, + }, + }, + } + ) + iterate_edges(pipeline, RuntimeTranslator(adapter)) + + # Verify the adapter set schema on the PyCallable + py_callable_step = adapter._steps[2] # HeaderFilter, Batch, PyCallable + assert py_callable_step.schema == "test-topic" + + # Swap source and sink for testing + source = PullSourceConfig.Test( + messages=[ + PyTestMessage(payload=b"span-0", headers={"item_type": b"1"}), + PyTestMessage(payload=b"span-1", headers={"item_type": b"1"}), + PyTestMessage(payload=b"span-2", headers={"item_type": b"2"}), # filtered out + PyTestMessage(payload=b"span-3", headers={"item_type": b"1"}), + PyTestMessage(payload=b"span-4", headers={"item_type": b"1"}), + ] + ) + mock_sink = PullOperator.MockSink() + + consumer = PullConsumer( + source=source, + steps=adapter._steps, + sink=mock_sink, + ) + consumer.run() + + # 4 matching messages, batch size 2 → 2 batches + # Each batch goes through _message_aware_transform (PyCallableStage with Message wrapping) + # MockSink captures the output + results = consumer.get_mock_sink_results() + assert len(results) == 2, f"Expected 2 batches, got {len(results)}" + + +def test_pull_consumer_e2e_complex_steps(): + """E2E test with real ComplexStep conversions: BatchParser and ParquetSerializer. + + Pipeline: Batch(2) → BatchParser[TraceItem] → Map(extract_org_id) → ParquetSerializer → MockSink + + This tests that: + - BatchParser.convert() produces a Map(batch_msg_parser) that works with Message wrapping + - batch_msg_parser uses msg.schema to find the codec and parses protobuf bytes + - ParquetSerializer.convert() produces a Map(serialize_to_parquet) that serializes to parquet + - The full chain works end-to-end through our pull pipeline + """ + from sentry_protos.snuba.v1.trace_item_pb2 import TraceItem as TraceItemProto + + from sentry_streams.pipeline.datatypes import Uint64 + from sentry_streams.pipeline.pipeline import BatchParser, ParquetSerializer + + # Create a simple processor that extracts org_id into a dict + def extract_org_id(msg): + """Map function: Sequence[TraceItem] → list[dict]""" + return [{"org_id": item.organization_id} for item in msg.payload] + + pipeline: Pipeline[bytes] = ( + streaming_source(name="kafka", stream_name="snuba-items") + .apply(Batch(name="batcher", batch_size=2)) + .apply(BatchParser[TraceItemProto]("parser")) + .apply(Map(name="processor", function=extract_org_id)) + .apply( + ParquetSerializer( + name="serializer", + schema_fields={"org_id": Uint64()}, + ) + ) + .sink( + GCSSink( + name="gcs_sink", + bucket="test-bucket", + object_generator=_dummy_object_generator, + ) + ) + ) + + adapter = PullBasedAdapter.build( + { + "steps_config": { + "kafka": { + "bootstrap_servers": ["localhost:9092"], + "auto_offset_reset": "earliest", + "consumer_group": "test-group", + "override_params": {}, + }, + }, + } + ) + iterate_edges(pipeline, RuntimeTranslator(adapter)) + + # Create test messages: real serialized TraceItem protobufs + item1 = TraceItemProto() + item1.organization_id = 42 + item1.trace_id = b"0123456789abcdef" + + item2 = TraceItemProto() + item2.organization_id = 99 + item2.trace_id = b"fedcba9876543210" + + item3 = TraceItemProto() + item3.organization_id = 7 + item3.trace_id = b"aaaaaaaaaaaaaaaa" + + item4 = TraceItemProto() + item4.organization_id = 123 + item4.trace_id = b"bbbbbbbbbbbbbbbb" + + source = PullSourceConfig.Test( + messages=[ + PyTestMessage(payload=item1.SerializeToString(), headers={}), + PyTestMessage(payload=item2.SerializeToString(), headers={}), + PyTestMessage(payload=item3.SerializeToString(), headers={}), + PyTestMessage(payload=item4.SerializeToString(), headers={}), + ] + ) + mock_sink = PullOperator.MockSink() + + consumer = PullConsumer( + source=source, + steps=adapter._steps, + sink=mock_sink, + ) + consumer.run() + + # 4 messages, batch size 2 → 2 batches + # Each batch: batch_msg_parser (protobuf decode) → extract_org_id → parquet serialize + # MockSink should capture 2 parquet byte blobs + results = consumer.get_mock_sink_results() + assert len(results) == 2, f"Expected 2 parquet outputs, got {len(results)}" + + # Verify the results are actual parquet bytes (magic number: PAR1) + for i, result_bytes in enumerate(results): + assert result_bytes[:4] == b"PAR1", f"Result {i} doesn't start with PAR1 magic" From 050b267250432b29652c6d151eabf94308da0b16 Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:23:24 -0700 Subject: [PATCH 5/9] handle cancellation --- sentry_streams/Cargo.lock | 1 + sentry_streams/Cargo.toml | 3 +- sentry_streams/src/pull/mod.rs | 101 ++++++++++++++++++++++- sentry_streams/src/pull/pull_consumer.rs | 49 ++++++++--- sentry_streams/src/pull/pull_source.rs | 24 ++---- 5 files changed, 147 insertions(+), 31 deletions(-) diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index 0664da83..dea4d911 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -2232,6 +2232,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tokio-util", "tracing", "tracing-subscriber", ] diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index 64cdcd85..c3207983 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -16,7 +16,8 @@ ctrlc = "3.4.6" rdkafka = { version = "0.37.0", features = ["cmake-build", "tracing"] } anyhow = "1.0.98" reqwest = "0.12.15" -tokio = { version = "1.45.0", features = ["macros", "rt-multi-thread"] } +tokio = { version = "1.45.0", features = ["macros", "rt-multi-thread", "signal"] } +tokio-util = "0.7" log = "0.4.27" serde_json = "1.0.141" clap = { version = "4.5.45", features = ["derive"], optional = true } diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index 101f9f71..1e9bf076 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -321,8 +321,9 @@ mod tests { use super::pull_consumer::PullConsumer; use super::pull_operator::PullOperator; - use super::pull_source::PullSource; use futures::stream::Stream; + use futures::StreamExt; + use sentry_arroyo::processing::stream::PullSource; use std::pin::Pin; /// Test source that drains its messages on first stream() call. @@ -352,6 +353,10 @@ mod tests { fn committer(&self) -> &dyn OffsetCommitter { self.committer.as_ref() } + + fn shutdown(&self) { + // No-op for test source + } } /// Helper to create a raw StageResult (not wrapped in PipelineValue). @@ -471,4 +476,98 @@ mod tests { let partition = Partition::new(Topic::new("test"), 0); assert_eq!(last.get(&partition), Some(&4)); } + + // ── Cancellation test ─────────────────────────────────────────── + + #[tokio::test] + async fn test_pull_consumer_shutdown() { + use tokio_util::sync::CancellationToken; + + /// A source that blocks until shutdown is called, then emits its messages. + struct BlockingSource { + messages: Mutex>>, + committer: Arc, + cancel: CancellationToken, + } + + impl BlockingSource { + fn new(messages: Vec>) -> (Self, Arc) { + let committer = Arc::new(MockCommitter::new()); + let source = Self { + messages: Mutex::new(messages), + committer: committer.clone(), + cancel: CancellationToken::new(), + }; + (source, committer) + } + } + + impl PullSource for BlockingSource { + fn stream(&self) -> Pin> + '_>> { + let messages: Vec<_> = self.messages.lock().unwrap().drain(..).collect(); + // Emit buffered messages, then block until shutdown + let msg_stream: Pin> + '_>> = + Box::pin(futures::stream::iter(messages)); + let pending: Pin> + '_>> = + Box::pin(futures::stream::pending()); + Box::pin( + msg_stream + .chain(pending) + .take_until(self.cancel.cancelled()), + ) + } + + fn committer(&self) -> &dyn OffsetCommitter { + self.committer.as_ref() + } + + fn shutdown(&self) { + self.cancel.cancel(); + } + } + + let messages = vec![ + make_raw_envelope_with_header(b"msg-0", 0, "item_type", 1), + make_raw_envelope_with_header(b"msg-1", 1, "item_type", 1), + ]; + + let (source, committer) = BlockingSource::new(messages); + + // Get a reference to trigger shutdown later + let source = Arc::new(source); + let source_for_shutdown = source.clone(); + + let stages = pyo3::Python::attach(|py| { + vec![PullOperator::HeaderFilter { + header_name: "item_type".into(), + expected_value: 1, + }] + .iter() + .map(|op| op.build_stage(py)) + .collect() + }); + + let consumer = PullConsumer { + source: source as Arc, + stages, + sink: None, + }; + + // Spawn shutdown after a short delay + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(50)).await; + source_for_shutdown.shutdown(); + }); + + // run_pipeline should return cleanly after shutdown + let result = consumer.run_pipeline().await; + assert!(result.is_ok(), "Pipeline should exit cleanly on shutdown"); + + // Messages should have been processed and committed + let committed = committer.committed(); + assert!( + !committed.is_empty(), + "Offsets should be committed on shutdown" + ); + } } diff --git a/sentry_streams/src/pull/pull_consumer.rs b/sentry_streams/src/pull/pull_consumer.rs index c7cdf23a..ad77c20b 100644 --- a/sentry_streams/src/pull/pull_consumer.rs +++ b/sentry_streams/src/pull/pull_consumer.rs @@ -1,24 +1,25 @@ use std::pin::Pin; +use std::sync::Arc; use std::time::Duration; use futures::stream::Stream; use futures::StreamExt; use pyo3::prelude::*; use sentry_arroyo::processing::strategies::offset_tracker::OffsetTracker; -use sentry_arroyo::processing::stream::{LogHandler, PipelineExt, StageResult}; +use sentry_arroyo::processing::stream::{LogHandler, PipelineExt, PullSource, StageResult}; use super::pipeline_sink::PipelineSink; use super::pipeline_stage::PipelineStage; use super::pipeline_value::PipelineValue; use super::pull_operator::PullOperator; -use super::pull_source::{PullSource, PullSourceConfig}; +use super::pull_source::PullSourceConfig; /// Pull-based pipeline consumer. Fully configured at construction time. -/// `run()` is parameterless. +/// `run()` is parameterless. Handles SIGINT/SIGTERM for graceful shutdown. #[pyclass] pub struct PullConsumer { - source: Box, - stages: Vec, + pub(crate) source: Arc, + pub(crate) stages: Vec, pub sink: Option, } @@ -31,7 +32,7 @@ impl PullConsumer { steps: Vec>, sink: Option>, ) -> Self { - let source = source.get().build(py); + let source: Arc = Arc::from(source.get().build(py)); let stages = steps.iter().map(|op| op.get().build_stage(py)).collect(); let sink = sink.map(|s| s.get().build_sink(py)); @@ -42,7 +43,7 @@ impl PullConsumer { } } - /// Run the pipeline. Blocks until completion or fatal error. + /// Run the pipeline. Blocks until completion, fatal error, or signal. fn run(&self) -> PyResult<()> { let rt = tokio::runtime::Runtime::new().map_err(|e| { pyo3::exceptions::PyRuntimeError::new_err(format!( @@ -54,8 +55,7 @@ impl PullConsumer { .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}"))) } - /// Get results captured by a MockSink. Returns a list of byte arrays. - /// Raises RuntimeError if the sink is not a MockSink. + /// Get results captured by a MockSink. fn get_mock_sink_results(&self) -> PyResult>> { match &self.sink { Some(PipelineSink::Mock(handler)) => Ok(handler.get_results()), @@ -74,13 +74,16 @@ impl PullConsumer { sink: Option, ) -> Self { Self { - source: Box::new(source), + source: Arc::new(source), stages, sink, } } pub async fn run_pipeline(&self) -> Result<(), Box> { + // Install signal handlers for graceful shutdown + self.install_signal_handlers(); + let committer = self.source.committer(); let error_handler = LogHandler; let commit_interval = Duration::from_secs(5); @@ -121,4 +124,30 @@ impl PullConsumer { stream.on_reject(&error_handler).commit(&mut tracker).await } } + + fn install_signal_handlers(&self) { + // SIGINT (Ctrl+C) + let source = self.source.clone(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + tracing::info!("Received SIGINT, shutting down..."); + source.shutdown(); + } + }); + + // SIGTERM (K8s pod termination) + #[cfg(unix)] + { + let source = self.source.clone(); + tokio::spawn(async move { + let mut sigterm = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("Failed to register SIGTERM handler"); + + sigterm.recv().await; + tracing::info!("Received SIGTERM, shutting down..."); + source.shutdown(); + }); + } + } } diff --git a/sentry_streams/src/pull/pull_source.rs b/sentry_streams/src/pull/pull_source.rs index f6bb0f69..6d61e561 100644 --- a/sentry_streams/src/pull/pull_source.rs +++ b/sentry_streams/src/pull/pull_source.rs @@ -7,28 +7,10 @@ use pyo3::prelude::*; use sentry_arroyo::backends::kafka::types::{Headers, KafkaPayload}; use sentry_arroyo::processing::strategies::offset_tracker::OffsetCommitter; use sentry_arroyo::processing::stream::source::KafkaSource; +use sentry_arroyo::processing::stream::PullSource; use sentry_arroyo::processing::stream::{MessageMetadata, PipelineEnvelope, StageResult}; use sentry_arroyo::types::{Partition, Topic}; -/// Trait for pipeline sources. Provides a stream of raw Kafka payloads -/// and an offset committer. -pub trait PullSource: Send + Sync { - fn stream(&self) -> Pin> + '_>>; - fn committer(&self) -> &dyn OffsetCommitter; -} - -impl PullSource for KafkaSource { - fn stream(&self) -> Pin> + '_>> { - Box::pin(KafkaSource::stream(self)) - } - - fn committer(&self) -> &dyn OffsetCommitter { - self - } -} - -// ── Test infrastructure ───────────────────────────────────────────── - /// A test message with payload bytes and optional headers. #[pyclass] #[derive(Clone)] @@ -129,6 +111,10 @@ impl PullSource for VecSource { fn committer(&self) -> &dyn OffsetCommitter { &self.committer } + + fn shutdown(&self) { + // No-op for test source + } } /// No-op committer for testing. From d91f41d7f7d2e7170ccddde85d59b012605c1ed7 Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:40:16 -0700 Subject: [PATCH 6/9] Rebuild pipeline on rebalance. --- sentry_streams/Cargo.lock | 23 ++++ .../sentry_streams/rust_streams.pyi | 3 - sentry_streams/src/pull/mod.rs | 71 ++---------- sentry_streams/src/pull/pipeline_sink.rs | 37 +------ sentry_streams/src/pull/pull_consumer.rs | 101 ++++++++++-------- sentry_streams/src/pull/pull_operator.rs | 8 +- sentry_streams/src/pull/pull_source.rs | 6 +- .../adapters/arroyo/test_pull_adapter.py | 62 ++++++----- 8 files changed, 121 insertions(+), 190 deletions(-) diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index dea4d911..cd9f432c 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -257,6 +257,28 @@ version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -2512,6 +2534,7 @@ dependencies = [ name = "sentry_arroyo" version = "2.42.0" dependencies = [ + "async-stream", "chrono", "coarsetime", "futures", diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index c9cbdfb4..3086ba07 100644 --- a/sentry_streams/sentry_streams/rust_streams.pyi +++ b/sentry_streams/sentry_streams/rust_streams.pyi @@ -221,8 +221,6 @@ class PullOperator: def PyCallable(cls, callable: Callable[..., Any], name: str, schema: str) -> Self: ... @classmethod def GcsSink(cls, bucket: str, object_generator: Callable[[], str]) -> Self: ... - @classmethod - def MockSink(cls) -> Self: ... class PyTestMessage: def __init__(self, payload: bytes, headers: Mapping[str, bytes] | None = None) -> None: ... @@ -241,4 +239,3 @@ class PullConsumer: sink: PullOperator | None = None, ) -> None: ... def run(self) -> None: ... - def get_mock_sink_results(self) -> list[bytes]: ... diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index 1e9bf076..c799ca15 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -18,10 +18,10 @@ mod tests { use futures::stream; use sentry_arroyo::backends::kafka::types::{Headers, KafkaPayload}; - use sentry_arroyo::processing::strategies::offset_tracker::{OffsetCommitter, OffsetTracker}; use sentry_arroyo::processing::stream::{ LogHandler, MessageMetadata, PipelineEnvelope, PipelineExt, Stage, StageResult, }; + use sentry_arroyo::processing::stream::{OffsetCommitter, OffsetTracker}; use sentry_arroyo::types::{Partition, Topic}; use super::pipeline_value::{PipelineValue, PipelineValueCaster}; @@ -319,6 +319,7 @@ mod tests { // ── PullConsumer e2e test ─────────────────────────────────────── + use super::pipeline_stage::PipelineStage; use super::pull_consumer::PullConsumer; use super::pull_operator::PullOperator; use futures::stream::Stream; @@ -401,7 +402,7 @@ mod tests { let (source, committer) = TestSource::new(messages); - let stages = pyo3::Python::attach(|py| { + let stages: Vec = pyo3::Python::attach(|py| { vec![ PullOperator::HeaderFilter { header_name: "item_type".into(), @@ -414,9 +415,9 @@ mod tests { .collect() }); - let consumer = PullConsumer::with_source(source, stages, None); + let source: Arc = Arc::new(source); - let result = consumer.run_pipeline().await; + let result = PullConsumer::run_pipeline(&source, &stages, None).await; assert!(result.is_ok()); let committed = committer.committed(); @@ -426,57 +427,6 @@ mod tests { assert_eq!(last.get(&partition), Some(&6)); } - // ── Test #3: e2e with sink ────────────────────────────────────── - - use super::pipeline_sink::{MockSinkHandler, PipelineSink}; - - #[tokio::test] - async fn test_pull_consumer_e2e_with_sink() { - // 4 messages with matching header → batch of 2 → 2 batches emitted → 2 sink calls - let messages = vec![ - make_raw_envelope_with_header(b"span-0", 0, "item_type", 1), - make_raw_envelope_with_header(b"span-1", 1, "item_type", 1), - make_raw_envelope_with_header(b"span-2", 2, "item_type", 1), - make_raw_envelope_with_header(b"span-3", 3, "item_type", 1), - ]; - - let (source, committer) = TestSource::new(messages); - - let stages = pyo3::Python::attach(|py| { - vec![ - PullOperator::HeaderFilter { - header_name: "item_type".into(), - expected_value: 1, - }, - PullOperator::Batch { max_batch_size: 2 }, - ] - .iter() - .map(|op| op.build_stage(py)) - .collect() - }); - - let mock_sink = MockSinkHandler::new(); - let consumer = - PullConsumer::with_source(source, stages, Some(PipelineSink::Mock(mock_sink))); - - let result = consumer.run_pipeline().await; - assert!(result.is_ok()); - - // Verify sink was called twice (2 batches of 2) - let results = match &consumer.sink { - Some(PipelineSink::Mock(h)) => h.get_results(), - _ => panic!("Expected MockSink"), - }; - assert_eq!(results.len(), 2, "Expected 2 sink calls (2 batches of 2)"); - - // Verify offsets committed - let committed = committer.committed(); - assert!(!committed.is_empty()); - let last = committed.last().unwrap(); - let partition = Partition::new(Topic::new("test"), 0); - assert_eq!(last.get(&partition), Some(&4)); - } - // ── Cancellation test ─────────────────────────────────────────── #[tokio::test] @@ -537,7 +487,7 @@ mod tests { let source = Arc::new(source); let source_for_shutdown = source.clone(); - let stages = pyo3::Python::attach(|py| { + let stages: Vec = pyo3::Python::attach(|py| { vec![PullOperator::HeaderFilter { header_name: "item_type".into(), expected_value: 1, @@ -547,12 +497,6 @@ mod tests { .collect() }); - let consumer = PullConsumer { - source: source as Arc, - stages, - sink: None, - }; - // Spawn shutdown after a short delay tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(50)).await; @@ -560,7 +504,8 @@ mod tests { }); // run_pipeline should return cleanly after shutdown - let result = consumer.run_pipeline().await; + let result = + PullConsumer::run_pipeline(&(source as Arc), &stages, None).await; assert!(result.is_ok(), "Pipeline should exit cleanly on shutdown"); // Messages should have been processed and committed diff --git a/sentry_streams/src/pull/pipeline_sink.rs b/sentry_streams/src/pull/pipeline_sink.rs index 30b97cd7..0a23def7 100644 --- a/sentry_streams/src/pull/pipeline_sink.rs +++ b/sentry_streams/src/pull/pipeline_sink.rs @@ -1,17 +1,12 @@ -use std::sync::{Arc, Mutex}; - use sentry_arroyo::processing::stream::handlers::next::NextHandler; use sentry_arroyo::processing::stream::PipelineEnvelope; use super::gcs_sink_handler::GcsSinkHandler; use super::pipeline_value::PipelineValue; -use super::pipeline_value_converter::PipelineValueConverter; -/// Enum dispatch for pipeline sinks. Mirrors PipelineStage pattern. +/// Enum dispatch for pipeline sinks. pub enum PipelineSink { Gcs(GcsSinkHandler), - /// Test sink that captures received payloads. - Mock(MockSinkHandler), } impl NextHandler for PipelineSink { @@ -21,36 +16,6 @@ impl NextHandler for PipelineSink { ) -> Result<(), Box> { match self { PipelineSink::Gcs(h) => h.handle(envelope).await, - PipelineSink::Mock(h) => h.handle(envelope).await, - } - } -} - -/// Mock sink handler that records what it receives as extracted bytes. -pub struct MockSinkHandler { - results: Arc>>>, -} - -impl MockSinkHandler { - pub fn new() -> Self { - Self { - results: Arc::new(Mutex::new(Vec::new())), } } - - pub fn get_results(&self) -> Vec> { - self.results.lock().unwrap().clone() - } -} - -impl NextHandler for MockSinkHandler { - async fn handle( - &self, - envelope: &PipelineEnvelope, - ) -> Result<(), Box> { - let bytes = PipelineValueConverter::extract_bytes(&envelope.payload) - .unwrap_or_else(|_| b"".to_vec()); - self.results.lock().unwrap().push(bytes); - Ok(()) - } } diff --git a/sentry_streams/src/pull/pull_consumer.rs b/sentry_streams/src/pull/pull_consumer.rs index ad77c20b..7a82da7c 100644 --- a/sentry_streams/src/pull/pull_consumer.rs +++ b/sentry_streams/src/pull/pull_consumer.rs @@ -5,8 +5,9 @@ use std::time::Duration; use futures::stream::Stream; use futures::StreamExt; use pyo3::prelude::*; -use sentry_arroyo::processing::strategies::offset_tracker::OffsetTracker; -use sentry_arroyo::processing::stream::{LogHandler, PipelineExt, PullSource, StageResult}; +use sentry_arroyo::processing::stream::{ + LogHandler, OffsetTracker, PipelineExit, PipelineExt, PullSource, StageResult, +}; use super::pipeline_sink::PipelineSink; use super::pipeline_stage::PipelineStage; @@ -14,13 +15,14 @@ use super::pipeline_value::PipelineValue; use super::pull_operator::PullOperator; use super::pull_source::PullSourceConfig; -/// Pull-based pipeline consumer. Fully configured at construction time. -/// `run()` is parameterless. Handles SIGINT/SIGTERM for graceful shutdown. +/// Pull-based pipeline consumer. Stores operator descriptions and +/// rebuilds fresh stages on each partition assignment (rebalance). +/// `run()` is parameterless — handles rebalance restart and signal shutdown. #[pyclass] pub struct PullConsumer { - pub(crate) source: Arc, - pub(crate) stages: Vec, - pub sink: Option, + source: Arc, + operators: Vec>, + sink_operator: Option>, } #[pymethods] @@ -33,17 +35,15 @@ impl PullConsumer { sink: Option>, ) -> Self { let source: Arc = Arc::from(source.get().build(py)); - let stages = steps.iter().map(|op| op.get().build_stage(py)).collect(); - let sink = sink.map(|s| s.get().build_sink(py)); - Self { source, - stages, - sink, + operators: steps, + sink_operator: sink, } } /// Run the pipeline. Blocks until completion, fatal error, or signal. + /// Rebuilds stages fresh on each rebalance. fn run(&self) -> PyResult<()> { let rt = tokio::runtime::Runtime::new().map_err(|e| { pyo3::exceptions::PyRuntimeError::new_err(format!( @@ -51,45 +51,53 @@ impl PullConsumer { )) })?; - rt.block_on(self.run_pipeline()) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}"))) - } - - /// Get results captured by a MockSink. - fn get_mock_sink_results(&self) -> PyResult>> { - match &self.sink { - Some(PipelineSink::Mock(handler)) => Ok(handler.get_results()), - _ => Err(pyo3::exceptions::PyRuntimeError::new_err( - "Sink is not a MockSink", - )), - } + rt.block_on(async { + self.install_signal_handlers(); + loop { + // Build fresh stages and sink for this assignment + let (stages, sink) = Python::attach(|py| { + let stages: Vec = self + .operators + .iter() + .map(|op| op.get().build_stage(py)) + .collect(); + let sink = self.sink_operator.as_ref().map(|s| s.get().build_sink(py)); + (stages, sink) + }); + + let exit = Self::run_pipeline(&self.source, &stages, sink.as_ref()) + .await + .map_err(|e| { + pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}")) + })?; + + match exit { + PipelineExit::Rebalance => { + tracing::info!("Rebalance detected, restarting pipeline"); + continue; + } + PipelineExit::Shutdown | PipelineExit::Complete => { + return Ok(()); + } + } + } + }) } } impl PullConsumer { - /// Construct from Rust with an injected source (for testing). - pub fn with_source( - source: impl PullSource + 'static, - stages: Vec, - sink: Option, - ) -> Self { - Self { - source: Arc::new(source), - stages, - sink, - } - } - - pub async fn run_pipeline(&self) -> Result<(), Box> { - // Install signal handlers for graceful shutdown - self.install_signal_handlers(); - - let committer = self.source.committer(); + /// Run a single pipeline iteration. + pub async fn run_pipeline( + source: &Arc, + stages: &[PipelineStage], + sink: Option<&PipelineSink>, + ) -> Result> { + let committer = source.committer(); let error_handler = LogHandler; let commit_interval = Duration::from_secs(5); let mut tracker = OffsetTracker::new(commit_interval, committer); - let source_stream = self.source.stream().map(|result| match result { + let source_stream = source.stream().map(|result| match result { StageResult::Emit(envelope) => { StageResult::Emit(envelope.map_payload(PipelineValue::Raw)) } @@ -105,16 +113,17 @@ impl PullConsumer { reason, }, StageResult::Fail(e) => StageResult::Fail(e), + StageResult::Exit(reason) => StageResult::Exit(reason), }); let mut stream: Pin> + '_>> = Box::pin(source_stream); - for stage in &self.stages { + for stage in stages { stream = Box::pin(stream.apply(stage)); } - if let Some(sink) = &self.sink { + if let Some(sink) = sink { stream .on_next(sink) .on_reject(&error_handler) @@ -126,7 +135,6 @@ impl PullConsumer { } fn install_signal_handlers(&self) { - // SIGINT (Ctrl+C) let source = self.source.clone(); tokio::spawn(async move { if tokio::signal::ctrl_c().await.is_ok() { @@ -135,7 +143,6 @@ impl PullConsumer { } }); - // SIGTERM (K8s pod termination) #[cfg(unix)] { let source = self.source.clone(); diff --git a/sentry_streams/src/pull/pull_operator.rs b/sentry_streams/src/pull/pull_operator.rs index e605b6dd..4dd66726 100644 --- a/sentry_streams/src/pull/pull_operator.rs +++ b/sentry_streams/src/pull/pull_operator.rs @@ -2,7 +2,7 @@ use pyo3::prelude::*; use super::gcs_client::GcsClient; use super::gcs_sink_handler::GcsSinkHandler; -use super::pipeline_sink::{MockSinkHandler, PipelineSink}; +use super::pipeline_sink::PipelineSink; use super::pipeline_stage::PipelineStage; use super::stages::batch::BatchAccumulatorStage; use super::stages::header_filter::HeaderFilterStage; @@ -32,9 +32,6 @@ pub enum PullOperator { bucket: String, object_generator: Py, }, - - #[pyo3(constructor = ())] - MockSink {}, } impl PullOperator { @@ -59,7 +56,7 @@ impl PullOperator { name.clone(), schema.clone(), )), - PullOperator::GcsSink { .. } | PullOperator::MockSink { .. } => { + PullOperator::GcsSink { .. } => { panic!("Sink operators are not stages — use build_sink()") } } @@ -74,7 +71,6 @@ impl PullOperator { let client = GcsClient::with_defaults(bucket.clone()); PipelineSink::Gcs(GcsSinkHandler::new(client, object_generator.clone_ref(py))) } - PullOperator::MockSink {} => PipelineSink::Mock(MockSinkHandler::new()), _ => panic!("build_sink() called on non-sink operator"), } } diff --git a/sentry_streams/src/pull/pull_source.rs b/sentry_streams/src/pull/pull_source.rs index 6d61e561..f16df0d7 100644 --- a/sentry_streams/src/pull/pull_source.rs +++ b/sentry_streams/src/pull/pull_source.rs @@ -5,10 +5,10 @@ use std::sync::{Arc, Mutex}; use futures::Stream; use pyo3::prelude::*; use sentry_arroyo::backends::kafka::types::{Headers, KafkaPayload}; -use sentry_arroyo::processing::strategies::offset_tracker::OffsetCommitter; use sentry_arroyo::processing::stream::source::KafkaSource; -use sentry_arroyo::processing::stream::PullSource; -use sentry_arroyo::processing::stream::{MessageMetadata, PipelineEnvelope, StageResult}; +use sentry_arroyo::processing::stream::{ + MessageMetadata, OffsetCommitter, PipelineEnvelope, PullSource, StageResult, +}; use sentry_arroyo::types::{Partition, Topic}; /// A test message with payload bytes and optional headers. diff --git a/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py b/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py index 6c03657e..87a9dd38 100644 --- a/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py +++ b/sentry_streams/tests/adapters/arroyo/test_pull_adapter.py @@ -192,23 +192,25 @@ def test_pull_adapter_batch_parser_converts_to_map(): # ── E2E: Python → Rust pipeline execution ─────────────────────── +# Shared capture list for e2e tests — populated by capturing Map steps. +_captured: list = [] -def _message_aware_transform(msg): - """Transform that reads msg.payload (like real pipeline callables do). - Appends '-processed' to each bytes item in the batch.""" - items = msg.payload # Message[Sequence[bytes]] -> Sequence[bytes] - return [item + b"-processed" if isinstance(item, bytes) else item for item in items] + +def _capture_and_passthrough(msg): + """Map function that captures msg.payload into _captured, then returns it.""" + _captured.append(msg.payload) + return msg.payload def test_pull_consumer_e2e_python_to_rust(): - """Full e2e: Python DSL → PullBasedAdapter → PullConsumer → stages → sink. + """Full e2e: Python DSL → PullBasedAdapter → PullConsumer → stages → verify output. - Pipeline: HeaderFilter → Batch(2) → PyCallable(transform) → MockSink + Pipeline: HeaderFilter → Batch(2) → PyCallable(capture) → no sink - The transform function receives a Message wrapper (because the adapter - sets schema on PyCallable operators), accesses .payload, and returns - the transformed data. PyCallableStage re-wraps the result in a new Message. + Uses a capturing Map step to verify data flows through the pipeline. + No mock sink needed — the capture happens in a regular pipeline stage. """ + _captured.clear() pipeline: Pipeline[bytes] = ( streaming_source(name="kafka", stream_name="test-topic") @@ -220,7 +222,7 @@ def test_pull_consumer_e2e_python_to_rust(): ) ) .apply(Batch(name="batcher", batch_size=2)) - .apply(Map(name="transformer", function=_message_aware_transform)) + .apply(Map(name="capture", function=_capture_and_passthrough)) .sink( GCSSink( name="gcs_sink", @@ -230,7 +232,6 @@ def test_pull_consumer_e2e_python_to_rust(): ) ) - # Translate DSL to PullOperator list via the adapter adapter = PullBasedAdapter.build( { "steps_config": { @@ -249,7 +250,6 @@ def test_pull_consumer_e2e_python_to_rust(): py_callable_step = adapter._steps[2] # HeaderFilter, Batch, PyCallable assert py_callable_step.schema == "test-topic" - # Swap source and sink for testing source = PullSourceConfig.Test( messages=[ PyTestMessage(payload=b"span-0", headers={"item_type": b"1"}), @@ -259,28 +259,24 @@ def test_pull_consumer_e2e_python_to_rust(): PyTestMessage(payload=b"span-4", headers={"item_type": b"1"}), ] ) - mock_sink = PullOperator.MockSink() consumer = PullConsumer( source=source, steps=adapter._steps, - sink=mock_sink, + sink=None, # no sink — capture step verifies output ) consumer.run() - # 4 matching messages, batch size 2 → 2 batches - # Each batch goes through _message_aware_transform (PyCallableStage with Message wrapping) - # MockSink captures the output - results = consumer.get_mock_sink_results() - assert len(results) == 2, f"Expected 2 batches, got {len(results)}" + # 4 matching messages, batch size 2 → 2 batches captured + assert len(_captured) == 2, f"Expected 2 batches, got {len(_captured)}" def test_pull_consumer_e2e_complex_steps(): """E2E test with real ComplexStep conversions: BatchParser and ParquetSerializer. - Pipeline: Batch(2) → BatchParser[TraceItem] → Map(extract_org_id) → ParquetSerializer → MockSink + Pipeline: Batch(2) → BatchParser[TraceItem] → Map(extract_org_id) → ParquetSerializer → Map(capture) - This tests that: + Tests that: - BatchParser.convert() produces a Map(batch_msg_parser) that works with Message wrapping - batch_msg_parser uses msg.schema to find the codec and parses protobuf bytes - ParquetSerializer.convert() produces a Map(serialize_to_parquet) that serializes to parquet @@ -291,11 +287,17 @@ def test_pull_consumer_e2e_complex_steps(): from sentry_streams.pipeline.datatypes import Uint64 from sentry_streams.pipeline.pipeline import BatchParser, ParquetSerializer - # Create a simple processor that extracts org_id into a dict + _captured.clear() + def extract_org_id(msg): """Map function: Sequence[TraceItem] → list[dict]""" return [{"org_id": item.organization_id} for item in msg.payload] + def capture_parquet(msg): + """Capture parquet bytes output.""" + _captured.append(msg.payload) + return msg.payload + pipeline: Pipeline[bytes] = ( streaming_source(name="kafka", stream_name="snuba-items") .apply(Batch(name="batcher", batch_size=2)) @@ -307,6 +309,7 @@ def extract_org_id(msg): schema_fields={"org_id": Uint64()}, ) ) + .apply(Map(name="capture", function=capture_parquet)) .sink( GCSSink( name="gcs_sink", @@ -330,7 +333,6 @@ def extract_org_id(msg): ) iterate_edges(pipeline, RuntimeTranslator(adapter)) - # Create test messages: real serialized TraceItem protobufs item1 = TraceItemProto() item1.organization_id = 42 item1.trace_id = b"0123456789abcdef" @@ -355,21 +357,17 @@ def extract_org_id(msg): PyTestMessage(payload=item4.SerializeToString(), headers={}), ] ) - mock_sink = PullOperator.MockSink() consumer = PullConsumer( source=source, steps=adapter._steps, - sink=mock_sink, + sink=None, ) consumer.run() - # 4 messages, batch size 2 → 2 batches - # Each batch: batch_msg_parser (protobuf decode) → extract_org_id → parquet serialize - # MockSink should capture 2 parquet byte blobs - results = consumer.get_mock_sink_results() - assert len(results) == 2, f"Expected 2 parquet outputs, got {len(results)}" + # 4 messages, batch size 2 → 2 batches → 2 parquet outputs captured + assert len(_captured) == 2, f"Expected 2 parquet outputs, got {len(_captured)}" # Verify the results are actual parquet bytes (magic number: PAR1) - for i, result_bytes in enumerate(results): + for i, result_bytes in enumerate(_captured): assert result_bytes[:4] == b"PAR1", f"Result {i} doesn't start with PAR1 magic" From 4cc453f9f916b8044aed4129ffa6d294baa86f8a Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:52:35 -0700 Subject: [PATCH 7/9] Cleanup. --- sentry_streams/src/pull/mod.rs | 18 ++++++++---------- sentry_streams/src/pull/pipeline_stage.rs | 2 +- sentry_streams/src/pull/pipeline_value.rs | 6 ------ sentry_streams/src/pull/stages/batch.rs | 2 +- .../src/pull/stages/header_filter.rs | 2 +- sentry_streams/src/pull/stages/py_callable.rs | 9 ++++----- 6 files changed, 15 insertions(+), 24 deletions(-) diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index c799ca15..a130117f 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -13,18 +13,24 @@ pub mod stages; #[cfg(test)] mod tests { use std::collections::HashMap; + use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::time::Duration; use futures::stream; + use futures::stream::Stream; + use futures::StreamExt; use sentry_arroyo::backends::kafka::types::{Headers, KafkaPayload}; use sentry_arroyo::processing::stream::{ - LogHandler, MessageMetadata, PipelineEnvelope, PipelineExt, Stage, StageResult, + LogHandler, MessageMetadata, OffsetCommitter, OffsetTracker, PipelineEnvelope, PipelineExt, + PullSource, Stage, StageResult, }; - use sentry_arroyo::processing::stream::{OffsetCommitter, OffsetTracker}; use sentry_arroyo::types::{Partition, Topic}; + use super::pipeline_stage::PipelineStage; use super::pipeline_value::{PipelineValue, PipelineValueCaster}; + use super::pull_consumer::PullConsumer; + use super::pull_operator::PullOperator; use super::stages::batch::BatchAccumulatorStage; use super::stages::header_filter::HeaderFilterStage; @@ -319,14 +325,6 @@ mod tests { // ── PullConsumer e2e test ─────────────────────────────────────── - use super::pipeline_stage::PipelineStage; - use super::pull_consumer::PullConsumer; - use super::pull_operator::PullOperator; - use futures::stream::Stream; - use futures::StreamExt; - use sentry_arroyo::processing::stream::PullSource; - use std::pin::Pin; - /// Test source that drains its messages on first stream() call. /// Wraps committer in Arc so it can be inspected after run. struct TestSource { diff --git a/sentry_streams/src/pull/pipeline_stage.rs b/sentry_streams/src/pull/pipeline_stage.rs index 50d81882..4397bd1d 100644 --- a/sentry_streams/src/pull/pipeline_stage.rs +++ b/sentry_streams/src/pull/pipeline_stage.rs @@ -31,7 +31,7 @@ impl Stage for PipelineStage { } } - fn name(&self) -> &'static str { + fn name(&self) -> &str { match self { PipelineStage::HeaderFilter(s) => s.name(), PipelineStage::Batch(s) => s.name(), diff --git a/sentry_streams/src/pull/pipeline_value.rs b/sentry_streams/src/pull/pipeline_value.rs index af28fe5a..5e4b7e23 100644 --- a/sentry_streams/src/pull/pipeline_value.rs +++ b/sentry_streams/src/pull/pipeline_value.rs @@ -36,12 +36,6 @@ impl fmt::Debug for PipelineValue { } } -// PipelineValue is Send because all variants are Send: -// - KafkaPayload is Send -// - Box is Send -// - Py is Send (PyO3 guarantees this) -unsafe impl Send for PipelineValue {} - /// Error returned when a stage receives an unexpected PipelineValue variant. #[derive(Debug)] pub struct DowncastError { diff --git a/sentry_streams/src/pull/stages/batch.rs b/sentry_streams/src/pull/stages/batch.rs index 085e4d31..f965e105 100644 --- a/sentry_streams/src/pull/stages/batch.rs +++ b/sentry_streams/src/pull/stages/batch.rs @@ -99,7 +99,7 @@ impl Stage for BatchAccumulatorStage { } } - fn name(&self) -> &'static str { + fn name(&self) -> &str { "batch_accumulator" } } diff --git a/sentry_streams/src/pull/stages/header_filter.rs b/sentry_streams/src/pull/stages/header_filter.rs index ef708629..7be15261 100644 --- a/sentry_streams/src/pull/stages/header_filter.rs +++ b/sentry_streams/src/pull/stages/header_filter.rs @@ -61,7 +61,7 @@ impl Stage for HeaderFilterStage { } } - fn name(&self) -> &'static str { + fn name(&self) -> &str { "header_filter" } } diff --git a/sentry_streams/src/pull/stages/py_callable.rs b/sentry_streams/src/pull/stages/py_callable.rs index 3d8c7820..13b73105 100644 --- a/sentry_streams/src/pull/stages/py_callable.rs +++ b/sentry_streams/src/pull/stages/py_callable.rs @@ -12,16 +12,15 @@ use crate::pull::pipeline_value_converter::PipelineValueConverter; /// `Message[T]` with `.payload`, `.headers`, `.timestamp`, `.schema`. pub struct PyCallableStage { callable: Py, - stage_name: &'static str, + stage_name: String, schema: String, } impl PyCallableStage { pub fn new(callable: Py, name: impl Into, schema: impl Into) -> Self { - let leaked: &'static str = Box::leak(name.into().into_boxed_str()); Self { callable, - stage_name: leaked, + stage_name: name.into(), schema: schema.into(), } } @@ -55,7 +54,7 @@ impl Stage for PyCallableStage { } } - fn name(&self) -> &'static str { - self.stage_name + fn name(&self) -> &str { + &self.stage_name } } From a43447618c5681bac9bbd560b19c48b637132b4c Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:05:08 -0700 Subject: [PATCH 8/9] Always signal drain on pipeline complete regardless of outcome. --- sentry_streams/src/pull/mod.rs | 6 ++++++ sentry_streams/src/pull/pull_consumer.rs | 14 +++++++++----- sentry_streams/src/pull/pull_source.rs | 4 ++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index a130117f..edcb367f 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -356,6 +356,10 @@ mod tests { fn shutdown(&self) { // No-op for test source } + + fn signal_drain_complete(&self) { + // No-op for test source + } } /// Helper to create a raw StageResult (not wrapped in PipelineValue). @@ -472,6 +476,8 @@ mod tests { fn shutdown(&self) { self.cancel.cancel(); } + + fn signal_drain_complete(&self) {} } let messages = vec![ diff --git a/sentry_streams/src/pull/pull_consumer.rs b/sentry_streams/src/pull/pull_consumer.rs index 7a82da7c..7ee93c27 100644 --- a/sentry_streams/src/pull/pull_consumer.rs +++ b/sentry_streams/src/pull/pull_consumer.rs @@ -65,11 +65,15 @@ impl PullConsumer { (stages, sink) }); - let exit = Self::run_pipeline(&self.source, &stages, sink.as_ref()) - .await - .map_err(|e| { - pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}")) - })?; + let exit = Self::run_pipeline(&self.source, &stages, sink.as_ref()).await; + + // Always signal drain complete — unblocks the rebalance callback + // if one is waiting. No-op if no rebalance is in progress. + self.source.signal_drain_complete(); + + let exit = exit.map_err(|e| { + pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}")) + })?; match exit { PipelineExit::Rebalance => { diff --git a/sentry_streams/src/pull/pull_source.rs b/sentry_streams/src/pull/pull_source.rs index f16df0d7..5d4d5ae5 100644 --- a/sentry_streams/src/pull/pull_source.rs +++ b/sentry_streams/src/pull/pull_source.rs @@ -115,6 +115,10 @@ impl PullSource for VecSource { fn shutdown(&self) { // No-op for test source } + + fn signal_drain_complete(&self) { + // No-op for test source + } } /// No-op committer for testing. From e4c8c4865fe0f0dffb104615b3a4888f136ad0d3 Mon Sep 17 00:00:00 2001 From: tryangul <11639460+tryangul@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:21:33 -0700 Subject: [PATCH 9/9] Revert draining on rebalance. --- sentry_streams/src/pull/mod.rs | 6 ------ sentry_streams/src/pull/pull_consumer.rs | 14 +++++--------- sentry_streams/src/pull/pull_source.rs | 4 ---- 3 files changed, 5 insertions(+), 19 deletions(-) diff --git a/sentry_streams/src/pull/mod.rs b/sentry_streams/src/pull/mod.rs index edcb367f..a130117f 100644 --- a/sentry_streams/src/pull/mod.rs +++ b/sentry_streams/src/pull/mod.rs @@ -356,10 +356,6 @@ mod tests { fn shutdown(&self) { // No-op for test source } - - fn signal_drain_complete(&self) { - // No-op for test source - } } /// Helper to create a raw StageResult (not wrapped in PipelineValue). @@ -476,8 +472,6 @@ mod tests { fn shutdown(&self) { self.cancel.cancel(); } - - fn signal_drain_complete(&self) {} } let messages = vec![ diff --git a/sentry_streams/src/pull/pull_consumer.rs b/sentry_streams/src/pull/pull_consumer.rs index 7ee93c27..7a82da7c 100644 --- a/sentry_streams/src/pull/pull_consumer.rs +++ b/sentry_streams/src/pull/pull_consumer.rs @@ -65,15 +65,11 @@ impl PullConsumer { (stages, sink) }); - let exit = Self::run_pipeline(&self.source, &stages, sink.as_ref()).await; - - // Always signal drain complete — unblocks the rebalance callback - // if one is waiting. No-op if no rebalance is in progress. - self.source.signal_drain_complete(); - - let exit = exit.map_err(|e| { - pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}")) - })?; + let exit = Self::run_pipeline(&self.source, &stages, sink.as_ref()) + .await + .map_err(|e| { + pyo3::exceptions::PyRuntimeError::new_err(format!("Pipeline failed: {e}")) + })?; match exit { PipelineExit::Rebalance => { diff --git a/sentry_streams/src/pull/pull_source.rs b/sentry_streams/src/pull/pull_source.rs index 5d4d5ae5..f16df0d7 100644 --- a/sentry_streams/src/pull/pull_source.rs +++ b/sentry_streams/src/pull/pull_source.rs @@ -115,10 +115,6 @@ impl PullSource for VecSource { fn shutdown(&self) { // No-op for test source } - - fn signal_drain_complete(&self) { - // No-op for test source - } } /// No-op committer for testing.