feat: serialize metadata response

Author: darthunix

src/backend.rs

@@ -1,20 +1,20 @@
 use anyhow::Result as AnyResult;
 use libc::c_long;
 use pgrx::pg_sys::{
-    error, fetch_search_path_array, get_namespace_oid, get_relname_relid, palloc0,
+    self, error, fetch_search_path_array, get_namespace_oid, get_relname_relid, palloc0,
     CustomExecMethods, CustomScan, CustomScanMethods, CustomScanState, EState, ExplainState,
     InvalidOid, List, ListCell, MyLatch, Node, NodeTag, Oid, ParamListInfo,
     RegisterCustomScanMethods, ResetLatch, TupleTableSlot, WaitLatch, PG_WAIT_EXTENSION,
     WL_LATCH_SET, WL_POSTMASTER_DEATH, WL_TIMEOUT,
 };
 use pgrx::{check_for_interrupts, pg_guard};
 use rmp::decode::{read_array_len, read_bin_len};
-use smallvec::{smallvec, SmallVec};
+use smallvec::SmallVec;
 use std::ffi::c_char;
 use std::time::Duration;
 
 use crate::error::FusionError;
-use crate::ipc::{my_slot, Bus, SlotHandler, SlotNumber, SlotStream, CURRENT_SLOT};
+use crate::ipc::{my_slot, Bus, SlotStream};
 use crate::protocol::{consume_header, read_error, send_params, send_query, Direction, Packet};
 
 thread_local! {

src/protocol.rs

@@ -7,9 +7,11 @@ use datafusion::scalar::ScalarValue;
 use datafusion_sql::TableReference;
 use pgrx::pg_sys::{Oid, ParamExternData, ProcSendSignal};
 use pgrx::prelude::*;
-use rmp::decode::{read_array_len, read_bin_len, read_pfix, read_str_len, read_u16};
+use pgrx::{pg_guard, PgRelation};
+use rmp::decode::{read_array_len, read_bin_len, read_pfix, read_str_len, read_u16, read_u8};
 use rmp::encode::{
-    write_array_len, write_bin_len, write_bool, write_pfix, write_str, write_u16, RmpWrite,
+    write_array_len, write_bin_len, write_bool, write_pfix, write_str, write_u16, write_u32,
+    write_u8, RmpWrite,
 };
 
 #[repr(u8)]
@@ -331,23 +333,71 @@ pub(crate) fn send_table_refs(
 }
 
 #[inline]
-pub(crate) fn write_column(
-    stream: &mut SlotStream,
-    column: &str,
-    is_null: bool,
-    etype: EncodedType,
+#[pg_guard]
+fn serialize_table(rel_oid: Oid, stream: &mut SlotStream) -> Result<()> {
+    // The destructor will release the lock.
+    let rel = unsafe { PgRelation::with_lock(rel_oid, pg_sys::AccessShareLock as i32) };
+    let tuple_desc = rel.tuple_desc();
+    let attr_num = u32::try_from(tuple_desc.iter().filter(|a| !a.is_dropped()).count())?;
+    write_u32(stream, rel_oid.as_u32())?;
+    write_array_len(stream, attr_num)?;
+    for attr in tuple_desc.iter() {
+        if attr.is_dropped() {
+            continue;
+        }
+        let etype = EncodedType::try_from(attr.type_oid().value())?;
+        let is_nullable = !attr.attnotnull;
+        let name = attr.name();
+        write_array_len(stream, 3)?;
+        write_str(stream, name)?;
+        write_u8(stream, etype as u8)?;
+        write_bool(stream, is_nullable)?;
+    }
+    Ok(())
+}
+
+pub(crate) fn prepare_metadata(rel_oids: &[Oid], stream: &mut SlotStream) -> Result<()> {
+    stream.reset();
+    // We don't know the length of the table metadata yet. So we write
+    // an invalid header to replace it with the correct one later.
+    write_header(stream, &Header::default())?;
+    let pos_init = stream.position();
+    write_array_len(stream, rel_oids.len() as u32)?;
+    for &rel_oid in rel_oids {
+        serialize_table(rel_oid, stream)?;
+    }
+    let pos_final = stream.position();
+    let length = u16::try_from(pos_final - pos_init)?;
+    let header = Header {
+        direction: Direction::ToWorker,
+        packet: Packet::Metadata,
+        length,
+        flag: Flag::Last,
+    };
+    stream.reset();
+    write_header(stream, &header)?;
+    stream.rewind(length as usize)?;
+    Ok(())
+}
+
+pub(crate) fn send_metadata(
+    slot_id: SlotNumber,
+    mut stream: SlotStream,
+    rel_oids: &[Oid],
 ) -> Result<()> {
-    write_str(stream, column)?;
-    write_bool(stream, is_null)?;
-    write_pfix(stream, etype as u8)?;
+    prepare_metadata(rel_oids, &mut stream)?;
+    // Unlock the slot after writing the metadata.
+    let _guard = Slot::from(stream);
+    signal(slot_id, Direction::ToWorker);
     Ok(())
 }
 
 #[cfg(any(test, feature = "pg_test"))]
 #[pg_schema]
 mod tests {
     use super::*;
-    use pgrx::pg_sys::Datum;
+    use pgrx::pg_sys::{Datum, Oid};
+    use rmp::decode::{read_bool, read_u32};
     use std::ptr::addr_of_mut;
 
     const SLOT_SIZE: usize = 8204;
@@ -482,4 +532,58 @@ mod tests {
         let t2 = stream.look_ahead(t2_len as usize).unwrap();
         assert_eq!(t2, b"table2\0");
     }
+
+    #[pg_test]
+    fn test_metadata_response() {
+        Spi::run("create table if not exists t1(a int not null, b text);").unwrap();
+        let t1_oid = Spi::get_one::<Oid>("select 't1'::regclass::oid;")
+            .unwrap()
+            .unwrap();
+
+        let mut slot_buf: [u8; SLOT_SIZE] = [1; SLOT_SIZE];
+        let ptr = addr_of_mut!(slot_buf) as *mut u8;
+        Slot::init(ptr, slot_buf.len());
+        let slot = Slot::from_bytes(ptr, slot_buf.len());
+        let mut stream: SlotStream = slot.into();
+
+        prepare_metadata(&[t1_oid], &mut stream).unwrap();
+        stream.reset();
+        let header = consume_header(&mut stream).unwrap();
+        assert_eq!(header.direction, Direction::ToWorker);
+        assert_eq!(header.packet, Packet::Metadata);
+        assert_eq!(header.flag, Flag::Last);
+
+        // Check table metadata deserialization
+        let table_num = read_array_len(&mut stream).unwrap();
+        assert_eq!(table_num, 1);
+        // t1
+        let oid = read_u32(&mut stream).unwrap();
+        assert_eq!(oid, t1_oid.as_u32());
+        let attr_num = read_array_len(&mut stream).unwrap();
+        assert_eq!(attr_num, 2);
+        // a
+        let elem_num = read_array_len(&mut stream).unwrap();
+        assert_eq!(elem_num, 3);
+        let name_len = read_str_len(&mut stream).unwrap();
+        assert_eq!(name_len, "a".len() as u32);
+        let name = stream.look_ahead(name_len as usize).unwrap();
+        assert_eq!(name, b"a");
+        stream.rewind(name_len as usize).unwrap();
+        let etype = read_u8(&mut stream).unwrap();
+        assert_eq!(etype, EncodedType::Int32 as u8);
+        let is_nullable = read_bool(&mut stream).unwrap();
+        assert!(!is_nullable);
+        // b
+        let elem_num = read_array_len(&mut stream).unwrap();
+        assert_eq!(elem_num, 3);
+        let name_len = read_str_len(&mut stream).unwrap();
+        assert_eq!(name_len, "b".len() as u32);
+        let name = stream.look_ahead(name_len as usize).unwrap();
+        assert_eq!(name, b"b");
+        stream.rewind(name_len as usize).unwrap();
+        let etype = read_u8(&mut stream).unwrap();
+        assert_eq!(etype, EncodedType::Utf8 as u8);
+        let is_nullable = read_bool(&mut stream).unwrap();
+        assert!(is_nullable);
+    }
 }