From 40691a5c86a3cb99f20319a03baf8da5bed2e73a Mon Sep 17 00:00:00 2001
From: Vipul Vaibhaw <vaibhaw.vipul@gmail.com>
Date: Thu, 8 Aug 2024 16:21:13 +0530
Subject: [PATCH] not using unwraps to avoid panic while acquiring locks

---
 src/api/csv.rs           | 6 ++++--
 src/api/duckdb.rs        | 6 ++++--
 src/api/parquet.rs       | 6 ++++--
 src/duckdb/connection.rs | 8 ++++++--
 src/fdw/trigger.rs       | 6 ++++--
 5 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/src/api/csv.rs b/src/api/csv.rs
index 3e2faf72..0fa84b2b 100644
--- a/src/api/csv.rs
+++ b/src/api/csv.rs
@@ -15,7 +15,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use duckdb::types::Value;
 use pgrx::*;
 
@@ -71,7 +71,9 @@ fn sniff_csv_impl(files: &str, sample_size: Option<i64>) -> Result<Vec<SniffCsvR
     .collect::<Vec<String>>()
     .join(", ");
     let conn = get_global_connection()?;
-    let conn = conn.lock().unwrap();
+    let conn = conn
+        .lock()
+        .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?;
     let query = format!("SELECT * FROM sniff_csv({schema_str})");
     let mut stmt = conn.prepare(&query)?;
 
diff --git a/src/api/duckdb.rs b/src/api/duckdb.rs
index 354329c1..96de0914 100644
--- a/src/api/duckdb.rs
+++ b/src/api/duckdb.rs
@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use pgrx::*;
 
 use crate::duckdb::connection;
@@ -38,7 +38,9 @@ pub fn duckdb_settings() -> iter::TableIterator<
 #[inline]
 fn duckdb_settings_impl() -> Result<Vec<DuckdbSettingsRow>> {
     let conn = get_global_connection()?;
-    let conn = conn.lock().unwrap();
+    let conn = conn
+        .lock()
+        .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?;
     let mut stmt = conn.prepare("SELECT * FROM duckdb_settings()")?;
 
     Ok(stmt
diff --git a/src/api/parquet.rs b/src/api/parquet.rs
index be222698..0a1f98e7 100644
--- a/src/api/parquet.rs
+++ b/src/api/parquet.rs
@@ -15,7 +15,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use pgrx::*;
 
 use crate::duckdb::utils;
@@ -89,7 +89,9 @@ pub fn parquet_schema(
 fn parquet_schema_impl(files: &str) -> Result<Vec<ParquetSchemaRow>> {
     let schema_str = utils::format_csv(files);
     let conn = get_global_connection()?;
-    let conn = conn.lock().unwrap();
+    let conn = conn
+        .lock()
+        .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?;
     let query = format!("SELECT * FROM parquet_schema({schema_str})");
     let mut stmt = conn.prepare(&query)?;
 
diff --git a/src/duckdb/connection.rs b/src/duckdb/connection.rs
index e473e1a0..2182883d 100644
--- a/src/duckdb/connection.rs
+++ b/src/duckdb/connection.rs
@@ -188,13 +188,17 @@ pub fn get_batches() -> Result<Vec<RecordBatch>> {
 
 pub fn execute<P: Params>(sql: &str, params: P) -> Result<usize> {
     let conn = get_global_connection()?;
-    let conn = conn.lock().unwrap();
+    let conn = conn
+        .lock()
+        .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?;
     conn.execute(sql, params).map_err(|err| anyhow!("{err}"))
 }
 
 pub fn drop_relation(table_name: &str, schema_name: &str) -> Result<()> {
     let conn = get_global_connection()?;
-    let conn = conn.lock().unwrap();
+    let conn = conn
+        .lock()
+        .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?;
     let mut statement = conn.prepare(format!("SELECT table_type from information_schema.tables WHERE table_schema = '{schema_name}' AND table_name = '{table_name}' LIMIT 1").as_str())?;
     if let Ok(Some(row)) = statement.query([])?.next() {
         let table_type: String = row.get(0)?;
diff --git a/src/fdw/trigger.rs b/src/fdw/trigger.rs
index 260c5b5c..11a8df30 100644
--- a/src/fdw/trigger.rs
+++ b/src/fdw/trigger.rs
@@ -15,7 +15,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program. If not, see <http://www.gnu.org/licenses/>.
 
-use anyhow::{bail, Result};
+use anyhow::{anyhow, bail, Result};
 use pgrx::*;
 use std::collections::HashMap;
 use std::ffi::CStr;
@@ -145,7 +145,9 @@ unsafe fn auto_create_schema_impl(fcinfo: pg_sys::FunctionCallInfo) -> Result<()
 
     // Get DuckDB schema
     let conn = get_global_connection()?;
-    let conn = conn.lock().unwrap();
+    let conn = conn
+        .lock()
+        .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?;
     let query = format!("DESCRIBE {schema_name}.{table_name}");
     let mut stmt = conn.prepare(&query)?;