migrate db to mongodb from diesel(postgres)

Cargo.toml

@@ -6,14 +6,15 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-chrono = "0.4.40"
+async-trait = "0.1.88"
+bson = { version = "2.14.0", features = ["chrono-0_4"] }
+chrono = { version = "0.4.40", features = ["serde"] }
 chrono-tz = "0.10.3"
-diesel = { version = "2.2.8", features = ["postgres", "chrono"] }
-diesel-async = { version = "0.5.2", features = ["bb8", "postgres"] }
-diesel-derive-enum = "2.1.0"
 dotenvy = "0.15.7"
 enum_stringify = "0.6.3"
 envconfig = "0.11.0"
+futures = "0.3.31"
+mongodb = "3.2.3"
 serde = { version = "1.0.219", features = ["derive", "serde_derive"] }
 teloxide = { version = "0.14.0", features = ["macros", "postgres-storage-nativetls"] }
 tokio = { version = "1.44.1", features = ["rt-multi-thread", "macros"] }

bacon.toml

@@ -0,0 +1,111 @@
+# This is a configuration file for the bacon tool
+#
+# Complete help on configuration: https://dystroy.org/bacon/config/
+# 
+# You may check the current default at
+#   https://github.com/Canop/bacon/blob/main/defaults/default-bacon.toml
+
+default_job = "check"
+env.CARGO_TERM_COLOR = "always"
+
+[jobs.check]
+command = ["cargo", "check"]
+need_stdout = false
+
+[jobs.check-all]
+command = ["cargo", "check", "--all-targets"]
+need_stdout = false
+
+# Run clippy on the default target
+[jobs.clippy]
+command = ["cargo", "clippy"]
+need_stdout = false
+
+# Run clippy on all targets
+# To disable some lints, you may change the job this way:
+#    [jobs.clippy-all]
+#    command = [
+#        "cargo", "clippy",
+#        "--all-targets",
+#    	 "--",
+#    	 "-A", "clippy::bool_to_int_with_if",
+#    	 "-A", "clippy::collapsible_if",
+#    	 "-A", "clippy::derive_partial_eq_without_eq",
+#    ]
+# need_stdout = false
+[jobs.clippy-all]
+command = ["cargo", "clippy", "--all-targets"]
+need_stdout = false
+
+# This job lets you run
+# - all tests: bacon test
+# - a specific test: bacon test -- config::test_default_files
+# - the tests of a package: bacon test -- -- -p config
+[jobs.test]
+command = ["cargo", "test"]
+need_stdout = true
+
+[jobs.nextest]
+command = [
+    "cargo", "nextest", "run",
+    "--hide-progress-bar", "--failure-output", "final"
+]
+need_stdout = true
+analyzer = "nextest"
+
+[jobs.doc]
+command = ["cargo", "doc", "--no-deps"]
+need_stdout = false
+
+# If the doc compiles, then it opens in your browser and bacon switches
+# to the previous job
+[jobs.doc-open]
+command = ["cargo", "doc", "--no-deps", "--open"]
+need_stdout = false
+on_success = "back" # so that we don't open the browser at each change
+
+# You can run your application and have the result displayed in bacon,
+# if it makes sense for this crate.
+[jobs.run]
+command = [
+    "cargo", "run",
+    # put launch parameters for your program behind a `--` separator
+]
+need_stdout = true
+allow_warnings = true
+background = true
+
+# Run your long-running application (eg server) and have the result displayed in bacon.
+# For programs that never stop (eg a server), `background` is set to false
+# to have the cargo run output immediately displayed instead of waiting for
+# program's end.
+# 'on_change_strategy' is set to `kill_then_restart` to have your program restart
+# on every change (an alternative would be to use the 'F5' key manually in bacon).
+# If you often use this job, it makes sense to override the 'r' key by adding
+# a binding `r = job:run-long` at the end of this file .
+[jobs.run-long]
+command = [
+    "cargo", "run",
+    # put launch parameters for your program behind a `--` separator
+]
+need_stdout = true
+allow_warnings = true
+background = false
+on_change_strategy = "kill_then_restart"
+
+# This parameterized job runs the example of your choice, as soon
+# as the code compiles.
+# Call it as
+#    bacon ex -- my-example
+[jobs.ex]
+command = ["cargo", "run", "--example"]
+need_stdout = true
+allow_warnings = true
+
+# You may define here keybindings that would be specific to
+# a project, for example a shortcut to launch a specific job.
+# Shortcuts to internal functions (scrolling, toggling, etc.)
+# should go in your personal global prefs.toml file instead.
+[keybindings]
+# alt-m = "job:my-job"
+c = "job:clippy-all" # comment this to have 'c' run clippy on only the default target

migrations/2025-04-18-100731_add_media_group_id/down.sql

@@ -0,0 +1 @@
+ALTER TABLE media DROP COLUMN media_group_id;

migrations/2025-04-18-100731_add_media_group_id/up.sql

@@ -0,0 +1,2 @@
+ALTER TABLE media ADD COLUMN media_group_id VARCHAR NOT NULL;
+

migrations/2025-04-18-101829_optional_media_group_id/down.sql

@@ -0,0 +1 @@
+ALTER TABLE media ALTER COLUMN media_group_id SET NOT NULL;

migrations/2025-04-18-101829_optional_media_group_id/up.sql

@@ -0,0 +1 @@
+ALTER TABLE media ALTER COLUMN media_group_id DROP NOT NULL;

src/admin.rs

@@ -3,7 +3,7 @@ use teloxide::{
     utils::{command::BotCommands, render::RenderMessageTextHelper},
 };
 
-use crate::db::DB;
+use crate::db::{CallDB, DB};
 use crate::LogMsg;
 
 // These are should not appear in /help

src/db.rs

@@ -1,266 +0,0 @@
-pub mod models;
-pub mod schema;
-use crate::Config;
-
-use self::models::*;
-
-use chrono::Utc;
-use diesel::prelude::*;
-use diesel_async::pooled_connection::bb8::Pool;
-use diesel_async::pooled_connection::AsyncDieselConnectionManager;
-use diesel_async::AsyncPgConnection;
-use diesel_async::RunQueryDsl;
-use enum_stringify::EnumStringify;
-
-#[derive(EnumStringify)]
-#[enum_stringify(case = "flat")]
-pub enum ReservationStatus {
-    Booked,
-    Paid,
-}
-
-pub trait GetReservationStatus {
-    fn get_status(&self) -> Option<ReservationStatus>;
-}
-
-impl GetReservationStatus for models::Reservation {
-    fn get_status(&self) -> Option<ReservationStatus> {
-        ReservationStatus::try_from(self.status.clone()).ok()
-    }
-}
-
-#[derive(Clone)]
-pub struct DB {
-    pool: diesel_async::pooled_connection::bb8::Pool<AsyncPgConnection>,
-}
-
-impl DB {
-    pub async fn new<S: Into<String>>(db_url: S) -> Self {
-        let config = AsyncDieselConnectionManager::<diesel_async::AsyncPgConnection>::new(db_url);
-        let pool = Pool::builder().build(config).await.unwrap();
-        DB { pool }
-    }
-
-    pub async fn get_users(&mut self) -> Vec<User> {
-        use self::schema::users::dsl::*;
-        let mut conn = self.pool.get().await.unwrap();
-        users
-            .filter(id.gt(0))
-            .load::<User>(&mut conn)
-            .await
-            .unwrap()
-    }
-
-    pub async fn set_admin(&mut self, userid: i64, isadmin: bool) {
-        use self::schema::users::dsl::*;
-        let connection = &mut self.pool.get().await.unwrap();
-        diesel::update(users)
-            .filter(id.eq(userid))
-            .set(is_admin.eq(isadmin))
-            .execute(connection)
-            .await
-            .unwrap();
-    }
-
-    pub async fn get_or_init_user(&mut self, userid: i64, firstname: &str) -> User {
-        use self::schema::users::dsl::*;
-        let connection = &mut self.pool.get().await.unwrap();
-
-        let user = users
-            .filter(id.eq(userid))
-            .first::<User>(connection)
-            .await
-            .optional()
-            .unwrap();
-
-        match user {
-            Some(existing_user) => existing_user,
-            None => diesel::insert_into(users)
-                .values((
-                    id.eq(userid as i64),
-                    is_admin.eq(false),
-                    first_name.eq(firstname),
-                ))
-                .get_result(connection)
-                .await
-                .unwrap(),
-        }
-    }
-
-    pub async fn get_message(
-        &mut self,
-        chatid: i64,
-        messageid: i32,
-    ) -> Result<Option<Message>, Box<dyn std::error::Error>> {
-        use self::schema::messages::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        let msg = messages
-            .filter(chat_id.eq(chatid))
-            .filter(message_id.eq(messageid as i64))
-            .first::<Message>(conn)
-            .await
-            .optional()?;
-
-        Ok(msg)
-    }
-
-    pub async fn get_message_literal(
-        &mut self,
-        chatid: i64,
-        messageid: i32,
-    ) -> Result<Option<String>, Box<dyn std::error::Error>> {
-        let msg = self.get_message(chatid, messageid).await?;
-        Ok(msg.map(|m| m.token))
-    }
-
-    pub async fn set_message_literal(
-        &mut self,
-        chatid: i64,
-        messageid: i32,
-        literal: &str,
-    ) -> Result<(), Box<dyn std::error::Error>> {
-        use self::schema::messages::dsl::*;
-        let conn = &mut self.pool.get().await?;
-
-        let msg = self.clone().get_message(chatid, messageid).await?;
-
-        match msg {
-            Some(msg) => {
-                diesel::update(messages)
-                    .filter(id.eq(msg.id))
-                    .set(token.eq(literal))
-                    .execute(conn)
-                    .await?;
-            }
-            None => {
-                diesel::insert_into(messages)
-                    .values((
-                        chat_id.eq(chatid),
-                        message_id.eq(messageid as i64),
-                        token.eq(literal),
-                    ))
-                    .execute(conn)
-                    .await?;
-            }
-        };
-
-        Ok(())
-    }
-
-    async fn get_literal(
-        &mut self,
-        literal: &str,
-    ) -> Result<Option<Literal>, Box<dyn std::error::Error>> {
-        use self::schema::literals::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        let literal = literals
-            .filter(token.eq(literal))
-            .first::<Literal>(conn)
-            .await
-            .optional()?;
-
-        Ok(literal)
-    }
-
-    pub async fn get_literal_value(
-        &mut self,
-        literal: &str,
-    ) -> Result<Option<String>, Box<dyn std::error::Error>> {
-        let literal = self.get_literal(literal).await?;
-
-        Ok(literal.map(|l| l.value))
-    }
-
-    pub async fn set_literal(
-        &mut self,
-        literal: &str,
-        valuestr: &str,
-    ) -> Result<(), Box<dyn std::error::Error>> {
-        use self::schema::literals::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        diesel::insert_into(literals)
-            .values((token.eq(literal), value.eq(valuestr)))
-            .on_conflict(token)
-            .do_update()
-            .set(value.eq(valuestr))
-            .execute(conn)
-            .await?;
-
-        Ok(())
-    }
-
-    pub async fn get_all_events(&mut self) -> Vec<Event> {
-        use self::schema::events::dsl::*;
-        let mut conn = self.pool.get().await.unwrap();
-        events
-            .filter(id.gt(0))
-            .load::<Event>(&mut conn)
-            .await
-            .unwrap()
-    }
-
-    pub async fn create_event(
-        &mut self,
-        event_datetime: chrono::DateTime<Utc>,
-    ) -> Result<Event, Box<dyn std::error::Error>> {
-        use self::schema::events::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        let new_event = diesel::insert_into(events)
-            .values((time.eq(event_datetime),))
-            .get_result::<Event>(conn)
-            .await?;
-
-        Ok(new_event)
-    }
-
-    pub async fn get_media(
-        &mut self,
-        literal: &str,
-    ) -> Result<Vec<Media>, Box<dyn std::error::Error>> {
-        use self::schema::media::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        let media_items = media.filter(token.eq(literal)).load::<Media>(conn).await?;
-
-        Ok(media_items)
-    }
-
-    pub async fn drop_media(&mut self, literal: &str) -> Result<usize, Box<dyn std::error::Error>> {
-        use self::schema::media::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        let deleted_count = diesel::delete(media.filter(token.eq(literal)))
-            .execute(conn)
-            .await?;
-
-        Ok(deleted_count)
-    }
-
-    pub async fn add_media(
-        &mut self,
-        literal: &str,
-        mediatype: &str,
-        fileid: &str,
-    ) -> Result<Media, Box<dyn std::error::Error>> {
-        use self::schema::media::dsl::*;
-        let conn = &mut self.pool.get().await.unwrap();
-
-        let new_media = diesel::insert_into(media)
-            .values((
-                token.eq(literal),
-                media_type.eq(mediatype),
-                file_id.eq(fileid),
-            ))
-            .get_result::<Media>(conn)
-            .await?;
-
-        Ok(new_media)
-    }
-}
-
-#[cfg(test)]
-mod tests;

src/db/mod.rs

@@ -0,0 +1,348 @@
+use async_trait::async_trait;
+use chrono::{DateTime, Utc};
+use enum_stringify::EnumStringify;
+use futures::stream::{StreamExt, TryStreamExt};
+
+use mongodb::Database;
+use mongodb::{
+    bson::doc,
+    options::{ClientOptions, ResolverConfig},
+    Client,
+};
+use serde::{Deserialize, Serialize};
+
+#[derive(EnumStringify)]
+#[enum_stringify(case = "flat")]
+pub enum ReservationStatus {
+    Booked,
+    Paid,
+}
+
+pub trait GetReservationStatus {
+    fn get_status(&self) -> Option<ReservationStatus>;
+}
+
+//impl GetReservationStatus for models::Reservation {
+//    fn get_status(&self) -> Option<ReservationStatus> {
+//        ReservationStatus::try_from(self.status.clone()).ok()
+//    }
+//}
+#[derive(Serialize, Deserialize, Default)]
+pub struct User {
+    pub _id: bson::oid::ObjectId,
+    pub id: i64,
+    pub is_admin: bool,
+    pub first_name: String,
+    pub last_name: Option<String>,
+    pub username: Option<String>,
+    pub language_code: Option<String>,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct Message {
+    pub _id: bson::oid::ObjectId,
+    pub chat_id: i64,
+    pub message_id: i64,
+    pub token: String,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct Literal {
+    pub _id: bson::oid::ObjectId,
+    pub token: String,
+    pub value: String,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct Event {
+    pub _id: bson::oid::ObjectId,
+    pub time: DateTime<Utc>,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct Media {
+    pub _id: bson::oid::ObjectId,
+    pub token: String,
+    pub media_type: String,
+    pub file_id: String,
+    pub media_group_id: Option<String>,
+}
+
+#[derive(Clone)]
+pub struct DB {
+    client: Client,
+}
+
+impl DB {
+    pub async fn new<S: Into<String>>(db_url: S) -> Self {
+        let options = ClientOptions::parse(db_url.into()).await.unwrap();
+        let client = Client::with_options(options).unwrap();
+
+        DB { client }
+    }
+}
+
+#[async_trait]
+impl CallDB for DB {
+    async fn get_database(&mut self) -> Database {
+        self.client.database("gongbot")
+    }
+}
+
+#[async_trait]
+pub trait CallDB {
+    //type C;
+    async fn get_database(&mut self) -> Database;
+    //async fn get_pool(&mut self) -> PooledConnection<'_, AsyncDieselConnectionManager<C>>;
+    async fn get_users(&mut self) -> Vec<User> {
+        let db = self.get_database().await;
+        let users = db.collection::<User>("users");
+        users
+            .find(doc! {})
+            .await
+            .unwrap()
+            .map(|u| u.unwrap())
+            .collect()
+            .await
+    }
+
+    async fn set_admin(&mut self, userid: i64, isadmin: bool) {
+        let db = self.get_database().await;
+        let users = db.collection::<User>("users");
+        users
+            .update_one(
+                doc! {
+                    "id": userid
+                },
+                doc! {
+                    "$set": { "is_admin": isadmin }
+                },
+            )
+            .await
+            .unwrap();
+    }
+
+    async fn get_or_init_user(&mut self, userid: i64, firstname: &str) -> User {
+        let db = self.get_database().await;
+        let users = db.collection::<User>("users");
+
+        users
+            .update_one(
+                doc! { "id": userid },
+                doc! { "$set": { "first_name": firstname } },
+            )
+            .upsert(true)
+            .await
+            .unwrap();
+
+        users
+            .find_one(doc! { "id": userid })
+            .await
+            .unwrap()
+            .expect("no such user created")
+    }
+
+    async fn get_message(
+        &mut self,
+        chatid: i64,
+        messageid: i32,
+    ) -> Result<Option<Message>, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let messages = db.collection::<Message>("messages");
+
+        let msg = messages
+            .find_one(doc! { "chat_id": chatid, "message_id": messageid as i64 })
+            .await?;
+
+        Ok(msg)
+    }
+
+    async fn get_message_literal(
+        &mut self,
+        chatid: i64,
+        messageid: i32,
+    ) -> Result<Option<String>, Box<dyn std::error::Error>> {
+        let msg = self.get_message(chatid, messageid).await?;
+        Ok(msg.map(|m| m.token))
+    }
+
+    async fn set_message_literal(
+        &mut self,
+        chatid: i64,
+        messageid: i32,
+        literal: &str,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let messages = db.collection::<Message>("messages");
+
+        messages
+            .update_one(
+                doc! {
+                    "chat_id": chatid,
+                    "message_id": messageid as i64
+                },
+                doc! {
+                    "$set": { "token": literal }
+                },
+            )
+            .upsert(true)
+            .await?;
+
+        Ok(())
+    }
+
+    async fn get_literal(
+        &mut self,
+        literal: &str,
+    ) -> Result<Option<Literal>, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let messages = db.collection::<Literal>("messages");
+
+        let literal = messages.find_one(doc! { "token": literal }).await?;
+
+        Ok(literal)
+    }
+
+    async fn get_literal_value(
+        &mut self,
+        literal: &str,
+    ) -> Result<Option<String>, Box<dyn std::error::Error>> {
+        let literal = self.get_literal(literal).await?;
+
+        Ok(literal.map(|l| l.value))
+    }
+
+    async fn set_literal(
+        &mut self,
+        literal: &str,
+        valuestr: &str,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let literals = db.collection::<Literal>("literals");
+
+        literals
+            .update_one(
+                doc! { "token": literal },
+                doc! { "$set": { "value": valuestr } },
+            )
+            .upsert(true)
+            .await?;
+
+        Ok(())
+    }
+
+    async fn get_all_events(&mut self) -> Vec<Event> {
+        let db = self.get_database().await;
+        let events = db.collection::<Event>("events");
+
+        events
+            .find(doc! {})
+            .await
+            .unwrap()
+            .map(|e| e.unwrap())
+            .collect()
+            .await
+    }
+
+    async fn create_event(
+        &mut self,
+        event_datetime: chrono::DateTime<Utc>,
+    ) -> Result<Event, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let events = db.collection::<Event>("events");
+
+        let new_event = Event {
+            _id: bson::oid::ObjectId::new(),
+            time: event_datetime,
+        };
+
+        events.insert_one(&new_event).await?;
+
+        Ok(new_event)
+    }
+
+    async fn get_media(&mut self, literal: &str) -> Result<Vec<Media>, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let media = db.collection::<Media>("media");
+
+        let media_items = media
+            .find(doc! { "token": literal })
+            .await?
+            .try_collect()
+            .await?;
+
+        Ok(media_items)
+    }
+
+    async fn is_media_group_exists(
+        &mut self,
+        media_group: &str,
+    ) -> Result<bool, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let media = db.collection::<Media>("media");
+
+        let is_exists = media
+            .count_documents(doc! { "media_group_id": media_group })
+            .await?
+            > 0;
+
+        Ok(is_exists)
+    }
+
+    async fn drop_media(&mut self, literal: &str) -> Result<usize, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let media = db.collection::<Media>("media");
+
+        let deleted_count = media
+            .delete_many(doc! { "token": literal })
+            .await?
+            .deleted_count;
+
+        Ok(deleted_count as usize)
+    }
+
+    async fn drop_media_except(
+        &mut self,
+        literal: &str,
+        except_group: &str,
+    ) -> Result<usize, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let media = db.collection::<Media>("media");
+
+        let deleted_count = media
+            .delete_many(doc! {
+                "token": literal,
+                "media_group_id": { "$ne": except_group }
+            })
+            .await?
+            .deleted_count;
+
+        Ok(deleted_count as usize)
+    }
+
+    async fn add_media(
+        &mut self,
+        literal: &str,
+        mediatype: &str,
+        fileid: &str,
+        media_group: Option<&str>,
+    ) -> Result<Media, Box<dyn std::error::Error>> {
+        let db = self.get_database().await;
+        let media = db.collection::<Media>("media");
+
+        let new_media = Media {
+            _id: bson::oid::ObjectId::new(),
+            token: literal.to_string(),
+            media_type: mediatype.to_string(),
+            file_id: fileid.to_string(),
+            media_group_id: media_group.map(|g| g.to_string()),
+        };
+
+        media.insert_one(&new_media).await?;
+
+        Ok(new_media)
+    }
+}
+
+#[cfg(test)]
+mod tests;

src/db/models.rs

@@ -31,6 +31,7 @@ pub struct Media {
     pub token: String,
     pub media_type: String,
     pub file_id: String,
+    pub media_group_id: Option<String>,
 }
 
 #[derive(Queryable, Debug, Identifiable)]

src/db/schema.rs

@@ -22,6 +22,7 @@ diesel::table! {
         token -> Varchar,
         media_type -> Varchar,
         file_id -> Varchar,
+        media_group_id -> Nullable<Varchar>,
     }
 }
 

src/db/tests.rs

@@ -1,7 +1,6 @@
-use diesel::Connection;
-use diesel_async::AsyncPgConnection;
 use dotenvy;
 
+use super::CallDB;
 use super::DB;
 
 async fn setup_db() -> DB {
@@ -16,31 +15,29 @@ async fn setup_db() -> DB {
 async fn test_get_media() {
     let mut db = setup_db().await;
 
-    let result = db.drop_media("test_get_media_literal").await;
+    let result = db.drop_media("test_get_media_literal").await.unwrap();
-    assert!(result.is_ok());
 
     let media_items = db.get_media("test_get_media_literal").await.unwrap();
     assert_eq!(media_items.len(), 0);
 
     let result = db
-        .add_media("test_get_media_literal", "photo", "file_id_1")
+        .add_media("test_get_media_literal", "photo", "file_id_1", None)
-        .await;
+        .await
-    assert!(result.is_ok());
+        .unwrap();
 
     let media_items = db.get_media("test_get_media_literal").await.unwrap();
     assert_eq!(media_items.len(), 1);
 
     let result = db
-        .add_media("test_get_media_literal", "video", "file_id_2")
+        .add_media("test_get_media_literal", "video", "file_id_2", None)
-        .await;
+        .await
-    assert!(result.is_ok());
+        .unwrap();
 
     let media_items = db.get_media("test_get_media_literal").await.unwrap();
     assert_eq!(media_items.len(), 2);
 
     // Clean up after test
-    let result = db.drop_media("test_get_media_literal").await;
+    let result = db.drop_media("test_get_media_literal").await.unwrap();
-    assert!(result.is_ok());
 }
 
 #[tokio::test]
@@ -51,11 +48,12 @@ async fn test_add_media() {
     let media_type = "photo";
     let file_id = "LjaldhAOh";
 
-    let result = db.drop_media(literal).await;
+    let result = db.drop_media(literal).await.unwrap();
-    assert!(result.is_ok());
 
-    let result = db.add_media(literal, media_type, file_id).await;
+    let result = db
-    assert!(result.is_ok());
+        .add_media(literal, media_type, file_id, None)
+        .await
+        .unwrap();
 
     // Verify that the media was added is correct
     let media_items = db.get_media(literal).await.unwrap();
@@ -65,8 +63,7 @@ async fn test_add_media() {
     assert_eq!(media_items[0].file_id, file_id);
 
     // Clean up after test
-    let result = db.drop_media(literal).await;
+    let result = db.drop_media(literal).await.unwrap();
-    assert!(result.is_ok());
 }
 
 #[tokio::test]
@@ -74,22 +71,104 @@ async fn test_drop_media() {
     let mut db = setup_db().await;
 
     let result = db
-        .add_media("test_drop_media_literal", "photo", "file_id_1")
+        .add_media("test_drop_media_literal", "photo", "file_id_1", None)
-        .await;
+        .await
-    assert!(result.is_ok());
+        .unwrap();
 
     // Verify that the media was added
     let media_items = db.get_media("test_drop_media_literal").await.unwrap();
     assert_eq!(media_items.len(), 1);
 
-    let result = db.drop_media("test_drop_media_literal").await;
+    let result = db.drop_media("test_drop_media_literal").await.unwrap();
-    assert!(result.is_ok());
 
     // Verify that the media has been dropped
     let media_items = db.get_media("test_drop_media_literal").await.unwrap();
     assert_eq!(media_items.len(), 0);
 
     // Clean up after test
-    let result = db.drop_media("test_drop_media_literal").await;
+    let result = db.drop_media("test_drop_media_literal").await.unwrap();
-    assert!(result.is_ok());
+}
+
+#[tokio::test]
+async fn test_is_media_group_exists() {
+    let mut db = setup_db().await;
+
+    let media_group = "test_media_group";
+    let literal = "test_media_group_literal";
+
+    let _ = db.drop_media(literal).await.unwrap();
+
+    let exists = db.is_media_group_exists(media_group).await.unwrap();
+    assert!(!exists);
+
+    let _ = db
+        .add_media(literal, "photo", "file_id_1", Some(media_group))
+        .await
+        .unwrap();
+
+    let exists = db.is_media_group_exists(media_group).await.unwrap();
+    assert!(exists);
+
+    let _ = db.drop_media(literal).await.unwrap();
+
+    let exists = db.is_media_group_exists(media_group).await.unwrap();
+    assert!(!exists);
+}
+
+#[tokio::test]
+async fn test_drop_media_except() {
+    let mut db = setup_db().await;
+
+    let media_group = "test_media_group_except";
+    let literal = "test_media_group_except_literal";
+    let _ = db.drop_media(literal).await.unwrap();
+
+    let _ = db
+        .add_media(literal, "photo", "file_id_2", None)
+        .await
+        .unwrap();
+    let _ = db
+        .add_media(literal, "photo", "file_id_3", None)
+        .await
+        .unwrap();
+
+    let media_items = db.get_media(literal).await.unwrap();
+    assert_eq!(media_items.len(), 2);
+
+    let deleted_count = db.drop_media_except(literal, media_group).await.unwrap();
+    let media_items = db.get_media(literal).await.unwrap();
+    assert_eq!(media_items.len(), 0);
+
+    let _ = db
+        .add_media(literal, "photo", "file_id_1", Some(media_group))
+        .await
+        .unwrap();
+    let _ = db
+        .add_media(literal, "photo", "file_id_2", None)
+        .await
+        .unwrap();
+    let _ = db
+        .add_media(literal, "photo", "file_id_3", None)
+        .await
+        .unwrap();
+
+    let deleted_count = db.drop_media_except(literal, media_group).await.unwrap();
+    let media_items = db.get_media(literal).await.unwrap();
+    assert_eq!(media_items.len(), 1);
+    let _ = db.drop_media(literal).await.unwrap();
+
+    let _ = db
+        .add_media(literal, "photo", "file_id_1", Some(media_group))
+        .await
+        .unwrap();
+    let _ = db
+        .add_media(literal, "photo", "file_id_2", Some(media_group))
+        .await
+        .unwrap();
+
+    let deleted_count = db.drop_media_except(literal, media_group).await.unwrap();
+    let media_items = db.get_media(literal).await.unwrap();
+    assert_eq!(media_items.len(), 2);
+
+    let _ = db.drop_media(literal).await.unwrap();
 }

src/main.rs

@@ -1,24 +1,25 @@
 pub mod admin;
 pub mod db;
 
+use std::time::Duration;
+
 use crate::admin::{admin_command_handler, AdminCommands};
 use crate::admin::{secret_command_handler, SecretCommands};
-use crate::db::DB;
+use crate::db::{CallDB, DB};
 
 use chrono::{DateTime, Utc};
 use chrono_tz::Asia;
-use db::schema::events;
 use envconfig::Envconfig;
 use serde::{Deserialize, Serialize};
 use teloxide::dispatching::dialogue::serializer::Json;
-use teloxide::dispatching::dialogue::{GetChatId, InMemStorage, PostgresStorage};
+use teloxide::dispatching::dialogue::{GetChatId, PostgresStorage};
 use teloxide::types::{
-    InlineKeyboardButton, InlineKeyboardMarkup, MediaKind, MessageKind, ReplyMarkup,
+    InlineKeyboardButton, InlineKeyboardMarkup, InputFile, InputMedia, MediaKind, MessageKind,
+    ParseMode, ReplyMarkup,
 };
 use teloxide::{
     payloads::SendMessageSetters,
     prelude::*,
-    types::InputFile,
     utils::{command::BotCommands, render::RenderMessageTextHelper},
 };
 
@@ -61,6 +62,7 @@ pub enum State {
     Edit {
         literal: String,
         lang: String,
+        is_caption_set: bool,
     },
 }
 
@@ -82,7 +84,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     for event in events {
         match db.clone().create_event(event).await {
-            Ok(e) => println!("Created event {}", e.id),
+            Ok(e) => println!("Created event {}", e._id),
             Err(err) => println!("Failed to create event, error: {}", err),
         }
     }
@@ -111,7 +113,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                         })
                         .endpoint(edit_msg_cmd_handler),
                 )
-                .branch(dptree::case![State::Edit { literal, lang }].endpoint(edit_msg_handler)),
+                .branch(
+                    dptree::case![State::Edit {
+                        literal,
+                        lang,
+                        is_caption_set
+                    }]
+                    .endpoint(edit_msg_handler),
+                ),
         )
         .branch(Update::filter_message().endpoint(echo));
 
@@ -178,7 +187,11 @@ async fn edit_msg_cmd_handler(
             // TODO: language selector will be implemented in future 😈
             let lang = "ru".to_string();
             dialogue
-                .update(State::Edit { literal, lang })
+                .update(State::Edit {
+                    literal,
+                    lang,
+                    is_caption_set: false,
+                })
                 .await
                 .unwrap();
             bot.send_message(
@@ -199,7 +212,7 @@ async fn edit_msg_handler(
     bot: Bot,
     mut db: DB,
     dialogue: BotDialogue,
-    (literal, lang): (String, String),
+    (literal, lang, is_caption_set): (String, String, bool),
     msg: Message,
 ) -> Result<(), teloxide::RequestError> {
     use teloxide::utils::render::Renderer;
@@ -215,12 +228,117 @@ async fn edit_msg_handler(
 
     match msg.media_kind {
         MediaKind::Text(text) => {
+            if is_caption_set {
+                return Ok(());
+            };
             let html_text = Renderer::new(&text.text, &text.entities).as_html();
             db.set_literal(&literal, &html_text).await.unwrap();
             bot.send_message(chat_id, "Updated text of message!")
                 .await?;
             dialogue.exit().await.unwrap();
         }
+        MediaKind::Photo(photo) => {
+            let group = photo.media_group_id;
+            if let Some(group) = group.clone() {
+                db.drop_media_except(&literal, &group).await.unwrap();
+            } else {
+                db.drop_media(&literal).await.unwrap();
+            }
+            let file_id = photo.photo[0].file.id.clone();
+            db.add_media(&literal, "photo", &file_id, group.as_deref())
+                .await
+                .unwrap();
+            match photo.caption {
+                Some(text) => {
+                    let html_text = Renderer::new(&text, &photo.caption_entities).as_html();
+                    db.set_literal(&literal, &html_text).await.unwrap();
+                    bot.send_message(chat_id, "Updated photo caption!").await?;
+                }
+                None => {
+                    // if it is a first message in group,
+                    // or just a photo without caption (unwrap_or case),
+                    // set text empty
+                    if !db
+                        .is_media_group_exists(group.as_deref().unwrap_or(""))
+                        .await
+                        .unwrap()
+                    {
+                        db.set_literal(&literal, "").await.unwrap();
+                        bot.send_message(chat_id, "Set photo without caption")
+                            .await?;
+                    };
+                }
+            }
+            // Some workaround because Telegram's group system
+            // is not easily and obviously handled with this
+            // code architecture, but probably there is a solution.
+            //
+            // So, this code will just wait for all media group
+            // updates to be processed
+            dialogue
+                .update(State::Edit {
+                    literal,
+                    lang,
+                    is_caption_set: true,
+                })
+                .await
+                .unwrap();
+            tokio::spawn(async move {
+                tokio::time::sleep(Duration::from_millis(200)).await;
+                dialogue.exit().await.unwrap_or(());
+            });
+        }
+        MediaKind::Video(video) => {
+            let group = video.media_group_id;
+            if let Some(group) = group.clone() {
+                db.drop_media_except(&literal, &group).await.unwrap();
+            } else {
+                db.drop_media(&literal).await.unwrap();
+            }
+            let file_id = video.video.file.id;
+            db.add_media(&literal, "video", &file_id, group.as_deref())
+                .await
+                .unwrap();
+            match video.caption {
+                Some(text) => {
+                    let html_text = Renderer::new(&text, &video.caption_entities).as_html();
+                    db.set_literal(&literal, &html_text).await.unwrap();
+                    bot.send_message(chat_id, "Updated video caption!").await?;
+                }
+                None => {
+                    // if it is a first message in group,
+                    // or just a video without caption (unwrap_or case),
+                    // set text empty
+                    if !db
+                        .is_media_group_exists(group.as_deref().unwrap_or(""))
+                        .await
+                        .unwrap()
+                    {
+                        db.set_literal(&literal, "").await.unwrap();
+                        bot.send_message(chat_id, "Set video without caption")
+                            .await?;
+                    };
+                }
+            }
+            // Some workaround because Telegram's group system
+            // is not easily and obviously handled with this
+            // code architecture, but probably there is a solution.
+            //
+            // So, this code will just wait for all media group
+            // updates to be processed
+            dialogue
+                .update(State::Edit {
+                    literal,
+                    lang,
+                    is_caption_set: true,
+                })
+                .await
+                .unwrap();
+            tokio::spawn(async move {
+                tokio::time::sleep(Duration::from_millis(200)).await;
+                dialogue.exit().await.unwrap_or(());
+            });
+        }
         _ => {
             bot.send_message(chat_id, "this type of message is not supported yet")
                 .await?;
@@ -307,15 +425,119 @@ async fn answer_message<RM: Into<ReplyMarkup>>(
         .await
         .unwrap()
         .unwrap_or("Please, set content of this message".into());
-    let msg = bot.send_message(ChatId(chat_id), text);
+    let media = db.get_media(&literal).await.unwrap();
-    let msg = match keyboard {
+    let (chat_id, msg_id) = match media.len() {
-        Some(kbd) => msg.reply_markup(kbd),
+        // just a text
-        None => msg,
+        0 => {
+            let msg = bot.send_message(ChatId(chat_id), text);
+            let msg = match keyboard {
+                Some(kbd) => msg.reply_markup(kbd),
+                None => msg,
+            };
+            let msg = msg.parse_mode(teloxide::types::ParseMode::Html);
+            println!("ENTS: {:?}", msg.entities);
+            let msg = msg.await?;
+
+            (msg.chat.id.0, msg.id.0)
+        }
+        // single media
+        1 => {
+            let media = &media[0]; // safe, cause we just checked len
+            match media.media_type.as_str() {
+                "photo" => {
+                    let msg = bot.send_photo(
+                        ChatId(chat_id),
+                        InputFile::file_id(media.file_id.to_string()),
+                    );
+                    let msg = match text.as_str() {
+                        "" => msg,
+                        text => msg.caption(text),
+                    };
+                    let msg = match keyboard {
+                        Some(kbd) => msg.reply_markup(kbd),
+                        None => msg,
+                    };
+
+                    let msg = msg.parse_mode(teloxide::types::ParseMode::Html);
+                    let msg = msg.await?;
+
+                    (msg.chat.id.0, msg.id.0)
+                }
+                "video" => {
+                    let msg = bot.send_video(
+                        ChatId(chat_id),
+                        InputFile::file_id(media.file_id.to_string()),
+                    );
+                    let msg = match text.as_str() {
+                        "" => msg,
+                        text => msg.caption(text),
+                    };
+                    let msg = match keyboard {
+                        Some(kbd) => msg.reply_markup(kbd),
+                        None => msg,
+                    };
+
+                    let msg = msg.parse_mode(teloxide::types::ParseMode::Html);
+                    let msg = msg.await?;
+
+                    (msg.chat.id.0, msg.id.0)
+                }
+                _ => {
+                    todo!()
+                }
+            }
+        }
+        // >= 2, should use media group
+        _ => {
+            let media: Vec<InputMedia> = media
+                .into_iter()
+                .enumerate()
+                .map(|(i, m)| {
+                    let ifile = InputFile::file_id(m.file_id);
+                    let caption = if i == 0 {
+                        match text.as_str() {
+                            "" => None,
+                            text => Some(text.to_string()),
+                        }
+                    } else {
+                        None
+                    };
+                    match m.media_type.as_str() {
+                        "photo" => InputMedia::Photo(teloxide::types::InputMediaPhoto {
+                            media: ifile,
+                            caption,
+                            parse_mode: Some(ParseMode::Html),
+                            caption_entities: None,
+                            has_spoiler: false,
+                            show_caption_above_media: false,
+                        }),
+                        "video" => InputMedia::Video(teloxide::types::InputMediaVideo {
+                            media: ifile,
+                            thumbnail: None,
+                            caption,
+                            parse_mode: Some(ParseMode::Html),
+                            caption_entities: None,
+                            show_caption_above_media: false,
+                            width: None,
+                            height: None,
+                            duration: None,
+                            supports_streaming: None,
+                            has_spoiler: false,
+                        }),
+                        _ => {
+                            todo!()
+                        }
+                    }
+                })
+                .collect();
+            let msg = bot.send_media_group(ChatId(chat_id), media);
+
+            let msg = msg.await?;
+
+            (msg[0].chat.id.0, msg[0].id.0)
+        }
     };
-    let msg = msg.parse_mode(teloxide::types::ParseMode::Html);
+    db.set_message_literal(chat_id, msg_id, literal)
-    println!("ENTS: {:?}", msg.entities);
-    let msg = msg.await?;
-    db.set_message_literal(msg.chat.id.0, msg.id.0, literal)
         .await
         .unwrap();
     Ok(())
@@ -329,7 +551,7 @@ async fn make_start_buttons(db: &mut DB) -> InlineKeyboardMarkup {
         .map(|e| {
             vec![InlineKeyboardButton::callback(
                 e.time.with_timezone(&Asia::Dubai).to_string(),
-                format!("event:{}", e.id),
+                format!("event:{}", e._id),
             )]
         })
         .collect();