Skip to content

Commit

Permalink
feat: add array cell type support
Browse files Browse the repository at this point in the history
  • Loading branch information
burmecia committed Sep 19, 2024
1 parent fd4a3c3 commit e89a375
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 4 deletions.
72 changes: 72 additions & 0 deletions supabase-wrappers/src/interface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,13 @@ pub enum Cell {
Timestamp(Timestamp),
Timestamptz(TimestampWithTimeZone),
Json(JsonB),
BoolArray(Vec<Option<bool>>),
I16Array(Vec<Option<i16>>),
I32Array(Vec<Option<i32>>),
I64Array(Vec<Option<i64>>),
F32Array(Vec<Option<f32>>),
F64Array(Vec<Option<f64>>),
StringArray(Vec<Option<String>>),
}

impl Clone for Cell {
Expand All @@ -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<T: std::fmt::Display>(
array: &[Option<T>],
f: &mut fmt::Formatter<'_>,
) -> fmt::Result {
let res = array
.iter()
.map(|e| match e {
Some(val) => format!("{}", val),
None => "null".to_owned(),
})
.collect::<Vec<String>>()
.join(",");
write!(f, "[{}]", res)
}

impl fmt::Display for Cell {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Expand Down Expand Up @@ -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),
}
}
}
Expand All @@ -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(),
}
}

Expand All @@ -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
}
}

Expand Down Expand Up @@ -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::<Option<bool>>::from_datum(datum, false).map(Cell::BoolArray)
}
PgOid::BuiltIn(PgBuiltInOids::INT2ARRAYOID) => {
Vec::<Option<i16>>::from_datum(datum, false).map(Cell::I16Array)
}
PgOid::BuiltIn(PgBuiltInOids::INT4ARRAYOID) => {
Vec::<Option<i32>>::from_datum(datum, false).map(Cell::I32Array)
}
PgOid::BuiltIn(PgBuiltInOids::INT8ARRAYOID) => {
Vec::<Option<i64>>::from_datum(datum, false).map(Cell::I64Array)
}
PgOid::BuiltIn(PgBuiltInOids::FLOAT4ARRAYOID) => {
Vec::<Option<f32>>::from_datum(datum, false).map(Cell::F32Array)
}
PgOid::BuiltIn(PgBuiltInOids::FLOAT8ARRAYOID) => {
Vec::<Option<f64>>::from_datum(datum, false).map(Cell::F64Array)
}
PgOid::BuiltIn(PgBuiltInOids::TEXTARRAYOID) => {
Vec::<Option<String>>::from_datum(datum, false).map(Cell::StringArray)
}
_ => None,
}
}
Expand All @@ -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
Expand Down
9 changes: 8 additions & 1 deletion wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
));
}
})
}
Expand Down Expand Up @@ -397,6 +399,11 @@ impl ForeignDataWrapper<BigQueryFdwError> 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(),
));
}
}
}
}
Expand Down
5 changes: 2 additions & 3 deletions wrappers/src/fdw/bigquery_fdw/mod.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand All @@ -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),
Expand Down
1 change: 1 addition & 0 deletions wrappers/src/fdw/wasm_fdw/bindings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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!(),
}
}
}
Expand Down

0 comments on commit e89a375

Please sign in to comment.