diff --git a/Cargo.lock b/Cargo.lock index b5a717a..536a528 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -137,6 +137,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "anyhow" +version = "1.0.101" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea" + [[package]] name = "apache-avro" version = "0.21.0" @@ -723,6 +729,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.0", +] + [[package]] name = "chrono" version = "0.4.43" @@ -893,6 +910,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -2119,6 +2145,20 @@ dependencies = [ "wasip2", ] +[[package]] +name = "getrandom" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "rand_core 0.10.0", + "wasip2", + "wasip3", +] + [[package]] name = "getset" version = "0.1.6" @@ -2335,6 +2375,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + [[package]] name = "ident_case" version = "1.0.1" @@ -2370,6 +2416,8 @@ checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" dependencies = [ "equivalent", "hashbrown 0.16.1", + "serde", + "serde_core", ] [[package]] @@ -2481,6 +2529,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + [[package]] name = "lexical-core" version = "1.0.6" @@ -2902,9 +2956,9 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.37.3" +version = "0.38.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6fcd410bc6990bd8d20b3fe3cd879a3c3ec250bdb1cb12537b528818823b02c9" +checksum = "89d5e5a60d3f6e40c91f6a2a7f8d09665e636272bd5611977253559b6651aabb" dependencies = [ "async-trait", "base64", @@ -2917,7 +2971,7 @@ dependencies = [ "md5", "pg_interval_2", "postgres-types", - "rand 0.9.2", + "rand 0.10.0", "ring", "rust_decimal", "rustls-pki-types", @@ -3216,6 +3270,17 @@ dependencies = [ "rand_core 0.9.3", ] +[[package]] +name = "rand" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" +dependencies = [ + "chacha20", + "getrandom 0.4.1", + "rand_core 0.10.0", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -3254,6 +3319,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" + [[package]] name = "rayon" version = "1.11.0" @@ -3578,7 +3649,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -4104,6 +4175,12 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + [[package]] name = "untrusted" version = "0.9.0" @@ -4174,7 +4251,16 @@ version = "1.0.1+wasi-0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" dependencies = [ - "wit-bindgen", + "wit-bindgen 0.46.0", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen 0.51.0", ] [[package]] @@ -4235,6 +4321,40 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags 2.10.0", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + [[package]] name = "web-sys" version = "0.3.83" @@ -4525,6 +4645,94 @@ version = "0.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck 0.5.0", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck 0.5.0", + "indexmap", + "prettyplease", + "syn 2.0.114", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.114", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags 2.10.0", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + [[package]] name = "wkb" version = "0.9.2" diff --git a/Cargo.toml b/Cargo.toml index 7d44f29..0f42ace 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ bytes = "1.11.1" chrono = { version = "0.4", features = ["std"] } datafusion = { version = "52", default-features = false } futures = "0.3" -pgwire = { version = "0.37.3", default-features = false } +pgwire = { version = "0.38", default-features = false } postgres-types = "0.2" rust_decimal = { version = "1.40", features = ["db-postgres"] } tokio = { version = "1", default-features = false } diff --git a/arrow-pg/src/encoder.rs b/arrow-pg/src/encoder.rs index 58dc89c..e848b3f 100644 --- a/arrow-pg/src/encoder.rs +++ b/arrow-pg/src/encoder.rs @@ -8,8 +8,10 @@ use chrono::{NaiveDate, NaiveDateTime}; #[cfg(feature = "datafusion")] use datafusion::arrow::{array::*, datatypes::*}; use pg_interval::Interval as PgInterval; -use pgwire::api::results::{DataRowEncoder, FieldInfo}; +use pgwire::api::results::{CopyEncoder, DataRowEncoder, FieldInfo}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; +use pgwire::messages::copy::CopyData; +use pgwire::messages::data::DataRow; use pgwire::types::ToSqlText; use postgres_types::ToSql; use rust_decimal::Decimal; @@ -22,12 +24,18 @@ use crate::list_encoder::encode_list; use crate::struct_encoder::encode_struct; pub trait Encoder { + type Item; + fn encode_field(&mut self, value: &T, pg_field: &FieldInfo) -> PgWireResult<()> where T: ToSql + ToSqlText + Sized; + + fn take_row(&mut self) -> Self::Item; } impl Encoder for DataRowEncoder { + type Item = DataRow; + fn encode_field(&mut self, value: &T, pg_field: &FieldInfo) -> PgWireResult<()> where T: ToSql + ToSqlText + Sized, @@ -39,6 +47,25 @@ impl Encoder for DataRowEncoder { pg_field.format_options(), ) } + + fn take_row(&mut self) -> Self::Item { + self.take_row() + } +} + +impl Encoder for CopyEncoder { + type Item = CopyData; + + fn encode_field(&mut self, value: &T, _pg_field: &FieldInfo) -> PgWireResult<()> + where + T: ToSql + ToSqlText + Sized, + { + self.encode_field(value) + } + + fn take_row(&mut self) -> Self::Item { + self.take_copy() + } } fn get_bool_value(arr: &Arc, idx: usize) -> Option { @@ -522,6 +549,8 @@ mod tests { } impl Encoder for MockEncoder { + type Item = String; + fn encode_field(&mut self, value: &T, pg_field: &FieldInfo) -> PgWireResult<()> where T: ToSql + ToSqlText + Sized, @@ -533,6 +562,10 @@ mod tests { self.encoded_value = string.unwrap(); Ok(()) } + + fn take_row(&mut self) -> Self::Item { + std::mem::take(&mut self.encoded_value) + } } let val = "~!@&$[]()@@!!"; diff --git a/arrow-pg/src/struct_encoder.rs b/arrow-pg/src/struct_encoder.rs index 9359973..6af1f4e 100644 --- a/arrow-pg/src/struct_encoder.rs +++ b/arrow-pg/src/struct_encoder.rs @@ -174,6 +174,8 @@ impl StructEncoder { } impl Encoder for StructEncoder { + type Item = BytesMut; + fn encode_field(&mut self, value: &T, pg_field: &FieldInfo) -> PgWireResult<()> where T: ToSql + ToSqlText + Sized, @@ -226,4 +228,8 @@ impl Encoder for StructEncoder { self.curr_col += 1; Ok(()) } + + fn take_row(&mut self) -> Self::Item { + std::mem::take(&mut self.row_buffer) + } }