From 916cd78db10ef6687c34146d573e7249c5f5c0da Mon Sep 17 00:00:00 2001
From: Daniel Harvey <danieljamesharvey@gmail.com>
Date: Wed, 1 Nov 2023 15:02:27 +0000
Subject: [PATCH] Count connection acquisition failures

---
 crates/query-engine/execution/src/execution.rs | 14 ++++++++++++--
 crates/query-engine/execution/src/metrics.rs   | 12 ++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/crates/query-engine/execution/src/execution.rs b/crates/query-engine/execution/src/execution.rs
index 96902023c..260fd3254 100644
--- a/crates/query-engine/execution/src/execution.rs
+++ b/crates/query-engine/execution/src/execution.rs
@@ -34,7 +34,12 @@ pub async fn execute(
         .acquire()
         .instrument(info_span!("Acquire connection"))
         .await;
-    let mut connection = acquisition_timer.complete_with(connection_result)?;
+    let mut connection = acquisition_timer
+        .complete_with(connection_result)
+        .map_err(|err| {
+            metrics.error_metrics.record_connection_acquisition_error();
+            err
+        })?;
 
     let query_timer = metrics.time_query_execution();
     let rows_result = execute_queries(&mut connection, database_info, query, plan.variables).await;
@@ -113,7 +118,12 @@ pub async fn explain(
             .acquire()
             .instrument(info_span!("Acquire connection"))
             .await;
-        let mut connection = acquisition_timer.complete_with(connection_result)?;
+        let mut connection = acquisition_timer
+            .complete_with(connection_result)
+            .map_err(|err| {
+                metrics.error_metrics.record_connection_acquisition_error();
+                err
+            })?;
 
         // run and fetch from the database
         sqlx_query
diff --git a/crates/query-engine/execution/src/metrics.rs b/crates/query-engine/execution/src/metrics.rs
index 0102034f9..cf4e33f34 100644
--- a/crates/query-engine/execution/src/metrics.rs
+++ b/crates/query-engine/execution/src/metrics.rs
@@ -303,6 +303,8 @@ pub struct ErrorMetrics {
     connector_error_total: IntCounter,
     /// the database emmited an error.
     database_error_total: IntCounter,
+    /// we failed to acquire a database connection from the pool
+    connection_acquisition_error_total: IntCounter,
 }
 
 impl ErrorMetrics {
@@ -338,12 +340,19 @@ impl ErrorMetrics {
             "Total number of requests failed due to a database error.",
         )?;
 
+        let connection_acquisition_error_total = add_int_counter_metric(
+            metrics_registry,
+            "ndc_postgres_error_connection_acquisition_error_total_count",
+            "Total number of failures to acquire a database connection.",
+        )?;
+
         Ok(ErrorMetrics {
             invalid_request_total,
             unsupported_capability_total,
             unsupported_feature_total,
             connector_error_total,
             database_error_total,
+            connection_acquisition_error_total,
         })
     }
 
@@ -362,4 +371,7 @@ impl ErrorMetrics {
     pub fn record_database_error(&self) {
         self.database_error_total.inc();
     }
+    pub fn record_connection_acquisition_error(&self) {
+        self.connection_acquisition_error_total.inc()
+    }
 }