address

Author: aykut-bozkurt

src/arrow_parquet/uri_utils.rs

@@ -1,4 +1,4 @@
-use std::{ffi::CStr, panic, sync::Arc};
+use std::{ffi::CStr, fs::File, os::fd::FromRawFd, panic, sync::Arc};
 
 use arrow::datatypes::SchemaRef;
 use object_store::{path::Path, ObjectStoreScheme};
@@ -15,8 +15,9 @@ use pgrx::{
     ereport,
     ffi::c_char,
     pg_sys::{
-        get_role_oid, has_privs_of_role, palloc0, superuser, AsPgCStr, DataDir, FileClose,
-        FilePathName, GetUserId, InvalidOid, OpenTemporaryFile, TempTablespacePath, MAXPGPATH,
+        get_role_oid, has_privs_of_role, palloc0, superuser, AsPgCStr, ClosePipeStream, DataDir,
+        FileClose, FilePathName, GetUserId, InvalidOid, OpenPipeStream, OpenTemporaryFile,
+        TempTablespacePath, __sFILE, MAXPGPATH, PG_BINARY_R, PG_BINARY_W,
     },
 };
 use url::Url;
@@ -39,22 +40,42 @@ pub(crate) struct ParsedUriInfo {
     pub(crate) bucket: Option<String>,
     pub(crate) path: Path,
     pub(crate) scheme: ObjectStoreScheme,
-    pub(crate) stdio_tmp_fd: Option<i32>,
-    pub(crate) program: Option<*mut c_char>,
+    // tmp_fd is used as intermediate file for copying data to/from stdin/out or program pipes
+    pub(crate) tmp_fd: Option<i32>,
+    // pipe_file is used to hold the pipe file descriptor for copying data to/from a program
+    // call open_program_pipe to open the pipe to a program
+    pub(crate) is_program: bool,
+    pub(crate) pipe_file: *mut __sFILE,
 }
 
 impl ParsedUriInfo {
     pub(crate) fn for_std_inout() -> Self {
-        Self::create_tmp_file()
+        Self::with_tmp_file()
     }
 
-    pub(crate) fn for_program(program: *mut c_char) -> Self {
-        let mut uri_info = Self::create_tmp_file();
-        uri_info.program = Some(program);
+    pub(crate) fn for_program() -> Self {
+        let mut uri_info = Self::with_tmp_file();
+        uri_info.is_program = true;
         uri_info
     }
 
-    fn create_tmp_file() -> Self {
+    pub(crate) fn open_program_pipe(&mut self, program: &str, copy_from: bool) -> File {
+        let pipe_mode = if copy_from { PG_BINARY_R } else { PG_BINARY_W };
+
+        let pipe_file = unsafe { OpenPipeStream(program.as_pg_cstr(), pipe_mode.as_ptr()) };
+
+        if pipe_file.is_null() {
+            panic!("Failed to open pipe stream for program: {}", program);
+        }
+
+        self.pipe_file = pipe_file;
+
+        let pipe_fd = (unsafe { *self.pipe_file })._file;
+
+        unsafe { File::from_raw_fd(pipe_fd as _) }
+    }
+
+    fn with_tmp_file() -> Self {
         // open temp postgres file, which is removed after transaction ends
         let tmp_path_fd = unsafe { OpenTemporaryFile(false) };
 
@@ -81,11 +102,15 @@ impl ParsedUriInfo {
 
         let mut parsed_uri = Self::try_from(tmp_path.as_str()).unwrap_or_else(|e| panic!("{}", e));
 
-        parsed_uri.stdio_tmp_fd = Some(tmp_path_fd);
+        parsed_uri.tmp_fd = Some(tmp_path_fd);
 
         parsed_uri
     }
 
+    fn is_std_inout(&self) -> bool {
+        self.tmp_fd.is_some() && !self.is_program
+    }
+
     fn try_parse_uri(uri: &str) -> Result<Url, String> {
         if !uri.contains("://") {
             // local file
@@ -139,17 +164,23 @@ impl TryFrom<&str> for ParsedUriInfo {
             bucket,
             path,
             scheme,
-            stdio_tmp_fd: None,
-            program: None,
+            tmp_fd: None,
+            is_program: false,
+            pipe_file: std::ptr::null_mut(),
         })
     }
 }
 
 impl Drop for ParsedUriInfo {
     fn drop(&mut self) {
-        if let Some(stdio_tmp_fd) = self.stdio_tmp_fd {
+        if let Some(tmp_fd) = self.tmp_fd {
             // close temp file, postgres api will remove it on close
-            unsafe { FileClose(stdio_tmp_fd) };
+            unsafe { FileClose(tmp_fd) };
+        }
+
+        if !self.pipe_file.is_null() {
+            // close pipe file, postgres api will remove it on close
+            unsafe { ClosePipeStream(self.pipe_file) };
         }
     }
 }
@@ -286,17 +317,15 @@ pub(crate) fn ensure_access_privilege_to_uri(uri_info: &ParsedUriInfo, copy_from
         return;
     }
 
-    let is_program = uri_info.program.is_some();
-
     // permission check is not needed for stdin/out
-    if uri_info.stdio_tmp_fd.is_some() && !is_program {
+    if uri_info.is_std_inout() {
         return;
     }
 
     let user_id = unsafe { GetUserId() };
     let is_file = uri_info.uri.scheme() == "file";
 
-    let required_role_name = if is_program {
+    let required_role_name = if uri_info.is_program {
         "pg_execute_server_program"
     } else if is_file {
         if copy_from {
@@ -317,7 +346,7 @@ pub(crate) fn ensure_access_privilege_to_uri(uri_info: &ParsedUriInfo, copy_from
         unsafe { get_role_oid(required_role_name.to_string().as_pg_cstr(), false) };
 
     let operation_str = if copy_from { "from" } else { "to" };
-    let object_type = if is_program {
+    let object_type = if uri_info.is_program {
         "program"
     } else if is_file {
         "file"

src/parquet_copy_hook/copy_from.rs

@@ -16,7 +16,8 @@ use crate::{
     parquet_copy_hook::{
         copy_from_program::copy_program_to_file,
         copy_utils::{
-            copy_from_stmt_create_option_list, copy_stmt_lock_mode, copy_stmt_relation_oid,
+            copy_from_stmt_create_option_list, copy_stmt_is_std_inout, copy_stmt_lock_mode,
+            copy_stmt_program, copy_stmt_relation_oid,
         },
     },
 };
@@ -120,7 +121,7 @@ pub(crate) fn execute_copy_from(
     p_stmt: &PgBox<PlannedStmt>,
     query_string: &CStr,
     query_env: &PgBox<QueryEnvironment>,
-    uri_info: &ParsedUriInfo,
+    mut uri_info: ParsedUriInfo,
 ) -> u64 {
     let rel_oid = copy_stmt_relation_oid(p_stmt);
 
@@ -149,15 +150,15 @@ pub(crate) fn execute_copy_from(
     let match_by = copy_from_stmt_match_by(p_stmt);
 
     unsafe {
-        if uri_info.program.is_some() {
-            copy_program_to_file(uri_info);
-        } else if uri_info.stdio_tmp_fd.is_some() {
+        if let Some(program) = copy_stmt_program(p_stmt) {
+            copy_program_to_file(&mut uri_info, &program);
+        } else if copy_stmt_is_std_inout(p_stmt) {
             let is_binary = true;
-            copy_stdin_to_file(uri_info, tupledesc.natts as _, is_binary);
+            copy_stdin_to_file(&uri_info, 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);
+        let parquet_reader_context = ParquetReaderContext::new(&uri_info, match_by, &tupledesc);
         push_parquet_reader_context(parquet_reader_context);
 
         // makes sure to set binary format

src/parquet_copy_hook/copy_from_program.rs

@@ -1,39 +1,13 @@
-use std::{ffi::CStr, io::pipe, process::Command};
-
 use crate::arrow_parquet::uri_utils::{uri_as_string, ParsedUriInfo};
 
-pub(crate) unsafe fn copy_program_to_file(uri_info: &ParsedUriInfo) {
-    let program = CStr::from_ptr(uri_info.program.expect("Program pointer is null"))
-        .to_str()
-        .unwrap_or_else(|e| panic!("Failed to convert program pointer to string: {e}"))
-        .to_string();
-
-    let (mut pipe_in, pipe_out) =
-        pipe().unwrap_or_else(|e| panic!("Failed to create command pipe: {e}"));
-
-    #[cfg(unix)]
-    let mut command = Command::new("/bin/sh")
-        .arg("-lc")
-        .arg(program)
-        .stdout(pipe_out)
-        .spawn()
-        .unwrap_or_else(|e| panic!("Failed to spawn command: {e}"));
-
-    #[cfg(windows)]
-    let mut command = Command::new("cmd")
-        .arg("/C")
-        .arg(program)
-        .stdout(pipe_out)
-        .spawn()
-        .unwrap_or_else(|e| panic!("Failed to spawn command: {e}"));
-
-    command
-        .wait()
-        .unwrap_or_else(|e| panic!("Failed to wait for command: {e}"));
-
-    // Write input pipe to the file
+pub(crate) unsafe fn copy_program_to_file(uri_info: &mut ParsedUriInfo, program: &str) {
+    // get tmp file
     let path = uri_as_string(&uri_info.uri);
 
+    // open and then get pipe file
+    let copy_from = true;
+    let mut pipe_file = uri_info.open_program_pipe(program, copy_from);
+
     // create or overwrite the local file
     let mut file = std::fs::OpenOptions::new()
         .write(true)
@@ -42,6 +16,7 @@ pub(crate) unsafe fn copy_program_to_file(uri_info: &ParsedUriInfo) {
         .open(&path)
         .unwrap_or_else(|e| panic!("{}", e));
 
-    std::io::copy(&mut pipe_in, &mut file)
+    // Write pipe file to temp file
+    std::io::copy(&mut pipe_file, &mut file)
         .unwrap_or_else(|e| panic!("Failed to copy command stdout to file: {e}"));
 }

src/parquet_copy_hook/copy_to_dest_receiver.rs

@@ -36,8 +36,8 @@ pub(crate) struct CopyToParquetDestReceiver {
     collected_tuple_column_sizes: *mut i64,
     target_batch_size: i64,
     uri: *const c_char,
-    is_to_stdout: bool,
     program: *mut c_char,
+    is_to_stdout: bool,
     copy_options: CopyToParquetOptions,
     copy_memory_context: MemoryContext,
     row_group_memory_context: MemoryContext,
@@ -156,10 +156,16 @@ impl CopyToParquetDestReceiver {
 
     fn copy_to_program(&self) {
         if !self.program.is_null() {
+            let program = unsafe {
+                CStr::from_ptr(self.program)
+                    .to_str()
+                    .expect("invalid program")
+            };
+
             let uri = unsafe { CStr::from_ptr(self.uri).to_str().expect("invalid uri") };
-            let uri_info = ParsedUriInfo::try_from(uri).expect("invalid uri");
+            let mut uri_info = ParsedUriInfo::try_from(uri).expect("invalid uri");
 
-            unsafe { copy_file_to_program(uri_info, self.program) };
+            unsafe { copy_file_to_program(&mut uri_info, program) };
         }
     }
 
@@ -370,8 +376,8 @@ fn tuple_column_sizes(tuple_datums: &[Option<Datum>], tupledesc: &PgTupleDesc) -
 #[pg_guard]
 pub(crate) extern "C-unwind" fn create_copy_to_parquet_dest_receiver(
     uri: *const c_char,
-    is_to_stdout: bool,
     program: *mut c_char,
+    is_to_stdout: bool,
     options: CopyToParquetOptions,
 ) -> *mut CopyToParquetDestReceiver {
     let row_group_memory_context = unsafe {
@@ -403,8 +409,8 @@ pub(crate) extern "C-unwind" fn create_copy_to_parquet_dest_receiver(
     parquet_dest.dest.rDestroy = Some(copy_destroy);
     parquet_dest.dest.mydest = CommandDest::DestCopyOut;
     parquet_dest.uri = uri;
-    parquet_dest.is_to_stdout = is_to_stdout;
     parquet_dest.program = program;
+    parquet_dest.is_to_stdout = is_to_stdout;
     parquet_dest.tupledesc = std::ptr::null_mut();
     parquet_dest.parquet_writer_context = std::ptr::null_mut();
     parquet_dest.natts = 0;

src/parquet_copy_hook/copy_to_program.rs

@@ -1,48 +1,20 @@
-use std::{
-    ffi::{c_char, CStr},
-    fs::File,
-    io::pipe,
-    process::Command,
-};
+use std::fs::File;
 
 use crate::arrow_parquet::uri_utils::{uri_as_string, ParsedUriInfo};
 
-pub(crate) unsafe fn copy_file_to_program(uri_info: ParsedUriInfo, program: *mut c_char) {
-    let (pipe_in, mut pipe_out) =
-        pipe().unwrap_or_else(|e| panic!("Failed to create command pipe: {e}"));
-
-    let program = CStr::from_ptr(program).to_string_lossy().into_owned();
-
-    #[cfg(unix)]
-    let mut command = Command::new("/bin/sh")
-        .arg("-lc")
-        .arg(program)
-        .stdin(pipe_in)
-        .spawn()
-        .unwrap_or_else(|e| panic!("Failed to spawn command: {e}"));
-
-    #[cfg(windows)]
-    let mut command = Command::new("cmd")
-        .arg("/C")
-        .arg(program)
-        .stdin(pipe_in)
-        .spawn()
-        .unwrap_or_else(|e| panic!("Failed to spawn command: {e}"));
-
-    // Write the file to output pipe
+pub(crate) unsafe fn copy_file_to_program(uri_info: &mut ParsedUriInfo, program: &str) {
+    // get tmp file
     let path = uri_as_string(&uri_info.uri);
 
     let mut file = File::open(path).unwrap_or_else(|e| {
         panic!("could not open temp file: {e}");
     });
 
-    std::io::copy(&mut file, &mut pipe_out)
-        .unwrap_or_else(|e| panic!("Failed to copy file to command stdin: {e}"));
-
-    // close output pipe to unblock the command
-    drop(pipe_out);
+    // open and then get pipe file
+    let copy_from = false;
+    let mut pipe_file = uri_info.open_program_pipe(program, copy_from);
 
-    command
-        .wait()
-        .unwrap_or_else(|e| panic!("Failed to wait for command: {e}"));
+    // Write temp file to pipe file
+    std::io::copy(&mut file, &mut pipe_file)
+        .unwrap_or_else(|e| panic!("Failed to copy file to command stdin: {e}"));
 }

src/parquet_copy_hook/copy_to_split_dest_receiver.rs

@@ -23,8 +23,8 @@ pub(crate) const INVALID_FILE_SIZE_BYTES: i64 = 0;
 struct CopyToParquetSplitDestReceiver {
     dest: DestReceiver,
     uri: *const c_char,
-    is_to_stdout: bool,
     program: *mut c_char,
+    is_to_stdout: bool,
     tupledesc: TupleDesc,
     operation: i32,
     options: CopyToParquetOptions,
@@ -50,8 +50,8 @@ impl CopyToParquetSplitDestReceiver {
         let child_uri = self.create_uri_for_child();
         self.current_child_receiver = create_copy_to_parquet_dest_receiver(
             child_uri,
-            self.is_to_stdout,
             self.program,
+            self.is_to_stdout,
             self.options,
         );
         self.current_child_id += 1;
@@ -212,8 +212,8 @@ extern "C-unwind" fn copy_split_destroy(_dest: *mut DestReceiver) {}
 #[allow(clippy::too_many_arguments)]
 pub extern "C-unwind" fn create_copy_to_parquet_split_dest_receiver(
     uri: *const c_char,
-    is_to_stdout: bool,
     program: *mut c_char,
+    is_to_stdout: bool,
     file_size_bytes: *const i64,
     field_ids: *const c_char,
     row_group_size: *const i64,