Add metrics endpoint behind feature

11 files changed, 84 insertions, 7 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 8acab35..83330b6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -545,7 +545,7 @@ dependencies = [
 
 [[package]]
 name = "hdt-api"
-version = "0.1.4"
+version = "0.1.5"
 dependencies = [
  "addr",
  "anyhow",
diff --git a/Cargo.toml b/Cargo.toml
index 0658f47..a4c96f6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "hdt-api"
-version = "0.1.4"
+version = "0.1.5"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/src/db/metrics.rs b/src/db/metrics.rs
new file mode 100644
index 0000000..7557d47
--- /dev/null
+++ b/src/db/metrics.rs
@@ -0,0 +1,15 @@
+use crate::db::models::Metrics;
+use crate::db::strings;
+use sqlx::{Pool, Postgres};
+
+
+pub async fn get_metrics(
+    pool: &Pool<Postgres>,
+) -> Result<Metrics, sqlx::Error> {
+    let mut transaction = pool.begin().await?;
+    let metrics = sqlx::query_as::<_, Metrics>(&strings::GET_METRICS)
+        .fetch_one(&mut transaction)
+        .await?;
+    transaction.commit().await?;
+    Ok(metrics)
+}
diff --git a/src/db/mod.rs b/src/db/mod.rs
index c88440f..f06dd3b 100644
--- a/src/db/mod.rs
+++ b/src/db/mod.rs
@@ -1,6 +1,8 @@
 pub mod records;
 pub mod users;
 pub mod zones;
+pub mod metrics;
 
 pub mod models;
 mod strings;
+
diff --git a/src/db/models.rs b/src/db/models.rs
index d9591d8..0d82018 100644
--- a/src/db/models.rs
+++ b/src/db/models.rs
@@ -40,3 +40,13 @@ pub struct Record {
     pub created_at: DateTime<Utc>,
     pub modified_at: DateTime<Utc>,
 }
+
+#[derive(Serialize, Deserialize, FromRow, Debug)]
+pub struct Metrics {
+    pub p50: f64,
+    pub p90: f64,
+    pub p95: f64,
+    pub p99: f64,
+    pub avg: f64,
+    pub count: i64,
+}
\ No newline at end of file
diff --git a/src/db/strings.rs b/src/db/strings.rs
index 7a7a5cf..dc6e44d 100644
--- a/src/db/strings.rs
+++ b/src/db/strings.rs
@@ -38,10 +38,21 @@ lazy_static! {
             FROM records WHERE zone_id = $1
     ";
     pub(crate) static ref GET_USER_FROM_API_KEY: &'static str = r"
-    SELECT users.id,email,password,display_name,users.created_at,modified_at,admin,enabled,totp_secret FROM api_keys
-        JOIN users
-            ON users.id = api_keys.owner_uuid
-        WHERE api_keys.token_hash = $1 
-            AND api_keys.expires_at > (now() AT TIME ZONE 'UTC');
+        SELECT users.id,email,password,display_name,users.created_at,modified_at,admin,enabled,totp_secret FROM api_keys
+            JOIN users
+                ON users.id = api_keys.owner_uuid
+            WHERE api_keys.token_hash = $1 
+                AND api_keys.expires_at > (now() AT TIME ZONE 'UTC');
     ";
+
+    pub(crate) static ref GET_METRICS: &'static str = r#"
+        SELECT
+            percentile_cont(0.50) WITHIN GROUP (ORDER BY queries.duration_us) AS p50,
+            percentile_cont(0.90) WITHIN GROUP (ORDER BY queries.duration_us) AS p90,
+            percentile_cont(0.95) WITHIN GROUP (ORDER BY queries.duration_us) AS p95,
+            percentile_cont(0.99) WITHIN GROUP (ORDER BY queries.duration_us) AS p99,
+            avg(queries.duration_us)::FLOAT8 AS avg,
+            count(queries.duration_us) as count
+        FROM queries WHERE "time" > (NOW() AT TIME ZONE 'UTC' - '1 day'::INTERVAL);
+    "#;
 }
diff --git a/src/features.rs b/src/features.rs
index 9adad1b..378bc46 100644
--- a/src/features.rs
+++ b/src/features.rs
@@ -6,4 +6,6 @@ lazy_static! {
         env::var("SIGNUPS_ENABLED").unwrap_or_else(|_| String::from("false")) == "true";
     pub static ref TOTP_ENABLED: bool =
         env::var("TOTP_ENABLED").unwrap_or_else(|_| String::from("false")) == "true";
+    pub static ref METRICS_ENABLED: bool =
+        env::var("METRICS_URL").is_ok() && env::var("METRICS_ENABLED").unwrap_or_else(|_| String::from("false")) == "true";
 }
diff --git a/src/main.rs b/src/main.rs
index e7e2489..40ef67a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -50,6 +50,17 @@ async fn main() {
             .unwrap(),
     );
     info!("Postgres pool initialized");
+    let metrics_pool = match env::var("METRICS_URL") {
+        Ok(url) => Some(Arc::new(
+            PgPoolOptions::new()
+                .max_connections(5)
+                .connect(&url)
+                .await
+                .unwrap(),
+        )),
+        Err(_) => None,
+    };
+    info!("Metrics pool (possibly) initialized");
 
     // Create our WhoIs client
     let whois_client = whois_rust::WhoIs::from_string(include_str!(concat!(
@@ -65,6 +76,7 @@ async fn main() {
                 "/v1",
                 Router::new()
                     .route("/features", get(routes::v1::features::get_features))
+                    .route("/metrics", get(routes::v1::metrics::get_metrics))
                     .nest(
                         "/users",
                         Router::new()
@@ -95,6 +107,7 @@ async fn main() {
         )
         .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()))
         .layer(Extension(pg_pool))
+        .layer(Extension(metrics_pool))
         .layer(Extension(whois_client));
 
     let addr = SocketAddr::from(([0, 0, 0, 0], 8000));
diff --git a/src/routes/v1/features.rs b/src/routes/v1/features.rs
index ea48abc..8385ec7 100644
--- a/src/routes/v1/features.rs
+++ b/src/routes/v1/features.rs
@@ -10,6 +10,7 @@ pub async fn get_features() -> impl IntoResponse {
             "version": option_env!("CARGO_PKG_VERSION").unwrap_or_else(|| "unknown"),
             "signup": *features::SIGNUPS_ENABLED,
             "totp": *features::TOTP_ENABLED,
+            "metrics": *features::METRICS_ENABLED,
         })),
     )
 }
diff --git a/src/routes/v1/metrics.rs b/src/routes/v1/metrics.rs
new file mode 100644
index 0000000..38485d0
--- /dev/null
+++ b/src/routes/v1/metrics.rs
@@ -0,0 +1,22 @@
+use crate::{db, features};
+use crate::extractors::{Json};
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use axum::Extension;
+use serde_json::json;
+use sqlx::{Pool, Postgres};
+use std::sync::Arc;
+
+pub async fn get_metrics(
+    Extension(metrics_pool): Extension<Option<Arc<Pool<Postgres>>>>,
+) -> impl IntoResponse {
+    if !*features::METRICS_ENABLED {
+        return (StatusCode::NOT_FOUND, Json(json!({"error": "Metrics not enabled"})));
+    }
+
+    let metrics = db::metrics::get_metrics(&metrics_pool.as_ref().unwrap()).await;
+    match metrics {
+        Ok(metrics) => (StatusCode::OK, Json(json!(metrics))),
+        Err(err) => (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": err.to_string()}))),
+    }
+}
diff --git a/src/routes/v1/mod.rs b/src/routes/v1/mod.rs
index f7ac287..df0880d 100644
--- a/src/routes/v1/mod.rs
+++ b/src/routes/v1/mod.rs
@@ -2,5 +2,6 @@ pub mod features;
 pub mod records;
 pub mod users;
 pub mod zones;
+pub mod metrics;
 
 mod requests;