Skip to content

Commit

Permalink
Try adding some diesel metrics
Browse files Browse the repository at this point in the history
  • Loading branch information
Eijebong committed May 18, 2024
1 parent 6925ef7 commit c0fe7a9
Show file tree
Hide file tree
Showing 6 changed files with 175 additions and 11 deletions.
107 changes: 99 additions & 8 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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" }
41 changes: 41 additions & 0 deletions src/db.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -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());
}
_ => {}
};
}
}
3 changes: 2 additions & 1 deletion src/diesel_uuid.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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)
}
}
10 changes: 10 additions & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use db::{DbInstrumentation, QUERY_HISTOGRAM};
use diesel::r2d2::Pool;
use diesel::sqlite::Sqlite;
use diesel::{r2d2::ConnectionManager, SqliteConnection};
Expand Down Expand Up @@ -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");
{
Expand Down Expand Up @@ -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())
Expand Down
21 changes: 19 additions & 2 deletions src/views/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit c0fe7a9

Please sign in to comment.