From 6255acb29d5311bc3ebfdc8eca38f9f78768b1d4 Mon Sep 17 00:00:00 2001 From: Kould Date: Sat, 21 Dec 2024 01:46:07 +0800 Subject: [PATCH] chore: fix server.rs --- src/bin/server.rs | 26 ++++++++++++++------------ src/db.rs | 2 +- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/src/bin/server.rs b/src/bin/server.rs index f45edd9..849206b 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -172,19 +172,21 @@ impl SimpleQueryHandler for SessionBackend { _ => { let mut guard = self.tx.lock(); - let iter = if let Some(transaction) = guard.as_mut() { - unsafe { transaction.as_mut().run(query) }.map(Box::new) - as Result, _> - } else { - self.inner.run(query).map(Box::new) - } - .map_err(|e| PgWireError::ApiError(Box::new(e)))?; - let mut tuples = Vec::new(); - for tuple in iter { - tuples.push(tuple.map_err(|e| PgWireError::ApiError(Box::new(e)))?); - } - Ok(vec![Response::Query(encode_tuples(iter.schema(), tuples)?)]) + let response = if let Some(transaction) = guard.as_mut() { + let mut iter = unsafe { transaction.as_mut().run(query) }.map_err(|e| PgWireError::ApiError(Box::new(e)))?; + for tuple in iter.by_ref() { + tuples.push(tuple.map_err(|e| PgWireError::ApiError(Box::new(e)))?); + } + encode_tuples(iter.schema(), tuples)? + } else { + let mut iter = self.inner.run(query).map_err(|e| PgWireError::ApiError(Box::new(e)))?; + for tuple in iter.by_ref() { + tuples.push(tuple.map_err(|e| PgWireError::ApiError(Box::new(e)))?); + } + encode_tuples(iter.schema(), tuples)? + }; + Ok(vec![Response::Query(response)]) } } } diff --git a/src/db.rs b/src/db.rs index a7379ad..2362848 100644 --- a/src/db.rs +++ b/src/db.rs @@ -328,7 +328,7 @@ impl Database { } } -pub trait ResultIter: Iterator> + Sized { +pub trait ResultIter: Iterator> { fn schema(&self) -> &SchemaRef; fn done(self) -> Result<(), DatabaseError>;