ryerraguntla commented on code in PR #4394: URL: https://github.com/apache/iggy/pull/4394#discussion_r4212463823
########## core/connectors/sinks/bigquery_sink/src/schema.rs: ########## @@ -0,0 +1,769 @@ +// 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. + +//! BigQuery table schema model and its mapping to Arrow. +//! +//! The schema comes from the `tables.get` REST call as JSON. It is parsed +//! with local serde types rather than a client library model so that every +//! BigQuery type (GEOGRAPHY, RANGE, ...) deserializes, and the unsupported +//! ones can be rejected with a clear message in `open()`. +//! +//! Arrow types follow the Storage Write API Arrow mapping: +//! +//! | BigQuery | Arrow | +//! |------------|--------------------------------| +//! | STRING | Utf8 | +//! | BYTES | Binary | +//! | INT64 | Int64 | +//! | FLOAT64 | Float64 | +//! | BOOL | Boolean | +//! | NUMERIC | Decimal128(38, 9) | +//! | BIGNUMERIC | Decimal256(76, 38) | +//! | TIMESTAMP | Timestamp(Microsecond, "UTC") | +//! | DATETIME | Timestamp(Microsecond, none) | +//! | DATE | Date32 | +//! | TIME | Time64(Microsecond) | +//! | GEOGRAPHY | Utf8 (WKT) | +//! | JSON | Utf8 (JSON text) | +//! | RECORD | Struct | +//! | REPEATED | List | + +use crate::{Settings, WriteMode}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields, TimeUnit}; +use serde::Deserialize; +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +pub(crate) const UTC: &str = "UTC"; +const LIST_ITEM: &str = "item"; + +/// One table column, recursively for RECORD types. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Column { + pub name: String, + pub bq_type: BqType, + pub mode: Mode, + pub has_default: bool, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum BqType { + String, + Bytes, + Int64, + Float64, + Bool, + Numeric, + BigNumeric, + Timestamp, + Datetime, + Date, + Time, + Geography, + Json, + Record(Vec<Column>), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Mode { + Nullable, + Required, + Repeated, +} + +/// Iggy metadata written next to the payload. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MetaColumn { + Stream, + Topic, + PartitionId, + Offset, + Timestamp, + Id, + Headers, +} + +/// What the `raw` mode payload column holds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RawKind { + Json, + String, + Bytes, +} + +/// The validated mapping from messages to the target table, built once in +/// `open()`. +#[derive(Debug)] +pub(crate) struct TableLayout { + /// Payload columns in table order. Empty in `raw` mode. + pub columns: Vec<Column>, + /// Arrow field for each entry of `columns`, same index. + pub fields: Vec<FieldRef>, + /// Column name to index into `columns`. + pub index: HashMap<String, usize>, + /// Enabled metadata columns with their Arrow fields. + pub metadata: Vec<(MetaColumn, FieldRef)>, + /// `raw` mode payload column. + pub raw: Option<(RawKind, FieldRef)>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SchemaError { + Malformed(String), + EmptySchema, + UnsupportedType { + column: String, + bq_type: String, + }, + MissingColumn { + column: String, + hint: String, + }, + WrongType { + column: String, + expected: String, + actual: String, + }, + UnfillableColumn(String), +} + +impl fmt::Display for SchemaError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + SchemaError::Malformed(reason) => write!(f, "malformed table schema: {reason}"), + SchemaError::EmptySchema => write!(f, "table has no columns"), + SchemaError::UnsupportedType { column, bq_type } => write!( + f, + "column '{column}' has type {bq_type}, which the BigQuery sink does not support" + ), + SchemaError::MissingColumn { column, hint } => { + write!(f, "column '{column}' is missing from the table. {hint}") + } + SchemaError::WrongType { + column, + expected, + actual, + } => write!(f, "column '{column}' must be {expected}, found {actual}"), + SchemaError::UnfillableColumn(column) => write!( + f, + "column '{column}' is REQUIRED without a default value and raw mode never writes it" + ), + } + } +} + +impl MetaColumn { + pub(crate) const ALL: [MetaColumn; 7] = [ + MetaColumn::Stream, + MetaColumn::Topic, + MetaColumn::PartitionId, + MetaColumn::Offset, + MetaColumn::Timestamp, + MetaColumn::Id, + MetaColumn::Headers, + ]; + + pub(crate) fn name(self) -> &'static str { + match self { + MetaColumn::Stream => "iggy_stream", + MetaColumn::Topic => "iggy_topic", + MetaColumn::PartitionId => "iggy_partition_id", + MetaColumn::Offset => "iggy_offset", + MetaColumn::Timestamp => "iggy_timestamp", + MetaColumn::Id => "iggy_id", + MetaColumn::Headers => "iggy_headers", + } + } + + fn ddl_type(self) -> &'static str { + match self { + MetaColumn::Stream | MetaColumn::Topic | MetaColumn::Id => "STRING", + MetaColumn::PartitionId | MetaColumn::Offset => "INT64", + MetaColumn::Timestamp => "TIMESTAMP", + MetaColumn::Headers => "JSON", + } + } + + fn accepts(self, bq_type: &BqType) -> bool { + match self { + MetaColumn::Stream | MetaColumn::Topic | MetaColumn::Id => { + matches!(bq_type, BqType::String) + } + MetaColumn::PartitionId | MetaColumn::Offset => matches!(bq_type, BqType::Int64), + MetaColumn::Timestamp => matches!(bq_type, BqType::Timestamp), + MetaColumn::Headers => matches!(bq_type, BqType::Json | BqType::String), + } + } + + fn enabled(self, settings: &Settings) -> bool { + match self { + MetaColumn::Headers => settings.include_headers, + _ => settings.include_metadata, + } + } +} + +/// Parse the body of a `tables.get` response. +pub(crate) fn parse_table_schema(body: &[u8]) -> Result<Vec<Column>, SchemaError> { + let resource: TableResource = + serde_json::from_slice(body).map_err(|e| SchemaError::Malformed(e.to_string()))?; + let fields = resource + .schema + .map(|schema| schema.fields) + .unwrap_or_default(); + if fields.is_empty() { + return Err(SchemaError::EmptySchema); + } + fields.iter().map(Column::try_from).collect() +} + +impl TableLayout { + pub(crate) fn build(columns: Vec<Column>, settings: &Settings) -> Result<Self, SchemaError> { + let mut by_name: HashMap<&str, &Column> = + columns.iter().map(|c| (c.name.as_str(), c)).collect(); + + let mut metadata = Vec::new(); + for meta in MetaColumn::ALL { + if !meta.enabled(settings) { + continue; + } + let column = by_name + .remove(meta.name()) + .ok_or_else(|| SchemaError::MissingColumn { + column: meta.name().to_owned(), + hint: format!( + "Add `{} {}` to the table or disable it in the connector config.", + meta.name(), + meta.ddl_type() + ), + })?; + if column.mode == Mode::Repeated || !meta.accepts(&column.bq_type) { + return Err(SchemaError::WrongType { + column: column.name.clone(), + expected: meta.ddl_type().to_owned(), + actual: column.type_label(), + }); + } + metadata.push((meta, Arc::new(meta.arrow_field(column.mode)))); + } + + match settings.mode { + WriteMode::Raw => Self::build_raw(&columns, by_name, metadata, settings), + WriteMode::Mapped => Ok(Self::build_mapped(&columns, &metadata)), + } + } + + fn build_raw( + columns: &[Column], + mut remaining: HashMap<&str, &Column>, + metadata: Vec<(MetaColumn, FieldRef)>, + settings: &Settings, + ) -> Result<Self, SchemaError> { + let name = settings.payload_column.as_str(); + let column = remaining + .remove(name) + .ok_or_else(|| SchemaError::MissingColumn { + column: name.to_owned(), + hint: "Add a JSON, STRING or BYTES column for the payload or set payload_column." + .to_owned(), + })?; + let kind = match (&column.bq_type, column.mode) { + (BqType::Json, Mode::Nullable | Mode::Required) => RawKind::Json, + (BqType::String, Mode::Nullable | Mode::Required) => RawKind::String, + (BqType::Bytes, Mode::Nullable | Mode::Required) => RawKind::Bytes, + _ => { + return Err(SchemaError::WrongType { + column: column.name.clone(), + expected: "JSON, STRING or BYTES".to_owned(), + actual: column.type_label(), + }); + } + }; + + // Every other column is left out of the writer schema, so BigQuery + // fills it from its default. A REQUIRED column without one would + // reject every row. + if let Some(unfillable) = columns.iter().find(|c| { + remaining.contains_key(c.name.as_str()) && c.mode == Mode::Required && !c.has_default + }) { + return Err(SchemaError::UnfillableColumn(unfillable.name.clone())); + } + + let field = Arc::new(column.arrow_field()); + Ok(TableLayout { + columns: Vec::new(), + fields: Vec::new(), + index: HashMap::new(), + metadata, + raw: Some((kind, field)), + }) + } + + fn build_mapped(columns: &[Column], metadata: &[(MetaColumn, FieldRef)]) -> Self { + let payload_columns: Vec<Column> = columns + .iter() + .filter(|c| !metadata.iter().any(|(meta, _)| meta.name() == c.name)) + .cloned() + .collect(); + let fields = payload_columns + .iter() + .map(|c| Arc::new(c.arrow_field())) + .collect(); + let index = payload_columns + .iter() + .enumerate() + .map(|(i, c)| (c.name.clone(), i)) + .collect(); + TableLayout { + columns: payload_columns, + fields, + index, + metadata: metadata.to_vec(), + raw: None, + } + } +} + +impl Column { + /// Arrow field as the Storage Write API expects it. REQUIRED maps to a + /// non-nullable field. REPEATED maps to a list whose items are + /// non-nullable, because BigQuery arrays cannot hold NULL. + pub(crate) fn arrow_field(&self) -> Field { + let element = self.bq_type.arrow_type(); + match self.mode { + Mode::Nullable => Field::new(&self.name, element, true), + Mode::Required => Field::new(&self.name, element, false), + Mode::Repeated => Field::new( + &self.name, + DataType::List(Arc::new(Field::new(LIST_ITEM, element, false))), + true, + ), + } + } + + fn type_label(&self) -> String { + let base = self.bq_type.label(); + match self.mode { + Mode::Repeated => format!("REPEATED {base}"), + _ => base.to_owned(), + } + } +} + +impl BqType { + fn arrow_type(&self) -> DataType { + match self { + BqType::String | BqType::Geography | BqType::Json => DataType::Utf8, + BqType::Bytes => DataType::Binary, + BqType::Int64 => DataType::Int64, + BqType::Float64 => DataType::Float64, + BqType::Bool => DataType::Boolean, + BqType::Numeric => DataType::Decimal128(38, 9), Review Comment: NUMERIC/BIGNUMERIC hardcoded to Decimal128(38,9)/Decimal256(76,38) regardless of the table's real declared scale; RestField never deserializes BigQuery's actual precision/scale from tables.get (zero matches, grepped). arrow-cast 58.4.0 (the actually-locked version, verified against parse.rs:914-924) silently truncates fractional digits ########## core/connectors/sinks/bigquery_sink/src/sink.rs: ########## @@ -0,0 +1,316 @@ +// 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::client::{AppendOutcome, BigQueryClient}; +use crate::encode::{self, Chunk, RunContext}; +use crate::error::{AppendError, TableError}; +use crate::schema::{TableLayout, parse_table_schema}; +use crate::{BigQuerySink, OpenState}; +use async_trait::async_trait; +use iggy_connector_sdk::retry::{RetryPolicy, retry_async}; +use iggy_connector_sdk::{ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata}; +use std::sync::atomic::Ordering; +use tracing::{debug, error, info, warn}; + +#[async_trait] +impl Sink for BigQuerySink { + async fn open(&mut self) -> Result<(), Error> { + info!( + "Opening BigQuery sink connector ID: {} for table {} (mode: {:?})", + self.id, self.target, self.settings.mode + ); + // The SDK turns a failed open into a bare status code, so this is the + // only place the reason reaches the runtime log. + self.try_open().await.inspect_err(|e| { + error!( + "Failed to open BigQuery sink connector ID: {} for table {}: {e}", + self.id, self.target + ); + }) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec<ConsumedMessage>, + ) -> Result<(), Error> { + if messages.is_empty() { + return Ok(()); + } + let Some(state) = &self.state else { + return Err(Error::InitError("BigQuery sink is not open".into())); + }; + let received = messages.len(); + debug!( + "BigQuery sink ID: {} received {received} messages from {}/{} partition {} current_offset {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + messages_metadata.current_offset + ); + + let ctx = RunContext { + layout: &state.layout, + topic: topic_metadata, + messages: &messages_metadata, + max_request_bytes: self.settings.max_request_bytes, + }; + let encoded = encode::encode(&ctx, messages).map_err(|e| { + self.counters + .rows_failed + .fetch_add(received as u64, Ordering::Relaxed); + error!( + "BigQuery sink ID: {} cannot encode {received} messages from {}/{} partition {}: {e}", + self.id, topic_metadata.stream, topic_metadata.topic, messages_metadata.partition_id + ); + Error::Serialization(e.to_string()) + })?; + + for rejected in &encoded.rejected { + warn!( + "BigQuery sink ID: {} dropped message stream={} topic={} partition={} offset={}: {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + rejected.offset, + rejected.reason + ); + } + self.counters + .rows_rejected + .fetch_add(encoded.rejected.len() as u64, Ordering::Relaxed); + + let mut written = 0u64; + let mut last_error = None; + for chunk in &encoded.chunks { + match self + .write_chunk(&state.client, chunk, topic_metadata, &messages_metadata) + .await + { + Ok(rows) => written += rows, + Err(e) => last_error = Some(e), + } + } + self.counters + .rows_written + .fetch_add(written, Ordering::Relaxed); + + if self.settings.verbose { + info!( + "BigQuery sink ID: {} wrote {written} of {received} messages from {}/{} partition {} to {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + self.target + ); + } else { + debug!( + "BigQuery sink ID: {} wrote {written} of {received} messages from {}/{} partition {} to {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + self.target + ); + } + + match last_error { + Some(e) => Err(e), + None => Ok(()), + } + } + + async fn close(&mut self) -> Result<(), Error> { + self.state = None; + info!( + "Closed BigQuery sink connector ID: {}, rows written: {}, rows rejected: {}, rows failed: {}", + self.id, + self.counters.rows_written.load(Ordering::Relaxed), + self.counters.rows_rejected.load(Ordering::Relaxed), + self.counters.rows_failed.load(Ordering::Relaxed) + ); + Ok(()) + } +} + +impl BigQuerySink { + async fn try_open(&mut self) -> Result<(), Error> { + self.config.validate()?; + + let mut client = BigQueryClient::connect(&self.config, &self.settings).await?; + let policy = self.retry_policy(); + + let context = format!("BigQuery sink ID: {} tables.get", self.id); + let body = retry_async(policy, &context, TableError::is_retryable, || { + client.fetch_table() + }) + .await + .map_err(|failure| Error::InitError(format!("cannot read the table schema: {failure}")))?; + + let columns = parse_table_schema(&body).map_err(|e| Error::from(TableError::Schema(e)))?; + let layout = TableLayout::build(columns, &self.settings) + .map_err(|e| Error::from(TableError::Schema(e)))?; + + let context = format!("BigQuery sink ID: {} write stream setup", self.id); + let stream = retry_async(policy, &context, AppendError::is_retryable, || { + client.resolve_write_stream() + }) + .await + .map_err(|failure| Error::InitError(format!("cannot open the write stream: {failure}")))?; + client.set_write_stream(stream); + + info!( + "Opened BigQuery sink connector ID: {} for table {}: {} payload column(s), {} metadata column(s)", + self.id, + self.target, + layout.raw.as_ref().map_or(layout.columns.len(), |_| 1), + layout.metadata.len() + ); + self.state = Some(OpenState { client, layout }); + Ok(()) + } + + fn retry_policy(&self) -> RetryPolicy { + RetryPolicy { + max_attempts: self.settings.max_retries, + base_delay: self.settings.retry_delay, + max_delay: self.settings.max_retry_delay, + } + } + + /// Append one chunk. When BigQuery reports row errors, nothing was + /// appended: drop the reported rows and append the rest once more. + /// Returns the number of rows written. + async fn write_chunk( + &self, + client: &BigQueryClient, + chunk: &Chunk, + topic: &TopicMetadata, + messages: &MessagesMetadata, + ) -> Result<u64, Error> { + let row_errors = match self + .append_with_retry(client, chunk, topic, messages) + .await? + { + AppendOutcome::Appended => return Ok(chunk.offsets.len() as u64), + AppendOutcome::RowErrors(row_errors) => row_errors, + }; + + let mut dropped = Vec::with_capacity(row_errors.len()); + for (row, reason) in row_errors { + let offset = chunk.offsets.get(row).copied(); + warn!( + "BigQuery sink ID: {} dropped row rejected by BigQuery stream={} topic={} partition={} offset={}: {reason}", + self.id, + topic.stream, + topic.topic, + messages.partition_id, + offset.map_or_else(|| format!("unknown (row {row})"), |o| o.to_string()) + ); + dropped.push(row); + } + self.counters + .rows_rejected + .fetch_add(dropped.len() as u64, Ordering::Relaxed); + + let Some(retry) = encode::without_rows(chunk, &dropped) + .map_err(|e| Error::Serialization(e.to_string()))? + else { + return Ok(0); + }; + match self + .append_with_retry(client, &retry, topic, messages) + .await? + { + AppendOutcome::Appended => Ok(retry.offsets.len() as u64), + AppendOutcome::RowErrors(again) => { + self.counters + .rows_failed + .fetch_add(retry.offsets.len() as u64, Ordering::Relaxed); + error!( + "BigQuery sink ID: {} dropped {} rows (offsets {}..={}) of {}/{} partition {}: BigQuery reported {} more row error(s) after the bad rows were removed", + self.id, + retry.offsets.len(), + retry.offsets.first().copied().unwrap_or_default(), + retry.offsets.last().copied().unwrap_or_default(), + topic.stream, + topic.topic, + messages.partition_id, + again.len() + ); + Err(Error::PermanentHttpError(format!( + "{} row error(s) persisted after removing rejected rows", + again.len() + ))) + } + } + } + + async fn append_with_retry( Review Comment: core/connectors/runtime/src/sink.rs:475-525 × core/connectors/runtime/src/manager/sink.rs:103-114 — exhausted-but-still-retryable AppendRows failures touch no counter (error! log only). Runtime hardcodes AutoCommit::When(AutoCommitWhen::PollingMessages) (commits the Iggy offset before consume() runs), consume()'s Err ends the whole spawned task, and set_error only flips ConnectorStatus::Error with no restart/redelivery anywhere. Net effect: silent, permanent, unmetered row loss, plus the entire partition stops being consumed until an operator notices and restarts. **Fix**: count these rows, make the failure an alertable metric not a log line, rewrite README.md:155-172 to lead with permanent-loss-and-halt, not just "may duplicate on retry." -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
