add sqlite database and message editing support

Cargo.toml

@@ -10,15 +10,15 @@ rss = "2.0"
 reqwest = {version = "0.11", default-features = false, features = ["rustls"]}
 tokio = {version = "1.28", features = ["rt-multi-thread", "signal"]}
 lazy-regex = "3.0.2"
-#time = {version = "0.3", features = ["local-offset" ,"macros", "parsing"]}
 chrono = "0.4.25"
 chrono-tz = "0.8.2"
 anyhow = "1.0.71"
 log = "0.4.18"
 env_logger = "0.10.0"
 dotenvy = "0.15.7"
-sled = { version = "0.34.7", features = ["compression"] }
-bincode = "1.3.3"
+diesel = { version = "2.1.3", features = ["sqlite", "chrono"] }
+diesel_migrations = "2.1.0"
+getset = "0.1.2"
 
 [dependencies.serenity]
 version = "0.11"

diesel.toml

@@ -0,0 +1,9 @@
+# For documentation on how to configure this file,
+# see https://diesel.rs/guides/configuring-diesel-cli
+
+[print_schema]
+file = "src/database/schema.rs"
+custom_type_derives = ["diesel::query_builder::QueryId"]
+
+[migrations_directory]
+dir = "migrations"

migrations/2023-11-05-222006_create_posts/down.sql

@@ -0,0 +1 @@
+DROP TABLE posts;

migrations/2023-11-05-222006_create_posts/up.sql

@@ -0,0 +1,7 @@
+CREATE TABLE posts (
+    announcement_id TEXT NOT NULL,
+    webhook_id UNSIGNED BIG INT NOT NULL,
+    message_id UNSIGNED BIG INT NOT NULL,
+    last_updated TIMESTAMP NOT NULL,
+    PRIMARY KEY (announcement_id, webhook_id)
+);

src/bin/disbahn.rs

@@ -0,0 +1,30 @@
+use disbahn::database::Database;
+use disbahn::DisbahnClient;
+use serenity::http::Http;
+use std::env;
+use std::env::VarError;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    dotenvy::dotenv()?;
+
+    env_logger::builder()
+        .filter_module(module_path!(), log::LevelFilter::Debug)
+        .init();
+
+    let database_url = match env::var("DATABASE_URL") {
+        Ok(url) => url,
+        Err(VarError::NotPresent) => "disbahn.db".to_string(),
+        Err(err) => return Err(err.into()),
+    };
+    let webhook_url = env::var("WEBHOOK_URL")?;
+    let feed_url = env::var("FEED_URL")?;
+
+    let database = Database::connect(&database_url)?;
+    let http = Http::new("");
+    let webhook = http.get_webhook_from_url(&webhook_url).await.unwrap();
+
+    let mut disbahn_client = DisbahnClient::new(database, webhook, http, feed_url);
+
+    disbahn_client.refresh().await
+}

src/database/mod.rs

@@ -0,0 +1,39 @@
+pub mod models;
+pub mod schema;
+
+use anyhow::format_err;
+use diesel::prelude::*;
+use diesel::SqliteConnection;
+use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
+use log::debug;
+
+pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
+
+pub struct Database {
+    conn: SqliteConnection,
+}
+
+impl Database {
+    fn new(conn: SqliteConnection) -> Self {
+        Self { conn }
+    }
+
+    pub fn connect(url: &str) -> anyhow::Result<Self> {
+        let mut connection = SqliteConnection::establish(url)
+            .map_err(|err| format_err!("Unable to connect to database at {url}: {err}"))?;
+        debug!("Established connection to SQLite database at {url}");
+
+        let migration_versions = connection
+            .run_pending_migrations(MIGRATIONS)
+            .map_err(|err| format_err!("Unable to run database migrations: {err}"))?;
+        if !migration_versions.is_empty() {
+            debug!("Ran database migrations for versions {migration_versions:?}");
+        }
+
+        Ok(Self::new(connection))
+    }
+
+    pub fn conn(&mut self) -> &mut SqliteConnection {
+        &mut self.conn
+    }
+}

src/database/models.rs

@@ -0,0 +1,51 @@
+use super::schema::*;
+use chrono::NaiveDateTime;
+use diesel::{Insertable, Queryable};
+use getset::Getters;
+use serenity::model::id::{MessageId, WebhookId};
+use std::borrow::Cow;
+
+#[derive(Queryable, Getters)]
+pub struct Post {
+    #[getset(get = "pub")]
+    announcement_id: String,
+    webhook_id: i64,
+    message_id: i64,
+    #[getset(get = "pub")]
+    last_updated: NaiveDateTime,
+}
+
+impl Post {
+    pub fn webhook_id(&self) -> WebhookId {
+        WebhookId(u64::from_le_bytes(self.webhook_id.to_le_bytes()))
+    }
+
+    pub fn message_id(&self) -> MessageId {
+        MessageId(u64::from_le_bytes(self.message_id.to_le_bytes()))
+    }
+}
+
+#[derive(Insertable)]
+#[diesel(table_name = posts)]
+pub struct NewPost<'a> {
+    announcement_id: Cow<'a, str>,
+    webhook_id: i64,
+    message_id: i64,
+    last_updated: NaiveDateTime,
+}
+
+impl<'a> NewPost<'a> {
+    pub fn new(
+        announcement_id: impl Into<Cow<'a, str>>,
+        webhook_id: WebhookId,
+        message_id: MessageId,
+        last_updated: NaiveDateTime,
+    ) -> Self {
+        Self {
+            announcement_id: announcement_id.into(),
+            webhook_id: i64::from_le_bytes(webhook_id.0.to_le_bytes()),
+            message_id: i64::from_le_bytes(message_id.0.to_le_bytes()),
+            last_updated,
+        }
+    }
+}

src/database/schema.rs

@@ -0,0 +1,10 @@
+// @generated automatically by Diesel CLI.
+
+diesel::table! {
+    posts (announcement_id, webhook_id) {
+        announcement_id -> Text,
+        webhook_id -> BigInt,
+        message_id -> BigInt,
+        last_updated -> Timestamp,
+    }
+}

src/lib.rs

@@ -1,5 +1,11 @@
+pub mod database;
+
+use crate::database::models::{NewPost, Post};
+use crate::database::schema::posts::dsl::posts;
+use crate::database::Database;
 use anyhow::{anyhow, Context};
-use chrono::{DateTime, NaiveDateTime, TimeZone, Utc};
+use chrono::{DateTime, NaiveDateTime, TimeZone};
+use diesel::{ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl};
 use lazy_regex::regex;
 use log::{debug, info};
 use reqwest::IntoUrl;
@@ -7,51 +13,21 @@ use serenity::http::Http;
 use serenity::json::Value;
 use serenity::model::channel::Embed;
 use serenity::model::webhook::Webhook;
-use std::borrow::Cow;
-use std::env;
 
-struct RefreshDb(sled::Db);
-
-impl RefreshDb {
-    fn get(&self, guid: &str) -> anyhow::Result<Option<DateTime<Utc>>> {
-        self.0
-            .get(guid)
-            .with_context(|| "db error: unable to get value")?
-            .map(|b| {
-                bincode::deserialize::<i64>(&b)
-                    .with_context(|| "db error: unable to deserialize i64")
-            })
-            .transpose()?
-            .map(|i| {
-                Utc.timestamp_opt(i, 0)
-                    .single()
-                    .with_context(|| "db error: unable to parse timestamp from i64")
-            })
-            .transpose()
-    }
-
-    fn insert(&self, guid: String, datetime: DateTime<Utc>) -> anyhow::Result<()> {
-        let timestamp = datetime.timestamp();
-        let bytes = bincode::serialize(&timestamp).with_context(|| "db error")?;
-        self.0.insert(guid, bytes).with_context(|| "db error")?;
-        Ok(())
-    }
-}
-
-struct DisbahnClient {
+pub struct DisbahnClient {
+    database: Database,
     webhook: Webhook,
     http: Http,
     rss_url: String,
-    known_guids: RefreshDb,
 }
 
 impl DisbahnClient {
-    fn new(webhook: Webhook, http: Http, rss_url: String, known_guids: RefreshDb) -> Self {
+    pub fn new(database: Database, webhook: Webhook, http: Http, rss_url: String) -> Self {
         Self {
+            database,
             webhook,
             http,
             rss_url,
-            known_guids,
         }
     }
 
@@ -93,7 +69,7 @@ impl DisbahnClient {
         }
     }
 
-    fn item_to_embed(item: &rss::Item, is_update: bool) -> anyhow::Result<Value> {
+    fn item_to_embed(item: &rss::Item) -> anyhow::Result<Value> {
         const COLOUR: u32 = 0x008d4f;
         const FOOTER_ICON_URL: &str = "https://www.zuginfo.nrw/img/customer/apple-touch-icon.png";
 
@@ -135,12 +111,9 @@ impl DisbahnClient {
 
         let pub_timestamp = pub_datetime.to_rfc3339_opts(chrono::SecondsFormat::Secs, true);
 
-        let embed = Embed::fake(|e| {
-            e.title(if is_update {
-                Cow::Owned(format!("UPDATE: {title}"))
-            } else {
-                Cow::Borrowed(title)
-            })
+        let embed =
+            Embed::fake(|e| {
+                e.title(title)
             .url(link)
             .thumbnail(icon_url)
             .colour(COLOUR)
@@ -158,7 +131,9 @@ impl DisbahnClient {
         Ok(embed)
     }
 
-    async fn refresh(&mut self) -> anyhow::Result<()> {
+    pub async fn refresh(&mut self) -> anyhow::Result<()> {
+        use crate::database::schema::posts::{self, dsl};
+
         debug!("Refreshing RSS feed ...");
         let channel = Self::get_rss_channel(&self.rss_url)
             .await
@@ -177,53 +152,54 @@ impl DisbahnClient {
                 .naive_utc()
                 .and_utc();
 
-            let is_update;
-            if let Some(last_pub_datetime) = self.known_guids.get(guid)? {
-                if last_pub_datetime < pub_datetime {
-                    is_update = true;
-                    info!("Updated item: {guid}");
-                } else {
-                    continue;
-                }
-            } else {
-                is_update = false;
-                info!("New item: {guid}");
-            }
+            let existing_post: Option<Post> = posts
+                .filter(dsl::webhook_id.eq(i64::from_le_bytes(self.webhook.id.0.to_le_bytes())))
+                .filter(dsl::announcement_id.eq(guid))
+                .first(self.database.conn())
+                .optional()
+                .with_context(|| "Error loading posts from database")?;
 
-            let embed = Self::item_to_embed(item, is_update)?;
+            if let Some(existing_post) = existing_post {
+                if existing_post.last_updated().and_utc() < pub_datetime {
+                    info!("Updated item: {guid}");
+                    let embed = Self::item_to_embed(item)?;
                     self.webhook
-                .execute(&self.http, false, |w| w.embeds(vec![embed]))
+                        .edit_message(&self.http, existing_post.message_id(), |w| {
+                            w.embeds(vec![embed])
+                        })
                         .await
-                .with_context(|| "Failed to execute webhook")?;
-            self.known_guids
-                .insert(guid.to_string(), pub_datetime)
-                .with_context(|| "Failed to write to database")?;
+                        .with_context(|| "Failed to edit message")?;
+
+                    diesel::update(
+                        posts.find((guid, i64::from_le_bytes(self.webhook.id.0.to_le_bytes()))),
+                    )
+                    .set(dsl::last_updated.eq(pub_datetime.naive_utc()))
+                    .execute(self.database.conn())
+                    .with_context(|| "Error updating post in database")?;
+                }
+            } else {
+                info!("New item: {guid}");
+                let embed = Self::item_to_embed(item)?;
+                let message = self
+                    .webhook
+                    .execute(&self.http, true, |w| w.embeds(vec![embed]))
+                    .await
+                    .with_context(|| "Failed to send message")?
+                    .with_context(|| "Discord did not return a message id")?;
+
+                diesel::insert_into(posts::table)
+                    .values(NewPost::new(
+                        guid,
+                        self.webhook.id,
+                        message.id,
+                        pub_datetime.naive_utc(),
+                    ))
+                    .execute(self.database.conn())
+                    .with_context(|| "Error inserting new post into database")?;
+            }
         }
 
         debug!("Done.");
         Ok(())
     }
 }
-
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
-    dotenvy::dotenv()?;
-
-    env_logger::builder()
-        .filter_module(module_path!(), log::LevelFilter::Debug)
-        .init();
-
-    let webhook_url = env::var("WEBHOOK_URL")?;
-    let feed_url = env::var("FEED_URL")?;
-    let known_guids_db = env::var("KNOWN_GUIDS_DB")?;
-
-    let known_guids_db =
-        RefreshDb(sled::open(known_guids_db).with_context(|| "unable to open database")?);
-
-    let http = Http::new("");
-    let webhook = http.get_webhook_from_url(&webhook_url).await.unwrap();
-
-    let mut disbahn_client = DisbahnClient::new(webhook, http, feed_url, known_guids_db);
-
-    disbahn_client.refresh().await
-}