From 53990672aea03037b90f186f8d58ce1a2161da09 Mon Sep 17 00:00:00 2001
From: Greg Heartsfield <scsibug@imap.cc>
Date: Wed, 23 Feb 2022 16:38:16 -0600
Subject: [PATCH] improvement: move db pool operations closer to query, do not
 panic on failure

---
 src/db.rs   | 62 +++++++++++++++++++++++++++++------------------------
 src/main.rs |  4 +---
 2 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index b157ec7..58e1861 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -503,7 +503,7 @@ fn query_from_sub(sub: &Subscription) -> (String, Vec<Box<dyn ToSql>>) {
 /// query is immediately aborted.
 pub async fn db_query(
     sub: Subscription,
-    conn: PooledConnection,
+    pool: SqlitePool,
     query_tx: tokio::sync::mpsc::Sender<QueryResult>,
     mut abandon_query_rx: tokio::sync::oneshot::Receiver<()>,
 ) {
@@ -514,36 +514,42 @@ pub async fn db_query(
         // generate SQL query
         let (q, p) = query_from_sub(&sub);
         debug!("SQL generated in {:?}", start.elapsed());
+        // show pool stats
+        debug!("DB pool stats: {:?}", pool.state());
         let start = Instant::now();
-        // execute the query. Don't cache, since queries vary so much.
-        let mut stmt = conn.prepare(&q)?;
-        let mut event_rows = stmt.query(rusqlite::params_from_iter(p))?;
-        let mut first_result = true;
-        while let Some(row) = event_rows.next()? {
-            if first_result {
-                debug!("time to first result: {:?}", start.elapsed());
-                first_result = false;
+        if let Ok(conn) = pool.get() {
+            // execute the query. Don't cache, since queries vary so much.
+            let mut stmt = conn.prepare(&q)?;
+            let mut event_rows = stmt.query(rusqlite::params_from_iter(p))?;
+            let mut first_result = true;
+            while let Some(row) = event_rows.next()? {
+                if first_result {
+                    debug!("time to first result: {:?}", start.elapsed());
+                    first_result = false;
+                }
+                // check if this is still active
+                // TODO:  check every N rows
+                if abandon_query_rx.try_recv().is_ok() {
+                    debug!("query aborted");
+                    return Ok(());
+                }
+                row_count += 1;
+                let event_json = row.get(0)?;
+                query_tx
+                    .blocking_send(QueryResult {
+                        sub_id: sub.get_id(),
+                        event: event_json,
+                    })
+                    .ok();
             }
-            // check if this is still active
-            // TODO:  check every N rows
-            if abandon_query_rx.try_recv().is_ok() {
-                debug!("query aborted");
-                return Ok(());
-            }
-            row_count += 1;
-            let event_json = row.get(0)?;
-            query_tx
-                .blocking_send(QueryResult {
-                    sub_id: sub.get_id(),
-                    event: event_json,
-                })
-                .ok();
+            debug!(
+                "query completed ({} rows) in {:?}",
+                row_count,
+                start.elapsed()
+            );
+        } else {
+            warn!("Could not get a database connection for querying");
         }
-        debug!(
-            "query completed ({} rows) in {:?}",
-            row_count,
-            start.elapsed()
-        );
         let ok: Result<()> = Ok(());
         ok
     });
diff --git a/src/main.rs b/src/main.rs
index 97f316b..3f8fd10 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -519,9 +519,7 @@ async fn nostr_server(
                                     previous_query.send(()).ok();
                                 }
                                 // start a database query
-                                // show pool stats
-                                debug!("DB pool stats: {:?}", pool.state());
-                                db::db_query(s, pool.get().expect("could not get connection"), query_tx.clone(), abandon_query_rx).await;
+                                db::db_query(s, pool.clone(), query_tx.clone(), abandon_query_rx).await;
                             },
                             Err(e) => {
                                 info!("Subscription error: {}", e);