feat: process metadata request (get table oids)

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,14 +1,19 @@
+use anyhow::Result as AnyResult;
 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;
 
+use crate::error::FusionError;
 use crate::ipc::{my_slot, Bus, SlotHandler, SlotNumber, SlotStream, CURRENT_SLOT};
 use crate::protocol::{consume_header, read_error, send_params, send_query, Direction, Packet};
 
@@ -145,6 +150,74 @@ unsafe extern "C" fn explain_df_scan(
     todo!()
 }
 
+// We expect that the header is already consumed and the packet type is `Packet::Metadata`.
+fn table_oids(stream: &mut SlotStream) -> AnyResult<SmallVec<[Oid; 16]>> {
+    let table_not_found = |c_table_name: &[u8]| -> Result<(), FusionError> {
+        assert!(!c_table_name.is_empty());
+        let table_name = c_table_name[..c_table_name.len() - 1].as_ref();
+        match std::str::from_utf8(table_name) {
+            Ok(name) => Err(FusionError::NotFound("Table".into(), name.into())),
+            Err(_) => Err(FusionError::NotFound(
+                "Table".into(),
+                format!("{:?}", table_name),
+            )),
+        }
+    };
+    let table_num = read_array_len(stream)?;
+    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)?;
+        match elem_num {
+            1 => {
+                let table_len = read_bin_len(stream)?;
+                let table_name = stream.look_ahead(table_len as usize)?;
+                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 {
+                    table_not_found(table_name)?;
+                }
+                stream.rewind(table_len as usize)?;
+            }
+            2 => {
+                let schema_len = read_bin_len(stream)?;
+                let schema_name = stream.look_ahead(schema_len as usize)?;
+                // 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)?;
+                let table_len = read_bin_len(stream)?;
+                let table_name = stream.look_ahead(table_len as usize)?;
+                let rel_oid =
+                    unsafe { get_relname_relid(table_name.as_ptr() as *const c_char, ns_oid) };
+                if rel_oid == InvalidOid {
+                    table_not_found(table_name)?;
+                }
+                stream.rewind(table_len as usize)?;
+                oids.push(rel_oid);
+            }
+            _ => {
+                return Err(FusionError::InvalidName(
+                    "Table".into(),
+                    "support only 'schema.table' format".into(),
+                ))?
+            }
+        }
+    }
+    Ok(oids)
+}
+
 fn wait_latch(timeout: Option<Duration>) {
     let timeout: c_long = timeout
         .map(|t| t.as_millis().try_into().unwrap())
@@ -215,9 +288,17 @@ fn list_nth(list: *mut List, n: i32) -> *mut ListCell {
 #[cfg(any(test, feature = "pg_test"))]
 #[pgrx::pg_schema]
 mod tests {
+    use crate::ipc::Slot;
+    use crate::protocol::prepare_table_refs;
+
     use super::*;
+    use datafusion_sql::TableReference;
     use pgrx::prelude::*;
+    use pgrx::spi::Spi;
     use std::ffi::c_void;
+    use std::ptr::addr_of_mut;
+
+    const SLOT_SIZE: usize = 8204;
 
     #[pg_test]
     fn test_node() {
@@ -231,4 +312,45 @@ mod tests {
             pg_sys::pfree(ptr);
         }
     }
+
+    #[pg_test]
+    fn test_table_oids() {
+        Spi::run("create table if not exists t1(a int, b text);").unwrap();
+        Spi::run("create schema if not exists s1;").unwrap();
+        Spi::run("create table if not exists s1.t2(a int, b text);").unwrap();
+        let t1_oid = Spi::get_one::<Oid>("select 't1'::regclass::oid;")
+            .unwrap()
+            .unwrap();
+        let t2_oid = Spi::get_one::<Oid>("select 's1.t2'::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();
+
+        // Check valid tables.
+        let t1 = TableReference::bare("t1");
+        let t2 = TableReference::partial("s1", "t2");
+        let tables = vec![&t1, &t2];
+        prepare_table_refs(&mut stream, &tables).unwrap();
+        stream.reset();
+        let _ = consume_header(&mut stream).unwrap();
+        let oids = table_oids(&mut stream).unwrap();
+        assert_eq!(oids.len(), 2);
+        assert_eq!(oids[0], t1_oid);
+        assert_eq!(oids[1], t2_oid);
+        stream.reset();
+
+        // Check invalid table.
+        let t3 = TableReference::bare("t3");
+        let tables = vec![&t3];
+        prepare_table_refs(&mut stream, &tables).unwrap();
+        stream.reset();
+        let _ = consume_header(&mut stream).unwrap();
+        let err = table_oids(&mut stream).unwrap_err();
+        assert_eq!(err.to_string(), "Table not found: t3");
+    }
 }

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,9 @@ pub enum FusionError {
     #[error("Payload is too large: {0} bytes")]
     PayloadTooLarge(usize),
     #[error("Failed to deserialize {0}: {1}")]
-    DeserializeU8(String, u8),
+    Deserialize(String, u64),
+    #[error("{0} not found: {1}")]
+    NotFound(String, String),
+    #[error("{0} name is not valid: {1}")]
+    InvalidName(String, String),
 }

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
@@ -276,7 +291,10 @@ fn write_table_ref(stream: &mut SlotStream, table: &TableReference) -> Result<()
     Ok(())
 }
 
-fn prepare_table_refs(stream: &mut SlotStream, tables: &[&TableReference]) -> Result<()> {
+pub(crate) fn prepare_table_refs(
+    stream: &mut SlotStream,
+    tables: &[&TableReference],
+) -> Result<()> {
     stream.reset();
     // We don't know the length of the tables yet. So we write an invalid header
     // to replace it with the correct one later.
@@ -312,6 +330,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 {