async-session-r2d2 0.1.3

/*!
 * Provides the SQLite implementation of a session store.
 */
use super::{PooledSessionStore, TableOperations};
use anyhow::Context;
use async_session::{chrono, serde_json, Session};
use r2d2_sqlite::{
    rusqlite, rusqlite::OpenFlags, rusqlite::OptionalExtension,
    SqliteConnectionManager as SessionManager,
};
use std::{path::PathBuf, str::FromStr};

/// A SQLite-backed session store.
pub type SessionStore = PooledSessionStore<SessionManager>;

impl TableOperations<SessionManager> for SessionStore {
    fn add_table(&self) -> async_session::Result<()> {
        self.connection()?
            .execute(
                &format!(
                    "
                    CREATE TABLE IF NOT EXISTS {:?} (
                        id TEXT PRIMARY KEY NOT NULL,
                        expires INTEGER NULL,
                        session TEXT NOT NULL
                    )
                    ",
                    self.table_name
                ),
                rusqlite::params![],
            )
            .context("Failed to create the table for handling sessions")
            .and(Ok(()))
    }

    fn clean_up(&self) -> async_session::Result<()> {
        self.connection()?
            .execute(
                &format!("DELETE from {} WHERE expires <= ?", self.table_name),
                rusqlite::params![chrono::Utc::now().timestamp()],
            )
            .and(Ok(()))
            .context("Failed to remove sessions that have expired.")
    }

    fn connection(&self) -> async_session::Result<r2d2::PooledConnection<SessionManager>> {
        self.pool
            .get()
            .context("Failed to obtain a connection to the SQLite database.")
    }

    fn count(&self) -> async_session::Result<usize> {
        self.connection()?
            .query_row(
                &format!("SELECT count(*) FROM {}", self.table_name),
                rusqlite::params![],
                |row| row.get("count(*)"),
            )
            .context("Could not get the count of sessions.")
    }
}

#[async_trait::async_trait]
impl async_session::SessionStore for SessionStore {
    async fn load_session(&self, cookie_value: String) -> async_session::Result<Option<Session>> {
        let id = Session::id_from_cookie_value(&cookie_value)?;
        let connection = self.connection()?;

        log::trace!("Looking for the session with the ID of {:?}", id);

        connection
            .query_row(
                &format!(
                    "
            SELECT
                session
            FROM {:?}
            WHERE
                id = :id AND
                (expires IS NULL OR expires > :expires)
            ",
                    self.table_name
                ),
                rusqlite::named_params![
                    ":id": id,
                    ":expires": chrono::Utc::now().timestamp()
                ],
                |row| row.get::<&str, String>("session"),
            )
            .optional()
            .context("Failed to get the column 'session' from the table")
            .and_then(|session_str_opt| {
                if let Some(session_string) = session_str_opt {
                    log::trace!("Found the session #{}", id);
                    serde_json::from_str(&session_string)
                        .context("Failed to deserialize session from JSON")
                } else {
                    log::trace!("No session was found with the ID {}", id);
                    Ok(None)
                }
            })
    }
    async fn store_session(&self, session: Session) -> async_session::Result<Option<String>> {
        let id = session.id();
        let string = serde_json::to_string(&session)?;
        let connection = self.connection()?;

        connection
            .execute(
                &format!(
                    "
            INSERT INTO {}
                (id, session, expires)
            VALUES (:id, :session, :expires)
            ON CONFLICT(id)
            DO UPDATE SET
                expires = excluded.expires,
                session = excluded.session
        ",
                    self.table_name
                ),
                rusqlite::named_params![
                    ":id": id,
                    ":session": string,
                    ":expires": session.expiry().map(|e| e.timestamp())
                ],
            )
            .context("Failed to add or update the session.")
            .and(Ok(session.into_cookie_value()))
    }
    async fn destroy_session(&self, session: Session) -> async_session::Result {
        let connection = self.connection()?;

        connection
            .execute(
                &format!("delete from {} where id = :id", self.table_name),
                rusqlite::named_params![":id": session.id()],
            )
            .and(Ok(()))
            .context("Failed to clear out the specified session.")
    }
    async fn clear_store(&self) -> async_session::Result {
        let connection = self.connection()?;

        connection
            .execute(
                &format!("delete from {}", self.table_name),
                rusqlite::params![],
            )
            .and(Ok(()))
            .context("Failed to clear out the sessions.")
    }
}

impl SessionStore {
    /// Creates a new SQLite-backed session store backed by memory.
    ///
    /// There's a bit of a catch in how `r2d2_sqlite` handles in-memory connections,
    /// which in turn, is a bit of a 'hack' around SQLite itself. For more information
    /// see <https://github.com/ivanceras/r2d2-sqlite/issues/39>.
    pub fn from_memory(table_name: String) -> async_session::Result<Self> {
        Self::from_path(
            PathBuf::from_str("file:memory:?mode=memory&cache=private").unwrap(),
            table_name,
        )
    }

    /// Creates a new SQLite-backed session store backed by a database on disk.
    pub fn from_path(
        session_database_path: PathBuf,
        table_name: String,
    ) -> async_session::Result<Self> {
        Self::from_pool(
            r2d2::Pool::builder()
                .test_on_check_out(true)
                .build(SessionManager::with_flags(
                    SessionManager::file(session_database_path),
                    OpenFlags::SQLITE_OPEN_URI
                        | OpenFlags::SQLITE_OPEN_READ_WRITE
                        | OpenFlags::SQLITE_OPEN_CREATE,
                ))
                .unwrap(),
            table_name,
        )
    }
}