fix: Postgres query metrics

src/handlers/register_client.rs

@@ -96,6 +96,7 @@ pub async fn handler(
                 token: body.token,
                 always_raw,
             },
+            state.metrics.as_ref(),
         )
         .await?;
 

src/metrics/mod.rs

@@ -1,4 +1,13 @@
-use wc::metrics::{otel::metrics::Counter, ServiceMetrics};
+use {
+    std::time::Instant,
+    wc::metrics::{
+        otel::{
+            metrics::{Counter, Histogram},
+            KeyValue,
+        },
+        ServiceMetrics,
+    },
+};
 
 #[derive(Clone)]
 pub struct Metrics {
@@ -16,6 +25,9 @@ pub struct Metrics {
 
     pub tenant_suspensions: Counter<u64>,
     pub client_suspensions: Counter<u64>,
+
+    postgres_queries: Counter<u64>,
+    postgres_query_latency: Histogram<u64>,
 }
 
 impl Default for Metrics {
@@ -84,6 +96,16 @@ impl Metrics {
             .with_description("The number of clients that have been suspended")
             .init();
 
+        let postgres_queries: Counter<u64> = meter
+            .u64_counter("postgres_queries")
+            .with_description("The number of Postgres queries executed")
+            .init();
+
+        let postgres_query_latency: Histogram<u64> = meter
+            .u64_histogram("postgres_query_latency")
+            .with_description("The latency Postgres queries")
+            .init();
+
         Metrics {
             registered_clients: clients_counter,
             received_notifications: received_notification_counter,
@@ -96,6 +118,17 @@ impl Metrics {
             tenant_fcm_v1_updates: tenant_fcm_v1_updates_counter,
             tenant_suspensions: tenant_suspensions_counter,
             client_suspensions: client_suspensions_counter,
+            postgres_queries,
+            postgres_query_latency,
         }
     }
+
+    pub fn postgres_query(&self, query_name: &'static str, start: Instant) {
+        let elapsed = start.elapsed();
+
+        let attributes = [KeyValue::new("name", query_name)];
+        self.postgres_queries.add(1, &attributes);
+        self.postgres_query_latency
+            .record(elapsed.as_millis() as u64, &attributes);
+    }
 }

src/stores/client.rs

@@ -1,10 +1,12 @@
 use {
     crate::{
+        metrics::Metrics,
         providers::ProviderKind,
         stores::{self, StoreError::NotFound},
     },
     async_trait::async_trait,
     sqlx::Executor,
+    std::time::Instant,
     tracing::{debug, instrument},
 };
 
@@ -19,15 +21,27 @@ pub struct Client {
 
 #[async_trait]
 pub trait ClientStore {
-    async fn create_client(&self, tenant_id: &str, id: &str, client: Client) -> stores::Result<()>;
+    async fn create_client(
+        &self,
+        tenant_id: &str,
+        id: &str,
+        client: Client,
+        metrics: Option<&Metrics>,
+    ) -> stores::Result<()>;
     async fn get_client(&self, tenant_id: &str, id: &str) -> stores::Result<Client>;
     async fn delete_client(&self, tenant_id: &str, id: &str) -> stores::Result<()>;
 }
 
 #[async_trait]
 impl ClientStore for sqlx::PgPool {
-    #[instrument(skip(self, client))]
-    async fn create_client(&self, tenant_id: &str, id: &str, client: Client) -> stores::Result<()> {
+    #[instrument(skip(self, client, metrics))]
+    async fn create_client(
+        &self,
+        tenant_id: &str,
+        id: &str,
+        client: Client,
+        metrics: Option<&Metrics>,
+    ) -> stores::Result<()> {
         debug!(
             "ClientStore::create_client tenant_id={tenant_id} id={id} token={} with locking",
             client.token
@@ -37,6 +51,7 @@ impl ClientStore for sqlx::PgPool {
 
         // Statement for locking based on the client id to prevent an issue #230
         // and locking based on the token to prevent an issue #292
+        let start = Instant::now();
         sqlx::query(
             "SELECT
                 pg_advisory_xact_lock(abs(hashtext($1::text))),
@@ -46,13 +61,21 @@ impl ClientStore for sqlx::PgPool {
         .bind(client.token.clone())
         .execute(&mut transaction)
         .await?;
+        if let Some(metrics) = metrics {
+            metrics.postgres_query("create_client_pg_advisory_xact_lock", start);
+        }
 
+        let start = Instant::now();
         sqlx::query("DELETE FROM public.clients WHERE id = $1 OR device_token = $2")
             .bind(id)
             .bind(client.token.clone())
             .execute(&mut transaction)
             .await?;
+        if let Some(metrics) = metrics {
+            metrics.postgres_query("create_client_delete", start);
+        }
 
+        let start = Instant::now();
         let mut insert_query = sqlx::QueryBuilder::new(
             "INSERT INTO public.clients (id, tenant_id, push_type, device_token, always_raw)",
         );
@@ -73,7 +96,15 @@ impl ClientStore for sqlx::PgPool {
             },
         );
         insert_query.build().execute(&mut transaction).await?;
+        if let Some(metrics) = metrics {
+            metrics.postgres_query("create_client_insert", start);
+        }
+
+        let start = Instant::now();
         transaction.commit().await?;
+        if let Some(metrics) = metrics {
+            metrics.postgres_query("create_client_commit", start);
+        }
 
         Ok(())
     }

tests/functional/stores/client.rs

@@ -24,6 +24,7 @@ async fn client_creation(ctx: &mut StoreContext) {
                 token,
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -46,6 +47,7 @@ async fn client_creation_fcm(ctx: &mut StoreContext) {
                 token,
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -68,6 +70,7 @@ async fn client_creation_apns(ctx: &mut StoreContext) {
                 token,
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -92,6 +95,7 @@ async fn client_upsert_token(ctx: &mut StoreContext) {
                 token: token.clone(),
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -133,6 +137,7 @@ async fn client_upsert_token(ctx: &mut StoreContext) {
                 token: updated_token.clone(),
                 always_raw: true,
             },
+            None,
         )
         .await
         .unwrap();
@@ -164,6 +169,7 @@ async fn client_upsert_id(ctx: &mut StoreContext) {
                 token: token.clone(),
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -205,6 +211,7 @@ async fn client_upsert_id(ctx: &mut StoreContext) {
                 token: token.clone(),
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -239,6 +246,7 @@ async fn client_create_same_id_and_token(ctx: &mut StoreContext) {
                 token: token.clone(),
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -280,6 +288,7 @@ async fn client_create_same_id_and_token(ctx: &mut StoreContext) {
                 token: token.clone(),
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -310,6 +319,7 @@ async fn client_deletion(ctx: &mut StoreContext) {
                 token,
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();
@@ -332,6 +342,7 @@ async fn client_fetch(ctx: &mut StoreContext) {
                 token: token.clone(),
                 always_raw: false,
             },
+            None,
         )
         .await
         .unwrap();

tests/functional/stores/notification.rs

@@ -24,6 +24,7 @@ pub async fn create_client(client_store: &ClientStoreArc) -> String {
                 token,
                 always_raw: false,
             },
+            None,
         )
         .await
         .expect("failed to create client for notification test");