feat(bigquery): Handle arrays with null values (#251)

Author: iambriccardo

etl-destinations/src/bigquery/client.rs

@@ -210,11 +210,12 @@ impl BigQueryClient {
         table_descriptor: &TableDescriptor,
         table_rows: Vec<TableRow>,
     ) -> EtlResult<()> {
-        // We have to map table rows into the new type due to the limitations of how Rust works.
+        // We map the table rows into `BigQueryTableRow`s instances, so that we also do a validation
+        // of the cells.
         let table_rows = table_rows
             .into_iter()
-            .map(BigQueryTableRow)
-            .collect::<Vec<_>>();
+            .map(BigQueryTableRow::try_from)
+            .collect::<EtlResult<Vec<_>>>()?;
 
         // We create a slice on table rows, which will be updated while the streaming progresses.
         //

etl-destinations/src/bigquery/encoding.rs

@@ -1,5 +1,6 @@
 use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
 use etl::config::BatchConfig;
+use etl::error::ErrorKind;
 use etl::state::table::TableReplicationPhaseType;
 use etl::test_utils::database::{spawn_source_database, test_table_name};
 use etl::test_utils::notify::NotifyingStore;
@@ -11,6 +12,8 @@ use etl_destinations::bigquery::install_crypto_provider_for_bigquery;
 use etl_telemetry::init_test_tracing;
 use rand::random;
 use std::str::FromStr;
+use std::time::Duration;
+use tokio::time::sleep;
 
 use crate::common::bigquery::{
     BigQueryOrder, BigQueryUser, NonNullableColsScalar, NullableColsArray, NullableColsScalar,
@@ -1474,3 +1477,151 @@ async fn table_non_nullable_array_columns() {
 
     pipeline.shutdown_and_wait().await.unwrap();
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn table_array_with_null_values() {
+    init_test_tracing();
+    install_crypto_provider_for_bigquery();
+
+    let database = spawn_source_database().await;
+    let bigquery_database = setup_bigquery_connection().await;
+    let table_name = test_table_name("array_with_nulls");
+    let table_id = database
+        .create_table(table_name.clone(), true, &[("int_array", "int4[]")])
+        .await
+        .unwrap();
+
+    let store = NotifyingStore::new();
+    let raw_destination = bigquery_database.build_destination(store.clone()).await;
+    let destination = TestDestinationWrapper::wrap(raw_destination);
+
+    let publication_name = "test_pub_array_nulls".to_string();
+    database
+        .create_publication(&publication_name, std::slice::from_ref(&table_name))
+        .await
+        .expect("Failed to create publication");
+
+    let pipeline_id: PipelineId = random();
+    let mut pipeline = create_pipeline(
+        &database.config,
+        pipeline_id,
+        publication_name.clone(),
+        store.clone(),
+        destination.clone(),
+    );
+
+    let table_sync_done_notification = store
+        .notify_on_table_state(table_id, TableReplicationPhaseType::SyncDone)
+        .await;
+
+    pipeline.start().await.unwrap();
+
+    table_sync_done_notification.notified().await;
+
+    // Insert array with null value
+    database
+        .client
+        .as_ref()
+        .unwrap()
+        .execute(
+            &format!(
+                "INSERT INTO {} (int_array) VALUES (ARRAY[1, NULL])",
+                table_name.as_quoted_identifier()
+            ),
+            &[],
+        )
+        .await
+        .unwrap();
+
+    // We sleep to wait for the event to be processed. This is not ideal, but if we wanted to do
+    // this better, we would have to also implement error handling within the apply worker to write
+    // in the state store.
+    sleep(Duration::from_secs(1)).await;
+
+    // Wait for the pipeline expecting an error to be returned.
+    let err = pipeline.shutdown_and_wait().await.err().unwrap();
+    assert_eq!(err.kinds().len(), 1);
+    assert_eq!(err.kinds()[0], ErrorKind::NullValuesNotSupportedInArray);
+
+    // Reset and try with valid array (no nulls)
+    database
+        .client
+        .as_ref()
+        .unwrap()
+        .execute(
+            &format!(
+                "DELETE FROM {} WHERE true",
+                table_name.as_quoted_identifier()
+            ),
+            &[],
+        )
+        .await
+        .unwrap();
+
+    // We have to reset the state of the table and copy it from scratch, otherwise the CDC will contain
+    // the inserts and deletes, failing again.
+    store.reset_table_state(table_id).await.unwrap();
+    // We also clear the events so that it's more idiomatic to wait for them, since we don't have
+    // the insert of before.
+    destination.clear_events().await;
+
+    // We recreate the pipeline and try again.
+    let mut pipeline = create_pipeline(
+        &database.config,
+        pipeline_id,
+        publication_name,
+        store.clone(),
+        destination.clone(),
+    );
+
+    let table_sync_done_notification = store
+        .notify_on_table_state(table_id, TableReplicationPhaseType::SyncDone)
+        .await;
+
+    pipeline.start().await.unwrap();
+
+    table_sync_done_notification.notified().await;
+
+    let event_notify = destination
+        .wait_for_events_count(vec![(EventType::Insert, 1)])
+        .await;
+
+    // Insert array without null values
+    database
+        .client
+        .as_ref()
+        .unwrap()
+        .execute(
+            &format!(
+                "INSERT INTO {} (int_array) VALUES (ARRAY[1, 2, 3])",
+                table_name.as_quoted_identifier()
+            ),
+            &[],
+        )
+        .await
+        .unwrap();
+
+    event_notify.notified().await;
+
+    pipeline.shutdown_and_wait().await.unwrap();
+
+    // Verify the valid array was successfully replicated to BigQuery
+    let table_rows = bigquery_database
+        .query_table(table_name.clone())
+        .await
+        .unwrap();
+
+    // Check that there is only the valid row in BigQuery
+    assert_eq!(table_rows.len(), 1);
+
+    // Check that the int array contains 3 elements, meaning it must be the second insert with all
+    // NON-NULL values
+    let row = &table_rows[0];
+    if let Some(columns) = &row.columns {
+        assert_eq!(columns.len(), 2);
+        assert_eq!(
+            columns[1].value.clone().unwrap().as_array().unwrap().len(),
+            3
+        );
+    }
+}

etl-destinations/tests/integration/bigquery_test.rs

@@ -3,6 +3,10 @@ use numeric::PgNumeric;
 use std::fmt::Debug;
 use uuid::Uuid;
 
+use crate::bail;
+use crate::error::ErrorKind;
+use crate::error::EtlError;
+
 pub mod bool;
 pub mod event;
 pub mod hex;
@@ -103,3 +107,172 @@ impl ArrayCell {
         }
     }
 }
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum CellNonOptional {
+    Null,
+    Bool(bool),
+    String(String),
+    I16(i16),
+    I32(i32),
+    U32(u32),
+    I64(i64),
+    F32(f32),
+    F64(f64),
+    Numeric(PgNumeric),
+    Date(NaiveDate),
+    Time(NaiveTime),
+    TimeStamp(NaiveDateTime),
+    TimeStampTz(DateTime<Utc>),
+    Uuid(Uuid),
+    Json(serde_json::Value),
+    Bytes(Vec<u8>),
+    Array(ArrayCellNonOptional),
+}
+
+impl TryFrom<Cell> for CellNonOptional {
+    type Error = EtlError;
+
+    fn try_from(cell: Cell) -> Result<Self, Self::Error> {
+        match cell {
+            Cell::Null => Ok(CellNonOptional::Null),
+            Cell::Bool(val) => Ok(CellNonOptional::Bool(val)),
+            Cell::String(val) => Ok(CellNonOptional::String(val)),
+            Cell::I16(val) => Ok(CellNonOptional::I16(val)),
+            Cell::I32(val) => Ok(CellNonOptional::I32(val)),
+            Cell::U32(val) => Ok(CellNonOptional::U32(val)),
+            Cell::I64(val) => Ok(CellNonOptional::I64(val)),
+            Cell::F32(val) => Ok(CellNonOptional::F32(val)),
+            Cell::F64(val) => Ok(CellNonOptional::F64(val)),
+            Cell::Numeric(val) => Ok(CellNonOptional::Numeric(val)),
+            Cell::Date(val) => Ok(CellNonOptional::Date(val)),
+            Cell::Time(val) => Ok(CellNonOptional::Time(val)),
+            Cell::TimeStamp(val) => Ok(CellNonOptional::TimeStamp(val)),
+            Cell::TimeStampTz(val) => Ok(CellNonOptional::TimeStampTz(val)),
+            Cell::Uuid(val) => Ok(CellNonOptional::Uuid(val)),
+            Cell::Json(val) => Ok(CellNonOptional::Json(val)),
+            Cell::Bytes(val) => Ok(CellNonOptional::Bytes(val)),
+            Cell::Array(array_cell) => {
+                let array_cell_non_optional = ArrayCellNonOptional::try_from(array_cell)?;
+                Ok(CellNonOptional::Array(array_cell_non_optional))
+            }
+        }
+    }
+}
+
+impl CellNonOptional {
+    pub fn clear(&mut self) {
+        match self {
+            CellNonOptional::Null => {}
+            CellNonOptional::Bool(b) => *b = false,
+            CellNonOptional::String(s) => s.clear(),
+            CellNonOptional::I16(i) => *i = 0,
+            CellNonOptional::I32(i) => *i = 0,
+            CellNonOptional::I64(i) => *i = 0,
+            CellNonOptional::F32(i) => *i = 0.,
+            CellNonOptional::F64(i) => *i = 0.,
+            CellNonOptional::Numeric(n) => *n = PgNumeric::default(),
+            CellNonOptional::Date(t) => *t = NaiveDate::default(),
+            CellNonOptional::Time(t) => *t = NaiveTime::default(),
+            CellNonOptional::TimeStamp(t) => *t = NaiveDateTime::default(),
+            CellNonOptional::TimeStampTz(t) => *t = DateTime::<Utc>::default(),
+            CellNonOptional::Uuid(u) => *u = Uuid::default(),
+            CellNonOptional::Json(j) => *j = serde_json::Value::default(),
+            CellNonOptional::U32(u) => *u = 0,
+            CellNonOptional::Bytes(b) => b.clear(),
+            CellNonOptional::Array(vec) => {
+                vec.clear();
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum ArrayCellNonOptional {
+    Null,
+    Bool(Vec<bool>),
+    String(Vec<String>),
+    I16(Vec<i16>),
+    I32(Vec<i32>),
+    U32(Vec<u32>),
+    I64(Vec<i64>),
+    F32(Vec<f32>),
+    F64(Vec<f64>),
+    Numeric(Vec<PgNumeric>),
+    Date(Vec<NaiveDate>),
+    Time(Vec<NaiveTime>),
+    TimeStamp(Vec<NaiveDateTime>),
+    TimeStampTz(Vec<DateTime<Utc>>),
+    Uuid(Vec<Uuid>),
+    Json(Vec<serde_json::Value>),
+    Bytes(Vec<Vec<u8>>),
+}
+
+macro_rules! convert_array_variant {
+    ($variant:ident, $vec:expr) => {
+        if $vec.iter().any(|v| v.is_none()) {
+            bail!(
+                ErrorKind::NullValuesNotSupportedInArray,
+                "NULL values in arrays are not supported",
+                format!(
+                    "Remove the NULL values from the array {:?} and try again",
+                    $vec
+                )
+            )
+        } else {
+            Ok(ArrayCellNonOptional::$variant(
+                $vec.into_iter().flatten().collect(),
+            ))
+        }
+    };
+}
+
+impl TryFrom<ArrayCell> for ArrayCellNonOptional {
+    type Error = EtlError;
+
+    fn try_from(array_cell: ArrayCell) -> Result<Self, Self::Error> {
+        match array_cell {
+            ArrayCell::Null => Ok(ArrayCellNonOptional::Null),
+            ArrayCell::Bool(vec) => convert_array_variant!(Bool, vec),
+            ArrayCell::String(vec) => convert_array_variant!(String, vec),
+            ArrayCell::I16(vec) => convert_array_variant!(I16, vec),
+            ArrayCell::I32(vec) => convert_array_variant!(I32, vec),
+            ArrayCell::U32(vec) => convert_array_variant!(U32, vec),
+            ArrayCell::I64(vec) => convert_array_variant!(I64, vec),
+            ArrayCell::F32(vec) => convert_array_variant!(F32, vec),
+            ArrayCell::F64(vec) => convert_array_variant!(F64, vec),
+            ArrayCell::Numeric(vec) => convert_array_variant!(Numeric, vec),
+            ArrayCell::Date(vec) => convert_array_variant!(Date, vec),
+            ArrayCell::Time(vec) => convert_array_variant!(Time, vec),
+            ArrayCell::TimeStamp(vec) => convert_array_variant!(TimeStamp, vec),
+            ArrayCell::TimeStampTz(vec) => convert_array_variant!(TimeStampTz, vec),
+            ArrayCell::Uuid(vec) => convert_array_variant!(Uuid, vec),
+            ArrayCell::Json(vec) => convert_array_variant!(Json, vec),
+            ArrayCell::Bytes(vec) => convert_array_variant!(Bytes, vec),
+        }
+    }
+}
+
+impl ArrayCellNonOptional {
+    fn clear(&mut self) {
+        match self {
+            ArrayCellNonOptional::Null => {}
+            ArrayCellNonOptional::Bool(vec) => vec.clear(),
+            ArrayCellNonOptional::String(vec) => vec.clear(),
+            ArrayCellNonOptional::I16(vec) => vec.clear(),
+            ArrayCellNonOptional::I32(vec) => vec.clear(),
+            ArrayCellNonOptional::U32(vec) => vec.clear(),
+            ArrayCellNonOptional::I64(vec) => vec.clear(),
+            ArrayCellNonOptional::F32(vec) => vec.clear(),
+            ArrayCellNonOptional::F64(vec) => vec.clear(),
+            ArrayCellNonOptional::Numeric(vec) => vec.clear(),
+            ArrayCellNonOptional::Date(vec) => vec.clear(),
+            ArrayCellNonOptional::Time(vec) => vec.clear(),
+            ArrayCellNonOptional::TimeStamp(vec) => vec.clear(),
+            ArrayCellNonOptional::TimeStampTz(vec) => vec.clear(),
+            ArrayCellNonOptional::Uuid(vec) => vec.clear(),
+            ArrayCellNonOptional::Json(vec) => vec.clear(),
+            ArrayCellNonOptional::Bytes(vec) => vec.clear(),
+        }
+    }
+}