Checkpoint

src/groups/moderation/arguments.rs

@@ -1,20 +1,14 @@
-use std::{
-    ops::{Deref, DerefMut},
-    time::Duration,
-};
+use std::ops::{Deref, DerefMut};
+use std::time::Duration;
 
 use scylla::client::session::Session;
 
-use crate::{
-    database::models::{Message, User},
-    extras::Extras,
-    framework::commands::{
-        arguments::{ArgumentParsingError, FromArguments, NeedyString},
-        context::CommandContext,
-    },
-    state::State,
-    utils::split_at_nth_from_end,
-};
+use crate::database::models::{Message, User};
+use crate::extras::Extras;
+use crate::framework::commands::arguments::{ArgumentParsingError, FromArguments, NeedyString};
+use crate::framework::commands::context::CommandContext;
+use crate::state::State;
+use crate::utils::split_at_nth_from_end;
 
 /// Fetch a user by their ID from the database.
 ///
@@ -110,7 +104,24 @@ impl FromArguments<State, Extras> for UserArgument {
                     .parse()
                     .map_err(|_| ArgumentParsingError::InvalidArgument)?;
 
-                user = get_user_by_id(user_id, &ctx.rt.state.db).await?;
+                match get_user_by_id(user_id, &ctx.rt.state.db).await {
+                    Ok(result) => user = result,
+                    Err(error) => {
+                        if matches!(error, ArgumentParsingError::InvalidArgument) {
+                            let message =
+                                Message::get_by_propagation_message_id(user_id, &ctx.rt.state.db)
+                                    .await?;
+
+                            if let Some(message) = message {
+                                user = get_user_by_id(message.author_id, &ctx.rt.state.db).await?;
+                            } else {
+                                return Err(ArgumentParsingError::InvalidArgument);
+                            }
+                        } else {
+                            return Err(error);
+                        }
+                    }
+                }
             } else {
                 // The argument is (probably) a username.