feedsignal/crates/db/src/lib.rs

mod schema;

use anyhow::Result;
use chrono::Utc;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use diesel_async::pooled_connection::bb8::Pool;
use diesel_async::pooled_connection::AsyncDieselConnectionManager;
use diesel_async::sync_connection_wrapper::SyncConnectionWrapper;
use diesel_async::RunQueryDsl;
use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
use feedsignal_core::{Article, ReadingEvent, ReadingOutcome, TopicAffinities};
use uuid::Uuid;

const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations");

/// SQLite has no native async driver, so `diesel-async` wraps a blocking
/// `SqliteConnection` and runs it on a blocking thread under the hood.
type AsyncSqliteConnection = SyncConnectionWrapper<SqliteConnection>;

#[derive(Clone)]
pub struct Db {
    pool: Pool<AsyncSqliteConnection>,
}

impl Db {
    /// `database_url` is a plain SQLite file path (e.g. "feedsignal.db"),
    /// created if it doesn't exist. Everything lives in one file, no
    /// separate database server to run.
    pub async fn connect(database_url: &str) -> Result<Self> {
        // `MigrationHarness` only works on a synchronous `Connection`, and
        // this only runs once at boot, so open a throwaway sync connection
        // just for it rather than pulling migrations through the async pool.
        // Each pending migration runs in its own transaction (SQLite
        // supports transactional DDL), and `down.sql` gives each one a
        // matching revert.
        let url = database_url.to_string();
        tokio::task::spawn_blocking(move || -> Result<()> {
            let mut conn = SqliteConnection::establish(&url)?;
            conn.run_pending_migrations(MIGRATIONS)
                .map_err(|e| anyhow::anyhow!(e))?;
            Ok(())
        })
        .await??;

        let manager = AsyncDieselConnectionManager::<AsyncSqliteConnection>::new(database_url);
        let pool = Pool::builder().max_size(5).build(manager).await?;
        Ok(Self { pool })
    }

    pub async fn upsert_feed(&self, url: &str, title: &str) -> Result<Uuid> {
        use schema::feeds::dsl;
        let mut conn = self.pool.get().await?;
        let new_id = Uuid::new_v4().to_string();
        diesel::insert_into(dsl::feeds)
            .values((dsl::id.eq(&new_id), dsl::url.eq(url), dsl::title.eq(title)))
            .on_conflict(dsl::url)
            .do_update()
            .set(dsl::title.eq(title))
            .execute(&mut conn)
            .await?;
        let id: String = dsl::feeds
            .filter(dsl::url.eq(url))
            .select(dsl::id)
            .first(&mut conn)
            .await?;
        Ok(Uuid::parse_str(&id)?)
    }

    pub async fn insert_article(&self, article: &Article) -> Result<()> {
        use schema::articles::dsl;
        let mut conn = self.pool.get().await?;
        let topics = serde_json::to_string(&article.topics)?;
        diesel::insert_into(dsl::articles)
            .values((
                dsl::id.eq(article.id.to_string()),
                dsl::feed_id.eq(article.feed_id.to_string()),
                dsl::url.eq(&article.url),
                dsl::title.eq(&article.title),
                dsl::summary.eq(&article.summary),
                dsl::fetched_at.eq(Utc::now().to_rfc3339()),
                dsl::topics.eq(topics),
                dsl::embedding_score.eq(article.embedding_score),
                dsl::llm_score.eq(article.llm_score),
                dsl::final_score.eq(article.final_score),
                dsl::estimated_read_seconds.eq(article.estimated_read_seconds.map(|v| v as i32)),
            ))
            .on_conflict_do_nothing()
            .execute(&mut conn)
            .await?;
        Ok(())
    }

    /// Articles above the embedding-similarity threshold that haven't been
    /// through the (slower) LLM scoring stage yet.
    pub async fn shortlist_for_llm_scoring(&self, threshold: f32, limit: i64) -> Result<Vec<Uuid>> {
        use schema::articles::dsl;
        let mut conn = self.pool.get().await?;
        let ids: Vec<String> = dsl::articles
            .filter(dsl::embedding_score.ge(threshold))
            .filter(dsl::llm_score.is_null())
            .order(dsl::embedding_score.desc())
            .limit(limit)
            .select(dsl::id)
            .load(&mut conn)
            .await?;
        ids.into_iter().map(|s| Ok(Uuid::parse_str(&s)?)).collect()
    }

    /// Title/summary/topics/embedding_score for one article, used to build
    /// the LLM-judging prompt for it.
    pub async fn article_scoring_fields(
        &self,
        article_id: Uuid,
    ) -> Result<Option<(String, String, Vec<String>, f32)>> {
        use schema::articles::dsl;
        let mut conn = self.pool.get().await?;
        let row: Option<(String, String, String, Option<f32>)> = dsl::articles
            .filter(dsl::id.eq(article_id.to_string()))
            .select((dsl::title, dsl::summary, dsl::topics, dsl::embedding_score))
            .first(&mut conn)
            .await
            .optional()?;
        Ok(match row {
            Some((title, summary, topics_json, score)) => {
                let topics: Vec<String> = serde_json::from_str(&topics_json).unwrap_or_default();
                Some((title, summary, topics, score.unwrap_or(0.0)))
            }
            None => None,
        })
    }

    pub async fn store_llm_result(
        &self,
        article_id: Uuid,
        llm_score: f32,
        rationale: &str,
        final_score: f32,
    ) -> Result<()> {
        use schema::articles::dsl;
        let mut conn = self.pool.get().await?;
        diesel::update(dsl::articles.filter(dsl::id.eq(article_id.to_string())))
            .set((
                dsl::llm_score.eq(llm_score),
                dsl::llm_rationale.eq(rationale),
                dsl::final_score.eq(final_score),
            ))
            .execute(&mut conn)
            .await?;
        Ok(())
    }

    /// Highest-ranked articles for display, most relevant first.
    pub async fn list_ranked_articles(
        &self,
        limit: i64,
    ) -> Result<Vec<(String, String, String, String, Vec<String>, Option<f32>)>> {
        use schema::articles::dsl;
        let mut conn = self.pool.get().await?;
        // SQLite sorts NULL before any value, so `DESC` already puts NULL
        // `final_score`s last — no separate NULLS LAST clause needed here.
        let rows: Vec<(String, String, String, String, String, Option<f32>)> = dsl::articles
            .order(dsl::final_score.desc())
            .limit(limit)
            .select((
                dsl::id,
                dsl::title,
                dsl::url,
                dsl::summary,
                dsl::topics,
                dsl::final_score,
            ))
            .load(&mut conn)
            .await?;
        Ok(rows
            .into_iter()
            .map(|(id, title, url, summary, topics_json, final_score)| {
                (
                    id,
                    title,
                    url,
                    summary,
                    serde_json::from_str(&topics_json).unwrap_or_default(),
                    final_score,
                )
            })
            .collect())
    }

    pub async fn record_event(&self, event: &ReadingEvent) -> Result<()> {
        use schema::reading_events::dsl;
        let mut conn = self.pool.get().await?;
        let (outcome, dwell) = match &event.outcome {
            ReadingOutcome::Impression => ("impression", None),
            ReadingOutcome::Opened { dwell_seconds } => ("opened", dwell_seconds.map(|v| v as i32)),
            ReadingOutcome::Starred => ("starred", None),
            ReadingOutcome::Dismissed => ("dismissed", None),
        };
        diesel::insert_into(dsl::reading_events)
            .values((
                dsl::id.eq(event.id.to_string()),
                dsl::article_id.eq(event.article_id.to_string()),
                dsl::occurred_at.eq(event.occurred_at.to_rfc3339()),
                dsl::outcome.eq(outcome),
                dsl::dwell_seconds.eq(dwell),
            ))
            .execute(&mut conn)
            .await?;
        Ok(())
    }

    pub async fn load_affinities(&self) -> Result<TopicAffinities> {
        use schema::topic_affinities::dsl;
        let mut conn = self.pool.get().await?;
        let row: Option<String> = dsl::topic_affinities
            .filter(dsl::user_id.eq("default"))
            .select(dsl::affinities)
            .first(&mut conn)
            .await
            .optional()?;
        Ok(match row {
            Some(json) => serde_json::from_str(&json)?,
            None => TopicAffinities::default(),
        })
    }

    pub async fn save_affinities(&self, affinities: &TopicAffinities) -> Result<()> {
        use schema::topic_affinities::dsl;
        let mut conn = self.pool.get().await?;
        let json = serde_json::to_string(affinities)?;
        let updated_at = Utc::now().to_rfc3339();
        diesel::insert_into(dsl::topic_affinities)
            .values((
                dsl::user_id.eq("default"),
                dsl::affinities.eq(&json),
                dsl::updated_at.eq(&updated_at),
            ))
            .on_conflict(dsl::user_id)
            .do_update()
            .set((dsl::affinities.eq(&json), dsl::updated_at.eq(&updated_at)))
            .execute(&mut conn)
            .await?;
        Ok(())
    }
}