wip: process metadata request

Author: darthunix

Cargo.toml

@@ -16,6 +16,7 @@ pg17 = ["pgrx/pg17", "pgrx-tests/pg17"]
 pg_test = []
 
 [dependencies]
+ahash = "0.8"
 anyhow = "1.0"
 datafusion = "44.0"
 datafusion-sql = "44.0"
@@ -24,6 +25,7 @@ libc = "0.2"
 pgrx = "0.12"
 rmp = "0.8"
 rust-fsm = { version = "0.7", features = ["diagram"] }
+smallvec = { version = "1.14", features = ["const_generics", "union"] }
 smol_str = "0.3"
 thiserror = "2.0"
 tokio = { version = "1.42", features = ["full"] }

src/backend.rs

@@ -1,11 +1,14 @@
 use libc::c_long;
 use pgrx::pg_sys::{
-    error, palloc0, CustomExecMethods, CustomScan, CustomScanMethods, CustomScanState, EState,
-    ExplainState, List, ListCell, MyLatch, Node, NodeTag, ParamListInfo, RegisterCustomScanMethods,
-    ResetLatch, TupleTableSlot, WaitLatch, PG_WAIT_EXTENSION, WL_LATCH_SET, WL_POSTMASTER_DEATH,
-    WL_TIMEOUT,
+    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 std::ffi::c_char;
 use std::time::Duration;
 
@@ -145,6 +148,79 @@ unsafe extern "C" fn explain_df_scan(
     todo!()
 }
 
+// We expect that the header is already consumed and the packet type is `Packet::Metadata`.
+#[pg_guard]
+fn table_oids(stream: &mut SlotStream) -> SmallVec<[Oid; 16]> {
+    let table_num = read_array_len(stream).expect("Failed to read the number of tables");
+    let mut oids: SmallVec<[Oid; 16]> = SmallVec::with_capacity(table_num as usize);
+    for _ in 0..table_num {
+        let elem_num = read_array_len(stream)
+            .expect("Failed to read the number of elements in the table name");
+        match elem_num {
+            1 => {
+                let table_len = read_bin_len(stream).expect("Failed to read the table name length");
+                let table_name = stream
+                    .look_ahead(table_len as usize)
+                    .expect("Failed to read the table name");
+                let mut search_path: [Oid; 16] = [InvalidOid; 16];
+                let path_len = unsafe {
+                    fetch_search_path_array(search_path.as_mut_ptr(), search_path.len() as i32)
+                };
+                let path = &search_path[..path_len as usize];
+                let mut rel_oid = InvalidOid;
+                for ns_oid in path {
+                    rel_oid =
+                        unsafe { get_relname_relid(table_name.as_ptr() as *const c_char, *ns_oid) };
+                    if rel_oid != InvalidOid {
+                        oids.push(rel_oid);
+                        break;
+                    }
+                }
+                if rel_oid == InvalidOid {
+                    match std::str::from_utf8(table_name) {
+                        Ok(name) => error!("Table not found: {}", name),
+                        Err(_) => error!("Table not found: {:?}", table_name),
+                    }
+                }
+                stream
+                    .rewind(table_len as usize)
+                    .expect("Failed to rewind the stream");
+            }
+            2 => {
+                let schema_len =
+                    read_bin_len(stream).expect("Failed to read the schema name length");
+                let schema_name = stream
+                    .look_ahead(schema_len as usize)
+                    .expect("Failed to read the schema name");
+                // Through an error if schema name not found.
+                let ns_oid =
+                    unsafe { get_namespace_oid(schema_name.as_ptr() as *const c_char, false) };
+                stream
+                    .rewind(schema_len as usize)
+                    .expect("Failed to rewind the stream");
+                let table_len = read_bin_len(stream).expect("Failed to read the table name length");
+                let table_name = stream
+                    .look_ahead(table_len as usize)
+                    .expect("Failed to read the table name");
+                let rel_oid =
+                    unsafe { get_relname_relid(table_name.as_ptr() as *const c_char, ns_oid) };
+                if rel_oid == InvalidOid {
+                    match std::str::from_utf8(table_name) {
+                        Ok(name) => error!("Table not found: {}", name),
+                        Err(_) => error!("Table not found: {:?}", table_name),
+                    }
+                }
+                stream
+                    .rewind(table_len as usize)
+                    .expect("Failed to rewind the stream");
+                oids.push(rel_oid);
+            }
+            _ => error!("Table name should consist either of the name or the schema and the name"),
+        }
+    }
+    oids
+}
+
 fn wait_latch(timeout: Option<Duration>) {
     let timeout: c_long = timeout
         .map(|t| t.as_millis().try_into().unwrap())

src/data_type.rs

@@ -93,7 +93,7 @@ fn type_to_oid(type_: &DataType) -> pg_sys::Oid {
 }
 
 #[repr(u8)]
-enum EncodedType {
+pub(crate) enum EncodedType {
     Boolean = 0,
     Utf8 = 1,
     Int16 = 2,
@@ -123,9 +123,33 @@ impl TryFrom<u8> for EncodedType {
             8 => Ok(EncodedType::Time64),
             9 => Ok(EncodedType::Timestamp),
             10 => Ok(EncodedType::Interval),
-            _ => Err(FusionError::DeserializeU8(
+            _ => Err(FusionError::Deserialize(
                 "encoded type".to_string(),
-                value,
+                value.into(),
+            )),
+        }
+    }
+}
+
+impl TryFrom<pg_sys::Oid> for EncodedType {
+    type Error = FusionError;
+
+    fn try_from(value: pg_sys::Oid) -> Result<Self, Self::Error> {
+        match value {
+            pg_sys::BOOLOID => Ok(EncodedType::Boolean),
+            pg_sys::TEXTOID => Ok(EncodedType::Utf8),
+            pg_sys::INT2OID => Ok(EncodedType::Int16),
+            pg_sys::INT4OID => Ok(EncodedType::Int32),
+            pg_sys::INT8OID => Ok(EncodedType::Int64),
+            pg_sys::FLOAT4OID => Ok(EncodedType::Float32),
+            pg_sys::FLOAT8OID => Ok(EncodedType::Float64),
+            pg_sys::DATEOID => Ok(EncodedType::Date32),
+            pg_sys::TIMEOID => Ok(EncodedType::Time64),
+            pg_sys::TIMESTAMPOID => Ok(EncodedType::Timestamp),
+            pg_sys::INTERVALOID => Ok(EncodedType::Interval),
+            _ => Err(FusionError::Deserialize(
+                "encoded type".to_string(),
+                value.as_u32().into(),
             )),
         }
     }

src/error.rs

@@ -10,5 +10,5 @@ pub enum FusionError {
     #[error("Payload is too large: {0} bytes")]
     PayloadTooLarge(usize),
     #[error("Failed to deserialize {0}: {1}")]
-    DeserializeU8(String, u8),
+    Deserialize(String, u64),
 }

src/protocol.rs

@@ -1,14 +1,16 @@
-use crate::data_type::{datum_to_scalar, read_scalar_value, write_scalar_value};
+use crate::data_type::{datum_to_scalar, read_scalar_value, write_scalar_value, EncodedType};
 use crate::error::FusionError;
 use crate::ipc::{Bus, Slot, SlotNumber, SlotStream, DATA_SIZE};
 use crate::worker::worker_id;
 use anyhow::Result;
 use datafusion::scalar::ScalarValue;
 use datafusion_sql::TableReference;
-use pgrx::pg_sys::{ParamExternData, ProcSendSignal};
+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 rmp::encode::{write_array_len, write_bin_len, write_pfix, write_str, write_u16, RmpWrite};
+use rmp::encode::{
+    write_array_len, write_bin_len, write_bool, write_pfix, write_str, write_u16, RmpWrite,
+};
 
 #[repr(u8)]
 #[derive(Clone, Debug, Default, PartialEq)]
@@ -26,7 +28,10 @@ impl TryFrom<u8> for Direction {
         match value {
             0 => Ok(Direction::ToWorker),
             1 => Ok(Direction::ToBackend),
-            _ => Err(FusionError::DeserializeU8("direction".to_string(), value)),
+            _ => Err(FusionError::Deserialize(
+                "direction".to_string(),
+                value.into(),
+            )),
         }
     }
 }
@@ -53,7 +58,7 @@ impl TryFrom<u8> for Packet {
             2 => Ok(Packet::Failure),
             3 => Ok(Packet::Metadata),
             4 => Ok(Packet::Parse),
-            _ => Err(FusionError::DeserializeU8("packet".to_string(), value)),
+            _ => Err(FusionError::Deserialize("packet".to_string(), value.into())),
         }
     }
 }
@@ -74,7 +79,7 @@ impl TryFrom<u8> for Flag {
         match value {
             0 => Ok(Flag::More),
             1 => Ok(Flag::Last),
-            _ => Err(FusionError::DeserializeU8("flag".to_string(), value)),
+            _ => Err(FusionError::Deserialize("flag".to_string(), value.into())),
         }
     }
 }
@@ -110,6 +115,8 @@ fn signal(slot_id: SlotNumber, direction: Direction) {
     }
 }
 
+// HEADER
+
 pub(crate) fn consume_header(stream: &mut SlotStream) -> Result<Header> {
     assert_eq!(stream.position(), 0);
     let direction = Direction::try_from(read_pfix(stream)?)?;
@@ -132,6 +139,8 @@ pub(crate) fn write_header(stream: &mut SlotStream, header: &Header) -> Result<(
     Ok(())
 }
 
+// PARSE
+
 /// Reads the query from the stream, but leaves the stream position at the beginning of the query.
 /// It is required to return the reference to the query bytes without copying them. It is the
 /// caller's responsibility to move the stream position to the end of the query.
@@ -170,6 +179,8 @@ pub(crate) fn send_query(slot_id: SlotNumber, mut stream: SlotStream, query: &st
     Ok(())
 }
 
+// BIND
+
 fn prepare_params(stream: &mut SlotStream, params: &[ParamExternData]) -> Result<()> {
     stream.reset();
     // We don't know the length of the parameters yet. So we write an invalid header
@@ -217,6 +228,8 @@ pub(crate) fn send_params(
     Ok(())
 }
 
+// FAILURE
+
 pub(crate) fn read_error(stream: &mut SlotStream) -> Result<String> {
     let len = read_str_len(stream)?;
     let buf = stream.look_ahead(len as usize)?;
@@ -255,6 +268,8 @@ fn write_c_str(stream: &mut SlotStream, s: &str) -> Result<()> {
     Ok(())
 }
 
+// METADATA
+
 /// Writes a table reference as null-terminated strings to
 /// the stream. It would be used by the Rust wrappers to the
 /// C code, so if we serialize the table and schema as
@@ -312,6 +327,19 @@ pub(crate) fn send_table_refs(
     Ok(())
 }
 
+#[inline]
+pub(crate) fn write_column(
+    stream: &mut SlotStream,
+    column: &str,
+    is_null: bool,
+    etype: EncodedType,
+) -> Result<()> {
+    write_str(stream, column)?;
+    write_bool(stream, is_null)?;
+    write_pfix(stream, etype as u8)?;
+    Ok(())
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pg_schema]
 mod tests {

src/sql.rs

@@ -2,6 +2,7 @@ use std::cell::OnceCell;
 use std::collections::HashMap;
 use std::sync::Arc;
 
+use ahash::AHashMap;
 use datafusion::arrow::datatypes::DataType;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::DFSchemaRef;
@@ -93,7 +94,7 @@ impl Builtin {
 
 struct Catalog {
     builtin: Arc<Builtin>,
-    tables: HashMap<SmolStr, Arc<dyn TableSource>>,
+    tables: AHashMap<SmolStr, Arc<dyn TableSource>>,
 }
 
 impl ContextProvider for Catalog {