Actually save things! Yay!

src/main.rs

@@ -1,50 +1,72 @@
 use std::{net::SocketAddr, str::FromStr};
 
 use axum::{
+    extract::State,
     http::{HeaderMap, StatusCode},
     routing::post,
     Json, Router,
 };
-use base64::{engine, alphabet, Engine};
+use base64::{alphabet, engine, Engine};
 use ed25519_dalek::{Signature, VerifyingKey};
 use serde::Deserialize;
+use sqlx::{postgres::PgPoolOptions, PgPool};
 use twilight_http::Client;
-use twilight_interactions::command::{CommandModel, CreateCommand, CommandInputData};
+use twilight_interactions::command::{CommandInputData, CommandModel, CreateCommand};
+use twilight_mention::{timestamp::{Timestamp, TimestampStyle}, Mention};
 use twilight_model::{
-    application::{
-        interaction::{Interaction, InteractionType, InteractionData},
-    },
-    http::interaction::{InteractionResponse, InteractionResponseType, InteractionResponseData}, id::Id,
+    application::interaction::{Interaction, InteractionData, InteractionType},
+    http::interaction::{InteractionResponse, InteractionResponseData, InteractionResponseType},
+    id::{Id, marker::{UserMarker, ChannelMarker, InteractionMarker}}, channel::message::MessageFlags,
 };
 
-
 #[derive(CommandModel, CreateCommand)]
-#[command(name="save_fact", desc="Quietly save a fact")]
-struct SaveFactCommand {
-    #[command(rename="name", desc="Fact name")]
+#[command(name = "set_fact", desc = "Quietly save a fact")]
+struct SetFactCommand {
+    #[command(rename = "name", desc = "Fact name")]
     fact_name: String,
-    #[command(rename="value", desc="Fact value")]
+    #[command(rename = "value", desc = "Fact value")]
     fact_value: String,
 }
 
+#[derive(CommandModel, CreateCommand)]
+#[command(name = "get_fact", desc = "Retrieve and display the value of a fact")]
+struct GetFactCommand {
+    #[command(rename = "name", desc = "Fact name")]
+    fact_name: String,
+    #[command(desc = "Should it be displayed publically, by default it won't be")]
+    public: Option<bool>,
+}
+
 #[tokio::main]
-async fn main() {
+async fn main() -> anyhow::Result<()> {
     let port = 4635;
-    let app = Router::new().route("/", post(post_interaction));
+    dotenvy::dotenv().ok();
+
+    let pg_pool = PgPoolOptions::new()
+        .max_connections(5)
+        .connect(database_url().as_str())
+        .await?;
+
+    sqlx::migrate!().run(&pg_pool).await?;
+
+    let app = Router::new()
+        .route("/", post(post_interaction))
+        .with_state(pg_pool);
 
     let addr = SocketAddr::from(([127, 0, 0, 1], port));
-    dotenvy::dotenv().ok();
-    register_command().await;
+
+    register_commands().await;
 
     axum::Server::bind(&addr)
         .serve(app.into_make_service())
-        .await
-        .unwrap();
+        .await?;
+
+    Ok(())
 }
 
 type InteractionResult = Result<(StatusCode, Json<InteractionResponse>), (StatusCode, String)>;
 
-async fn post_interaction(headers: HeaderMap, body: String) -> InteractionResult {
+fn validate_request(headers: HeaderMap, body: String) -> Result<Interaction, (StatusCode, String)> {
     let Ok(interaction): Result<Interaction, _> = serde_json::from_str(&body) else {
         return Err((StatusCode::BAD_REQUEST, "request contained invalid json".to_string()))
     };
@@ -68,6 +90,119 @@ async fn post_interaction(headers: HeaderMap, body: String) -> InteractionResult
         return Err((StatusCode::UNAUTHORIZED, "interaction failed signature verification".to_string()))
     };
 
+    return Ok(interaction);
+}
+
+async fn set_fact(
+    interaction_id: Id<InteractionMarker>,
+    channel_id: Option<Id<ChannelMarker>>,
+    author_id: Id<UserMarker>,
+    command_data: SetFactCommand,
+    pg_pool: &PgPool,
+) -> Result<InteractionResponse, (StatusCode, String)> {
+    let Ok(rows) = sqlx::query!("
+INSERT INTO facts (\"last_interaction_id\", \"channel_id\", \"author_id\", \"name\", \"value\") 
+VALUES ($1, $2, $3, $4, $5)
+ON CONFLICT ON CONSTRAINT facts_origin_key DO UPDATE SET value = $5, version = facts.version + 1
+    ",
+        interaction_id.to_string(),
+        channel_id.map(|cid| cid.to_string()),
+        author_id.to_string(),
+        command_data.fact_name,
+        command_data.fact_value,
+    ).execute(pg_pool).await.and_then(|rows| Ok(rows.rows_affected())) else {
+        return Err((StatusCode::INTERNAL_SERVER_ERROR, "Error saving fact.".to_string()));
+    };
+
+    if rows != 1 {
+        return Err((StatusCode::INTERNAL_SERVER_ERROR, "Error saving fact".to_string()));
+    }
+
+    Ok(InteractionResponse {
+        kind: InteractionResponseType::ChannelMessageWithSource,
+        data: Some(InteractionResponseData {
+            content: Some(format!(
+                "Set {0} to {1}",
+                command_data.fact_name, command_data.fact_value
+            )),
+            flags: Some(MessageFlags::EPHEMERAL),
+            ..Default::default()
+        }),
+    })
+}
+
+struct FactResponse {
+    value: String,
+    version: i32,
+    created_at: time::OffsetDateTime,
+    updated_at: time::OffsetDateTime,
+}
+
+async fn get_fact(
+    channel_id: Option<Id<ChannelMarker>>,
+    author_id: Id<UserMarker>,
+    command_data: GetFactCommand,
+    pg_pool: &PgPool,
+) -> Result<InteractionResponse, (StatusCode, String)> {
+    let Ok(facts) = sqlx::query_as!(FactResponse, 
+        "
+SELECT \"value\", \"version\", \"created_at\", \"updated_at\"
+FROM facts
+WHERE
+    channel_id IS NOT DISTINCT FROM $1 AND
+    author_id = $2 AND
+    name = $3
+        ", channel_id.map(|cid| cid.to_string()), author_id.to_string(), command_data.fact_name).fetch_all(pg_pool).await else {
+            return Err((StatusCode::INTERNAL_SERVER_ERROR, "Querying facts failed".to_string()));
+        };
+    if facts.len() == 0 {        
+        return Ok(InteractionResponse { 
+            kind: InteractionResponseType::ChannelMessageWithSource, 
+            data: Some(InteractionResponseData { 
+                content: Some(format!("Fact {0} in channel {1} by you, {2}, was not found.",
+                    command_data.fact_name,
+                    channel_id.map_or("<none>".to_string(), |cid| cid.mention().to_string()),
+                    author_id.mention().to_string(),
+                )),
+                flags: match command_data.public { Some(true) => None, _ => Some(MessageFlags::EPHEMERAL) },
+                ..Default::default()
+            }),
+        });
+    }
+    if facts.len() > 1 {
+        return Err((StatusCode::INTERNAL_SERVER_ERROR, "Too many facts found, wtf, impossible".to_string()));
+    }
+
+    let fact = &facts[0];
+
+    Ok(InteractionResponse {
+        kind: InteractionResponseType::ChannelMessageWithSource,
+        data: Some(InteractionResponseData {
+            content: Some(format!(
+                "Fact **{0}** was set to **{1}** by {2} at {3}, and was reset {4} times in total since {5}.",
+                command_data.fact_name, 
+                fact.value, 
+                author_id.mention().to_string(),
+                Timestamp::new(fact.updated_at.unix_timestamp().try_into().unwrap(), Some(TimestampStyle::RelativeTime)).mention(),
+                fact.version,
+                Timestamp::new(fact.created_at.unix_timestamp().try_into().unwrap(), Some(TimestampStyle::ShortDateTime)).mention(),
+            )),
+            flags: match command_data.public { Some(true) => None, _ => Some(MessageFlags::EPHEMERAL) },
+            ..Default::default()
+        }),
+    })
+}
+
+async fn post_interaction(
+    headers: HeaderMap,
+    State(pg_pool): State<PgPool>,
+    body: String,
+) -> InteractionResult {
+    let interaction = match validate_request(headers, body) {
+        Ok(interaction) => interaction,
+        Err(error) => return Err(error),
+    };
+
     match interaction.kind {
         InteractionType::Ping => {
             let pong = InteractionResponse {
@@ -77,25 +212,38 @@ async fn post_interaction(headers: HeaderMap, body: String) -> InteractionResult
             Ok((StatusCode::OK, Json(pong)))
         }
         InteractionType::ApplicationCommand => {
+            let author_id = interaction.author_id();
             let Some(InteractionData::ApplicationCommand(data)) = interaction.data else {
                 return not_found();
             };
             let command_input_data = CommandInputData::from(*data.clone());
             match &*data.name {
-                "save_fact" => {
-                    let Ok(command_data) = SaveFactCommand::from_interaction(command_input_data) else {
-                        return Err((StatusCode::BAD_REQUEST, "invalid save fact command".to_string()));
+                "set_fact" => {
+                    let Ok(command_data) = SetFactCommand::from_interaction(command_input_data) else {
+                        return Err((StatusCode::BAD_REQUEST, "invalid set fact command".to_string()));
                     };
-                    let reply = InteractionResponse {
-                        kind: InteractionResponseType::ChannelMessageWithSource,
-                        data: Some(InteractionResponseData {
-                            content: Some(format!("Set {0} to {1}", command_data.fact_name, command_data.fact_value)),
-                            ..Default::default()
-                        })
+                    let Some(author_id) = author_id else {
+                        return Err((StatusCode::BAD_REQUEST, "save_fact requires a user".to_string()));
                     };
-                    Ok((StatusCode::OK, Json(reply)))
-                }
-                _ =>  not_found(),
+                    match set_fact(interaction.id, interaction.channel_id, author_id, command_data, &pg_pool).await {
+                        Ok(response) => Ok((StatusCode::OK, Json(response))),
+                        Err(err) => Err(err),
+                    }
+                },
+                "get_fact" => {
+                    let Ok(command_data) = GetFactCommand::from_interaction(command_input_data) else {
+                        return Err((StatusCode::BAD_REQUEST, "invalid get fact command".to_string()));
+                    };
+                    let Some(author_id) = author_id else {
+                        return Err((StatusCode::BAD_REQUEST, "get_fact requires a user".to_string()));
+                    };
+                    
+                    match get_fact(interaction.channel_id, author_id, command_data, &pg_pool).await {
+                        Ok(response) => Ok((StatusCode::OK, Json(response))),
+                        Err(err) => Err(err),
+                    }
+                },
+                _ => not_found(),
             }
         }
         _ => not_found(),
@@ -118,10 +266,13 @@ fn discord_pub_key() -> VerifyingKey {
     VerifyingKey::from_bytes(&pub_key_bytes).unwrap()
 }
 
-async fn register_command() {
+async fn register_commands() {
     discord_client()
         .interaction(Id::from_str(&discord_client_id()).unwrap())
-        .set_global_commands(&[SaveFactCommand::create_command().into()])
+        .set_global_commands(&[
+            GetFactCommand::create_command().into(),
+            SetFactCommand::create_command().into(),
+            ])
         .await
         .unwrap();
 }
@@ -133,11 +284,7 @@ struct ClientCredentialsResponse {
 
 fn authorization() -> String {
     let engine = engine::GeneralPurpose::new(&alphabet::STANDARD, engine::general_purpose::PAD);
-    let auth = format!(
-        "{}:{}",
-        discord_client_id(),
-        discord_client_secret(),
-    );
+    let auth = format!("{}:{}", discord_client_id(), discord_client_secret(),);
     engine.encode(auth)
 }
 
@@ -147,7 +294,10 @@ fn client_credentials_grant() -> ClientCredentialsResponse {
         .send_form(&[
             ("grant_type", "client_credentials"),
             ("scope", "applications.commands.update"),
-        ]).unwrap().into_json().unwrap()
+        ])
+        .unwrap()
+        .into_json()
+        .unwrap()
 }
 
 fn discord_client_id() -> String {
@@ -162,3 +312,7 @@ fn discord_client() -> Client {
     let token = client_credentials_grant().access_token;
     Client::new(format!("Bearer {token}"))
 }
+
+fn database_url() -> String {
+    std::env::var("DATABASE_URL").unwrap()
+}