factoriauth/src/db.rs

use std::str::FromStr;

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

use crate::secrets::UserToken;

// TODO: allow configuring this via envar
const DB_URI_DEFAULT: &str = "sqlite://sqlite.db";

const TABLE_USER_TOKENS: &str = "user_tokens";

pub enum UserTokenEntry {
    Valid(UserToken, time::Time, time::Time),
    Invalid(UserToken, time::Time, time::Time),
}

pub trait Database {
    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>;
}

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

impl SqliteDatabase {
    #[instrument]
    pub async fn open() -> Self {
        let options = SqliteConnectOptions::from_str(DB_URI_DEFAULT)
            .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(&format!(
            "CREATE TABLE IF NOT EXISTS {TABLE_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");
    }
}

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

        Ok(row.map(|row| match row {
            (token_str, true, created, last_used) => {
                UserTokenEntry::Valid(UserToken::from(token_str), created, last_used)
            }

            (token_str, false, created, last_used) => {
                UserTokenEntry::Invalid(UserToken::from(token_str), created, last_used)
            }
        }))
    }

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

        Ok(())
    }
}