Try adding some diesel metrics

Author: Eijebong

Cargo.lock

@@ -15,7 +15,10 @@ diesel_migrations = { version = "2.1.0", features = ["sqlite"] }
 dotenvy = "0.15.0"
 either = "1.9.0"
 encoding_rs = "0.8.33"
+headers = "0.3.0"
 itertools = "0.12.1"
+once_cell = "1.19.0"
+prometheus = { version = "0.13.4", default-features = false }
 reqwest = { version = "0.11.24", features = ["multipart", "json"] }
 rocket = { version = "0.5.0", features = ["uuid", "secrets"] }
 rocket_oauth2 = "0.5.0"
@@ -41,3 +44,4 @@ walkdir = "2.5.0"
 
 [patch.crates-io]
 rocket_prometheus = { git = "https://github.com/Eijebong/rocket_prometheus.git", branch = "add-request-filter" }
+diesel = { git = "https://github.com/diesel-rs/diesel.git" }

Cargo.toml

@@ -1,14 +1,18 @@
 use std::collections::HashMap;
+use std::time::Instant;
 
 use crate::diesel_uuid::Uuid as DieselUuid;
 use crate::error::Result;
 use crate::schema::{discord_users, rooms, yamls};
 use crate::Context;
 
 use chrono::NaiveDateTime;
+use diesel::connection::Instrumentation;
 use diesel::dsl::{exists, now, AsSelect, SqlTypeOf};
 use diesel::prelude::*;
 use diesel::sqlite::Sqlite;
+use once_cell::sync::Lazy;
+use prometheus::{HistogramOpts, HistogramVec};
 use rocket::State;
 use uuid::Uuid;
 
@@ -344,3 +348,40 @@ impl RoomFilter {
         self
     }
 }
+
+#[derive(Default)]
+pub struct DbInstrumentation {
+    query_start: Option<Instant>,
+}
+
+pub(crate) static QUERY_HISTOGRAM: Lazy<HistogramVec> = Lazy::new(|| {
+    HistogramVec::new(
+        HistogramOpts::new("diesel_query_seconds", "SQL query duration").buckets(vec![
+            0.000005, 0.00001, 0.00005, 0.0001, 0.0005, 0.001, 0.005, 0.01, 0.1, 1.0,
+        ]),
+        &["query"],
+    )
+    .expect("Failed to create query histogram")
+});
+
+impl Instrumentation for DbInstrumentation {
+    fn on_connection_event(&mut self, event: diesel::connection::InstrumentationEvent<'_>) {
+        match event {
+            diesel::connection::InstrumentationEvent::StartQuery { .. } => {
+                self.query_start = Some(Instant::now());
+            }
+            diesel::connection::InstrumentationEvent::FinishQuery { query, .. } => {
+                let Some(query_start) = self.query_start else {
+                    return;
+                };
+                let elapsed = Instant::now().duration_since(query_start);
+                let query = query.to_string().replace('\n', " ");
+                let query = query.split("--").next().unwrap().trim();
+                QUERY_HISTOGRAM
+                    .with_label_values(&[query])
+                    .observe(elapsed.as_secs_f64());
+            }
+            _ => {}
+        };
+    }
+}

src/db.rs

@@ -1,6 +1,7 @@
 use diesel::backend::Backend;
 use diesel::deserialize::FromSqlRow;
 use diesel::expression::AsExpression;
+use diesel::serialize::ToSql;
 use diesel::sql_types::{Binary, SqlType};
 use std::fmt;
 use std::fmt::{Display, Formatter};
@@ -48,6 +49,6 @@ where
         out: &mut diesel::serialize::Output<'b, '_, DB>,
     ) -> diesel::serialize::Result {
         // XXX: This is probably wrong
-        self.0.as_bytes().to_sql(out)
+        <[u8; 16] as ToSql<Binary, DB>>::to_sql(self.0.as_bytes(), out)
     }
 }

src/diesel_uuid.rs

@@ -1,3 +1,4 @@
+use db::{DbInstrumentation, QUERY_HISTOGRAM};
 use diesel::r2d2::Pool;
 use diesel::sqlite::Sqlite;
 use diesel::{r2d2::ConnectionManager, SqliteConnection};
@@ -116,6 +117,11 @@ fn rocket() -> _ {
     let admin_token =
         AdminToken(std::env::var("ADMIN_TOKEN").expect("Plox provide a ADMIN_TOKEN env variable"));
 
+    diesel::connection::set_default_instrumentation(|| {
+        Some(Box::new(DbInstrumentation::default()))
+    })
+    .expect("Failed to set diesel instrumentation");
+
     let manager = ConnectionManager::<SqliteConnection>::new(db_url);
     let db_pool = Pool::new(manager).expect("Failed to create database pool, aborting");
     {
@@ -145,6 +151,10 @@ fn rocket() -> _ {
     let figment = rocket::Config::figment().merge(("limits", limits));
     let prometheus =
         PrometheusMetrics::new().with_request_filter(|request| request.uri().path() != "/metrics");
+    prometheus
+        .registry()
+        .register(Box::new(QUERY_HISTOGRAM.clone()))
+        .expect("Failed to register query histogram");
 
     rocket::custom(figment.clone())
         .attach(prometheus.clone())

src/main.rs

@@ -3,6 +3,7 @@ use std::str::FromStr;
 use crate::error::Result;
 use crate::{AdminToken, Context, Discord};
 use anyhow::anyhow;
+use headers::authorization::{Basic, Credentials};
 use reqwest::header::HeaderValue;
 use reqwest::Url;
 use rocket::figment::{Figment, Profile, Provider};
@@ -57,8 +58,24 @@ impl LoggedInSession {
 impl Session {
     pub fn from_request_sync(request: &Request) -> Self {
         let admin_token = request.rocket().state::<AdminToken>();
-        let x_api_key = request.headers().get("X-Api-Key").next();
-        if x_api_key == admin_token.map(|t| t.0.as_str()) {
+        let admin_token = admin_token.map(|t| t.0.as_str());
+
+        let authorization = request.headers().get_one("Authorization");
+        if let Some(authorization) = authorization {
+            let creds = Basic::decode(&HeaderValue::from_str(authorization).unwrap());
+            if let Some(creds) = creds {
+                if creds.username() == "admin" && Some(creds.password()) == admin_token {
+                    return Session {
+                        is_admin: true,
+                        is_logged_in: true,
+                        ..Default::default()
+                    };
+                }
+            }
+        }
+
+        let x_api_key = request.headers().get_one("X-Api-Key");
+        if x_api_key == admin_token {
             return Session {
                 is_admin: true,
                 is_logged_in: true,