From e89a375cc13e083a36ae36d3e7ed506e2f19e88a Mon Sep 17 00:00:00 2001 From: Bo Lu Date: Thu, 19 Sep 2024 15:55:28 +1000 Subject: [PATCH 1/3] feat: add array cell type support --- supabase-wrappers/src/interface.rs | 72 +++++++++++++++++++ wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs | 9 ++- wrappers/src/fdw/bigquery_fdw/mod.rs | 5 +- wrappers/src/fdw/wasm_fdw/bindings.rs | 1 + 4 files changed, 83 insertions(+), 4 deletions(-) diff --git a/supabase-wrappers/src/interface.rs b/supabase-wrappers/src/interface.rs index 420001237..e273b2605 100644 --- a/supabase-wrappers/src/interface.rs +++ b/supabase-wrappers/src/interface.rs @@ -47,6 +47,13 @@ pub enum Cell { Timestamp(Timestamp), Timestamptz(TimestampWithTimeZone), Json(JsonB), + BoolArray(Vec>), + I16Array(Vec>), + I32Array(Vec>), + I64Array(Vec>), + F32Array(Vec>), + F64Array(Vec>), + StringArray(Vec>), } impl Clone for Cell { @@ -65,10 +72,32 @@ impl Clone for Cell { Cell::Timestamp(v) => Cell::Timestamp(*v), Cell::Timestamptz(v) => Cell::Timestamptz(*v), Cell::Json(v) => Cell::Json(JsonB(v.0.clone())), + Cell::BoolArray(v) => Cell::BoolArray(v.clone()), + Cell::I16Array(v) => Cell::I16Array(v.clone()), + Cell::I32Array(v) => Cell::I32Array(v.clone()), + Cell::I64Array(v) => Cell::I64Array(v.clone()), + Cell::F32Array(v) => Cell::F32Array(v.clone()), + Cell::F64Array(v) => Cell::F64Array(v.clone()), + Cell::StringArray(v) => Cell::StringArray(v.clone()), } } } +fn write_array( + array: &[Option], + f: &mut fmt::Formatter<'_>, +) -> fmt::Result { + let res = array + .iter() + .map(|e| match e { + Some(val) => format!("{}", val), + None => "null".to_owned(), + }) + .collect::>() + .join(","); + write!(f, "[{}]", res) +} + impl fmt::Display for Cell { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -123,6 +152,13 @@ impl fmt::Display for Cell { ) }, Cell::Json(v) => write!(f, "{:?}", v), + Cell::BoolArray(v) => write_array(v, f), + Cell::I16Array(v) => write_array(v, f), + Cell::I32Array(v) => write_array(v, f), + Cell::I64Array(v) => write_array(v, f), + Cell::F32Array(v) => write_array(v, f), + Cell::F64Array(v) => write_array(v, f), + Cell::StringArray(v) => write_array(v, f), } } } @@ -143,6 +179,13 @@ impl IntoDatum for Cell { Cell::Timestamp(v) => v.into_datum(), Cell::Timestamptz(v) => v.into_datum(), Cell::Json(v) => v.into_datum(), + Cell::BoolArray(v) => v.into_datum(), + Cell::I16Array(v) => v.into_datum(), + Cell::I32Array(v) => v.into_datum(), + Cell::I64Array(v) => v.into_datum(), + Cell::F32Array(v) => v.into_datum(), + Cell::F64Array(v) => v.into_datum(), + Cell::StringArray(v) => v.into_datum(), } } @@ -165,6 +208,13 @@ impl IntoDatum for Cell { || other == pg_sys::TIMESTAMPOID || other == pg_sys::TIMESTAMPTZOID || other == pg_sys::JSONBOID + || other == pg_sys::BOOLARRAYOID + || other == pg_sys::TEXTARRAYOID + || other == pg_sys::INT2ARRAYOID + || other == pg_sys::INT4ARRAYOID + || other == pg_sys::INT8ARRAYOID + || other == pg_sys::FLOAT4ARRAYOID + || other == pg_sys::FLOAT8ARRAYOID } } @@ -212,6 +262,27 @@ impl FromDatum for Cell { PgOid::BuiltIn(PgBuiltInOids::JSONBOID) => { JsonB::from_datum(datum, is_null).map(Cell::Json) } + PgOid::BuiltIn(PgBuiltInOids::BOOLARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::BoolArray) + } + PgOid::BuiltIn(PgBuiltInOids::INT2ARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::I16Array) + } + PgOid::BuiltIn(PgBuiltInOids::INT4ARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::I32Array) + } + PgOid::BuiltIn(PgBuiltInOids::INT8ARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::I64Array) + } + PgOid::BuiltIn(PgBuiltInOids::FLOAT4ARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::F32Array) + } + PgOid::BuiltIn(PgBuiltInOids::FLOAT8ARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::F64Array) + } + PgOid::BuiltIn(PgBuiltInOids::TEXTARRAYOID) => { + Vec::>::from_datum(datum, false).map(Cell::StringArray) + } _ => None, } } @@ -234,6 +305,7 @@ impl CellFormatter for DefaultFormatter { format!("{}", cell) } } + /// A data row in a table /// /// The row contains a column name list and cell list with same number of diff --git a/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs b/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs index fbd9dc869..532e5862a 100644 --- a/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs +++ b/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs @@ -42,7 +42,9 @@ fn field_to_cell(rs: &ResultSet, field: &TableFieldSchema) -> BigQueryFdwResult< Cell::Timestamp(ts.to_utc()) }), _ => { - return Err(BigQueryFdwError::UnsupportedFieldType(field.r#type.clone())); + return Err(BigQueryFdwError::UnsupportedFieldType( + field.name.to_owned(), + )); } }) } @@ -397,6 +399,11 @@ impl ForeignDataWrapper for BigQueryFdw { Cell::Timestamp(v) => row_json[col_name] = json!(v), Cell::Timestamptz(v) => row_json[col_name] = json!(v), Cell::Json(v) => row_json[col_name] = json!(v), + _ => { + return Err(BigQueryFdwError::UnsupportedFieldType( + col_name.to_owned(), + )); + } } } } diff --git a/wrappers/src/fdw/bigquery_fdw/mod.rs b/wrappers/src/fdw/bigquery_fdw/mod.rs index 1e98caeb5..4f37bf791 100644 --- a/wrappers/src/fdw/bigquery_fdw/mod.rs +++ b/wrappers/src/fdw/bigquery_fdw/mod.rs @@ -1,7 +1,6 @@ #![allow(clippy::module_inception)] use gcp_bigquery_client::error::BQError; -use gcp_bigquery_client::model::field_type::FieldType; use pgrx::pg_sys::panic::ErrorReport; use pgrx::{DateTimeConversionError, PgSqlErrorCode}; use supabase_wrappers::prelude::{CreateRuntimeError, OptionsError}; @@ -21,8 +20,8 @@ enum BigQueryFdwError { #[error("big query error: {0}")] BigQueryError(#[from] BQError), - #[error("field type {0:?} not supported")] - UnsupportedFieldType(FieldType), + #[error("field {0} type not supported")] + UnsupportedFieldType(String), #[error("{0}")] NumericConversionError(#[from] pgrx::numeric::Error), diff --git a/wrappers/src/fdw/wasm_fdw/bindings.rs b/wrappers/src/fdw/wasm_fdw/bindings.rs index 8debc6208..f0a7a350d 100644 --- a/wrappers/src/fdw/wasm_fdw/bindings.rs +++ b/wrappers/src/fdw/wasm_fdw/bindings.rs @@ -79,6 +79,7 @@ impl From<&HostCell> for GuestCell { Self::Timestamptz(v.into_inner() + 946_684_800_000_000) } HostCell::Json(v) => Self::Json(v.0.to_string()), + _ => todo!(), } } } From a3e0b61852c022cc9f8b8abc01b045d18a9ded29 Mon Sep 17 00:00:00 2001 From: Bo Lu Date: Thu, 19 Sep 2024 15:59:56 +1000 Subject: [PATCH 2/3] adjust type match ordering --- supabase-wrappers/src/interface.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supabase-wrappers/src/interface.rs b/supabase-wrappers/src/interface.rs index e273b2605..67b088468 100644 --- a/supabase-wrappers/src/interface.rs +++ b/supabase-wrappers/src/interface.rs @@ -209,12 +209,12 @@ impl IntoDatum for Cell { || other == pg_sys::TIMESTAMPTZOID || other == pg_sys::JSONBOID || other == pg_sys::BOOLARRAYOID - || other == pg_sys::TEXTARRAYOID || other == pg_sys::INT2ARRAYOID || other == pg_sys::INT4ARRAYOID || other == pg_sys::INT8ARRAYOID || other == pg_sys::FLOAT4ARRAYOID || other == pg_sys::FLOAT8ARRAYOID + || other == pg_sys::TEXTARRAYOID } } From 22c77269e6ad42afc7d903d1b724e722cf6c0b34 Mon Sep 17 00:00:00 2001 From: Bo Lu Date: Fri, 20 Sep 2024 16:05:11 +1000 Subject: [PATCH 3/3] add todo message --- wrappers/src/fdw/wasm_fdw/bindings.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wrappers/src/fdw/wasm_fdw/bindings.rs b/wrappers/src/fdw/wasm_fdw/bindings.rs index f0a7a350d..c6ad81ea0 100644 --- a/wrappers/src/fdw/wasm_fdw/bindings.rs +++ b/wrappers/src/fdw/wasm_fdw/bindings.rs @@ -79,7 +79,7 @@ impl From<&HostCell> for GuestCell { Self::Timestamptz(v.into_inner() + 946_684_800_000_000) } HostCell::Json(v) => Self::Json(v.0.to_string()), - _ => todo!(), + _ => todo!("Add array type support for Wasm FDW"), } } }