From e2ad784af9e4433ea878cf63a8ae3a4d0d8d7f1e Mon Sep 17 00:00:00 2001
From: Garrett Thornburg <film42@gmail.com>
Date: Sat, 23 Oct 2021 22:28:38 -0600
Subject: [PATCH] Add example for wrapping PostgresConnectionManager to add
 state

With this example we show how easy it is to wrap the postgres connection
pool to prepare several queries when a connection is created. Then, once
a connection is checked out, this examples shows how easy it is to use
the custom state to pull a prepared statement.

Inspired by #110 which requests this kind of behavior in bb8 directly,
but for now it's easy to extend any connection manager with your custom
needs.
---
 postgres/examples/custom_state.rs | 133 ++++++++++++++++++++++++++++++
 1 file changed, 133 insertions(+)
 create mode 100644 postgres/examples/custom_state.rs

diff --git a/postgres/examples/custom_state.rs b/postgres/examples/custom_state.rs
new file mode 100644
index 0000000..d4218ca
--- /dev/null
+++ b/postgres/examples/custom_state.rs
@@ -0,0 +1,133 @@
+use bb8::{CustomizeConnection, Pool};
+use bb8_postgres::PostgresConnectionManager;
+use tokio_postgres::config::Config;
+use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
+use tokio_postgres::{Client, Error, Socket, Statement};
+
+use async_trait::async_trait;
+use std::collections::BTreeMap;
+use std::ops::Deref;
+use std::str::FromStr;
+
+#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
+enum QueryName {
+    BasicSelect,
+    Addition,
+}
+
+struct CustomPostgresConnection {
+    inner: Client,
+    custom_state: BTreeMap<QueryName, Statement>,
+}
+
+impl CustomPostgresConnection {
+    fn new(inner: Client) -> Self {
+        Self {
+            inner,
+            custom_state: Default::default(),
+        }
+    }
+}
+
+impl Deref for CustomPostgresConnection {
+    type Target = Client;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+struct CustomPostgresConnectionManager<Tls>
+where
+    Tls: MakeTlsConnect<Socket>,
+{
+    inner: PostgresConnectionManager<Tls>,
+}
+
+impl<Tls> CustomPostgresConnectionManager<Tls>
+where
+    Tls: MakeTlsConnect<Socket>,
+{
+    pub fn new(config: Config, tls: Tls) -> Self {
+        Self {
+            inner: PostgresConnectionManager::new(config, tls),
+        }
+    }
+}
+
+#[async_trait]
+impl<Tls> bb8::ManageConnection for CustomPostgresConnectionManager<Tls>
+where
+    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
+    <Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
+    <Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
+    <<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
+{
+    type Connection = CustomPostgresConnection;
+    type Error = Error;
+
+    async fn connect(&self) -> Result<Self::Connection, Self::Error> {
+        let conn = self.inner.connect().await?;
+        Ok(CustomPostgresConnection::new(conn))
+    }
+
+    async fn is_valid(
+        &self,
+        conn: &mut bb8::PooledConnection<'_, Self>,
+    ) -> Result<(), Self::Error> {
+        conn.simple_query("").await.map(|_| ())
+    }
+
+    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
+        self.inner.has_broken(&mut conn.inner)
+    }
+}
+
+#[derive(Debug)]
+struct Customizer;
+
+#[async_trait]
+impl<'a> CustomizeConnection<CustomPostgresConnection, Error> for Customizer {
+    async fn on_acquire(&self, conn: &mut CustomPostgresConnection) -> Result<(), Error> {
+        conn.custom_state
+            .insert(QueryName::BasicSelect, conn.prepare("SELECT 1").await?);
+
+        conn.custom_state
+            .insert(QueryName::Addition, conn.prepare("SELECT 1 + 1 + 1").await?);
+
+        Ok(())
+    }
+}
+
+// Select some static data from a Postgres DB
+//
+// The simplest way to start the db is using Docker:
+// docker run --name gotham-middleware-postgres -e POSTGRES_PASSWORD=mysecretpassword -p 5432:5432 -d postgres
+#[tokio::main]
+async fn main() {
+    let config =
+        tokio_postgres::config::Config::from_str("postgresql://postgres:docker@localhost:5432")
+            .unwrap();
+    let pg_mgr = CustomPostgresConnectionManager::new(config, tokio_postgres::NoTls);
+
+    let pool = Pool::builder()
+        .connection_customizer(Box::new(Customizer))
+        .build(pg_mgr)
+        .await
+        .expect("build error");
+
+    let connection = pool.get().await.expect("pool error");
+
+    let row = connection
+        .query_one(
+            connection
+                .custom_state
+                .get(&QueryName::Addition)
+                .expect("statement not predefined"),
+            &[],
+        )
+        .await
+        .expect("query failed");
+
+    println!("result: {}", row.get::<usize, i32>(0));
+}