Add Postgres COPY FROM/TO support #3951

weiznich · 2024-03-01T18:09:28Z

This commit adds support for PostgreSQL COPY commands.

For COPY FROM we expose a variant that allows users to configure the stream manually and write directly to the stream. We also support a variant that takes essentially a Vec<Insertable> (or equivalent batch insert containers) and uses the binary format to perform a streamed batch insert.

For COPY TO we expose again a variant that allows the user to configure the stream manually and read directly from the stream. We also support loading the results directly via the binary protocol into a iterator of Queryable structs.

This is not finished yet, I open the PR nevertheless to gather some feedback on the API.

Things that need to be addressed before merging:

Documentation
Write some more tests (especially compile tests)
Go over the naming of everything again
Have a look if we could use Selectable for the loading part to avoid running into nasty mismatch errors there

JosueMolinaMorales

This is some cool code. Just some questions (sorry if there are any dumb questions)

diesel/src/pg/connection/copy.rs

diesel/src/pg/query_builder/copy/copy_in.rs

JosueMolinaMorales · 2024-03-04T15:37:58Z

diesel/src/pg/query_builder/copy/mod.rs

+}
+
+#[derive(Default, Debug)]
+struct CommonOptions {


Is there a reason for not including HEADER, FORCE_QUOTE, ENCODING, etc?

ENCODING is currently always assumed to be the same as for the connection (so UTF-8). It might be worth to expose that at some point.
HEADER is currently only exposed for copy_in, but after reading the documentation again it seems to be useful to expose it also for copy_out without allowing MATCH
FORCE_QUOTE support requires figuring out passing down a list of columns to that function. That's not impossible, it's just something that I did not want to do for the first version.

diesel_tests/tests/copy.rs

diesel/src/pg/connection/copy.rs

pppp24 · 2024-03-07T06:51:11Z

diesel/src/pg/connection/raw.rs

+            .as_ref()
+            .map(|l| l.as_ptr())
+            .unwrap_or(std::ptr::null());
+        let ret = unsafe { pq_sys::PQputCopyEnd(self.internal_connection.as_ptr(), error) };


@weiznich what would happen in the event this line does not get called as a result of the '?' on line 164 ?

Good catch, addressed in dbc92f2

looks good 👍

pppp24 · 2024-03-07T07:31:25Z

diesel/src/pg/query_builder/copy/copy_out.rs

+
+    fn get<'b, I>(&'b self, idx: I) -> Option<Self::Field<'b>>
+    where
+        'a: 'b,


@weiznich does this mean that 'a must be a subtype of 'b or is it the other way around ?

https://doc.rust-lang.org/reference/trait-bounds.html#lifetime-bounds

So it's usually 'a needs to outlive 'b.

ah yep makes sense, thanks for the link 👍

peasee

A couple questions from me to help get myself up to speed 😄

peasee · 2024-04-04T07:38:53Z

diesel/src/pg/connection/copy.rs

+
+impl<'conn> Write for CopyFromSink<'conn> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.conn.put_copy_data(buf).unwrap();


question: Should we .unwrap() here if we're already returning a Result<_>?

Good catch, that shouldn't use unwrap. We don't want to use unwrap anywhere other than in tests.

peasee · 2024-04-04T07:46:12Z

diesel/src/pg/connection/copy.rs

+    #[allow(unsafe_code)] // ffi code
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe { pq_sys::PQfreemem(self.ptr as *mut ffi::c_void) };


question: This is my first time seeing CFFI code before. Let me see if I understand this right:

We have a function in pq_sys, PQfreemem which is a binding to the libpq library and drops (something) from memory?

The PQfreemem is equivalent to doing a drop() in Rust, and because we're using the C library we need to implement Drop to call PQfreemem so that we can drop(impl CopyToBuffer).

I have a gap in my knowledge here though, around the purpose of self.ptr as *mut ffi::c_void? Does this equate to C like void* self.ptr, so the library knows it's receiving a pointer but doesn't know what data type it is? The library then knows the data at that pointer is what it has to clear, but doesn't care what that data is?

PQfreemem is the postgresql specific variant of the C dealloc/free function. This drop impl essentially ensures that we always free the buffer this pointer points to. The cast is required as the c function takes a void pointer as argument.
Not calling that function would potentially leak memory, which would be a bug.

peasee · 2024-04-04T07:49:24Z

diesel/src/pg/connection/copy.rs

+                if !self.ptr.is_null() {
+                    pq_sys::PQfreemem(self.ptr as *mut ffi::c_void);
+                    self.ptr = std::ptr::null_mut();
+                }
+                let len =
+                    pq_sys::PQgetCopyData(self.conn.internal_connection.as_ptr(), &mut self.ptr, 0);


question: Okay, I think it's starting to click more for me. We clear the pointer from any previous library call, then do another FFI call supplying our (now empty) pointer.

If I'm following correctly, does this mean the library call will mutate our underlying pointer to make it point to the results of this library call, and this function returns the length of data that exists at the pointer location?

Yes this is correct.

peasee · 2024-04-04T07:59:00Z

diesel/src/pg/connection/mod.rs

+        Ok(res)
+    }
+
+    pub(crate) fn copy_to<T>(&mut self, command: CopyToCommand<T>) -> QueryResult<CopyToBuffer<'_>>


question: This might be a styling thing I need to get familiar with, but why do we do command: CopyToCommand<T> vs the copy_from function above which does target: S [...] where S: CopyInExpression<T>?

They are functionally different: command: CopyToCommand<T> is a struct with a generic type argument. target: S [...] where S: CopyInExpression <T> is a complete generic type that implements a certain trait (CopyInExpression in that case). The struct essentially holds a set of common values that is the same for all possible values of the generic type, while the trait variant allows fundamentally different types there. I've just used what seemed to fit the needs in each location.

yep totally right, tired brain was thinking there should've been a type T/S somewhere in the trait which isn't always true 🤦

peasee

A couple more comments from me here. I also noticed there's quite a few .expect() in the copy_to.rs changes, among other places. I'm not sure if that's a pattern we accept here, but for those that are in a Result<_> it might be suitable to change to errors instead of panicking.

Also, please let me know if my comments are too nitpick-y. I'm testing the waters to see how deep I go with my reviews 😄

peasee · 2024-04-05T08:34:38Z

diesel/src/pg/connection/mod.rs

+                let mut copy_in = CopyFromSink::new(&mut conn.raw_connection);
+                let r = source.target.callback(&mut copy_in);
+                copy_in.finish(r.as_ref().err().map(|e| e.to_string()))?;
+                let next_res = conn.raw_connection.get_next_result()?.expect("exists");


question: Should we give these .expect()'s a better error message? Also, curious why we wouldn't use something like .ok_or_else(|| SomeError())? here, as this function already returns a Result<_>?

peasee · 2024-04-05T08:36:12Z

diesel/src/pg/connection/mod.rs

+        ) -> Result<R, E>,
+    ) -> Result<R, E>


question: Why the type change from QueryResult to a Result?

The copy_from method allows to configure a callback that is then used to pull data into the database. That callback can return custom error, that we want to propagate back to the user. This callback is called somewhere in the callback passed to with_prepared_query, therefore we need to change the return type there to allow that.

peasee · 2024-04-05T08:41:22Z

diesel/src/pg/connection/raw.rs

+                CString::new("Error message contains a \\0 byte")
+                    .expect("Does not contain a null byte")


question: Should we make this .expect() more true to the error, in that it failed to initialize our new CString error? I would've expected to see the other .expect() in .map(CString::new) if it were to be somewhere.

I usually write expect() messages as factual statement. In this case we know at compile time (by looking at the string above) that is does not contain a null byte. This cannot ever be an error.

peasee · 2024-04-05T08:45:32Z

diesel/src/pg/query_builder/copy/copy_from.rs

+        self.target.walk_target(pass.reborrow())?;
+        pass.push_sql(" FROM STDIN");
+        self.target.options().walk_ast(pass.reborrow())?;
+        // todo: where?


question: Should this TODO be addressed before merge?

No, that's for now not implemented. I do not plan to support that for now, so I've removed the todo.

peasee · 2024-04-05T08:47:51Z

diesel/src/pg/query_builder/copy/copy_from.rs

+    type Error = E;
+
+    fn callback(&mut self, copy: &mut CopyFromSink<'_>) -> Result<(), Self::Error> {
+        (self.copy_callback)(copy)


question: Why do we wrap self.copy_callback in brackets if we already know it's a function? Couldn't we just self.copy_callback(copy)?

That's required by rust to call a function stored in a field. See this playground for a minimal example: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=f5a4af872196441cb75487ee05af5aca

peasee · 2024-04-05T08:52:15Z

diesel/src/pg/query_builder/copy/copy_from.rs

+        let values = self
+            .0
+            .take()
+            .expect("We only call this callback once")


question: .expect() in a Result<_>. Ditto for above, could we map this to an error?

That's one of the cases where I personally prefer using expect to get a panic if that's ever hit as that would be clearly a bug in diesel. My reasoning there is to panic for cases that are considered to be a problem with our (== diesel's) code to get reports on that. On the other hand anything that could be triggered and handled by a user should rather be a normal error.

okay that's fair enough, I'll keep that in mind for future reviews! 😁

peasee · 2024-04-05T08:59:47Z

diesel/src/pg/query_builder/copy/mod.rs

+        &'b self,
+        mut pass: crate::query_builder::AstPass<'_, 'b, Pg>,
+        comma: &mut &'static str,
+    ) -> crate::QueryResult<()> {


question: Why is this QueryResult<_> if we don't ? anything here?

peasee

Thanks for your hard work! 🚢

This commit adds support for PostgreSQL `COPY` commands. For `COPY FROM` we expose a variant that allows users to configure the stream manually and write directly to the stream. We also support a variant that takes essentially a `Vec<Insertable>` (or equivalent batch insert containers) and uses the binary format to perform a streamed batch insert. For `COPY TO` we expose again a variant that allows the user to configure the stream manually and read directly from the stream. We also support loading the results directly via the binary protocol into a iterator of `Queryable` structs.

* Renaming everything to follow `COPY FROM`/`COPY TO` nameing * Require `Selectable` for loading (to prevent errors) * Add instrumentation support + tests

rust features

…res versions

weiznich requested a review from a team March 1, 2024 18:09

JosueMolinaMorales reviewed Mar 4, 2024

View reviewed changes

ppsmol24 mentioned this pull request Mar 7, 2024

More code reviewers wanted! #1186

Open

pppp24 reviewed Mar 7, 2024

View reviewed changes

weiznich mentioned this pull request Mar 8, 2024

Fix most type-defs for #[auto_type] by cleaning up all the type defs #3783

Merged

weiznich force-pushed the feature/copy_in_out branch from 6b8f727 to dbc92f2 Compare March 22, 2024 14:03

peasee reviewed Apr 4, 2024

View reviewed changes

weiznich force-pushed the feature/copy_in_out branch from dbc92f2 to a740c1a Compare April 5, 2024 06:30

weiznich marked this pull request as ready for review April 5, 2024 06:30

weiznich force-pushed the feature/copy_in_out branch from 32b47f1 to fd49b3f Compare April 5, 2024 07:49

peasee reviewed Apr 5, 2024

View reviewed changes

peasee approved these changes Apr 9, 2024

View reviewed changes

weiznich enabled auto-merge April 10, 2024 05:01

weiznich force-pushed the feature/copy_in_out branch from 57ba5c0 to 3f32a02 Compare April 10, 2024 05:32

weiznich added 15 commits April 10, 2024 09:41

Address code review comments and other improvements

8e36638

* Renaming everything to follow `COPY FROM`/`COPY TO` nameing * Require `Selectable` for loading (to prevent errors) * Add instrumentation support + tests

Documentation & more renaming

f1dd012

Update nightly version used for compile tests to allow compiling new

aab0300

rust features

Add compile tests for copy_from/copy_to

49ceb38

Bump minimal supported rust version

e3ba94c

Add a changelog entry

e96a1dd

Fix minor issues

5ba0d6a

Clippy fixes

0df7ab4

Restore default features

42fa0be

CI fixes

2b3b607

Feature gate fixes for postgres_backend

2ff4c71

Address more review comments

a0098c1

Fix another compile test

77782ea

Use a different option as that seems not to be supported on old postg…

c3b6dd1

…res versions

weiznich force-pushed the feature/copy_in_out branch from 3f32a02 to c3b6dd1 Compare April 10, 2024 09:41

weiznich added this pull request to the merge queue Apr 10, 2024

Merged via the queue into diesel-rs:master with commit 8afe715 Apr 10, 2024
48 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Postgres COPY FROM/TO support #3951

Add Postgres COPY FROM/TO support #3951

weiznich commented Mar 1, 2024 •

edited

Loading

JosueMolinaMorales left a comment

JosueMolinaMorales Mar 4, 2024

weiznich Mar 5, 2024

pppp24 Mar 7, 2024

weiznich Mar 22, 2024

pppp24 Mar 26, 2024

pppp24 Mar 7, 2024

weiznich Mar 22, 2024 •

edited

Loading

pppp24 Mar 26, 2024

peasee left a comment

peasee Apr 4, 2024

weiznich Apr 4, 2024

peasee Apr 4, 2024

weiznich Apr 4, 2024

peasee Apr 4, 2024

weiznich Apr 4, 2024

peasee Apr 4, 2024

weiznich Apr 4, 2024

peasee Apr 4, 2024 •

edited

Loading

peasee left a comment

peasee Apr 5, 2024

peasee Apr 5, 2024

weiznich Apr 5, 2024

peasee Apr 5, 2024

weiznich Apr 5, 2024

peasee Apr 5, 2024

weiznich Apr 5, 2024

peasee Apr 5, 2024

weiznich Apr 5, 2024

peasee Apr 5, 2024

weiznich Apr 5, 2024

peasee Apr 5, 2024

peasee Apr 5, 2024

peasee left a comment

		CString::new("Error message contains a \\0 byte")
		.expect("Does not contain a null byte")

Add Postgres COPY FROM/TO support #3951

Add Postgres COPY FROM/TO support #3951

Conversation

weiznich commented Mar 1, 2024 • edited Loading

JosueMolinaMorales left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

weiznich Mar 22, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peasee left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peasee Apr 4, 2024 • edited Loading

Choose a reason for hiding this comment

peasee left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peasee left a comment

Choose a reason for hiding this comment

weiznich commented Mar 1, 2024 •

edited

Loading

weiznich Mar 22, 2024 •

edited

Loading

peasee Apr 4, 2024 •

edited

Loading