copy to/from stdout/stdin with (format parquet)

Author: aykut-bozkurt

.devcontainer/.env

@@ -28,4 +28,7 @@ GOOGLE_SERVICE_ACCOUNT_KEY='{"gcs_base_url": "http://localhost:4443","disable_oa
 GOOGLE_SERVICE_ENDPOINT=http://localhost:4443
 
 # Others
+# run pgrx tests with a single thread to avoid race conditions
 RUST_TEST_THREADS=1
+# pgrx runs test on the port base_port + pg_version
+PGRX_TEST_PG_BASE_PORT=5454

.github/workflows/ci.yml

@@ -94,7 +94,8 @@ jobs:
       - name: Install and configure pgrx
         run: |
           cargo install --locked [email protected]
-          cargo pgrx init --pg${{ env.PG_MAJOR }} /usr/lib/postgresql/${{ env.PG_MAJOR }}/bin/pg_config
+          cargo pgrx init --pg${{ env.PG_MAJOR }} /usr/lib/postgresql/${{ env.PG_MAJOR }}/bin/pg_config \
+                          --base-testing-port $PGRX_TEST_PG_BASE_PORT
 
       - name: Install cargo-llvm-cov for coverage report
         run: cargo install --locked [email protected]

Cargo.lock

@@ -46,6 +46,7 @@ serde = "1"
 serde_json = "1"
 tokio = {version = "1", default-features = false, features = ["rt", "time", "macros"]}
 url = "2"
+uuid = { version = "1", default-features = false, features = ["v4"] }
 
 [dev-dependencies]
 pgrx-tests = "=0.13.1"

Cargo.toml

@@ -130,7 +130,7 @@ impl ParquetWriterContext {
     }
 
     // finalize flushes the in progress rows to a new row group and finally writes metadata to the file.
-    fn finalize(&mut self) {
+    pub(crate) fn finalize(&mut self) {
         PG_BACKEND_TOKIO_RUNTIME
             .block_on(self.parquet_writer.finish())
             .unwrap_or_else(|e| panic!("failed to finish parquet writer: {}", e));
@@ -156,9 +156,3 @@ impl ParquetWriterContext {
         RecordBatch::try_new(schema, attribute_arrays).expect("Expected record batch")
     }
 }
-
-impl Drop for ParquetWriterContext {
-    fn drop(&mut self) {
-        self.finalize();
-    }
-}

src/arrow_parquet/parquet_writer.rs

@@ -1,4 +1,4 @@
-use std::{panic, sync::Arc};
+use std::{env::temp_dir, panic, sync::Arc};
 
 use arrow::datatypes::SchemaRef;
 use object_store::{path::Path, ObjectStoreScheme};
@@ -16,6 +16,7 @@ use pgrx::{
     pg_sys::{get_role_oid, has_privs_of_role, superuser, AsPgCStr, GetUserId},
 };
 use url::Url;
+use uuid::Uuid;
 
 use crate::{
     object_store::{
@@ -35,9 +36,21 @@ pub(crate) struct ParsedUriInfo {
     pub(crate) bucket: Option<String>,
     pub(crate) path: Path,
     pub(crate) scheme: ObjectStoreScheme,
+    pub(crate) is_stdio: bool,
 }
 
 impl ParsedUriInfo {
+    fn for_stdout() -> Self {
+        let path = temp_dir().join(format!("pg_parquet_{}", Uuid::new_v4()));
+
+        let mut parsed_uri = Self::try_from(path.to_str().expect("invalid temp path"))
+            .unwrap_or_else(|e| panic!("{}", e));
+
+        parsed_uri.is_stdio = true;
+
+        parsed_uri
+    }
+
     fn try_parse_uri(uri: &str) -> Result<Url, String> {
         if !uri.contains("://") {
             // local file
@@ -81,6 +94,10 @@ impl TryFrom<&str> for ParsedUriInfo {
     type Error = String;
 
     fn try_from(uri: &str) -> Result<Self, Self::Error> {
+        if uri == "std" {
+            return Ok(Self::for_stdout());
+        }
+
         let uri = Self::try_parse_uri(uri)?;
 
         let (scheme, path) = Self::try_parse_scheme(&uri)?;
@@ -92,6 +109,7 @@ impl TryFrom<&str> for ParsedUriInfo {
             bucket,
             path,
             scheme,
+            is_stdio: false,
         })
     }
 }

src/arrow_parquet/uri_utils.rs

@@ -53,6 +53,7 @@ impl ObjectStoreCache {
             path,
             scheme,
             bucket,
+            ..
         } = uri_info;
 
         // no need to cache local files

src/object_store/object_store_cache.rs

@@ -1,7 +1,9 @@
 pub(crate) mod copy_from;
+pub(crate) mod copy_from_stdin;
 pub(crate) mod copy_to;
 pub(crate) mod copy_to_dest_receiver;
 pub(crate) mod copy_to_split_dest_receiver;
+pub(crate) mod copy_to_stdout;
 pub(crate) mod copy_utils;
 pub(crate) mod hook;
 pub(crate) mod pg_compat;

src/parquet_copy_hook.rs

@@ -18,9 +18,12 @@ use crate::{
     },
 };
 
-use super::copy_utils::{
-    copy_from_stmt_match_by, copy_stmt_attribute_list, copy_stmt_create_namespace_item,
-    copy_stmt_create_parse_state, create_filtered_tupledesc_for_relation,
+use super::{
+    copy_from_stdin::copy_stdin_to_file,
+    copy_utils::{
+        copy_from_stmt_match_by, copy_stmt_attribute_list, copy_stmt_create_namespace_item,
+        copy_stmt_create_parse_state, create_filtered_tupledesc_for_relation,
+    },
 };
 
 // stack to store parquet reader contexts for COPY FROM.
@@ -140,6 +143,11 @@ pub(crate) fn execute_copy_from(
     let match_by = copy_from_stmt_match_by(p_stmt);
 
     unsafe {
+        if uri_info.is_stdio {
+            let is_binary = true;
+            copy_stdin_to_file(uri_info.clone(), tupledesc.natts as _, is_binary);
+        }
+
         // parquet reader context is used throughout the COPY FROM operation.
         let parquet_reader_context = ParquetReaderContext::new(uri_info, match_by, &tupledesc);
         push_parquet_reader_context(parquet_reader_context);

src/parquet_copy_hook/copy_from.rs

@@ -0,0 +1,216 @@
+use std::{ffi::CStr, io::Write};
+
+use pgrx::{
+    ffi::c_char,
+    pg_sys::{
+        makeStringInfo, pq_beginmessage, pq_copymsgbytes, pq_endmessage, pq_getmsgstring,
+        pq_sendbyte, pq_sendint16, QueryCancelHoldoffCount, StringInfo,
+    },
+};
+
+use crate::arrow_parquet::uri_utils::{uri_as_string, ParsedUriInfo};
+
+/*
+ * CopyFromStdinState is a simplified version of CopyFromState
+ * in PostgreSQL to use in ReceiveDataFromClient with names
+ * preserved.
+ */
+struct CopyFromStdinState {
+    /* buffer in which we store incoming bytes */
+    fe_msgbuf: StringInfo,
+
+    /* whether we reached the end-of-file */
+    raw_reached_eof: bool,
+}
+
+const MAX_READ_SIZE: usize = 65536;
+
+/*
+ * CopyInputToFile copies data from the socket to the given file.
+ * We request the client send a specific column count.
+ */
+pub(crate) unsafe fn copy_stdin_to_file(uri_info: ParsedUriInfo, natts: i16, is_binary: bool) {
+    let mut cstate = CopyFromStdinState {
+        fe_msgbuf: makeStringInfo(),
+        raw_reached_eof: false,
+    };
+
+    /* open the destination file for writing */
+    let path = uri_as_string(&uri_info.uri);
+
+    // create or overwrite the local file
+    let mut file = std::fs::OpenOptions::new()
+        .write(true)
+        .truncate(true)
+        .create(true)
+        .open(&path)
+        .unwrap_or_else(|e| panic!("{}", e));
+
+    /* tell the client we are ready for data */
+    send_copy_in_begin(natts, is_binary);
+
+    /* allocate on the heap since it's quite big */
+    let mut receive_buffer = vec![0u8; MAX_READ_SIZE];
+
+    while !cstate.raw_reached_eof {
+        /* copy some bytes from the client into fe_msgbuf */
+        let bytes_read = receive_data_from_client(&mut cstate, &mut receive_buffer);
+
+        if bytes_read == 0 {
+            break;
+        }
+
+        if bytes_read > 0 {
+            /* copy bytes from fe_msgbuf to the destination file */
+            file.write_all(&receive_buffer[..bytes_read])
+                .unwrap_or_else(|e| {
+                    panic!("could not write to file: {}", e);
+                });
+        }
+    }
+}
+
+/*
+ * send_copy_in_begin sends the CopyInResponse message that the client
+ * expects after a COPY .. FROM STDIN.
+ *
+ * This code is adapted from ReceiveCopyBegin in PostgreSQL.
+ */
+unsafe fn send_copy_in_begin(natts: i16, is_binary: bool) {
+    let buf = makeStringInfo();
+
+    pq_beginmessage(buf, 'G' as _);
+
+    let copy_format = if is_binary { 1 } else { 0 };
+    pq_sendbyte(buf, copy_format);
+
+    pq_sendint16(buf, natts as _);
+    for _ in 0..natts {
+        /* use the same format for all columns */
+        pq_sendint16(buf, copy_format as _);
+    }
+
+    pq_endmessage(buf);
+    ((*PqCommMethods).flush)();
+}
+
+const PQ_LARGE_MESSAGE_LIMIT: i32 = 1024 * 1024 * 1024 - 3;
+const PQ_SMALL_MESSAGE_LIMIT: i32 = 10000;
+
+unsafe fn receive_data_from_client(
+    cstate: &mut CopyFromStdinState,
+    receive_buffer: &mut [u8],
+) -> usize {
+    let mut databuf = receive_buffer;
+
+    let minread = 1;
+    let mut maxread = MAX_READ_SIZE;
+
+    let mut bytesread = 0;
+
+    while maxread > 0 && bytesread < minread && !cstate.raw_reached_eof {
+        let mut avail;
+        let mut flushed = false;
+
+        while flushed || (*cstate.fe_msgbuf).cursor >= (*cstate.fe_msgbuf).len {
+            /* Try to receive another message */
+
+            QueryCancelHoldoffCount += 1;
+
+            pq_startmsgread();
+
+            let mtype = pq_getbyte();
+            if mtype == -1 {
+                panic!("unexpected EOF on client connection with an open transaction");
+            }
+
+            /* Validate message type and set packet size limit */
+            let maxmsglen = match mtype as u8 as char {
+                'd' =>
+                /* CopyData */
+                {
+                    PQ_LARGE_MESSAGE_LIMIT
+                }
+                'c' | 'f' | 'H' | 'S' =>
+                /* CopyDone, CopyFail, Flush, Sync */
+                {
+                    PQ_SMALL_MESSAGE_LIMIT
+                }
+                _ => {
+                    panic!(
+                        "unexpected message type 0x{:02X} during COPY from stdin",
+                        mtype
+                    );
+                }
+            };
+
+            /* Now collect the message body */
+            if pq_getmessage(cstate.fe_msgbuf, maxmsglen) != 0 {
+                panic!("unexpected EOF on client connection with an open transaction");
+            }
+
+            QueryCancelHoldoffCount -= 1;
+
+            /* ... and process it */
+            match mtype as u8 as char {
+                'd' => {
+                    /* CopyData */
+                    break;
+                }
+                'c' => {
+                    /* CopyDone */
+                    cstate.raw_reached_eof = true;
+                    return bytesread;
+                }
+                'f' => {
+                    /* CopyFail */
+                    let msg = pq_getmsgstring(cstate.fe_msgbuf);
+                    let msg = CStr::from_ptr(msg).to_str().expect("invalid CStr");
+                    panic!("COPY from stdin failed: {msg}");
+                }
+                'H' | 'S' => {
+                    /* Flush, Sync */
+                    flushed = true;
+                    continue;
+                }
+                _ => {
+                    panic!(
+                        "unexpected message type 0x{:02X} during COPY from stdin",
+                        mtype
+                    );
+                }
+            }
+        }
+
+        avail = ((*cstate.fe_msgbuf).len - (*cstate.fe_msgbuf).cursor) as _;
+        if avail > maxread {
+            avail = maxread;
+        }
+
+        pq_copymsgbytes(cstate.fe_msgbuf, databuf.as_mut_ptr() as _, avail as _);
+        databuf = &mut databuf[avail..];
+        maxread -= avail;
+        bytesread += avail;
+    }
+
+    bytesread
+}
+
+// todo: move to pgrx (include libpq.h)
+#[repr(C)]
+struct PQcommMethods {
+    comm_reset: unsafe extern "C" fn(),
+    flush: unsafe extern "C" fn() -> i32,
+    flush_if_writable: unsafe extern "C" fn() -> i32,
+    is_send_pending: unsafe extern "C" fn() -> bool,
+    putmessage: unsafe extern "C" fn(msgtype: u32, s: *const c_char, len: usize) -> i32,
+    putmessage_noblock: unsafe extern "C" fn(msgtype: u32, s: *const c_char, len: usize),
+}
+
+unsafe extern "C" {
+    fn pq_startmsgread();
+    fn pq_getmessage(s: StringInfo, maxlen: i32) -> i32;
+    fn pq_getbyte() -> i32;
+
+    static PqCommMethods: *mut PQcommMethods;
+}