From 6c6a266c1449a5d754d1b3771df58ad5dbbf7b02 Mon Sep 17 00:00:00 2001
From: Orion Kindel <cakekindel@gmail.com>
Date: Thu, 6 Jul 2023 22:29:28 -0500
Subject: [PATCH] fix: dx changes

---
 postgres-types/src/lib.rs          |   4 +-
 postgres/src/generic_client.rs     | 278 +++++++++++++++++++----------
 postgres/src/lib.rs                |   4 +-
 tokio-postgres/src/lib.rs          |   2 +-
 tokio-postgres/src/statement.rs    |   2 +-
 tokio-postgres/src/to_statement.rs |  58 +++---
 6 files changed, 212 insertions(+), 136 deletions(-)

diff --git a/postgres-types/src/lib.rs b/postgres-types/src/lib.rs
index edd72397..b095c1cd 100644
--- a/postgres-types/src/lib.rs
+++ b/postgres-types/src/lib.rs
@@ -848,9 +848,7 @@ pub trait ToSql: fmt::Debug {
     /// The return value indicates if this value should be represented as
     /// `NULL`. If this is the case, implementations **must not** write
     /// anything to `out`.
-    fn to_sql(&self, ty: &Type, out: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>>
-    where
-        Self: Sized;
+    fn to_sql(&self, ty: &Type, out: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>>;
 
     /// Determines if a value of this type can be converted to the specified
     /// Postgres `Type`.
diff --git a/postgres/src/generic_client.rs b/postgres/src/generic_client.rs
index 12f07465..b11cb795 100644
--- a/postgres/src/generic_client.rs
+++ b/postgres/src/generic_client.rs
@@ -1,29 +1,29 @@
+#![allow(missing_docs)]
+
+use tokio_postgres::Column;
+use tokio_postgres::row::RowIndex;
+use tokio_postgres::types::FromSql;
+
 use crate::types::{BorrowToSql, ToSql, Type};
 use crate::{
     Client, CopyInWriter, CopyOutReader, Error, Row, RowIter, SimpleQueryMessage, Statement,
-    ToStatement, Transaction,
+    ToStatement, Transaction
 };
 
-mod private {
-    pub trait Sealed {}
-}
-
-/// A trait allowing abstraction over connections and transactions.
-///
-/// This trait is "sealed", and cannot be implemented outside of this crate.
-pub trait GenericClient: private::Sealed {
+macro_rules! common {
+  ($transaction:ty, $err:ty) => {
     /// Like `Client::execute`.
-    fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, Error>
+    fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, $err>
     where
         T: ?Sized + ToStatement;
 
     /// Like `Client::query`.
-    fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Row>, Error>
+    fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Self::Row>, $err>
     where
         T: ?Sized + ToStatement;
 
     /// Like `Client::query_one`.
-    fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Row, Error>
+    fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Self::Row, $err>
     where
         T: ?Sized + ToStatement;
 
@@ -32,12 +32,12 @@ pub trait GenericClient: private::Sealed {
         &mut self,
         query: &T,
         params: &[&(dyn ToSql + Sync)],
-    ) -> Result<Option<Row>, Error>
+    ) -> Result<Option<Self::Row>, $err>
     where
         T: ?Sized + ToStatement;
 
     /// Like `Client::query_raw`.
-    fn query_raw<T, P, I>(&mut self, query: &T, params: I) -> Result<RowIter<'_>, Error>
+    fn query_raw<T, P, I>(&mut self, query: &T, params: I) -> Result<RowIter<'_>, $err>
     where
         T: ?Sized + ToStatement,
         P: BorrowToSql,
@@ -45,34 +45,74 @@ pub trait GenericClient: private::Sealed {
         I::IntoIter: ExactSizeIterator;
 
     /// Like `Client::prepare`.
-    fn prepare(&mut self, query: &str) -> Result<Statement, Error>;
+    fn prepare(&mut self, query: &str) -> Result<Statement, $err>;
 
     /// Like `Client::prepare_typed`.
-    fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error>;
+    fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, $err>;
 
     /// Like `Client::copy_in`.
-    fn copy_in<T>(&mut self, query: &T) -> Result<CopyInWriter<'_>, Error>
+    fn copy_in<T>(&mut self, query: &T) -> Result<CopyInWriter<'_>, $err>
     where
         T: ?Sized + ToStatement;
 
     /// Like `Client::copy_out`.
-    fn copy_out<T>(&mut self, query: &T) -> Result<CopyOutReader<'_>, Error>
+    fn copy_out<T>(&mut self, query: &T) -> Result<CopyOutReader<'_>, $err>
     where
         T: ?Sized + ToStatement;
 
     /// Like `Client::simple_query`.
-    fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error>;
+    fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, $err>;
 
     /// Like `Client::batch_execute`.
-    fn batch_execute(&mut self, query: &str) -> Result<(), Error>;
+    fn batch_execute(&mut self, query: &str) -> Result<(), $err>;
 
     /// Like `Client::transaction`.
-    fn transaction(&mut self) -> Result<Transaction<'_>, Error>;
+    fn transaction(&mut self) -> Result<$transaction, $err>;
+  };
 }
 
-impl private::Sealed for Client {}
+pub trait GenericRow {
+    type Error: core::fmt::Debug;
+
+    fn columns(&self) -> &[Column];
+    fn len(&self) -> usize;
+
+    fn get<'a, I, T>(&'a self, idx: I) -> T
+    where
+        I: RowIndex + core::fmt::Display,
+        T: FromSql<'a>;
+
+    fn try_get<'a, I, T>(&'a self, idx: I) -> Result<T, Self::Error>
+    where
+        I: RowIndex + core::fmt::Display,
+        T: FromSql<'a>;
+}
+
+pub trait GenericTransaction<'a> where Self: Sized {
+  type Error: core::fmt::Debug;
+  type NestedTransaction<'b> where Self: 'b;
+  type Row: GenericRow<Error = Self::Error>;
+
+  fn commit(self) -> Result<(), Self::Error>;
+  fn rollback(self) -> Result<(), Self::Error>;
+
+  common!(Self::NestedTransaction<'_>, Self::Error);
+}
+
+/// A trait allowing abstraction over connections and transactions.
+pub trait GenericClient {
+    type Error: core::fmt::Debug;
+    type Transaction<'a>: GenericTransaction<'a> where Self: 'a;
+    type Row: GenericRow<Error = Self::Error>;
+
+    common!(Self::Transaction<'_>, Self::Error);
+}
 
 impl GenericClient for Client {
+    type Error = Error;
+    type Transaction<'a> = Transaction<'a>;
+    type Row = Row;
+
     fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, Error>
     where
         T: ?Sized + ToStatement,
@@ -80,14 +120,14 @@ impl GenericClient for Client {
         self.execute(query, params)
     }
 
-    fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Row>, Error>
+    fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Self::Row>, Error>
     where
         T: ?Sized + ToStatement,
     {
         self.query(query, params)
     }
 
-    fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Row, Error>
+    fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Self::Row, Error>
     where
         T: ?Sized + ToStatement,
     {
@@ -98,7 +138,7 @@ impl GenericClient for Client {
         &mut self,
         query: &T,
         params: &[&(dyn ToSql + Sync)],
-    ) -> Result<Option<Row>, Error>
+    ) -> Result<Option<Self::Row>, Error>
     where
         T: ?Sized + ToStatement,
     {
@@ -145,87 +185,131 @@ impl GenericClient for Client {
         self.batch_execute(query)
     }
 
-    fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
+    fn transaction(&mut self) -> Result<Self::Transaction<'_>, Error> {
         self.transaction()
     }
 }
 
-impl private::Sealed for Transaction<'_> {}
+macro_rules! transaction_common {
+  ($transaction:ty) => {
+    fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.execute(query, params)
+    }
+
+    fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Self::Row>, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.query(query, params)
+    }
+
+    fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Self::Row, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.query_one(query, params)
+    }
+
+    fn query_opt<T>(
+        &mut self,
+        query: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Option<Self::Row>, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.query_opt(query, params)
+    }
+
+    fn query_raw<T, P, I>(&mut self, query: &T, params: I) -> Result<RowIter<'_>, Error>
+    where
+        T: ?Sized + ToStatement,
+        P: BorrowToSql,
+        I: IntoIterator<Item = P>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        self.query_raw(query, params)
+    }
+
+    fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
+        self.prepare(query)
+    }
+
+    fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error> {
+        self.prepare_typed(query, types)
+    }
+
+    fn copy_in<T>(&mut self, query: &T) -> Result<CopyInWriter<'_>, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.copy_in(query)
+    }
+
+    fn copy_out<T>(&mut self, query: &T) -> Result<CopyOutReader<'_>, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.copy_out(query)
+    }
+
+    fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
+        self.simple_query(query)
+    }
+
+    fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
+        self.batch_execute(query)
+    }
+
+    fn transaction(&mut self) -> Result<$transaction, Error> {
+        self.transaction()
+    }
+  };
+}
+
+impl<'a> GenericTransaction<'a> for Transaction<'a> {
+  type Error = Error;
+  type NestedTransaction<'b> = Transaction<'b> where Self: 'b;
+  type Row = Row;
+
+  fn commit(self) -> Result<(), Error> {Transaction::commit(self)}
+  fn rollback(self) -> Result<(), Error> {Transaction::rollback(self)}
+
+  transaction_common!(Self::NestedTransaction<'_>);
+}
 
 impl GenericClient for Transaction<'_> {
-    fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, Error>
+    type Error = Error;
+    type Transaction<'a> = Transaction<'a> where Self: 'a;
+    type Row = Row;
+    transaction_common!(Self::Transaction<'_>);
+}
+
+impl GenericRow for Row {
+    type Error = Error;
+
+    fn columns(&self) -> &[Column] {
+      self.columns()
+    }
+
+    fn len(&self) -> usize {
+      self.len()
+    }
+
+    fn get<'a, I, T>(&'a self, idx: I) -> T
     where
-        T: ?Sized + ToStatement,
-    {
-        self.execute(query, params)
+        I: RowIndex + core::fmt::Display,
+        T: FromSql<'a> {
+      self.get(idx)
     }
 
-    fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Row>, Error>
+    fn try_get<'a, I, T>(&'a self, idx: I) -> Result<T, Error>
     where
-        T: ?Sized + ToStatement,
-    {
-        self.query(query, params)
-    }
-
-    fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Row, Error>
-    where
-        T: ?Sized + ToStatement,
-    {
-        self.query_one(query, params)
-    }
-
-    fn query_opt<T>(
-        &mut self,
-        query: &T,
-        params: &[&(dyn ToSql + Sync)],
-    ) -> Result<Option<Row>, Error>
-    where
-        T: ?Sized + ToStatement,
-    {
-        self.query_opt(query, params)
-    }
-
-    fn query_raw<T, P, I>(&mut self, query: &T, params: I) -> Result<RowIter<'_>, Error>
-    where
-        T: ?Sized + ToStatement,
-        P: BorrowToSql,
-        I: IntoIterator<Item = P>,
-        I::IntoIter: ExactSizeIterator,
-    {
-        self.query_raw(query, params)
-    }
-
-    fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
-        self.prepare(query)
-    }
-
-    fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error> {
-        self.prepare_typed(query, types)
-    }
-
-    fn copy_in<T>(&mut self, query: &T) -> Result<CopyInWriter<'_>, Error>
-    where
-        T: ?Sized + ToStatement,
-    {
-        self.copy_in(query)
-    }
-
-    fn copy_out<T>(&mut self, query: &T) -> Result<CopyOutReader<'_>, Error>
-    where
-        T: ?Sized + ToStatement,
-    {
-        self.copy_out(query)
-    }
-
-    fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
-        self.simple_query(query)
-    }
-
-    fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
-        self.batch_execute(query)
-    }
-
-    fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
-        self.transaction()
+        I: RowIndex + core::fmt::Display,
+        T: FromSql<'a> {
+      self.try_get(idx)
     }
 }
diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs
index fbe85cbd..5d407eba 100644
--- a/postgres/src/lib.rs
+++ b/postgres/src/lib.rs
@@ -69,7 +69,7 @@
 pub use fallible_iterator;
 pub use tokio_postgres::{
     error, row, tls, types, Column, IsolationLevel, Notification, Portal, SimpleQueryMessage,
-    Socket, Statement, ToStatement,
+    Socket, Statement, ToStatement, StatementOrString,
 };
 
 pub use crate::cancel_token::CancelToken;
@@ -79,7 +79,7 @@ pub use crate::copy_in_writer::CopyInWriter;
 pub use crate::copy_out_reader::CopyOutReader;
 #[doc(no_inline)]
 pub use crate::error::Error;
-pub use crate::generic_client::GenericClient;
+pub use crate::generic_client::{GenericClient, GenericTransaction, GenericRow};
 #[doc(inline)]
 pub use crate::notifications::Notifications;
 #[doc(no_inline)]
diff --git a/tokio-postgres/src/lib.rs b/tokio-postgres/src/lib.rs
index 2bb41018..7ccadfa1 100644
--- a/tokio-postgres/src/lib.rs
+++ b/tokio-postgres/src/lib.rs
@@ -138,7 +138,7 @@ pub use crate::statement::{Column, Statement};
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 pub use crate::tls::NoTls;
-pub use crate::to_statement::ToStatement;
+pub use crate::to_statement::{ToStatement, StatementOrString};
 pub use crate::transaction::Transaction;
 pub use crate::transaction_builder::{IsolationLevel, TransactionBuilder};
 use crate::types::ToSql;
diff --git a/tokio-postgres/src/statement.rs b/tokio-postgres/src/statement.rs
index 97561a8e..eeabd719 100644
--- a/tokio-postgres/src/statement.rs
+++ b/tokio-postgres/src/statement.rs
@@ -71,7 +71,7 @@ pub struct Column {
 }
 
 impl Column {
-    pub(crate) fn new(name: String, type_: Type) -> Column {
+    pub fn new(name: String, type_: Type) -> Column {
         Column { name, type_ }
     }
 
diff --git a/tokio-postgres/src/to_statement.rs b/tokio-postgres/src/to_statement.rs
index 427f77dd..11e2adc3 100644
--- a/tokio-postgres/src/to_statement.rs
+++ b/tokio-postgres/src/to_statement.rs
@@ -1,23 +1,26 @@
-use crate::to_statement::private::{Sealed, ToStatementType};
+#![allow(missing_docs)]
+
 use crate::Statement;
+use crate::{Client, Error};
 
-mod private {
-    use crate::{Client, Error, Statement};
+pub enum StatementOrString<'a> {
+    Statement(&'a Statement),
+    Query(&'a str),
+}
 
-    pub trait Sealed {}
-
-    pub enum ToStatementType<'a> {
-        Statement(&'a Statement),
-        Query(&'a str),
+impl<'a> StatementOrString<'a> {
+    pub async fn into_statement(self, client: &Client) -> Result<Statement, Error> {
+        match self {
+            StatementOrString::Statement(s) => Ok(s.clone()),
+            StatementOrString::Query(s) => client.prepare(s).await,
+        }
     }
 
-    impl<'a> ToStatementType<'a> {
-        pub async fn into_statement(self, client: &Client) -> Result<Statement, Error> {
-            match self {
-                ToStatementType::Statement(s) => Ok(s.clone()),
-                ToStatementType::Query(s) => client.prepare(s).await,
-            }
-        }
+    pub fn unwrap_str(self) -> &'a str {
+      match self {
+        StatementOrString::Query(s) => s,
+        _ => panic!(),
+      }
     }
 }
 
@@ -25,33 +28,24 @@ mod private {
 ///
 /// Many methods are generic over this bound, so that they support both a raw query string as well as a statement which
 /// was prepared previously.
-///
-/// This trait is "sealed" and cannot be implemented by anything outside this crate.
-pub trait ToStatement: Sealed {
-    #[doc(hidden)]
-    fn __convert(&self) -> ToStatementType<'_>;
+pub trait ToStatement {
+  fn __convert(&self) -> StatementOrString<'_>;
 }
 
 impl ToStatement for Statement {
-    fn __convert(&self) -> ToStatementType<'_> {
-        ToStatementType::Statement(self)
+    fn __convert(&self) -> StatementOrString<'_> {
+        StatementOrString::Statement(self)
     }
 }
 
-impl Sealed for Statement {}
-
 impl ToStatement for str {
-    fn __convert(&self) -> ToStatementType<'_> {
-        ToStatementType::Query(self)
+    fn __convert(&self) -> StatementOrString<'_> {
+        StatementOrString::Query(self)
     }
 }
 
-impl Sealed for str {}
-
 impl ToStatement for String {
-    fn __convert(&self) -> ToStatementType<'_> {
-        ToStatementType::Query(self)
+    fn __convert(&self) -> StatementOrString<'_> {
+        StatementOrString::Query(self)
     }
 }
-
-impl Sealed for String {}