factoriauth/src/db.rs

use std::{fmt::Debug, str::FromStr};

use axum::async_trait;
use secrecy::ExposeSecret;
use sqlx::{query, query_as, sqlite::SqliteConnectOptions, Connection, SqliteConnection};
use tracing::instrument;

use crate::secrets::UserToken;

// TODO: check again if it's possible to pass this as a parameter to query!
// const TABLE_USER_TOKENS: &str = "user_tokens";

pub enum UserTokenEntry {
    Valid(
        UserToken,
        time::OffsetDateTime,
        Option<time::OffsetDateTime>,
    ),
    Invalid(
        UserToken,
        time::OffsetDateTime,
        Option<time::OffsetDateTime>,
    ),
}

#[async_trait]
pub trait Database: Debug {
    async fn get_token(&mut self, username: &str) -> Result<Option<UserTokenEntry>, sqlx::Error>;

    async fn save_token(&mut self, username: &str, token: &UserToken) -> Result<(), sqlx::Error>;

    async fn update_token_last_used(
        &mut self,
        username: &str,
        token: &UserToken,
    ) -> Result<(), sqlx::Error>;
}

#[derive(Debug)]
pub struct SqliteDatabase {
    conn: SqliteConnection,
}

impl SqliteDatabase {
    #[instrument]
    pub async fn open(connection_string: &str) -> Self {
        let options = SqliteConnectOptions::from_str(connection_string)
            .expect("Invalid database URI")
            .create_if_missing(true);

        let mut db = Self {
            conn: SqliteConnection::connect_with(&options)
                .await
                .expect("Failed to open SQLite database"),
        };

        db.init().await;

        db
    }

    #[instrument]
    pub async fn init(&mut self) {
        query!(
            "CREATE TABLE IF NOT EXISTS user_tokens (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                username    VARCHAR(255) NOT NULL,
                token       NCHAR(30) NOT NULL,
                valid       BOOLEAN NOT NULL,
                created     DATETIME NOT NULL,
                last_used   DATETIME
            )"
        )
        .execute(&mut self.conn)
        .await
        .expect("Failed to initialize table user_tokens");
    }
}

#[async_trait]
impl Database for SqliteDatabase {
    #[instrument]
    async fn get_token(&mut self, username: &str) -> Result<Option<UserTokenEntry>, sqlx::Error> {
        struct TokenRow {
            token: String,
            valid: bool,
            created: time::OffsetDateTime,
            last_used: Option<time::OffsetDateTime>,
        }
        let row = query_as!(
            TokenRow,
            "SELECT token, valid, created, last_used
            FROM user_tokens
            WHERE username = $1
            ORDER BY created DESC",
            username
        )
        .fetch_optional(&mut self.conn)
        .await?;

        Ok(row.map(
            |TokenRow {
                 token,
                 valid,
                 created,
                 last_used,
             }| match valid {
                true => UserTokenEntry::Valid(UserToken::from(token), created, last_used),
                false => UserTokenEntry::Invalid(UserToken::from(token), created, last_used),
            },
        ))
    }

    #[instrument]
    async fn save_token(&mut self, username: &str, token: &UserToken) -> Result<(), sqlx::Error> {
        let token_inner = token.0.expose_secret();
        query!(
            "INSERT INTO user_tokens
                (username, token, created, valid)
            VALUES
                ($1, $2, DATETIME('NOW'), TRUE)
            ",
            username,
            token_inner
        )
        .execute(&mut self.conn)
        .await?;

        Ok(())
    }

    #[instrument]
    async fn update_token_last_used(
        &mut self,
        username: &str,
        token: &UserToken,
    ) -> Result<(), sqlx::Error> {
        let token_inner = token.0.expose_secret();
        query!(
            "UPDATE user_tokens
            SET last_used = DATETIME('NOW')
            WHERE username = $1 AND token = $2",
            username,
            token_inner
        )
        .execute(&mut self.conn)
        .await?;

        Ok(())
    }
}