Add support for Postgres store with tokio-postgres along with bb8-postgres. Implementation attached.

[premkumart]

I have implemented PostgresStore with tokio-postgres and bb8-postgres libraries. I hope you can consider this for future implementation. Given below the code:

use axum::async_trait;
use bb8::{Pool, RunError};
use bb8_postgres::PostgresConnectionManager;
use time::OffsetDateTime;
use tokio_postgres::NoTls;
use tower_sessions::{session, session::Id, session::Session, ExpiredDeletion, SessionStore};

// Alias to represent a database pool connections
pub type DBPool = Pool<PostgresConnectionManager<NoTls>>;

// Alias to error
pub type PostgresConnectionError = RunError<tokio_postgres::error::Error>;

#[derive(thiserror::Error, Debug)]
pub enum PostgresStoreError {
    /// A variant to map session errors.
    #[error(transparent)]
    Session(#[from] session::Error),

    /// A Db Connection Error
    #[error("Db connection error: {0}")]
    ConnectionError(#[from] PostgresConnectionError),

    /// A Db Query Error
    #[error("Db error: {0}")]
    PostgresError(#[from] tokio_postgres::Error),

    /// A variant to map `serde_json` errors.
    #[error("JSON serialization/deserialization error: {0}")]
    SerdeJson(#[from] serde_json::Error),

    /// A variant to map `rmp_serde` encode errors.
    #[error("Rust MsgPack encode error: {0}")]
    RmpSerdeEncode(#[from] rmp_serde::encode::Error),

    /// A variant to map `rmp_serde` decode errors.
    #[error("Rust MsgPack decode error: {0}")]
    RmpSerdeDecode(#[from] rmp_serde::decode::Error),
}

#[derive(Clone, Debug)]
pub struct PostgresStore {
    pool: DBPool,
    schema_name: String,
    table_name: String,
}

impl PostgresStore {
    pub async fn new(pool: &DBPool) -> Self {
        Self {
            pool: pool.clone(),
            schema_name: "tower_sessions".to_string(),
            table_name: "session".to_string(),
        }
    }

    pub async fn migrate(&self) -> Result<(), PostgresStoreError> {
        let create_schema_query = format!(
            r#"create schema if not exists "{schema_name}""#,
            schema_name = self.schema_name,
        );

        let mut conn = self
            .pool
            .get()
            .await
            .map_err(PostgresStoreError::ConnectionError)?;
        let tx = conn.transaction().await?;

        if let Err(err) = tx.execute(&create_schema_query, &[]).await {
            if !err
                .to_string()
                .contains("duplicate key value violates unique constraint")
            {
                return Err(PostgresStoreError::PostgresError(err));
            }

            return Ok(());
        }

        let create_table_query = format!(
            r#"
            create table if not exists "{schema_name}"."{table_name}"
            (
                id text primary key not null,
                data bytea not null,
                expiry_date timestamptz not null
            )
            "#,
            schema_name = self.schema_name,
            table_name = self.table_name
        );

        let _ = tx.execute(&create_table_query, &[]).await;

        let _ = tx.commit().await;

        Ok(())
    }
}

#[async_trait]
impl ExpiredDeletion for PostgresStore {
    async fn delete_expired(&self) -> Result<(), Self::Error> {
        let query = format!(
            r#"
            delete from "{schema_name}"."{table_name}"
            where expiry_date < (now() at time zone 'utc')
            "#,
            schema_name = self.schema_name,
            table_name = self.table_name
        );
        let conn = self
            .pool
            .get()
            .await
            .map_err(PostgresStoreError::ConnectionError)?;
        let _ = conn.execute(&query, &[]).await;
        Ok(())
    }
}

#[async_trait]
impl SessionStore for PostgresStore {
    type Error = PostgresStoreError;

    async fn save(&self, session: &Session) -> Result<(), Self::Error> {
        let query = format!(
            r#"
            insert into "{schema_name}"."{table_name}" (id, data, expiry_date)
            values ($1, $2, $3)
            on conflict (id) do update
            set
              data = excluded.data,
              expiry_date = excluded.expiry_date
            "#,
            schema_name = self.schema_name,
            table_name = self.table_name
        );
        let data = rmp_serde::to_vec(&session)?;
        let conn = self
            .pool
            .get()
            .await
            .map_err(PostgresStoreError::ConnectionError)?;
        let _ = conn
            .execute(
                &query,
                &[&session.id().to_string(), &data, &session.expiry_date()],
            )
            .await;
        Ok(())
    }

    async fn load(&self, session_id: &Id) -> Result<Option<Session>, Self::Error> {
        let query = format!(
            r#"
            select data from "{schema_name}"."{table_name}"
            where id = $1 and expiry_date > $2
            "#,
            schema_name = self.schema_name,
            table_name = self.table_name
        );
        let cur_date = OffsetDateTime::now_utc();
        let conn = self
            .pool
            .get()
            .await
            .map_err(PostgresStoreError::ConnectionError)?;
        let rows = conn
            .query(&query, &[&session_id.to_string(), &cur_date])
            .await?;

        if let Some(row) = rows.get(0) {
            let data: Vec<u8> = row.get("data");
            let session: Session = rmp_serde::from_slice(&data)?;

            return Ok(Some(session));
        }

        Ok(None)
    }

    async fn delete(&self, session_id: &Id) -> Result<(), Self::Error> {
        let query = format!(
            r#"delete from "{schema_name}"."{table_name}" where id = $1"#,
            schema_name = self.schema_name,
            table_name = self.table_name
        );
        let conn = self
            .pool
            .get()
            .await
            .map_err(PostgresStoreError::ConnectionError)?;
        let _ = conn.execute(&query, &[&session_id.to_string()]).await;
        Ok(())
    }
}

[maxcountryman]

Thanks for this! We very much want folks to write their own SessionStore implementations. We probably will not aim to expand the built in implementations given that, but it's very helpful to have examples like this. Would you want to add this an example of how folks can implement their own SessionStore?

[maxcountryman]

Hi again, if you publish this as a crate please let us know and we'll add it to the list of available session stores.