From f6d836a2cad4d57c6abd214cbe82dfb12be421ac Mon Sep 17 00:00:00 2001 From: avirajkhare00 Date: Tue, 21 Jul 2026 11:21:09 +0530 Subject: [PATCH 1/6] feat(connectors): add Apache Airflow trigger sink Add a sink plugin that consumes Iggy messages and creates Airflow DAG runs via the stable REST API, with basic/bearer auth, deterministic dag_run_id for idempotent retries, and WireMock integration tests. Closes #3715 --- Cargo.lock | 24 + Cargo.toml | 1 + .../connectors/airflow_sink.toml | 53 + core/connectors/sinks/README.md | 1 + core/connectors/sinks/airflow_sink/Cargo.toml | 58 ++ core/connectors/sinks/airflow_sink/README.md | 127 +++ .../connectors/sinks/airflow_sink/config.toml | 77 ++ core/connectors/sinks/airflow_sink/src/lib.rs | 984 ++++++++++++++++++ .../tests/connectors/airflow/airflow_sink.rs | 106 ++ .../tests/connectors/airflow/mod.rs | 20 + .../tests/connectors/airflow/sink.toml | 20 + .../wiremock/mappings/accept-dag-runs.json | 16 + .../wiremock/mappings/accept-version.json | 16 + .../connectors/fixtures/airflow/container.rs | 197 ++++ .../tests/connectors/fixtures/airflow/mod.rs | 21 + .../tests/connectors/fixtures/airflow/sink.rs | 81 ++ .../tests/connectors/fixtures/mod.rs | 2 + core/integration/tests/connectors/mod.rs | 1 + 18 files changed, 1805 insertions(+) create mode 100644 core/connectors/runtime/example_config/connectors/airflow_sink.toml create mode 100644 core/connectors/sinks/airflow_sink/Cargo.toml create mode 100644 core/connectors/sinks/airflow_sink/README.md create mode 100644 core/connectors/sinks/airflow_sink/config.toml create mode 100644 core/connectors/sinks/airflow_sink/src/lib.rs create mode 100644 core/integration/tests/connectors/airflow/airflow_sink.rs create mode 100644 core/integration/tests/connectors/airflow/mod.rs create mode 100644 core/integration/tests/connectors/airflow/sink.toml create mode 100644 core/integration/tests/connectors/airflow/wiremock/mappings/accept-dag-runs.json create mode 100644 core/integration/tests/connectors/airflow/wiremock/mappings/accept-version.json create mode 100644 core/integration/tests/connectors/fixtures/airflow/container.rs create mode 100644 core/integration/tests/connectors/fixtures/airflow/mod.rs create mode 100644 core/integration/tests/connectors/fixtures/airflow/sink.rs diff --git a/Cargo.lock b/Cargo.lock index 523f27c9db..b34875a9af 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6845,6 +6845,30 @@ dependencies = [ "uuid", ] +[[package]] +name = "iggy_connector_airflow_sink" +version = "0.4.1-edge.1" +dependencies = [ + "async-trait", + "base64", + "bytes", + "dashmap", + "humantime", + "iggy_common", + "iggy_connector_sdk", + "reqwest 0.13.4", + "reqwest-middleware", + "reqwest-retry", + "reqwest-tracing", + "secrecy", + "serde", + "serde_json", + "simd-json", + "tokio", + "toml 1.1.2+spec-1.1.0", + "tracing", +] + [[package]] name = "iggy_connector_clickhouse_sink" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 87a96add3e..6c5ae17f20 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ members = [ "core/configs_derive", "core/connectors/runtime", "core/connectors/sdk", + "core/connectors/sinks/airflow_sink", "core/connectors/sinks/clickhouse_sink", "core/connectors/sinks/delta_sink", "core/connectors/sinks/doris_sink", diff --git a/core/connectors/runtime/example_config/connectors/airflow_sink.toml b/core/connectors/runtime/example_config/connectors/airflow_sink.toml new file mode 100644 index 0000000000..84ed3e9b89 --- /dev/null +++ b/core/connectors/runtime/example_config/connectors/airflow_sink.toml @@ -0,0 +1,53 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +type = "sink" +key = "airflow" +enabled = true +version = 0 +name = "Airflow trigger" +path = "target/release/libiggy_connector_airflow_sink" +verbose = false + +[[streams]] +stream = "example_stream" +topics = ["example_topic"] +schema = "json" +batch_length = 50 +poll_interval = "100ms" +consumer_group = "airflow_sink_connector" + +[plugin_config] +base_url = "http://localhost:8080" +dag_id = "example_dag" +api_prefix = "/api/v1" +auth = "basic" +username = "admin" +password = "admin" +# dag_id_header = "airflow_dag_id" +conf_mode = "payload" +include_iggy_metadata_in_conf = false +health_check_enabled = true +health_path = "/api/v1/version" +timeout = "30s" +max_retries = 3 +retry_delay = "1s" +retry_backoff_multiplier = 2 +max_retry_delay = "30s" +tls_danger_accept_invalid_certs = false +max_connections = 10 +verbose_logging = false diff --git a/core/connectors/sinks/README.md b/core/connectors/sinks/README.md index 57ea055490..8abc0ccf8c 100644 --- a/core/connectors/sinks/README.md +++ b/core/connectors/sinks/README.md @@ -8,6 +8,7 @@ Sink connectors are responsible for writing data from Iggy streams to external s | Sink | Description | | ---- | ----------- | +| **airflow_sink** | Triggers Apache Airflow DAG runs via the REST API (one run per message) | | **doris_sink** | Loads JSON messages into Apache Doris tables via the Stream Load HTTP API | | **elasticsearch_sink** | Sends messages to Elasticsearch indices for full-text search and analytics | | **iceberg_sink** | Writes data to Apache Iceberg tables via REST catalog with S3/GCS/Azure storage | diff --git a/core/connectors/sinks/airflow_sink/Cargo.toml b/core/connectors/sinks/airflow_sink/Cargo.toml new file mode 100644 index 0000000000..6176394268 --- /dev/null +++ b/core/connectors/sinks/airflow_sink/Cargo.toml @@ -0,0 +1,58 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "iggy_connector_airflow_sink" +version = "0.4.1-edge.1" +description = "Iggy Apache Airflow sink connector for triggering DAG runs from stream messages via the Airflow REST API." +edition = "2024" +license = "Apache-2.0" +keywords = ["iggy", "messaging", "streaming", "airflow", "sink"] +categories = ["command-line-utilities", "database", "network-programming"] +homepage = "https://iggy.apache.org" +documentation = "https://iggy.apache.org/docs" +repository = "https://github.com/apache/iggy" +readme = "../../README.md" +publish = false + +[package.metadata.cargo-machete] +ignored = ["dashmap"] + +[lib] +crate-type = ["cdylib", "lib"] + +[dependencies] +async-trait = { workspace = true } +base64 = { workspace = true } +bytes = { workspace = true } +dashmap = { workspace = true } +humantime = { workspace = true } +iggy_common = { workspace = true } +iggy_connector_sdk = { workspace = true } +reqwest = { workspace = true } +reqwest-middleware = { workspace = true } +reqwest-retry = { workspace = true } +reqwest-tracing = { workspace = true } +secrecy = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +simd-json = { workspace = true } +toml = { workspace = true } diff --git a/core/connectors/sinks/airflow_sink/README.md b/core/connectors/sinks/airflow_sink/README.md new file mode 100644 index 0000000000..01abc8e4e8 --- /dev/null +++ b/core/connectors/sinks/airflow_sink/README.md @@ -0,0 +1,127 @@ +# Apache Airflow Sink Connector + +Consumes messages from Iggy streams and triggers Apache Airflow DAG runs via the +stable REST API (`POST /api/v1/dags/{dag_id}/dagRuns`). + +## Overview + +| | | +|---|---| +| **Type** | Sink (trigger) | +| **Direction** | Iggy → Airflow | +| **API** | Airflow REST (`api_prefix` default `/api/v1`) | +| **Idempotency** | Deterministic `dag_run_id`; HTTP **409** treated as success | + +Each message becomes one DAG run. The message payload is sent as the run's +`conf` object. Retries re-use the same `dag_run_id` so redelivery does not +create duplicate runs. + +## Configuration + +```toml +type = "sink" +key = "airflow" +enabled = true +version = 0 +name = "Airflow trigger" +path = "target/release/libiggy_connector_airflow_sink" + +[[streams]] +stream = "events" +topics = ["orders"] +schema = "json" +batch_length = 50 +poll_interval = "100ms" +consumer_group = "airflow_sink_group" + +[plugin_config] +base_url = "http://localhost:8080" +dag_id = "example_dag" +api_prefix = "/api/v1" +auth = "basic" # none | basic | bearer +username = "admin" +password = "admin" +# token = "..." # when auth = bearer +# dag_id_header = "airflow_dag_id" +conf_mode = "payload" +include_iggy_metadata_in_conf = false +health_check_enabled = true +health_path = "/api/v1/version" +timeout = "30s" +max_retries = 3 +retry_delay = "1s" +retry_backoff_multiplier = 2 +max_retry_delay = "30s" +``` + +### Plugin fields + +| Field | Required | Default | Description | +|-------|:--------:|---------|-------------| +| `base_url` | yes | — | Airflow webserver base URL | +| `dag_id` | yes* | — | Default DAG id (*optional if every message sets `dag_id_header`) | +| `api_prefix` | no | `/api/v1` | REST prefix for version differences | +| `auth` | no | `none` | `none`, `basic`, or `bearer` | +| `username` / `password` | basic | — | Basic auth credentials (`password` is secret) | +| `token` | bearer | — | Bearer/JWT token (secret) | +| `dag_id_header` | no | unset | Message header that overrides `dag_id` | +| `conf_mode` | no | `payload` | Whole payload → `conf` | +| `include_iggy_metadata_in_conf` | no | `false` | Nest stream/topic/offset under `conf.iggy` | +| `health_check_enabled` | no | `true` | `GET` health path in `open()` | +| `health_path` | no | `/api/v1/version` | Path relative to `base_url` | +| `timeout` | no | `30s` | HTTP timeout | +| `max_retries` | no | `3` | Total attempts including the first | +| `retry_delay` / `max_retry_delay` | no | `1s` / `30s` | Exponential backoff bounds | +| `tls_danger_accept_invalid_certs` | no | `false` | Dev only | +| `verbose_logging` | no | `false` | Extra per-trigger logs | + +Credentials use `SecretString` and are redacted in logs and `/stats` serialization. + +## Request shape + +```http +POST {base_url}{api_prefix}/dags/{dag_id}/dagRuns +Content-Type: application/json + +{ + "dag_run_id": "iggy-{partition}-{offset}-{message_id_hex}", + "conf": { ... message payload ... } +} +``` + +### Status handling + +| Status | Behavior | +|--------|----------| +| 2xx | Success | +| 409 | Success (run already exists — idempotent replay) | +| 400 / 401 / 403 / 404 / 422 | Permanent — drop message, do not retry | +| 429 / 5xx | Transient — retry with exponential backoff | +| Network errors | Transient — retry | + +## Out of scope (v1) + +- Waiting for DAG completion +- Airflow source (task/DAG state → Iggy) +- Batching multiple messages into one DAG run +- Airflow provider package on the Airflow side + +## Build and test + +```bash +cargo build -p iggy_connector_airflow_sink +cargo test -p iggy_connector_airflow_sink +cargo clippy -p iggy_connector_airflow_sink --all-targets -- -D warnings +``` + +Integration (Docker + WireMock): + +```bash +cargo build -p iggy_connector_airflow_sink +cargo test -p integration -- connectors::airflow +``` + +## Related + +- Issue: https://github.com/apache/iggy/issues/3715 +- Roadmap: https://github.com/apache/iggy/issues/2753 diff --git a/core/connectors/sinks/airflow_sink/config.toml b/core/connectors/sinks/airflow_sink/config.toml new file mode 100644 index 0000000000..53844f33e1 --- /dev/null +++ b/core/connectors/sinks/airflow_sink/config.toml @@ -0,0 +1,77 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +type = "sink" +key = "airflow" +enabled = true +version = 0 +name = "Airflow trigger" +path = "../../target/release/libiggy_connector_airflow_sink" +verbose = false + +[[streams]] +stream = "my_stream" +topics = ["my_topic"] +schema = "json" +batch_length = 50 +poll_interval = "100ms" +consumer_group = "airflow_sink_group" + +[plugin_config] +# Required — Airflow webserver base URL (no trailing slash required). +base_url = "http://localhost:8080" + +# Required — default DAG to trigger. Override per message via dag_id_header when set. +dag_id = "example_dag" + +# API path prefix (default: /api/v1). Use for Airflow version path differences. +api_prefix = "/api/v1" + +# Auth: "none" | "basic" | "bearer" (default: none). +auth = "basic" +username = "admin" +password = "admin" +# token = "eyJ..." # when auth = "bearer" + +# Optional message header that overrides dag_id for a single trigger. +# dag_id_header = "airflow_dag_id" + +# Put whole JSON payload into TriggerDAGRunPostBody.conf (default: payload). +conf_mode = "payload" + +# Nest Iggy stream/topic/partition/offset under conf.iggy (default: false). +include_iggy_metadata_in_conf = false + +# Health check on open (default: true). Path is relative to base_url. +health_check_enabled = true +health_path = "/api/v1/version" + +# Retry configuration (max_retries = total attempts including first try). +timeout = "30s" +max_retries = 3 +retry_delay = "1s" +retry_backoff_multiplier = 2 +max_retry_delay = "30s" + +# TLS — accept invalid certs (default: false). Use only for development. +tls_danger_accept_invalid_certs = false + +# Connection pool — max idle connections per host (default: 10). +max_connections = 10 + +# Verbose request/response logging (default: false). +verbose_logging = false diff --git a/core/connectors/sinks/airflow_sink/src/lib.rs b/core/connectors/sinks/airflow_sink/src/lib.rs new file mode 100644 index 0000000000..711be5bdf2 --- /dev/null +++ b/core/connectors/sinks/airflow_sink/src/lib.rs @@ -0,0 +1,984 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Apache Airflow trigger sink: consume Iggy messages and create DAG runs via REST. + +use async_trait::async_trait; +use base64::Engine; +use base64::engine::general_purpose; +use bytes::Bytes; +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Payload, Sink, TopicMetadata, + convert::owned_value_to_serde_json, sink_connector, +}; +use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue}; +use reqwest_middleware::{ClientBuilder, ClientWithMiddleware}; +use reqwest_retry::{ + RetryTransientMiddleware, Retryable, RetryableStrategy, policies::ExponentialBackoff, +}; +use reqwest_tracing::TracingMiddleware; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use std::str::FromStr; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use tracing::{debug, error, info, warn}; + +sink_connector!(AirflowSink); + +const CONNECTOR_NAME: &str = "Airflow sink"; +const DEFAULT_TIMEOUT: &str = "30s"; +const DEFAULT_RETRY_DELAY: &str = "1s"; +const DEFAULT_MAX_RETRY_DELAY: &str = "30s"; +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_BACKOFF_MULTIPLIER: u32 = 2; +const DEFAULT_MAX_CONNECTIONS: usize = 10; +const DEFAULT_API_PREFIX: &str = "/api/v1"; +const DEFAULT_HEALTH_PATH: &str = "/api/v1/version"; +const DEFAULT_TCP_KEEPALIVE_SECS: u64 = 30; +const DEFAULT_POOL_IDLE_TIMEOUT_SECS: u64 = 90; +const MAX_CONSECUTIVE_FAILURES: u32 = 3; +const MAX_RESPONSE_LOG_BYTES: usize = 500; + +/// Authentication mode for the Airflow REST API. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AuthMode { + #[default] + None, + Basic, + Bearer, +} + +/// How message payload maps into the DAG-run `conf` body. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ConfMode { + /// Entire JSON (or text) payload becomes `conf`. + #[default] + Payload, +} + +/// Plugin config from `[plugin_config]`. +#[derive(Debug, Serialize, Deserialize)] +pub struct AirflowSinkConfig { + /// Airflow webserver base URL (required), e.g. `http://localhost:8080`. + pub base_url: String, + /// Default DAG id to trigger (required). + pub dag_id: String, + /// REST path prefix (default: `/api/v1`). + pub api_prefix: Option, + /// Auth mode: `none` | `basic` | `bearer` (default: `none`). + pub auth: Option, + /// Basic-auth username (required when `auth = basic`). + pub username: Option, + #[serde( + default, + serialize_with = "iggy_common::serde_secret::serialize_optional_secret" + )] + pub password: Option, + #[serde( + default, + serialize_with = "iggy_common::serde_secret::serialize_optional_secret" + )] + pub token: Option, + /// Message header that overrides `dag_id` when present. + pub dag_id_header: Option, + /// How payload becomes `conf` (default: `payload`). + pub conf_mode: Option, + /// Nest Iggy metadata under `conf.iggy` (default: false). + pub include_iggy_metadata_in_conf: Option, + /// Run connectivity check in `open()` (default: true). + pub health_check_enabled: Option, + /// Path relative to `base_url` for the health check (default: `/api/v1/version`). + pub health_path: Option, + pub timeout: Option, + pub max_retries: Option, + pub retry_delay: Option, + pub retry_backoff_multiplier: Option, + pub max_retry_delay: Option, + pub tls_danger_accept_invalid_certs: Option, + pub max_connections: Option, + pub verbose_logging: Option, +} + +/// Trigger sink: one Airflow DAG run per consumed message. +#[derive(Debug)] +pub struct AirflowSink { + id: u32, + base_url: String, + log_url: String, + dag_id: String, + api_prefix: String, + auth: AuthMode, + username: Option, + password: Option, + token: Option, + dag_id_header: Option, + conf_mode: ConfMode, + include_iggy_metadata_in_conf: bool, + health_check_enabled: bool, + health_path: String, + timeout: Duration, + max_retries: u32, + retry_delay: Duration, + retry_backoff_multiplier: u32, + max_retry_delay: Duration, + tls_danger_accept_invalid_certs: bool, + max_connections: usize, + verbose: bool, + request_headers: Option, + client: Option, + trigger_attempts: AtomicU64, + messages_triggered: AtomicU64, + errors_count: AtomicU64, +} + +impl AirflowSink { + pub fn new(id: u32, config: AirflowSinkConfig) -> Self { + let base_url = config.base_url.trim_end_matches('/').to_string(); + let log_url = sanitize_url_for_log(&base_url); + let dag_id = config.dag_id; + let api_prefix = + normalize_path_prefix(config.api_prefix.as_deref().unwrap_or(DEFAULT_API_PREFIX)); + let auth = config.auth.unwrap_or_default(); + let conf_mode = config.conf_mode.unwrap_or_default(); + let include_iggy_metadata_in_conf = config.include_iggy_metadata_in_conf.unwrap_or(false); + let health_check_enabled = config.health_check_enabled.unwrap_or(true); + let health_path = + normalize_path(config.health_path.as_deref().unwrap_or(DEFAULT_HEALTH_PATH)); + let timeout = parse_duration(config.timeout.as_deref(), DEFAULT_TIMEOUT); + let max_retries = config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES); + let mut retry_delay = parse_duration(config.retry_delay.as_deref(), DEFAULT_RETRY_DELAY); + let retry_backoff_multiplier = config + .retry_backoff_multiplier + .unwrap_or(DEFAULT_BACKOFF_MULTIPLIER) + .max(1); + let mut max_retry_delay = + parse_duration(config.max_retry_delay.as_deref(), DEFAULT_MAX_RETRY_DELAY); + let tls_danger_accept_invalid_certs = + config.tls_danger_accept_invalid_certs.unwrap_or(false); + let max_connections = config.max_connections.unwrap_or(DEFAULT_MAX_CONNECTIONS); + let verbose = config.verbose_logging.unwrap_or(false); + + if retry_delay > max_retry_delay { + warn!( + "{CONNECTOR_NAME} ID: {id} — retry_delay ({retry_delay:?}) exceeds \ + max_retry_delay ({max_retry_delay:?}). Swapping values." + ); + std::mem::swap(&mut retry_delay, &mut max_retry_delay); + } + + if tls_danger_accept_invalid_certs { + warn!( + "{CONNECTOR_NAME} ID: {id} — tls_danger_accept_invalid_certs is enabled. \ + TLS certificate validation is DISABLED." + ); + } + + Self { + id, + base_url, + log_url, + dag_id, + api_prefix, + auth, + username: config.username, + password: config.password, + token: config.token, + dag_id_header: config.dag_id_header.filter(|h| !h.is_empty()), + conf_mode, + include_iggy_metadata_in_conf, + health_check_enabled, + health_path, + timeout, + max_retries, + retry_delay, + retry_backoff_multiplier, + max_retry_delay, + tls_danger_accept_invalid_certs, + max_connections, + verbose, + request_headers: None, + client: None, + trigger_attempts: AtomicU64::new(0), + messages_triggered: AtomicU64::new(0), + errors_count: AtomicU64::new(0), + } + } + + fn build_client(&self) -> Result { + let raw_client = reqwest::Client::builder() + .timeout(self.timeout) + .pool_max_idle_per_host(self.max_connections) + .pool_idle_timeout(Duration::from_secs(DEFAULT_POOL_IDLE_TIMEOUT_SECS)) + .tcp_keepalive(Duration::from_secs(DEFAULT_TCP_KEEPALIVE_SECS)) + .danger_accept_invalid_certs(self.tls_danger_accept_invalid_certs) + .build() + .map_err(|e| Error::InitError(format!("Failed to build HTTP client: {e}")))?; + + let retry_policy = ExponentialBackoff::builder() + .retry_bounds(self.retry_delay, self.max_retry_delay) + .base(self.retry_backoff_multiplier) + .build_with_max_retries(self.max_retries); + + let retry_middleware = RetryTransientMiddleware::new_with_policy_and_strategy( + retry_policy, + AirflowRetryStrategy, + ); + + Ok(ClientBuilder::new(raw_client) + .with(TracingMiddleware::default()) + .with(retry_middleware) + .build()) + } + + fn client(&self) -> Result<&ClientWithMiddleware, Error> { + self.client.as_ref().ok_or_else(|| { + Error::InitError(format!( + "{CONNECTOR_NAME} client not initialized — was open() called?" + )) + }) + } + + fn build_auth_headers(&self) -> Result { + let mut headers = HeaderMap::new(); + match self.auth { + AuthMode::None => {} + AuthMode::Basic => { + let username = self.username.as_deref().ok_or_else(|| { + Error::InvalidConfigValue("username is required when auth = basic".to_string()) + })?; + let password = self.password.as_ref().ok_or_else(|| { + Error::InvalidConfigValue("password is required when auth = basic".to_string()) + })?; + let credentials = format!("{username}:{}", password.expose_secret()); + let encoded = general_purpose::STANDARD.encode(credentials.as_bytes()); + let mut value = + HeaderValue::from_str(&format!("Basic {encoded}")).map_err(|e| { + Error::InitError(format!("Invalid basic auth header value: {e}")) + })?; + value.set_sensitive(true); + headers.insert(AUTHORIZATION, value); + } + AuthMode::Bearer => { + let token = self.token.as_ref().ok_or_else(|| { + Error::InvalidConfigValue("token is required when auth = bearer".to_string()) + })?; + let mut value = HeaderValue::from_str(&format!("Bearer {}", token.expose_secret())) + .map_err(|e| { + Error::InitError(format!("Invalid bearer auth header value: {e}")) + })?; + value.set_sensitive(true); + headers.insert(AUTHORIZATION, value); + } + } + Ok(headers) + } + + fn dag_runs_url(&self, dag_id: &str) -> String { + format!( + "{}{}/dags/{}/dagRuns", + self.base_url, + self.api_prefix, + encode_path_segment(dag_id) + ) + } + + fn health_url(&self) -> String { + format!("{}{}", self.base_url, self.health_path) + } + + fn resolve_dag_id(&self, message: &ConsumedMessage) -> Result { + if let Some(header_name) = self.dag_id_header.as_deref() + && let Some(headers) = message.headers.as_ref() + { + for (key, value) in headers { + if key.to_string_value().eq_ignore_ascii_case(header_name) { + let override_id = value.to_string_value(); + if override_id.is_empty() { + return Err(Error::InvalidRecordValue(format!( + "header '{header_name}' is present but empty" + ))); + } + return Ok(override_id); + } + } + } + if self.dag_id.is_empty() { + return Err(Error::InvalidConfigValue( + "dag_id is required when no per-message header override is present".to_string(), + )); + } + Ok(self.dag_id.clone()) + } + + fn build_conf( + &self, + message: &ConsumedMessage, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + payload: Payload, + ) -> Result { + let mut conf = match self.conf_mode { + ConfMode::Payload => match payload { + Payload::Json(value) => owned_value_to_serde_json(&value), + Payload::Text(text) => serde_json::json!({ "payload": text }), + Payload::Raw(bytes) | Payload::FlatBuffer(bytes) | Payload::Avro(bytes) => { + // Binary payloads are base64 so Airflow conf stays valid JSON. + serde_json::json!({ + "payload": general_purpose::STANDARD.encode(&bytes), + "iggy_payload_encoding": "base64", + }) + } + Payload::Proto(proto) => serde_json::json!({ + "payload": general_purpose::STANDARD.encode(proto.as_bytes()), + "iggy_payload_encoding": "base64", + }), + }, + }; + + if self.include_iggy_metadata_in_conf { + let conf_obj = conf.as_object_mut().ok_or_else(|| { + Error::InvalidRecordValue( + "conf must be a JSON object to nest iggy metadata".to_string(), + ) + })?; + conf_obj.insert( + "iggy".to_string(), + serde_json::json!({ + "stream": topic_metadata.stream, + "topic": topic_metadata.topic, + "partition_id": messages_metadata.partition_id, + "offset": message.offset, + "id": format_u128_as_hex(message.id), + "timestamp": message.timestamp, + }), + ); + } + + Ok(conf) + } + + async fn trigger_one( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + mut message: ConsumedMessage, + ) -> Result<(), Error> { + let dag_id = self.resolve_dag_id(&message)?; + let payload = std::mem::replace(&mut message.payload, Payload::Raw(vec![])); + let conf = self.build_conf(&message, topic_metadata, messages_metadata, payload)?; + let dag_run_id = + build_dag_run_id(messages_metadata.partition_id, message.offset, message.id); + + let body = serde_json::json!({ + "dag_run_id": dag_run_id, + "conf": conf, + }); + let body_bytes = serde_json::to_vec(&body) + .map_err(|e| Error::Serialization(format!("dag run body: {e}")))?; + + let url = self.dag_runs_url(&dag_id); + let client = self.client()?; + let headers = self.request_headers.as_ref().ok_or_else(|| { + Error::InitError("request headers not initialized — was open() called?".to_string()) + })?; + + if self.verbose { + debug!( + "{CONNECTOR_NAME} ID: {} — triggering DAG '{}' at {} (dag_run_id={}, offset={})", + self.id, dag_id, self.log_url, dag_run_id, message.offset + ); + } + + self.trigger_attempts.fetch_add(1, Ordering::Relaxed); + + let response = client + .post(&url) + .headers(headers.clone()) + .header(CONTENT_TYPE, "application/json") + .body(Bytes::from(body_bytes)) + .send() + .await + .map_err(|e| { + self.errors_count.fetch_add(1, Ordering::Relaxed); + error!( + "{CONNECTOR_NAME} ID: {} — trigger request failed after retries: {e:#}", + self.id + ); + Error::HttpRequestFailed(format!("Airflow trigger {url}: {e}")) + })?; + + let status = response.status(); + let status_code = status.as_u16(); + + // 2xx created / accepted; 409 means this dag_run_id already exists (idempotent replay). + if status.is_success() || status_code == 409 { + if self.verbose { + debug!( + "{CONNECTOR_NAME} ID: {} — DAG '{}' trigger ok (status {})", + self.id, dag_id, status_code + ); + } + self.messages_triggered.fetch_add(1, Ordering::Relaxed); + return Ok(()); + } + + let response_body = match response.text().await { + Ok(body) => body, + Err(e) => format!(""), + }; + let truncated = truncate_response(&response_body, MAX_RESPONSE_LOG_BYTES); + + self.errors_count.fetch_add(1, Ordering::Relaxed); + + if is_permanent_status(status_code) { + error!( + "{CONNECTOR_NAME} ID: {} — permanent trigger failure for DAG '{}' \ + (status {}). Response: {truncated}", + self.id, dag_id, status_code + ); + return Err(Error::PermanentHttpError(format!( + "Airflow trigger status {status_code}: {truncated}" + ))); + } + + error!( + "{CONNECTOR_NAME} ID: {} — trigger failure for DAG '{}' (status {}). Response: {truncated}", + self.id, dag_id, status_code + ); + Err(Error::HttpRequestFailed(format!( + "Airflow trigger status {status_code}: {truncated}" + ))) + } +} + +#[async_trait] +impl Sink for AirflowSink { + async fn open(&mut self) -> Result<(), Error> { + if self.base_url.is_empty() { + return Err(Error::InitError( + "base_url is required in [plugin_config]".to_string(), + )); + } + match reqwest::Url::parse(&self.base_url) { + Ok(parsed) => { + let scheme = parsed.scheme(); + if scheme != "http" && scheme != "https" { + return Err(Error::InitError(format!( + "base_url scheme '{scheme}' is not allowed — only http/https \ + (url: '{}')", + self.log_url + ))); + } + } + Err(e) => { + return Err(Error::InitError(format!( + "base_url '{}' is not a valid URL: {e}", + self.log_url + ))); + } + } + + if self.dag_id.is_empty() && self.dag_id_header.is_none() { + return Err(Error::InitError( + "dag_id is required unless dag_id_header is set for per-message overrides" + .to_string(), + )); + } + + // Validate auth credentials early. + self.request_headers = Some(self.build_auth_headers()?); + self.client = Some(self.build_client()?); + + if self.health_check_enabled { + let client = self.client.as_ref().expect("client just built"); + let headers = self + .request_headers + .as_ref() + .expect("request_headers just built"); + let health_url = self.health_url(); + let response = client + .get(&health_url) + .headers(headers.clone()) + .send() + .await + .map_err(|e| { + Error::Connection(format!("Health check failed for '{}': {e}", self.log_url)) + })?; + let status = response.status(); + if !status.is_success() { + return Err(Error::Connection(format!( + "Health check returned status {} for '{}{}'", + status.as_u16(), + self.log_url, + self.health_path + ))); + } + info!( + "{CONNECTOR_NAME} ID: {} — health check passed (status {})", + self.id, + status.as_u16() + ); + } + + info!( + "Opened {CONNECTOR_NAME} connector ID: {} for URL: {} (dag_id: {}, auth: {:?}, \ + max_retries: {})", + self.id, self.log_url, self.dag_id, self.auth, self.max_retries + ); + Ok(()) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec, + ) -> Result<(), Error> { + if messages.is_empty() { + return Ok(()); + } + + if self.verbose { + info!( + "{CONNECTOR_NAME} ID: {} consuming {} messages from stream: {}, topic: {}, \ + partition: {}, offset: {}", + self.id, + messages.len(), + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + messages_metadata.current_offset + ); + } else { + debug!( + "{CONNECTOR_NAME} ID: {} consuming {} messages", + self.id, + messages.len() + ); + } + + let total = messages.len(); + let mut triggered = 0u64; + let mut failures = 0u64; + let mut consecutive_failures = 0u32; + let mut last_error: Option = None; + + for message in messages { + let offset = message.offset; + match self + .trigger_one(topic_metadata, &messages_metadata, message) + .await + { + Ok(()) => { + triggered += 1; + consecutive_failures = 0; + } + Err(Error::PermanentHttpError(message)) => { + // Drop permanently bad records; do not abort the batch. + error!( + "{CONNECTOR_NAME} ID: {} dropping message at offset {} (permanent): {message}", + self.id, offset + ); + failures += 1; + consecutive_failures = 0; + } + Err(error) => { + error!( + "{CONNECTOR_NAME} ID: {} failed to trigger at offset {}: {error}", + self.id, offset + ); + failures += 1; + consecutive_failures += 1; + last_error = Some(error); + if consecutive_failures >= MAX_CONSECUTIVE_FAILURES { + let processed = triggered + failures; + let skipped = (total as u64).saturating_sub(processed); + error!( + "{CONNECTOR_NAME} ID: {} aborting batch after {} consecutive failures \ + ({} remaining messages skipped)", + self.id, consecutive_failures, skipped + ); + self.errors_count.fetch_add(skipped, Ordering::Relaxed); + break; + } + } + } + } + + match last_error { + Some(error) => { + error!( + "{CONNECTOR_NAME} ID: {} partial delivery: {}/{} triggered, {} failures", + self.id, triggered, total, failures + ); + Err(error) + } + None => Ok(()), + } + } + + async fn close(&mut self) -> Result<(), Error> { + let attempts = self.trigger_attempts.load(Ordering::Relaxed); + let triggered = self.messages_triggered.load(Ordering::Relaxed); + let errors = self.errors_count.load(Ordering::Relaxed); + info!( + "Closed {CONNECTOR_NAME} connector ID: {}, trigger_attempts: {}, \ + messages_triggered: {}, errors: {}", + self.id, attempts, triggered, errors + ); + self.request_headers = None; + self.client = None; + Ok(()) + } +} + +/// Retry 429/5xx and network errors. Treat 409 as success (no retry). +struct AirflowRetryStrategy; + +impl RetryableStrategy for AirflowRetryStrategy { + fn handle(&self, res: &reqwest_middleware::Result) -> Option { + match res { + Ok(response) => { + let status = response.status().as_u16(); + if (200..300).contains(&status) || status == 409 { + return None; + } + match status { + 429 | 500 | 502 | 503 | 504 => Some(Retryable::Transient), + _ => Some(Retryable::Fatal), + } + } + Err(_) => Some(Retryable::Transient), + } + } +} + +fn is_permanent_status(status: u16) -> bool { + matches!(status, 400 | 401 | 403 | 404 | 405 | 422) +} + +/// Deterministic DAG run id for idempotent redelivery. +fn build_dag_run_id(partition_id: u32, offset: u64, message_id: u128) -> String { + format!( + "iggy-{partition_id}-{offset}-{}", + format_u128_as_hex(message_id) + ) +} + +fn format_u128_as_hex(id: u128) -> String { + format!("{id:032x}") +} + +fn parse_duration(input: Option<&str>, default: &str) -> Duration { + let raw = input.unwrap_or(default); + HumanDuration::from_str(raw) + .map(|d| *d) + .unwrap_or_else(|e| { + warn!("Invalid duration '{raw}': {e}, using default '{default}'"); + *HumanDuration::from_str(default).expect("default duration must be valid") + }) +} + +fn normalize_path_prefix(path: &str) -> String { + let trimmed = path.trim().trim_end_matches('/'); + if trimmed.is_empty() { + return DEFAULT_API_PREFIX.to_string(); + } + if trimmed.starts_with('/') { + trimmed.to_string() + } else { + format!("/{trimmed}") + } +} + +fn normalize_path(path: &str) -> String { + let trimmed = path.trim(); + if trimmed.is_empty() { + return DEFAULT_HEALTH_PATH.to_string(); + } + if trimmed.starts_with('/') { + trimmed.to_string() + } else { + format!("/{trimmed}") + } +} + +/// Minimal path-segment encoding for DAG ids (slashes and spaces). +fn encode_path_segment(value: &str) -> String { + let mut out = String::with_capacity(value.len()); + for byte in value.bytes() { + match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { + out.push(byte as char); + } + _ => { + out.push('%'); + out.push_str(&format!("{byte:02X}")); + } + } + } + out +} + +fn truncate_response(body: &str, max_len: usize) -> &str { + if body.len() <= max_len { + body + } else { + let end = body.floor_char_boundary(max_len); + &body[..end] + } +} + +fn sanitize_url_for_log(url: &str) -> String { + match reqwest::Url::parse(url) { + Ok(parsed) if parsed.username().is_empty() && parsed.password().is_none() => { + url.to_string() + } + Ok(mut parsed) => { + let _ = parsed.set_username(""); + let _ = parsed.set_password(None); + parsed.to_string() + } + Err(_) => { + if let Some(scheme_end) = url.find("://") { + let after_scheme = &url[scheme_end + 3..]; + if let Some(at_pos) = after_scheme.find('@') { + let slash_pos = after_scheme.find('/').unwrap_or(after_scheme.len()); + if at_pos < slash_pos { + return format!( + "{}{}", + &url[..scheme_end + 3], + &after_scheme[at_pos + 1..] + ); + } + } + } + url.to_string() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use iggy_connector_sdk::Schema; + use std::collections::BTreeMap; + + fn test_config() -> AirflowSinkConfig { + AirflowSinkConfig { + base_url: "http://localhost:8080".to_string(), + dag_id: "example_dag".to_string(), + api_prefix: None, + auth: Some(AuthMode::None), + username: None, + password: None, + token: None, + dag_id_header: None, + conf_mode: None, + include_iggy_metadata_in_conf: None, + health_check_enabled: Some(false), + health_path: None, + timeout: None, + max_retries: None, + retry_delay: None, + retry_backoff_multiplier: None, + max_retry_delay: None, + tls_danger_accept_invalid_certs: None, + max_connections: None, + verbose_logging: None, + } + } + + fn sample_message(offset: u64, payload: Payload) -> ConsumedMessage { + ConsumedMessage { + id: 42, + offset, + timestamp: 1_700_000_000_000_000, + origin_timestamp: 1_700_000_000_000_000, + checksum: 0, + headers: None, + payload, + } + } + + #[test] + fn given_all_none_optional_fields_when_new_should_apply_defaults() { + let sink = AirflowSink::new(1, test_config()); + assert_eq!(sink.api_prefix, "/api/v1"); + assert_eq!(sink.auth, AuthMode::None); + assert_eq!(sink.conf_mode, ConfMode::Payload); + assert!(!sink.include_iggy_metadata_in_conf); + assert!(!sink.health_check_enabled); + assert_eq!(sink.health_path, DEFAULT_HEALTH_PATH); + assert_eq!(sink.timeout, Duration::from_secs(30)); + assert_eq!(sink.max_retries, DEFAULT_MAX_RETRIES); + assert_eq!(sink.retry_delay, Duration::from_secs(1)); + assert_eq!(sink.max_retry_delay, Duration::from_secs(30)); + assert_eq!(sink.max_connections, DEFAULT_MAX_CONNECTIONS); + assert!(!sink.verbose); + assert_eq!(sink.base_url, "http://localhost:8080"); + } + + #[test] + fn given_trailing_slash_base_url_when_new_should_strip_it() { + let mut config = test_config(); + config.base_url = "http://localhost:8080/".to_string(); + let sink = AirflowSink::new(1, config); + assert_eq!(sink.base_url, "http://localhost:8080"); + } + + #[test] + fn given_partition_offset_id_when_build_dag_run_id_should_be_deterministic() { + let first = build_dag_run_id(0, 7, 42); + let second = build_dag_run_id(0, 7, 42); + assert_eq!(first, second); + assert_eq!(first, format!("iggy-0-7-{}", format_u128_as_hex(42))); + } + + #[test] + fn given_special_dag_id_when_encode_path_segment_should_percent_encode() { + assert_eq!(encode_path_segment("my dag"), "my%20dag"); + assert_eq!(encode_path_segment("a/b"), "a%2Fb"); + assert_eq!(encode_path_segment("plain_dag-1.0"), "plain_dag-1.0"); + } + + #[test] + fn given_status_codes_when_classify_should_mark_permanent() { + assert!(is_permanent_status(400)); + assert!(is_permanent_status(401)); + assert!(is_permanent_status(403)); + assert!(is_permanent_status(404)); + assert!(is_permanent_status(422)); + assert!(!is_permanent_status(409)); + assert!(!is_permanent_status(500)); + assert!(!is_permanent_status(429)); + } + + fn json_payload(raw: &str) -> Payload { + let mut bytes = raw.as_bytes().to_vec(); + Payload::Json(simd_json::to_owned_value(&mut bytes).expect("valid JSON")) + } + + #[test] + fn given_json_payload_when_build_conf_should_use_object() { + let sink = AirflowSink::new(1, test_config()); + let message = sample_message(3, json_payload(r#"{"order_id":1}"#)); + let topic = TopicMetadata { + stream: "orders".to_string(), + topic: "created".to_string(), + }; + let meta = MessagesMetadata { + partition_id: 0, + current_offset: 3, + schema: Schema::Json, + }; + let conf = sink + .build_conf(&message, &topic, &meta, json_payload(r#"{"order_id":1}"#)) + .expect("conf"); + assert_eq!(conf["order_id"], 1); + assert!(conf.get("iggy").is_none()); + } + + #[test] + fn given_include_iggy_metadata_when_build_conf_should_nest_iggy_object() { + let mut config = test_config(); + config.include_iggy_metadata_in_conf = Some(true); + let sink = AirflowSink::new(1, config); + let message = sample_message(3, json_payload(r#"{"order_id":1}"#)); + let topic = TopicMetadata { + stream: "orders".to_string(), + topic: "created".to_string(), + }; + let meta = MessagesMetadata { + partition_id: 2, + current_offset: 3, + schema: Schema::Json, + }; + let conf = sink + .build_conf(&message, &topic, &meta, json_payload(r#"{"order_id":1}"#)) + .expect("conf"); + assert_eq!(conf["order_id"], 1); + assert_eq!(conf["iggy"]["stream"], "orders"); + assert_eq!(conf["iggy"]["topic"], "created"); + assert_eq!(conf["iggy"]["partition_id"], 2); + assert_eq!(conf["iggy"]["offset"], 3); + } + + #[test] + fn given_dag_id_header_when_resolve_should_override_config() { + use iggy_common::{HeaderKey, HeaderValue}; + + let mut config = test_config(); + config.dag_id_header = Some("airflow_dag_id".to_string()); + let sink = AirflowSink::new(1, config); + + let mut headers = BTreeMap::new(); + headers.insert( + HeaderKey::try_from("airflow_dag_id").unwrap(), + HeaderValue::try_from("override_dag").unwrap(), + ); + let mut message = sample_message(0, Payload::Text("x".into())); + message.headers = Some(headers); + + let resolved = sink.resolve_dag_id(&message).expect("dag id"); + assert_eq!(resolved, "override_dag"); + } + + #[test] + fn given_config_toml_when_deserialize_should_parse_plugin_fields() { + let raw = r#" +base_url = "http://airflow:8080" +dag_id = "demo" +auth = "basic" +username = "admin" +password = "secret" +"#; + let config: AirflowSinkConfig = toml::from_str(raw).expect("toml"); + assert_eq!(config.base_url, "http://airflow:8080"); + assert_eq!(config.dag_id, "demo"); + assert_eq!(config.auth, Some(AuthMode::Basic)); + assert_eq!(config.username.as_deref(), Some("admin")); + assert_eq!( + config + .password + .as_ref() + .map(|p| p.expose_secret().to_string()), + Some("secret".to_string()) + ); + } + + #[test] + fn given_url_with_userinfo_when_sanitize_should_strip_credentials() { + let sanitized = sanitize_url_for_log("http://user:pass@localhost:8080/path"); + assert!(!sanitized.contains("user")); + assert!(!sanitized.contains("pass")); + assert!(sanitized.contains("localhost:8080")); + } + + #[test] + fn given_dag_id_when_dag_runs_url_should_include_api_prefix() { + let sink = AirflowSink::new(1, test_config()); + assert_eq!( + sink.dag_runs_url("example_dag"), + "http://localhost:8080/api/v1/dags/example_dag/dagRuns" + ); + } +} diff --git a/core/integration/tests/connectors/airflow/airflow_sink.rs b/core/integration/tests/connectors/airflow/airflow_sink.rs new file mode 100644 index 0000000000..650567a690 --- /dev/null +++ b/core/integration/tests/connectors/airflow/airflow_sink.rs @@ -0,0 +1,106 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Airflow sink connector integration tests (WireMock as Airflow REST). + +use super::TEST_MESSAGE_COUNT; +use crate::connectors::fixtures::AirflowSinkFixture; +use bytes::Bytes; +use iggy_common::{Identifier, IggyMessage, MessageClient, Partitioning}; +use integration::harness::seeds; +use integration::iggy_harness; + +/// Publishes JSON messages and asserts each becomes a DAG-run POST with conf + dag_run_id. +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/airflow/sink.toml")), + seed = seeds::connector_stream +)] +async fn given_json_messages_when_consumed_should_post_dag_runs( + harness: &TestHarness, + fixture: AirflowSinkFixture, +) { + let client = harness.root_client().await.unwrap(); + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"order_id": 1, "status": "new"}), + serde_json::json!({"order_id": 2, "status": "paid"}), + serde_json::json!({"order_id": 3, "status": "shipped"}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("serialize"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("send messages"); + + let requests = fixture + .container() + .wait_for_dag_run_requests(TEST_MESSAGE_COUNT) + .await + .expect("WireMock did not receive expected DAG-run POSTs"); + + assert_eq!(requests.len(), TEST_MESSAGE_COUNT); + + for req in &requests { + assert_eq!(req.method, "POST"); + assert!( + req.url.contains("/api/v1/dags/example_dag/dagRuns"), + "unexpected url: {}", + req.url + ); + + let body = req.body_as_json().expect("JSON body"); + assert!( + body.get("dag_run_id") + .and_then(|v| v.as_str()) + .is_some_and(|id| id.starts_with("iggy-")), + "missing deterministic dag_run_id: {body}" + ); + let conf = body + .get("conf") + .expect("conf field required") + .as_object() + .expect("conf object"); + assert!( + conf.contains_key("order_id"), + "conf should keep payload: {body}" + ); + assert!( + conf.get("iggy").is_some(), + "include_iggy_metadata_in_conf should nest conf.iggy: {body}" + ); + } +} diff --git a/core/integration/tests/connectors/airflow/mod.rs b/core/integration/tests/connectors/airflow/mod.rs new file mode 100644 index 0000000000..58ff1b85d0 --- /dev/null +++ b/core/integration/tests/connectors/airflow/mod.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod airflow_sink; + +const TEST_MESSAGE_COUNT: usize = 3; diff --git a/core/integration/tests/connectors/airflow/sink.toml b/core/integration/tests/connectors/airflow/sink.toml new file mode 100644 index 0000000000..09814e82b8 --- /dev/null +++ b/core/integration/tests/connectors/airflow/sink.toml @@ -0,0 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[connectors] +config_type = "local" +config_dir = "../connectors/sinks/airflow_sink" diff --git a/core/integration/tests/connectors/airflow/wiremock/mappings/accept-dag-runs.json b/core/integration/tests/connectors/airflow/wiremock/mappings/accept-dag-runs.json new file mode 100644 index 0000000000..08dcfae322 --- /dev/null +++ b/core/integration/tests/connectors/airflow/wiremock/mappings/accept-dag-runs.json @@ -0,0 +1,16 @@ +{ + "request": { + "method": "POST", + "urlPathPattern": "/api/v1/dags/.*/dagRuns" + }, + "response": { + "status": 200, + "headers": { + "Content-Type": "application/json" + }, + "jsonBody": { + "dag_run_id": "mock-run", + "state": "queued" + } + } +} diff --git a/core/integration/tests/connectors/airflow/wiremock/mappings/accept-version.json b/core/integration/tests/connectors/airflow/wiremock/mappings/accept-version.json new file mode 100644 index 0000000000..c801a6c36c --- /dev/null +++ b/core/integration/tests/connectors/airflow/wiremock/mappings/accept-version.json @@ -0,0 +1,16 @@ +{ + "request": { + "method": "GET", + "urlPath": "/api/v1/version" + }, + "response": { + "status": 200, + "headers": { + "Content-Type": "application/json" + }, + "jsonBody": { + "version": "2.10.0", + "git_version": "mock" + } + } +} diff --git a/core/integration/tests/connectors/fixtures/airflow/container.rs b/core/integration/tests/connectors/fixtures/airflow/container.rs new file mode 100644 index 0000000000..d6c605a778 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/airflow/container.rs @@ -0,0 +1,197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::connectors::fixtures; +use integration::harness::TestBinaryError; +use std::time::Duration; +use testcontainers_modules::testcontainers::core::WaitFor::Healthcheck; +use testcontainers_modules::testcontainers::core::wait::HealthWaitStrategy; +use testcontainers_modules::testcontainers::core::{IntoContainerPort, Mount}; +use testcontainers_modules::testcontainers::runners::AsyncRunner; +use testcontainers_modules::testcontainers::{ContainerAsync, GenericImage, ImageExt}; +use tokio::time::sleep; +use tracing::info; + +const WIREMOCK_IMAGE: &str = "docker.io/wiremock/wiremock"; +const WIREMOCK_TAG: &str = "3.13.2"; +const WIREMOCK_PORT: u16 = 8080; + +pub(super) const DEFAULT_TEST_STREAM: &str = "test_stream"; +pub(super) const DEFAULT_TEST_TOPIC: &str = "test_topic"; +pub(super) const DEFAULT_POLL_ATTEMPTS: usize = 100; +pub(super) const DEFAULT_POLL_INTERVAL_MS: u64 = 100; + +pub(super) const ENV_SINK_PATH: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_PATH"; +pub(super) const ENV_SINK_STREAMS_0_STREAM: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_STREAMS_0_STREAM"; +pub(super) const ENV_SINK_STREAMS_0_TOPICS: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_STREAMS_0_TOPICS"; +pub(super) const ENV_SINK_STREAMS_0_SCHEMA: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_STREAMS_0_SCHEMA"; +pub(super) const ENV_SINK_STREAMS_0_CONSUMER_GROUP: &str = + "IGGY_CONNECTORS_SINK_AIRFLOW_STREAMS_0_CONSUMER_GROUP"; + +pub(super) const ENV_SINK_BASE_URL: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_BASE_URL"; +pub(super) const ENV_SINK_DAG_ID: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_DAG_ID"; +pub(super) const ENV_SINK_AUTH: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_AUTH"; +pub(super) const ENV_SINK_TIMEOUT: &str = "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_TIMEOUT"; +pub(super) const ENV_SINK_MAX_RETRIES: &str = + "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_MAX_RETRIES"; +pub(super) const ENV_SINK_RETRY_DELAY: &str = + "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_RETRY_DELAY"; +pub(super) const ENV_SINK_HEALTH_CHECK: &str = + "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_HEALTH_CHECK_ENABLED"; +pub(super) const ENV_SINK_VERBOSE_LOGGING: &str = + "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_VERBOSE_LOGGING"; +pub(super) const ENV_SINK_INCLUDE_IGGY_META: &str = + "IGGY_CONNECTORS_SINK_AIRFLOW_PLUGIN_CONFIG_INCLUDE_IGGY_METADATA_IN_CONF"; + +/// WireMock stand-in for Airflow REST API endpoints used by the sink. +pub struct AirflowWireMockContainer { + #[allow(dead_code)] + container: ContainerAsync, + pub(super) base_url: String, +} + +impl AirflowWireMockContainer { + pub(super) async fn start() -> Result { + let current_dir = std::env::current_dir().map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "AirflowWireMockContainer".to_string(), + message: format!("Failed to get current dir: {e}"), + })?; + + let container = GenericImage::new(WIREMOCK_IMAGE, WIREMOCK_TAG) + .with_exposed_port(WIREMOCK_PORT.tcp()) + .with_wait_for(Healthcheck(HealthWaitStrategy::default())) + .with_mount(Mount::bind_mount( + current_dir + .join("tests/connectors/airflow/wiremock/mappings") + .to_string_lossy() + .to_string(), + "/home/wiremock/mappings", + )) + .with_container_name(fixtures::unique_container_name("wiremock-airflow")) + .start() + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "AirflowWireMockContainer".to_string(), + message: format!("Failed to start container: {e}"), + })?; + + let host = container + .get_host() + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "AirflowWireMockContainer".to_string(), + message: format!("Failed to get host: {e}"), + })?; + + let host_port = container + .get_host_port_ipv4(WIREMOCK_PORT) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "AirflowWireMockContainer".to_string(), + message: format!("Failed to get port: {e}"), + })?; + + let base_url = format!("http://{host}:{host_port}"); + info!("Airflow sink WireMock container available at {base_url}"); + + Ok(Self { + container, + base_url, + }) + } + + pub async fn get_received_requests(&self) -> Result, TestBinaryError> { + let url = format!("{}/__admin/requests", self.base_url); + let response = reqwest::get(&url) + .await + .map_err(|e| TestBinaryError::InvalidState { + message: format!("Failed to query WireMock admin API: {e}"), + })?; + + let body: serde_json::Value = + response + .json() + .await + .map_err(|e| TestBinaryError::InvalidState { + message: format!("Failed to parse WireMock admin response: {e}"), + })?; + + let empty = vec![]; + let requests = body["requests"] + .as_array() + .unwrap_or(&empty) + .iter() + .map(|r| WireMockRequest { + method: r["request"]["method"].as_str().unwrap_or("").to_string(), + url: r["request"]["url"].as_str().unwrap_or("").to_string(), + body: r["request"]["body"].as_str().unwrap_or("").to_string(), + }) + .collect(); + + Ok(requests) + } + + pub async fn wait_for_dag_run_requests( + &self, + expected: usize, + ) -> Result, TestBinaryError> { + for _ in 0..DEFAULT_POLL_ATTEMPTS { + let requests = self.get_received_requests().await?; + let dag_runs: Vec<_> = requests + .into_iter() + .filter(|r| r.method == "POST" && r.url.contains("/dagRuns")) + .collect(); + if dag_runs.len() >= expected { + info!( + "WireMock received {} DAG-run POSTs (expected {})", + dag_runs.len(), + expected + ); + return Ok(dag_runs); + } + sleep(Duration::from_millis(DEFAULT_POLL_INTERVAL_MS)).await; + } + + let actual = self + .get_received_requests() + .await? + .into_iter() + .filter(|r| r.method == "POST" && r.url.contains("/dagRuns")) + .count(); + Err(TestBinaryError::InvalidState { + message: format!( + "Expected at least {expected} DAG-run POSTs after {} attempts, got {actual}", + DEFAULT_POLL_ATTEMPTS + ), + }) + } +} + +#[derive(Debug, Clone)] +pub struct WireMockRequest { + pub method: String, + pub url: String, + pub body: String, +} + +impl WireMockRequest { + pub fn body_as_json(&self) -> Result { + serde_json::from_str(&self.body).map_err(|e| TestBinaryError::InvalidState { + message: format!("Failed to parse request body as JSON: {e}"), + }) + } +} diff --git a/core/integration/tests/connectors/fixtures/airflow/mod.rs b/core/integration/tests/connectors/fixtures/airflow/mod.rs new file mode 100644 index 0000000000..e762021f2f --- /dev/null +++ b/core/integration/tests/connectors/fixtures/airflow/mod.rs @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod container; +mod sink; + +pub use sink::AirflowSinkFixture; diff --git a/core/integration/tests/connectors/fixtures/airflow/sink.rs b/core/integration/tests/connectors/fixtures/airflow/sink.rs new file mode 100644 index 0000000000..dca2061e4f --- /dev/null +++ b/core/integration/tests/connectors/fixtures/airflow/sink.rs @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::container::{ + AirflowWireMockContainer, DEFAULT_TEST_STREAM, DEFAULT_TEST_TOPIC, ENV_SINK_AUTH, + ENV_SINK_BASE_URL, ENV_SINK_DAG_ID, ENV_SINK_HEALTH_CHECK, ENV_SINK_INCLUDE_IGGY_META, + ENV_SINK_MAX_RETRIES, ENV_SINK_PATH, ENV_SINK_RETRY_DELAY, ENV_SINK_STREAMS_0_CONSUMER_GROUP, + ENV_SINK_STREAMS_0_SCHEMA, ENV_SINK_STREAMS_0_STREAM, ENV_SINK_STREAMS_0_TOPICS, + ENV_SINK_TIMEOUT, ENV_SINK_VERBOSE_LOGGING, +}; +use async_trait::async_trait; +use integration::harness::{TestBinaryError, TestFixture}; +use std::collections::HashMap; + +/// Default Airflow sink fixture: unauthenticated WireMock backend, fixed DAG id. +pub struct AirflowSinkFixture { + container: AirflowWireMockContainer, +} + +impl AirflowSinkFixture { + pub fn container(&self) -> &AirflowWireMockContainer { + &self.container + } + + fn base_envs(container: &AirflowWireMockContainer) -> HashMap { + let mut envs = HashMap::new(); + envs.insert(ENV_SINK_BASE_URL.to_string(), container.base_url.clone()); + envs.insert(ENV_SINK_DAG_ID.to_string(), "example_dag".to_string()); + envs.insert(ENV_SINK_AUTH.to_string(), "none".to_string()); + envs.insert(ENV_SINK_TIMEOUT.to_string(), "10s".to_string()); + envs.insert(ENV_SINK_MAX_RETRIES.to_string(), "1".to_string()); + envs.insert(ENV_SINK_RETRY_DELAY.to_string(), "100ms".to_string()); + envs.insert(ENV_SINK_HEALTH_CHECK.to_string(), "true".to_string()); + envs.insert(ENV_SINK_VERBOSE_LOGGING.to_string(), "true".to_string()); + envs.insert(ENV_SINK_INCLUDE_IGGY_META.to_string(), "true".to_string()); + envs.insert( + ENV_SINK_STREAMS_0_STREAM.to_string(), + DEFAULT_TEST_STREAM.to_string(), + ); + envs.insert( + ENV_SINK_STREAMS_0_TOPICS.to_string(), + format!("[{DEFAULT_TEST_TOPIC}]"), + ); + envs.insert(ENV_SINK_STREAMS_0_SCHEMA.to_string(), "json".to_string()); + envs.insert( + ENV_SINK_STREAMS_0_CONSUMER_GROUP.to_string(), + "airflow_sink_cg".to_string(), + ); + envs.insert( + ENV_SINK_PATH.to_string(), + "../../target/debug/libiggy_connector_airflow_sink".to_string(), + ); + envs + } +} + +#[async_trait] +impl TestFixture for AirflowSinkFixture { + async fn setup() -> Result { + let container = AirflowWireMockContainer::start().await?; + Ok(Self { container }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + Self::base_envs(&self.container) + } +} diff --git a/core/integration/tests/connectors/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index 885867de96..4c04454fb0 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -17,6 +17,7 @@ use uuid::Uuid; +mod airflow; mod clickhouse; mod delta; mod doris; @@ -47,6 +48,7 @@ pub(crate) fn unique_container_name(service: &str) -> String { ) } +pub use airflow::AirflowSinkFixture; pub use clickhouse::{ ClickHouseSinkFixture, ClickHouseSinkRowBinaryFixture, ClickHouseSinkStringFixture, }; diff --git a/core/integration/tests/connectors/mod.rs b/core/integration/tests/connectors/mod.rs index a1433160b9..6f527da2c4 100644 --- a/core/integration/tests/connectors/mod.rs +++ b/core/integration/tests/connectors/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod airflow; mod api; mod clickhouse; mod delta; From 66b02f7cbbf64495e19e042a7e2d707195d2060c Mon Sep 17 00:00:00 2001 From: avirajkhare00 Date: Tue, 21 Jul 2026 11:23:24 +0530 Subject: [PATCH 2/6] docs(connectors): fix markdownlint for airflow sink README Adjust table separators and link bare URLs so MD060/MD034 pass. --- core/connectors/sinks/airflow_sink/README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/core/connectors/sinks/airflow_sink/README.md b/core/connectors/sinks/airflow_sink/README.md index 01abc8e4e8..1b3863bfc8 100644 --- a/core/connectors/sinks/airflow_sink/README.md +++ b/core/connectors/sinks/airflow_sink/README.md @@ -6,7 +6,7 @@ stable REST API (`POST /api/v1/dags/{dag_id}/dagRuns`). ## Overview | | | -|---|---| +| --- | --- | | **Type** | Sink (trigger) | | **Direction** | Iggy → Airflow | | **API** | Airflow REST (`api_prefix` default `/api/v1`) | @@ -57,7 +57,7 @@ max_retry_delay = "30s" ### Plugin fields | Field | Required | Default | Description | -|-------|:--------:|---------|-------------| +| --- | :---: | --- | --- | | `base_url` | yes | — | Airflow webserver base URL | | `dag_id` | yes* | — | Default DAG id (*optional if every message sets `dag_id_header`) | | `api_prefix` | no | `/api/v1` | REST prefix for version differences | @@ -92,7 +92,7 @@ Content-Type: application/json ### Status handling | Status | Behavior | -|--------|----------| +| --- | --- | | 2xx | Success | | 409 | Success (run already exists — idempotent replay) | | 400 / 401 / 403 / 404 / 422 | Permanent — drop message, do not retry | @@ -123,5 +123,5 @@ cargo test -p integration -- connectors::airflow ## Related -- Issue: https://github.com/apache/iggy/issues/3715 -- Roadmap: https://github.com/apache/iggy/issues/2753 +- Issue: [#3715](https://github.com/apache/iggy/issues/3715) +- Roadmap: [#2753](https://github.com/apache/iggy/issues/2753) From ad139d1d2cde7516d997d751893d8bae8c65e778 Mon Sep 17 00:00:00 2001 From: avirajkhare00 Date: Tue, 21 Jul 2026 11:28:55 +0530 Subject: [PATCH 3/6] fix(connectors): drop unused tokio dep from airflow sink cargo-machete failed CI because tokio was declared but not used. --- core/connectors/sinks/airflow_sink/Cargo.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/core/connectors/sinks/airflow_sink/Cargo.toml b/core/connectors/sinks/airflow_sink/Cargo.toml index 6176394268..3333eebf81 100644 --- a/core/connectors/sinks/airflow_sink/Cargo.toml +++ b/core/connectors/sinks/airflow_sink/Cargo.toml @@ -50,7 +50,6 @@ reqwest-tracing = { workspace = true } secrecy = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true } tracing = { workspace = true } [dev-dependencies] From 27793f6c6926a212fde9e359b3de9e4a02ba4636 Mon Sep 17 00:00:00 2001 From: avirajkhare00 Date: Tue, 21 Jul 2026 11:29:07 +0530 Subject: [PATCH 4/6] chore: update Cargo.lock after removing airflow sink tokio dep --- Cargo.lock | 1 - 1 file changed, 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index b34875a9af..83f9785c61 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6864,7 +6864,6 @@ dependencies = [ "serde", "serde_json", "simd-json", - "tokio", "toml 1.1.2+spec-1.1.0", "tracing", ] From daa2ec6d26752153fff97edf8743a63d984d33af Mon Sep 17 00:00:00 2001 From: avirajkhare00 Date: Tue, 21 Jul 2026 11:47:42 +0530 Subject: [PATCH 5/6] chore: retrigger CI after unrelated C++ e2e flake finalize_pr failed solely because LowLevelE2E_Client.GetClientsReflectsSessionRemovalAfterDisconnect failed (278 passed, 1 failed). No C++ changes in this PR. From 734d54a52a9859afdb24550d7ae1e40d41e5afd7 Mon Sep 17 00:00:00 2001 From: avirajkhare00 Date: Fri, 24 Jul 2026 04:17:46 +0530 Subject: [PATCH 6/6] fix(connectors): trigger Airflow by poll batch, not per message Airflow runs are jobs; one DAG run per message overwhelms the scheduler on busy topics. Create one run per consume batch with conf.messages, a deterministic batch dag_run_id, and fail the batch on permanent HTTP errors instead of silent drops. Addresses review feedback on #3716. --- .../connectors/airflow_sink.toml | 3 +- core/connectors/sinks/README.md | 2 +- core/connectors/sinks/airflow_sink/README.md | 72 ++- .../connectors/sinks/airflow_sink/config.toml | 9 +- core/connectors/sinks/airflow_sink/src/lib.rs | 438 +++++++++++------- .../tests/connectors/airflow/airflow_sink.rs | 78 ++-- 6 files changed, 390 insertions(+), 212 deletions(-) diff --git a/core/connectors/runtime/example_config/connectors/airflow_sink.toml b/core/connectors/runtime/example_config/connectors/airflow_sink.toml index 84ed3e9b89..1d0dc0df13 100644 --- a/core/connectors/runtime/example_config/connectors/airflow_sink.toml +++ b/core/connectors/runtime/example_config/connectors/airflow_sink.toml @@ -38,8 +38,9 @@ api_prefix = "/api/v1" auth = "basic" username = "admin" password = "admin" -# dag_id_header = "airflow_dag_id" +# dag_id_header = "airflow_dag_id" # optional: group batch by header DAG id conf_mode = "payload" +# One DAG run per poll batch; conf.messages holds the batch payloads. include_iggy_metadata_in_conf = false health_check_enabled = true health_path = "/api/v1/version" diff --git a/core/connectors/sinks/README.md b/core/connectors/sinks/README.md index 8abc0ccf8c..4f9f1c843a 100644 --- a/core/connectors/sinks/README.md +++ b/core/connectors/sinks/README.md @@ -8,7 +8,7 @@ Sink connectors are responsible for writing data from Iggy streams to external s | Sink | Description | | ---- | ----------- | -| **airflow_sink** | Triggers Apache Airflow DAG runs via the REST API (one run per message) | +| **airflow_sink** | Triggers Apache Airflow DAG runs via the REST API (one run per poll batch) | | **doris_sink** | Loads JSON messages into Apache Doris tables via the Stream Load HTTP API | | **elasticsearch_sink** | Sends messages to Elasticsearch indices for full-text search and analytics | | **iceberg_sink** | Writes data to Apache Iceberg tables via REST catalog with S3/GCS/Azure storage | diff --git a/core/connectors/sinks/airflow_sink/README.md b/core/connectors/sinks/airflow_sink/README.md index 1b3863bfc8..4c7f594556 100644 --- a/core/connectors/sinks/airflow_sink/README.md +++ b/core/connectors/sinks/airflow_sink/README.md @@ -7,14 +7,20 @@ stable REST API (`POST /api/v1/dags/{dag_id}/dagRuns`). | | | | --- | --- | -| **Type** | Sink (trigger) | +| **Type** | Sink (batch trigger) | | **Direction** | Iggy → Airflow | | **API** | Airflow REST (`api_prefix` default `/api/v1`) | -| **Idempotency** | Deterministic `dag_run_id`; HTTP **409** treated as success | +| **Unit of work** | One DAG run per poll batch (not per message) | +| **Idempotency** | Deterministic batch `dag_run_id`; HTTP **409** treated as success | -Each message becomes one DAG run. The message payload is sent as the run's -`conf` object. Retries re-use the same `dag_run_id` so redelivery does not -create duplicate runs. +Airflow runs are jobs, not stream events. Each connector poll becomes **one DAG +run** whose `conf.messages` holds the batch. Tune volume with stream +`batch_length` / `poll_interval`. Retries re-use the same `dag_run_id` so +redelivery does not create duplicate runs. + +If `dag_id_header` is set and messages resolve to different DAG ids, the poll is +split into **one run per DAG id group** (still batches of messages, never one +run per message by default). ## Configuration @@ -64,9 +70,9 @@ max_retry_delay = "30s" | `auth` | no | `none` | `none`, `basic`, or `bearer` | | `username` / `password` | basic | — | Basic auth credentials (`password` is secret) | | `token` | bearer | — | Bearer/JWT token (secret) | -| `dag_id_header` | no | unset | Message header that overrides `dag_id` | -| `conf_mode` | no | `payload` | Whole payload → `conf` | -| `include_iggy_metadata_in_conf` | no | `false` | Nest stream/topic/offset under `conf.iggy` | +| `dag_id_header` | no | unset | Message header that overrides `dag_id` (groups batch by value) | +| `conf_mode` | no | `payload` | How each message body is placed under `conf.messages[].payload` | +| `include_iggy_metadata_in_conf` | no | `false` | Nest batch stream/topic/offset range under `conf.iggy` | | `health_check_enabled` | no | `true` | `GET` health path in `open()` | | `health_path` | no | `/api/v1/version` | Path relative to `base_url` | | `timeout` | no | `30s` | HTTP timeout | @@ -84,27 +90,69 @@ POST {base_url}{api_prefix}/dags/{dag_id}/dagRuns Content-Type: application/json { - "dag_run_id": "iggy-{partition}-{offset}-{message_id_hex}", - "conf": { ... message payload ... } + "dag_run_id": "iggy-{partition}-{first_offset}-{last_offset}-{message_count}", + "conf": { + "messages": [ + { + "offset": 0, + "id": "0000...0001", + "timestamp": 1700000000000000, + "payload": { "order_id": 1 } + }, + { + "offset": 1, + "id": "0000...0002", + "timestamp": 1700000000000001, + "payload": { "order_id": 2 } + } + ], + "iggy": { + "stream": "events", + "topic": "orders", + "partition_id": 0, + "first_offset": 0, + "last_offset": 1, + "message_count": 2 + } + } } ``` +`conf.iggy` is present only when `include_iggy_metadata_in_conf = true`. + +The sink always sets `dag_run_id` explicitly. Airflow does **not** fall back to +`manual__timestamp` when the field is provided. Replaying the same batch yields +the same id; Airflow responds **409**, which this sink treats as success. + ### Status handling | Status | Behavior | | --- | --- | | 2xx | Success | | 409 | Success (run already exists — idempotent replay) | -| 400 / 401 / 403 / 404 / 422 | Permanent — drop message, do not retry | +| 400 / 401 / 403 / 404 / 422 | Permanent — fail the batch (do not advance as if triggered) | | 429 / 5xx | Transient — retry with exponential backoff | | Network errors | Transient — retry | +Permanent failures return an error for the whole batch so consumer offsets are +not committed past an untriggered poll. Fix config/auth/DAG availability and +retry. + +### Throughput guidance + +- Prefer larger `batch_length` so one Airflow run processes many messages. +- Point this sink at **work-sized** streams (export ready, file landed, window + closed), not raw high-volume domain firehoses. +- For a single-message command topic, set `batch_length = 1` so each poll is + still one batch run with one entry in `conf.messages`. + ## Out of scope (v1) - Waiting for DAG completion - Airflow source (task/DAG state → Iggy) -- Batching multiple messages into one DAG run - Airflow provider package on the Airflow side +- In-sink content filters (`apply_function`-style); use a dedicated topic or a + future transform if you need selective triggering ## Build and test diff --git a/core/connectors/sinks/airflow_sink/config.toml b/core/connectors/sinks/airflow_sink/config.toml index 53844f33e1..debb48bdea 100644 --- a/core/connectors/sinks/airflow_sink/config.toml +++ b/core/connectors/sinks/airflow_sink/config.toml @@ -35,7 +35,7 @@ consumer_group = "airflow_sink_group" # Required — Airflow webserver base URL (no trailing slash required). base_url = "http://localhost:8080" -# Required — default DAG to trigger. Override per message via dag_id_header when set. +# Required — default DAG to trigger. One DAG run per poll batch (see batch_length). dag_id = "example_dag" # API path prefix (default: /api/v1). Use for Airflow version path differences. @@ -47,13 +47,14 @@ username = "admin" password = "admin" # token = "eyJ..." # when auth = "bearer" -# Optional message header that overrides dag_id for a single trigger. +# Optional message header that overrides dag_id. Different values split the poll +# into one batch run per DAG id (still not one run per message). # dag_id_header = "airflow_dag_id" -# Put whole JSON payload into TriggerDAGRunPostBody.conf (default: payload). +# How each message body is placed under conf.messages[].payload (default: payload). conf_mode = "payload" -# Nest Iggy stream/topic/partition/offset under conf.iggy (default: false). +# Nest batch stream/topic/partition/offset range under conf.iggy (default: false). include_iggy_metadata_in_conf = false # Health check on open (default: true). Path is relative to base_url. diff --git a/core/connectors/sinks/airflow_sink/src/lib.rs b/core/connectors/sinks/airflow_sink/src/lib.rs index 711be5bdf2..e889ca1b8c 100644 --- a/core/connectors/sinks/airflow_sink/src/lib.rs +++ b/core/connectors/sinks/airflow_sink/src/lib.rs @@ -15,7 +15,11 @@ // specific language governing permissions and limitations // under the License. -//! Apache Airflow trigger sink: consume Iggy messages and create DAG runs via REST. +//! Apache Airflow trigger sink: consume Iggy message batches and create one DAG run per batch. +//! +//! Airflow runs are jobs, not events. Each `consume()` poll becomes (by default) a single DAG +//! run whose `conf.messages` holds the batch. Redelivery reuses a deterministic `dag_run_id` +//! and treats HTTP 409 as success. use async_trait::async_trait; use base64::Engine; @@ -34,6 +38,7 @@ use reqwest_retry::{ use reqwest_tracing::TracingMiddleware; use secrecy::{ExposeSecret, SecretString}; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::str::FromStr; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; @@ -52,7 +57,6 @@ const DEFAULT_API_PREFIX: &str = "/api/v1"; const DEFAULT_HEALTH_PATH: &str = "/api/v1/version"; const DEFAULT_TCP_KEEPALIVE_SECS: u64 = 30; const DEFAULT_POOL_IDLE_TIMEOUT_SECS: u64 = 90; -const MAX_CONSECUTIVE_FAILURES: u32 = 3; const MAX_RESPONSE_LOG_BYTES: usize = 500; /// Authentication mode for the Airflow REST API. @@ -65,11 +69,11 @@ pub enum AuthMode { Bearer, } -/// How message payload maps into the DAG-run `conf` body. +/// How each message payload is placed under `conf.messages[].payload`. #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum ConfMode { - /// Entire JSON (or text) payload becomes `conf`. + /// Message body becomes `payload` (JSON object as-is; text/raw wrapped as needed). #[default] Payload, } @@ -79,7 +83,7 @@ pub enum ConfMode { pub struct AirflowSinkConfig { /// Airflow webserver base URL (required), e.g. `http://localhost:8080`. pub base_url: String, - /// Default DAG id to trigger (required). + /// Default DAG id to trigger (required unless every message sets `dag_id_header`). pub dag_id: String, /// REST path prefix (default: `/api/v1`). pub api_prefix: Option, @@ -97,11 +101,12 @@ pub struct AirflowSinkConfig { serialize_with = "iggy_common::serde_secret::serialize_optional_secret" )] pub token: Option, - /// Message header that overrides `dag_id` when present. + /// Message header that overrides `dag_id`. Messages with different values are + /// grouped; each group becomes one DAG run (still batch, not per-message). pub dag_id_header: Option, - /// How payload becomes `conf` (default: `payload`). + /// How each payload is encoded under `conf.messages` (default: `payload`). pub conf_mode: Option, - /// Nest Iggy metadata under `conf.iggy` (default: false). + /// Nest batch Iggy metadata under `conf.iggy` (default: false). pub include_iggy_metadata_in_conf: Option, /// Run connectivity check in `open()` (default: true). pub health_check_enabled: Option, @@ -117,7 +122,7 @@ pub struct AirflowSinkConfig { pub verbose_logging: Option, } -/// Trigger sink: one Airflow DAG run per consumed message. +/// Trigger sink: one Airflow DAG run per consumed poll batch (per DAG id group). #[derive(Debug)] pub struct AirflowSink { id: u32, @@ -328,64 +333,146 @@ impl AirflowSink { Ok(self.dag_id.clone()) } - fn build_conf( + /// Group messages by resolved DAG id, preserving first-seen order of groups and + /// original order of messages within each group. + fn group_by_dag_id( &self, - message: &ConsumedMessage, - topic_metadata: &TopicMetadata, - messages_metadata: &MessagesMetadata, - payload: Payload, - ) -> Result { - let mut conf = match self.conf_mode { + messages: Vec, + ) -> Result)>, Error> { + let mut groups: Vec<(String, Vec)> = Vec::new(); + let mut index_by_dag: HashMap = HashMap::new(); + + for message in messages { + let dag_id = match self.resolve_dag_id(&message) { + Ok(id) => id, + Err(e) => { + error!( + "{CONNECTOR_NAME} ID: {} — cannot resolve dag_id at offset {}: {e}", + self.id, message.offset + ); + self.errors_count.fetch_add(1, Ordering::Relaxed); + return Err(e); + } + }; + if let Some(&idx) = index_by_dag.get(&dag_id) { + groups[idx].1.push(message); + } else { + index_by_dag.insert(dag_id.clone(), groups.len()); + groups.push((dag_id, vec![message])); + } + } + + Ok(groups) + } + + fn payload_value(&self, payload: Payload) -> Result { + match self.conf_mode { ConfMode::Payload => match payload { - Payload::Json(value) => owned_value_to_serde_json(&value), - Payload::Text(text) => serde_json::json!({ "payload": text }), + Payload::Json(value) => Ok(owned_value_to_serde_json(&value)), + Payload::Text(text) => Ok(serde_json::Value::String(text)), Payload::Raw(bytes) | Payload::FlatBuffer(bytes) | Payload::Avro(bytes) => { - // Binary payloads are base64 so Airflow conf stays valid JSON. - serde_json::json!({ - "payload": general_purpose::STANDARD.encode(&bytes), - "iggy_payload_encoding": "base64", - }) + Ok(serde_json::json!({ + "data": general_purpose::STANDARD.encode(&bytes), + "encoding": "base64", + })) } - Payload::Proto(proto) => serde_json::json!({ - "payload": general_purpose::STANDARD.encode(proto.as_bytes()), - "iggy_payload_encoding": "base64", - }), + Payload::Proto(proto) => Ok(serde_json::json!({ + "data": general_purpose::STANDARD.encode(proto.as_bytes()), + "encoding": "base64", + })), }, - }; + } + } + + fn build_batch_conf( + &self, + messages: &mut [ConsumedMessage], + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + ) -> Result<(serde_json::Value, u64, u64, usize), Error> { + if messages.is_empty() { + return Err(Error::InvalidRecordValue( + "cannot build conf for empty message batch".to_string(), + )); + } + + let first_offset = messages.first().map(|m| m.offset).unwrap_or(0); + let last_offset = messages.last().map(|m| m.offset).unwrap_or(first_offset); + let mut entries = Vec::with_capacity(messages.len()); + let mut skipped = 0u64; + + for message in messages.iter_mut() { + let offset = message.offset; + let id = message.id; + let timestamp = message.timestamp; + let payload = std::mem::replace(&mut message.payload, Payload::Raw(vec![])); + match self.payload_value(payload) { + Ok(payload_value) => { + entries.push(serde_json::json!({ + "offset": offset, + "id": format_u128_as_hex(id), + "timestamp": timestamp, + "payload": payload_value, + })); + } + Err(e) => { + error!( + "{CONNECTOR_NAME} ID: {} — skipping message at offset {} \ + (payload conversion failed): {e}", + self.id, offset + ); + self.errors_count.fetch_add(1, Ordering::Relaxed); + skipped += 1; + } + } + } + + if entries.is_empty() { + return Err(Error::InvalidRecordValue(format!( + "all {skipped} messages in batch failed payload conversion" + ))); + } + + let message_count = entries.len(); + let mut conf = serde_json::json!({ + "messages": entries, + }); if self.include_iggy_metadata_in_conf { - let conf_obj = conf.as_object_mut().ok_or_else(|| { - Error::InvalidRecordValue( - "conf must be a JSON object to nest iggy metadata".to_string(), - ) - })?; - conf_obj.insert( + conf.as_object_mut().expect("conf is object").insert( "iggy".to_string(), serde_json::json!({ "stream": topic_metadata.stream, "topic": topic_metadata.topic, "partition_id": messages_metadata.partition_id, - "offset": message.offset, - "id": format_u128_as_hex(message.id), - "timestamp": message.timestamp, + "first_offset": first_offset, + "last_offset": last_offset, + "message_count": message_count, }), ); } - Ok(conf) + Ok((conf, first_offset, last_offset, message_count)) } - async fn trigger_one( + async fn trigger_batch( &self, + dag_id: &str, topic_metadata: &TopicMetadata, messages_metadata: &MessagesMetadata, - mut message: ConsumedMessage, + mut messages: Vec, ) -> Result<(), Error> { - let dag_id = self.resolve_dag_id(&message)?; - let payload = std::mem::replace(&mut message.payload, Payload::Raw(vec![])); - let conf = self.build_conf(&message, topic_metadata, messages_metadata, payload)?; - let dag_run_id = - build_dag_run_id(messages_metadata.partition_id, message.offset, message.id); + let (conf, first_offset, last_offset, message_count) = + self.build_batch_conf(&mut messages, topic_metadata, messages_metadata)?; + + // Deterministic across redelivery of the same poll range — not Airflow's + // auto `manual__timestamp` form. 409 on create means this batch was already triggered. + let dag_run_id = build_batch_dag_run_id( + messages_metadata.partition_id, + first_offset, + last_offset, + message_count, + ); let body = serde_json::json!({ "dag_run_id": dag_run_id, @@ -394,7 +481,7 @@ impl AirflowSink { let body_bytes = serde_json::to_vec(&body) .map_err(|e| Error::Serialization(format!("dag run body: {e}")))?; - let url = self.dag_runs_url(&dag_id); + let url = self.dag_runs_url(dag_id); let client = self.client()?; let headers = self.request_headers.as_ref().ok_or_else(|| { Error::InitError("request headers not initialized — was open() called?".to_string()) @@ -402,8 +489,9 @@ impl AirflowSink { if self.verbose { debug!( - "{CONNECTOR_NAME} ID: {} — triggering DAG '{}' at {} (dag_run_id={}, offset={})", - self.id, dag_id, self.log_url, dag_run_id, message.offset + "{CONNECTOR_NAME} ID: {} — triggering DAG '{}' at {} \ + (dag_run_id={}, messages={}, offsets={}..{})", + self.id, dag_id, self.log_url, dag_run_id, message_count, first_offset, last_offset ); } @@ -417,12 +505,13 @@ impl AirflowSink { .send() .await .map_err(|e| { - self.errors_count.fetch_add(1, Ordering::Relaxed); + self.errors_count + .fetch_add(message_count as u64, Ordering::Relaxed); error!( - "{CONNECTOR_NAME} ID: {} — trigger request failed after retries: {e:#}", + "{CONNECTOR_NAME} ID: {} — batch trigger request failed after retries: {e:#}", self.id ); - Error::HttpRequestFailed(format!("Airflow trigger {url}: {e}")) + Error::HttpRequestFailed(format!("Airflow batch trigger {url}: {e}")) })?; let status = response.status(); @@ -432,11 +521,12 @@ impl AirflowSink { if status.is_success() || status_code == 409 { if self.verbose { debug!( - "{CONNECTOR_NAME} ID: {} — DAG '{}' trigger ok (status {})", - self.id, dag_id, status_code + "{CONNECTOR_NAME} ID: {} — DAG '{}' batch trigger ok (status {}, messages={})", + self.id, dag_id, status_code, message_count ); } - self.messages_triggered.fetch_add(1, Ordering::Relaxed); + self.messages_triggered + .fetch_add(message_count as u64, Ordering::Relaxed); return Ok(()); } @@ -446,25 +536,29 @@ impl AirflowSink { }; let truncated = truncate_response(&response_body, MAX_RESPONSE_LOG_BYTES); - self.errors_count.fetch_add(1, Ordering::Relaxed); + self.errors_count + .fetch_add(message_count as u64, Ordering::Relaxed); if is_permanent_status(status_code) { + // Auth / missing DAG / bad request: fail the consume so offsets do not + // advance past an untriggered batch (do not silent-drop). error!( - "{CONNECTOR_NAME} ID: {} — permanent trigger failure for DAG '{}' \ + "{CONNECTOR_NAME} ID: {} — permanent batch trigger failure for DAG '{}' \ (status {}). Response: {truncated}", self.id, dag_id, status_code ); return Err(Error::PermanentHttpError(format!( - "Airflow trigger status {status_code}: {truncated}" + "Airflow batch trigger status {status_code}: {truncated}" ))); } error!( - "{CONNECTOR_NAME} ID: {} — trigger failure for DAG '{}' (status {}). Response: {truncated}", + "{CONNECTOR_NAME} ID: {} — batch trigger failure for DAG '{}' (status {}). \ + Response: {truncated}", self.id, dag_id, status_code ); Err(Error::HttpRequestFailed(format!( - "Airflow trigger status {status_code}: {truncated}" + "Airflow batch trigger status {status_code}: {truncated}" ))) } } @@ -540,7 +634,7 @@ impl Sink for AirflowSink { info!( "Opened {CONNECTOR_NAME} connector ID: {} for URL: {} (dag_id: {}, auth: {:?}, \ - max_retries: {})", + mode: batch, max_retries: {})", self.id, self.log_url, self.dag_id, self.auth, self.max_retries ); Ok(()) @@ -556,12 +650,13 @@ impl Sink for AirflowSink { return Ok(()); } + let total = messages.len(); if self.verbose { info!( "{CONNECTOR_NAME} ID: {} consuming {} messages from stream: {}, topic: {}, \ - partition: {}, offset: {}", + partition: {}, offset: {} (one DAG run per DAG-id group)", self.id, - messages.len(), + total, topic_metadata.stream, topic_metadata.topic, messages_metadata.partition_id, @@ -569,70 +664,30 @@ impl Sink for AirflowSink { ); } else { debug!( - "{CONNECTOR_NAME} ID: {} consuming {} messages", - self.id, - messages.len() + "{CONNECTOR_NAME} ID: {} consuming {} messages as batch trigger(s)", + self.id, total ); } - let total = messages.len(); - let mut triggered = 0u64; - let mut failures = 0u64; - let mut consecutive_failures = 0u32; - let mut last_error: Option = None; + let groups = self.group_by_dag_id(messages)?; + let group_count = groups.len(); - for message in messages { - let offset = message.offset; - match self - .trigger_one(topic_metadata, &messages_metadata, message) + for (dag_id, group_messages) in groups { + let group_size = group_messages.len(); + if let Err(error) = self + .trigger_batch(&dag_id, topic_metadata, &messages_metadata, group_messages) .await { - Ok(()) => { - triggered += 1; - consecutive_failures = 0; - } - Err(Error::PermanentHttpError(message)) => { - // Drop permanently bad records; do not abort the batch. - error!( - "{CONNECTOR_NAME} ID: {} dropping message at offset {} (permanent): {message}", - self.id, offset - ); - failures += 1; - consecutive_failures = 0; - } - Err(error) => { - error!( - "{CONNECTOR_NAME} ID: {} failed to trigger at offset {}: {error}", - self.id, offset - ); - failures += 1; - consecutive_failures += 1; - last_error = Some(error); - if consecutive_failures >= MAX_CONSECUTIVE_FAILURES { - let processed = triggered + failures; - let skipped = (total as u64).saturating_sub(processed); - error!( - "{CONNECTOR_NAME} ID: {} aborting batch after {} consecutive failures \ - ({} remaining messages skipped)", - self.id, consecutive_failures, skipped - ); - self.errors_count.fetch_add(skipped, Ordering::Relaxed); - break; - } - } - } - } - - match last_error { - Some(error) => { error!( - "{CONNECTOR_NAME} ID: {} partial delivery: {}/{} triggered, {} failures", - self.id, triggered, total, failures + "{CONNECTOR_NAME} ID: {} failed batch trigger for DAG '{}' \ + ({} messages in group, {} groups total): {error}", + self.id, dag_id, group_size, group_count ); - Err(error) + return Err(error); } - None => Ok(()), } + + Ok(()) } async fn close(&mut self) -> Result<(), Error> { @@ -675,12 +730,17 @@ fn is_permanent_status(status: u16) -> bool { matches!(status, 400 | 401 | 403 | 404 | 405 | 422) } -/// Deterministic DAG run id for idempotent redelivery. -fn build_dag_run_id(partition_id: u32, offset: u64, message_id: u128) -> String { - format!( - "iggy-{partition_id}-{offset}-{}", - format_u128_as_hex(message_id) - ) +/// Deterministic DAG run id for a message batch (idempotent redelivery). +/// +/// Explicitly set on create — Airflow does **not** auto-assign `manual__timestamp` +/// when `dag_run_id` is provided. +fn build_batch_dag_run_id( + partition_id: u32, + first_offset: u64, + last_offset: u64, + message_count: usize, +) -> String { + format!("iggy-{partition_id}-{first_offset}-{last_offset}-{message_count}") } fn format_u128_as_hex(id: u128) -> String { @@ -809,7 +869,7 @@ mod tests { fn sample_message(offset: u64, payload: Payload) -> ConsumedMessage { ConsumedMessage { - id: 42, + id: 42 + offset as u128, offset, timestamp: 1_700_000_000_000_000, origin_timestamp: 1_700_000_000_000_000, @@ -819,6 +879,25 @@ mod tests { } } + fn json_payload(raw: &str) -> Payload { + let mut bytes = raw.as_bytes().to_vec(); + Payload::Json(simd_json::to_owned_value(&mut bytes).expect("valid JSON")) + } + + fn topic_and_meta() -> (TopicMetadata, MessagesMetadata) { + ( + TopicMetadata { + stream: "orders".to_string(), + topic: "created".to_string(), + }, + MessagesMetadata { + partition_id: 0, + current_offset: 3, + schema: Schema::Json, + }, + ) + } + #[test] fn given_all_none_optional_fields_when_new_should_apply_defaults() { let sink = AirflowSink::new(1, test_config()); @@ -846,11 +925,11 @@ mod tests { } #[test] - fn given_partition_offset_id_when_build_dag_run_id_should_be_deterministic() { - let first = build_dag_run_id(0, 7, 42); - let second = build_dag_run_id(0, 7, 42); + fn given_batch_range_when_build_dag_run_id_should_be_deterministic() { + let first = build_batch_dag_run_id(0, 7, 12, 6); + let second = build_batch_dag_run_id(0, 7, 12, 6); assert_eq!(first, second); - assert_eq!(first, format!("iggy-0-7-{}", format_u128_as_hex(42))); + assert_eq!(first, "iggy-0-7-12-6"); } #[test] @@ -872,54 +951,91 @@ mod tests { assert!(!is_permanent_status(429)); } - fn json_payload(raw: &str) -> Payload { - let mut bytes = raw.as_bytes().to_vec(); - Payload::Json(simd_json::to_owned_value(&mut bytes).expect("valid JSON")) - } - #[test] - fn given_json_payload_when_build_conf_should_use_object() { + fn given_json_batch_when_build_conf_should_wrap_messages_array() { let sink = AirflowSink::new(1, test_config()); - let message = sample_message(3, json_payload(r#"{"order_id":1}"#)); - let topic = TopicMetadata { - stream: "orders".to_string(), - topic: "created".to_string(), - }; - let meta = MessagesMetadata { - partition_id: 0, - current_offset: 3, - schema: Schema::Json, - }; - let conf = sink - .build_conf(&message, &topic, &meta, json_payload(r#"{"order_id":1}"#)) + let (topic, meta) = topic_and_meta(); + let mut messages = vec![ + sample_message(1, json_payload(r#"{"order_id":1}"#)), + sample_message(2, json_payload(r#"{"order_id":2}"#)), + ]; + let (conf, first, last, count) = sink + .build_batch_conf(&mut messages, &topic, &meta) .expect("conf"); - assert_eq!(conf["order_id"], 1); + assert_eq!(first, 1); + assert_eq!(last, 2); + assert_eq!(count, 2); + let msgs = conf["messages"].as_array().expect("messages array"); + assert_eq!(msgs.len(), 2); + assert_eq!(msgs[0]["offset"], 1); + assert_eq!(msgs[0]["payload"]["order_id"], 1); + assert_eq!(msgs[1]["payload"]["order_id"], 2); assert!(conf.get("iggy").is_none()); } #[test] - fn given_include_iggy_metadata_when_build_conf_should_nest_iggy_object() { + fn given_include_iggy_metadata_when_build_conf_should_nest_batch_iggy_object() { let mut config = test_config(); config.include_iggy_metadata_in_conf = Some(true); let sink = AirflowSink::new(1, config); - let message = sample_message(3, json_payload(r#"{"order_id":1}"#)); - let topic = TopicMetadata { - stream: "orders".to_string(), - topic: "created".to_string(), - }; - let meta = MessagesMetadata { - partition_id: 2, - current_offset: 3, - schema: Schema::Json, - }; - let conf = sink - .build_conf(&message, &topic, &meta, json_payload(r#"{"order_id":1}"#)) + let (topic, mut meta) = topic_and_meta(); + meta.partition_id = 2; + let mut messages = vec![ + sample_message(3, json_payload(r#"{"order_id":1}"#)), + sample_message(4, json_payload(r#"{"order_id":2}"#)), + ]; + let (conf, _, _, _) = sink + .build_batch_conf(&mut messages, &topic, &meta) .expect("conf"); - assert_eq!(conf["order_id"], 1); assert_eq!(conf["iggy"]["stream"], "orders"); assert_eq!(conf["iggy"]["topic"], "created"); assert_eq!(conf["iggy"]["partition_id"], 2); - assert_eq!(conf["iggy"]["offset"], 3); + assert_eq!(conf["iggy"]["first_offset"], 3); + assert_eq!(conf["iggy"]["last_offset"], 4); + assert_eq!(conf["iggy"]["message_count"], 2); + assert_eq!(conf["messages"].as_array().unwrap().len(), 2); + } + + #[test] + fn given_dag_id_header_when_group_should_split_batches_by_dag() { + use iggy_common::{HeaderKey, HeaderValue}; + + let mut config = test_config(); + config.dag_id_header = Some("airflow_dag_id".to_string()); + let sink = AirflowSink::new(1, config); + + let mut a = sample_message(0, Payload::Text("x".into())); + let mut headers_a = BTreeMap::new(); + headers_a.insert( + HeaderKey::try_from("airflow_dag_id").unwrap(), + HeaderValue::try_from("dag_a").unwrap(), + ); + a.headers = Some(headers_a); + + let mut b = sample_message(1, Payload::Text("y".into())); + let mut headers_b = BTreeMap::new(); + headers_b.insert( + HeaderKey::try_from("airflow_dag_id").unwrap(), + HeaderValue::try_from("dag_b").unwrap(), + ); + b.headers = Some(headers_b); + + let mut c = sample_message(2, Payload::Text("z".into())); + let mut headers_c = BTreeMap::new(); + headers_c.insert( + HeaderKey::try_from("airflow_dag_id").unwrap(), + HeaderValue::try_from("dag_a").unwrap(), + ); + c.headers = Some(headers_c); + + let groups = sink.group_by_dag_id(vec![a, b, c]).expect("groups"); + assert_eq!(groups.len(), 2); + assert_eq!(groups[0].0, "dag_a"); + assert_eq!(groups[0].1.len(), 2); + assert_eq!(groups[0].1[0].offset, 0); + assert_eq!(groups[0].1[1].offset, 2); + assert_eq!(groups[1].0, "dag_b"); + assert_eq!(groups[1].1.len(), 1); } #[test] diff --git a/core/integration/tests/connectors/airflow/airflow_sink.rs b/core/integration/tests/connectors/airflow/airflow_sink.rs index 650567a690..cca362e9ff 100644 --- a/core/integration/tests/connectors/airflow/airflow_sink.rs +++ b/core/integration/tests/connectors/airflow/airflow_sink.rs @@ -24,12 +24,12 @@ use iggy_common::{Identifier, IggyMessage, MessageClient, Partitioning}; use integration::harness::seeds; use integration::iggy_harness; -/// Publishes JSON messages and asserts each becomes a DAG-run POST with conf + dag_run_id. +/// Publishes JSON messages and asserts the poll becomes one DAG-run POST with conf.messages. #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/airflow/sink.toml")), seed = seeds::connector_stream )] -async fn given_json_messages_when_consumed_should_post_dag_runs( +async fn given_json_messages_when_consumed_should_post_one_batch_dag_run( harness: &TestHarness, fixture: AirflowSinkFixture, ) { @@ -66,41 +66,53 @@ async fn given_json_messages_when_consumed_should_post_dag_runs( .await .expect("send messages"); + // Batch design: N messages in one poll => one DAG-run POST, not N. let requests = fixture .container() - .wait_for_dag_run_requests(TEST_MESSAGE_COUNT) + .wait_for_dag_run_requests(1) .await - .expect("WireMock did not receive expected DAG-run POSTs"); + .expect("WireMock did not receive expected DAG-run POST"); - assert_eq!(requests.len(), TEST_MESSAGE_COUNT); + assert_eq!(requests.len(), 1); - for req in &requests { - assert_eq!(req.method, "POST"); - assert!( - req.url.contains("/api/v1/dags/example_dag/dagRuns"), - "unexpected url: {}", - req.url - ); + let req = &requests[0]; + assert_eq!(req.method, "POST"); + assert!( + req.url.contains("/api/v1/dags/example_dag/dagRuns"), + "unexpected url: {}", + req.url + ); - let body = req.body_as_json().expect("JSON body"); - assert!( - body.get("dag_run_id") - .and_then(|v| v.as_str()) - .is_some_and(|id| id.starts_with("iggy-")), - "missing deterministic dag_run_id: {body}" - ); - let conf = body - .get("conf") - .expect("conf field required") - .as_object() - .expect("conf object"); - assert!( - conf.contains_key("order_id"), - "conf should keep payload: {body}" - ); - assert!( - conf.get("iggy").is_some(), - "include_iggy_metadata_in_conf should nest conf.iggy: {body}" - ); - } + let body = req.body_as_json().expect("JSON body"); + assert!( + body.get("dag_run_id") + .and_then(|v| v.as_str()) + .is_some_and(|id| id.starts_with("iggy-")), + "missing deterministic batch dag_run_id: {body}" + ); + + let conf = body + .get("conf") + .expect("conf field required") + .as_object() + .expect("conf object"); + + let batch = conf + .get("messages") + .and_then(|v| v.as_array()) + .expect("conf.messages array required for batch trigger"); + assert_eq!( + batch.len(), + TEST_MESSAGE_COUNT, + "expected all polled messages in one DAG run conf: {body}" + ); + assert_eq!(batch[0]["payload"]["order_id"], 1); + assert_eq!(batch[1]["payload"]["order_id"], 2); + assert_eq!(batch[2]["payload"]["order_id"], 3); + + assert!( + conf.get("iggy").is_some(), + "include_iggy_metadata_in_conf should nest conf.iggy: {body}" + ); + assert_eq!(conf["iggy"]["message_count"], TEST_MESSAGE_COUNT as u64); }