From 41b6401cd45097359fb1e519b1b7ac6530381012 Mon Sep 17 00:00:00 2001
From: fakeshadow <24548779@qq.com>
Date: Wed, 2 Oct 2024 17:53:56 +0800
Subject: [PATCH] add more methods to AsyncLendingIteratorExt

---
 postgres/CHANGES.md          |  4 +-
 postgres/src/iter.rs         | 89 ++++++++++++++++++++++++++++++++++--
 postgres/src/query/stream.rs |  3 +-
 3 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/postgres/CHANGES.md b/postgres/CHANGES.md
index ee91504e..abf2714e 100644
--- a/postgres/CHANGES.md
+++ b/postgres/CHANGES.md
@@ -1,4 +1,6 @@
-# unreleased 0.2.1
+# 0.2.1
+## Fix
+- relax lifetime bound on various query types 
 
 # 0.2.0
 ## Remove
diff --git a/postgres/src/iter.rs b/postgres/src/iter.rs
index 280f5bb7..caa50673 100644
--- a/postgres/src/iter.rs
+++ b/postgres/src/iter.rs
@@ -15,6 +15,8 @@ pub trait AsyncLendingIterator {
     }
 }
 
+impl<I> AsyncLendingIteratorExt for I where I: AsyncLendingIterator {}
+
 pub trait AsyncLendingIteratorExt: AsyncLendingIterator {
     fn map_ok<F, O>(self, func: F) -> MapOk<Self, F>
     where
@@ -24,6 +26,22 @@ pub trait AsyncLendingIteratorExt: AsyncLendingIterator {
         MapOk { iter: self, func }
     }
 
+    fn map_err<F, O>(self, func: F) -> MapErr<Self, F>
+    where
+        F: Fn(Self::Err) -> O,
+        Self: Sized,
+    {
+        MapErr { iter: self, func }
+    }
+
+    fn try_map<F, T, E>(self, func: F) -> Map<Self, F>
+    where
+        F: Fn(Result<Self::Ok<'_>, Self::Err>) -> Result<T, E>,
+        Self: Sized,
+    {
+        Map { iter: self, func }
+    }
+
     #[inline]
     fn try_collect<T>(self) -> impl Future<Output = Result<T, Self::Err>> + Send
     where
@@ -75,8 +93,73 @@ where
     }
 }
 
-impl<I> AsyncLendingIteratorExt for I where I: AsyncLendingIterator {}
+pub struct MapErr<I, F> {
+    iter: I,
+    func: F,
+}
+
+impl<I, F, O> AsyncLendingIterator for MapErr<I, F>
+where
+    I: AsyncLendingIterator + Send,
+    F: Fn(I::Err) -> O + Send,
+    O: Send,
+{
+    type Ok<'i>
+        = I::Ok<'i>
+    where
+        Self: 'i;
+    type Err = O;
+
+    async fn try_next(&mut self) -> Result<Option<Self::Ok<'_>>, Self::Err> {
+        self.iter.try_next().await.map_err(&self.func)
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+pub struct Map<I, F> {
+    iter: I,
+    func: F,
+}
+
+impl<I, F, T, E> AsyncLendingIterator for Map<I, F>
+where
+    I: AsyncLendingIterator + Send,
+    F: Fn(Result<I::Ok<'_>, I::Err>) -> Result<T, E> + Send,
+    T: Send,
+    E: Send,
+{
+    type Ok<'i>
+        = T
+    where
+        Self: 'i;
+    type Err = E;
+
+    async fn try_next(&mut self) -> Result<Option<Self::Ok<'_>>, Self::Err> {
+        match self.iter.try_next().await {
+            Ok(Some(t)) => (self.func)(Ok(t)).map(Some),
+            Ok(None) => Ok(None),
+            Err(e) => (self.func)(Err(e)).map(Some),
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+async fn _map_ok_err_try_collect(stream: crate::RowStreamOwned) -> Result<Vec<String>, crate::error::Error> {
+    stream
+        .map_ok(|row| row.get(0))
+        .map_err(|e| dbg!(e))
+        .try_collect::<Vec<_>>()
+        .await
+}
 
-async fn _try_collect_test(stream: crate::RowStreamOwned) -> Result<Vec<String>, crate::error::Error> {
-    stream.map_ok(|row| row.get(0)).try_collect::<Vec<_>>().await
+async fn _try_map_try_collect(stream: crate::RowStreamOwned) -> Result<Vec<String>, crate::error::Error> {
+    stream.try_map(|row| row?.try_get(0)).try_collect::<Vec<_>>().await
 }
diff --git a/postgres/src/query/stream.rs b/postgres/src/query/stream.rs
index b9678d4d..bc78b529 100644
--- a/postgres/src/query/stream.rs
+++ b/postgres/src/query/stream.rs
@@ -93,7 +93,8 @@ async fn try_next<'r>(
 /// // then collecting all user name to a collection
 /// let mut strings = Vec::new();
 /// while let Some(row) = stream.try_next().await? {
-///     strings.push(row.get::<String>("name"));
+///     let name = row.get::<String>("name");
+///     strings.push(name);
 /// }
 ///
 /// // the same operation with owned row stream can be simplified a bit: