ImageApi/src/database/user_ai_prefs_dao.rs

use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use std::ops::DerefMut;
use std::sync::{Arc, Mutex};

use crate::database::models::{UpsertUserAiPrefs, UserAiPrefs};
use crate::database::schema;
use crate::database::{DbError, DbErrorKind, connect};
use crate::otel::trace_db_call;

/// Generic single-row table that passively mirrors the latest client AI
/// request parameters (voice, timezone, library). Read by the nightly
/// pre-generation scheduler (Section D) to pick up user preferences.
pub trait UserAiPrefsDao: Sync + Send {
    /// Read the single row; `None` when it hasn't been populated yet.
    fn get_prefs(
        &mut self,
        context: &opentelemetry::Context,
    ) -> Result<Option<UserAiPrefs>, DbError>;

    /// Upsert the single row (id is always 1).
    #[allow(dead_code)]
    fn upsert_prefs(
        &mut self,
        context: &opentelemetry::Context,
        prefs: &UpsertUserAiPrefs,
    ) -> Result<(), DbError>;
}

pub struct SqliteUserAiPrefsDao {
    connection: Arc<Mutex<SqliteConnection>>,
}

impl Default for SqliteUserAiPrefsDao {
    fn default() -> Self {
        Self::new()
    }
}

impl SqliteUserAiPrefsDao {
    pub fn new() -> Self {
        Self {
            connection: Arc::new(Mutex::new(connect())),
        }
    }

    #[cfg(test)]
    pub fn from_connection(conn: Arc<Mutex<SqliteConnection>>) -> Self {
        Self { connection: conn }
    }
}

impl UserAiPrefsDao for SqliteUserAiPrefsDao {
    fn get_prefs(
        &mut self,
        context: &opentelemetry::Context,
    ) -> Result<Option<UserAiPrefs>, DbError> {
        trace_db_call(context, "query", "get_prefs", |_span| {
            use schema::user_ai_prefs::dsl;

            let mut connection = self
                .connection
                .lock()
                .expect("Unable to lock UserAiPrefsDao");

            dsl::user_ai_prefs
                .first::<UserAiPrefs>(connection.deref_mut())
                .optional()
                .map_err(|e| anyhow::anyhow!("Failed to get prefs: {}", e))
        })
        .map_err(|e| DbError::log(DbErrorKind::QueryError, e))
    }

    fn upsert_prefs(
        &mut self,
        context: &opentelemetry::Context,
        prefs: &UpsertUserAiPrefs,
    ) -> Result<(), DbError> {
        trace_db_call(context, "upsert", "upsert_prefs", |_span| {
            use schema::user_ai_prefs::dsl;

            let mut connection = self
                .connection
                .lock()
                .expect("Unable to lock UserAiPrefsDao");

            // Single-row table (id=1): one atomic upsert. The explicit id=1
            // makes the conflict target deterministic so the second call
            // updates in place rather than tripping the CHECK(id=1) constraint,
            // and real insert errors surface instead of being swallowed into a
            // separate update branch. The columns are set explicitly (rather
            // than via AsChangeset) so a None field overwrites to NULL — the
            // row mirrors the latest request exactly, not a merge of past ones.
            diesel::insert_into(dsl::user_ai_prefs)
                .values((dsl::id.eq(1), prefs))
                .on_conflict(dsl::id)
                .do_update()
                .set((
                    dsl::voice.eq(&prefs.voice),
                    dsl::tz_offset_minutes.eq(&prefs.tz_offset_minutes),
                    dsl::library.eq(&prefs.library),
                    dsl::display_tz.eq(&prefs.display_tz),
                    dsl::default_model.eq(&prefs.default_model),
                    dsl::default_backend.eq(&prefs.default_backend),
                    dsl::default_persona_id.eq(&prefs.default_persona_id),
                    dsl::reel_span.eq(&prefs.reel_span),
                    dsl::reel_max_segments.eq(&prefs.reel_max_segments),
                    dsl::updated_at.eq(&prefs.updated_at),
                ))
                .execute(connection.deref_mut())
                .map_err(|e| anyhow::anyhow!("Failed to upsert prefs: {}", e))?;
            Ok(())
        })
        .map_err(|e| DbError::log(DbErrorKind::InsertError, e))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use diesel::Connection;
    use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};

    const DB_MIGRATIONS: EmbeddedMigrations = embed_migrations!();

    fn setup_dao() -> SqliteUserAiPrefsDao {
        let mut conn = SqliteConnection::establish(":memory:")
            .expect("Unable to create in-memory db connection");
        conn.run_pending_migrations(DB_MIGRATIONS)
            .expect("Failure running DB migrations");
        SqliteUserAiPrefsDao::from_connection(Arc::new(Mutex::new(conn)))
    }

    fn ctx() -> opentelemetry::Context {
        opentelemetry::Context::new()
    }

    #[test]
    fn get_prefs_returns_none_when_empty() {
        let mut dao = setup_dao();
        let result = dao.get_prefs(&ctx()).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn upsert_prefs_inserts_row() {
        let mut dao = setup_dao();
        let now = 1_700_000_000i64;
        let prefs = UpsertUserAiPrefs {
            voice: Some("grandma".to_string()),
            tz_offset_minutes: Some(-480),
            library: Some("1".to_string()),
            updated_at: now,
            display_tz: Some("America/New_York".to_string()),
            default_model: Some("llama3.2:3b".to_string()),
            default_backend: Some("hybrid".to_string()),
            default_persona_id: Some("cabin".to_string()),
            reel_span: Some("week".to_string()),
            reel_max_segments: Some(24),
        };
        dao.upsert_prefs(&ctx(), &prefs).unwrap();

        let row = dao.get_prefs(&ctx()).unwrap().unwrap();
        assert_eq!(row.id, 1);
        assert_eq!(row.voice, Some("grandma".to_string()));
        assert_eq!(row.tz_offset_minutes, Some(-480));
        assert_eq!(row.library, Some("1".to_string()));
        assert_eq!(row.display_tz, Some("America/New_York".to_string()));
        assert_eq!(row.default_model, Some("llama3.2:3b".to_string()));
        assert_eq!(row.default_backend, Some("hybrid".to_string()));
        assert_eq!(row.default_persona_id, Some("cabin".to_string()));
        assert_eq!(row.reel_span, Some("week".to_string()));
        assert_eq!(row.reel_max_segments, Some(24));
        assert_eq!(row.updated_at, now);
    }

    #[test]
    fn upsert_prefs_replaces_existing() {
        let mut dao = setup_dao();
        let now1 = 1_700_000_000i64;
        let now2 = 1_800_000_000i64;

        let prefs1 = UpsertUserAiPrefs {
            voice: Some("grandma".to_string()),
            tz_offset_minutes: Some(-480),
            library: Some("1".to_string()),
            updated_at: now1,
            display_tz: Some("America/New_York".to_string()),
            default_model: Some("llama3.2:3b".to_string()),
            default_backend: Some("local".to_string()),
            default_persona_id: None,
            reel_span: Some("week".to_string()),
            reel_max_segments: Some(24),
        };
        dao.upsert_prefs(&ctx(), &prefs1).unwrap();

        let prefs2 = UpsertUserAiPrefs {
            voice: Some("dad".to_string()),
            tz_offset_minutes: Some(-300),
            library: None,
            updated_at: now2,
            display_tz: None,
            default_model: None,
            default_backend: None,
            default_persona_id: None,
            reel_span: None,
            reel_max_segments: None,
        };
        dao.upsert_prefs(&ctx(), &prefs2).unwrap();

        let row = dao.get_prefs(&ctx()).unwrap().unwrap();
        assert_eq!(row.voice, Some("dad".to_string()));
        assert_eq!(row.tz_offset_minutes, Some(-300));
        assert!(row.library.is_none());
        assert!(row.display_tz.is_none());
        // Owner-default columns are overwritten to NULL like the rest.
        assert!(row.default_model.is_none());
        assert!(row.default_backend.is_none());
        assert!(row.reel_span.is_none());
        assert!(row.reel_max_segments.is_none());
        assert_eq!(row.updated_at, now2);
    }

    #[test]
    fn upsert_partial_fields() {
        let mut dao = setup_dao();
        let now = 1_700_000_000i64;

        let prefs = UpsertUserAiPrefs {
            voice: None,
            tz_offset_minutes: Some(-480),
            library: None,
            updated_at: now,
            display_tz: Some("Europe/Berlin".to_string()),
            default_model: None,
            default_backend: None,
            default_persona_id: Some("default".to_string()),
            reel_span: Some("month".to_string()),
            reel_max_segments: None,
        };
        dao.upsert_prefs(&ctx(), &prefs).unwrap();

        let row = dao.get_prefs(&ctx()).unwrap().unwrap();
        assert_eq!(row.tz_offset_minutes, Some(-480));
        assert!(row.voice.is_none());
        assert!(row.library.is_none());
        assert_eq!(row.display_tz, Some("Europe/Berlin".to_string()));
        assert!(row.default_model.is_none());
        assert_eq!(row.default_persona_id, Some("default".to_string()));
        assert_eq!(row.reel_span, Some("month".to_string()));
        assert!(row.reel_max_segments.is_none());
    }
}