From dd00eefc508c46c516d16c01690c405fa52c9f44 Mon Sep 17 00:00:00 2001
From: psteinroe <philipp@steinroetter.com>
Date: Wed, 16 Apr 2025 09:50:59 +0200
Subject: [PATCH 1/2] fix: report db connection errors

---
 crates/pgt_typecheck/src/lib.rs              | 22 +++++++++-----------
 crates/pgt_typecheck/tests/diagnostics.rs    |  2 +-
 crates/pgt_workspace/src/workspace/server.rs | 19 ++++++++++-------
 3 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/crates/pgt_typecheck/src/lib.rs b/crates/pgt_typecheck/src/lib.rs
index 9311bb8e..bf5a6e76 100644
--- a/crates/pgt_typecheck/src/lib.rs
+++ b/crates/pgt_typecheck/src/lib.rs
@@ -3,10 +3,9 @@ mod diagnostics;
 pub use diagnostics::TypecheckDiagnostic;
 use diagnostics::create_type_error;
 use pgt_text_size::TextRange;
-use sqlx::Executor;
-use sqlx::PgPool;
 use sqlx::postgres::PgDatabaseError;
 pub use sqlx::postgres::PgSeverity;
+use sqlx::{Executor, PgPool};
 
 #[derive(Debug)]
 pub struct TypecheckParams<'a> {
@@ -29,7 +28,9 @@ pub struct TypeError {
     pub constraint: Option<String>,
 }
 
-pub async fn check_sql(params: TypecheckParams<'_>) -> Option<TypecheckDiagnostic> {
+pub async fn check_sql(
+    params: TypecheckParams<'_>,
+) -> Result<Option<TypecheckDiagnostic>, sqlx::Error> {
     // Check if the AST is not a supported statement type
     if !matches!(
         params.ast,
@@ -39,13 +40,10 @@ pub async fn check_sql(params: TypecheckParams<'_>) -> Option<TypecheckDiagnosti
             | pgt_query_ext::NodeEnum::DeleteStmt(_)
             | pgt_query_ext::NodeEnum::CommonTableExpr(_)
     ) {
-        return None;
+        return Ok(None);
     }
 
-    let mut conn = match params.conn.acquire().await {
-        Ok(c) => c,
-        Err(_) => return None,
-    };
+    let mut conn = params.conn.acquire().await?;
 
     // Postgres caches prepared statements within the current DB session (connection).
     // This can cause issues if the underlying table schema changes while statements
@@ -53,14 +51,14 @@ pub async fn check_sql(params: TypecheckParams<'_>) -> Option<TypecheckDiagnosti
     // each typecheck operation.
     conn.close_on_drop();
 
-    let res = conn.prepare(params.sql).await;
+    let res = params.conn.prepare(params.sql).await;
 
     match res {
-        Ok(_) => None,
+        Ok(_) => Ok(None),
         Err(sqlx::Error::Database(err)) => {
             let pg_err = err.downcast_ref::<PgDatabaseError>();
-            Some(create_type_error(pg_err, params.tree))
+            Ok(Some(create_type_error(pg_err, params.tree)))
         }
-        Err(_) => None,
+        Err(err) => Err(err),
     }
 }
diff --git a/crates/pgt_typecheck/tests/diagnostics.rs b/crates/pgt_typecheck/tests/diagnostics.rs
index 46daa8a1..4c780d74 100644
--- a/crates/pgt_typecheck/tests/diagnostics.rs
+++ b/crates/pgt_typecheck/tests/diagnostics.rs
@@ -37,7 +37,7 @@ async fn test(name: &str, query: &str, setup: &str) {
 
     Formatter::new(&mut writer)
         .write_markup(markup! {
-            {PrintDiagnostic::simple(&result.unwrap())}
+            {PrintDiagnostic::simple(&result.unwrap().unwrap())}
         })
         .unwrap();
 
diff --git a/crates/pgt_workspace/src/workspace/server.rs b/crates/pgt_workspace/src/workspace/server.rs
index 2ad119f5..3bf540cc 100644
--- a/crates/pgt_workspace/src/workspace/server.rs
+++ b/crates/pgt_workspace/src/workspace/server.rs
@@ -360,8 +360,6 @@ impl Workspace for WorkspaceServer {
 
         let mut diagnostics: Vec<SDiagnostic> = parser.document_diagnostics().to_vec();
 
-        // TODO: run this in parallel with rayon based on rayon.count()
-
         if let Some(pool) = self
             .connection
             .read()
@@ -385,13 +383,15 @@ impl Workspace for WorkspaceServer {
                                 })
                                 .await
                                 .map(|d| {
-                                    let r = d.location().span.map(|span| span + range.start());
+                                    d.map(|d| {
+                                        let r = d.location().span.map(|span| span + range.start());
 
-                                    d.with_file_path(path.as_path().display().to_string())
-                                        .with_file_span(r.unwrap_or(range))
+                                        d.with_file_path(path.as_path().display().to_string())
+                                            .with_file_span(r.unwrap_or(range))
+                                    })
                                 })
                             } else {
-                                None
+                                Ok(None)
                             }
                         }
                     })
@@ -400,8 +400,11 @@ impl Workspace for WorkspaceServer {
                     .await
             })?;
 
-            for result in async_results.into_iter().flatten() {
-                diagnostics.push(SDiagnostic::new(result));
+            for result in async_results.into_iter() {
+                let result = result?;
+                if let Some(diag) = result {
+                    diagnostics.push(SDiagnostic::new(diag));
+                }
             }
         }
 

From 5511c90fd96395828e74412469e9093ac37f84aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Steinr=C3=B6tter?= <philipp@steinroetter.com>
Date: Wed, 16 Apr 2025 14:08:08 +0200
Subject: [PATCH 2/2] Update lib.rs

---
 crates/pgt_typecheck/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/pgt_typecheck/src/lib.rs b/crates/pgt_typecheck/src/lib.rs
index bf5a6e76..f741c0e6 100644
--- a/crates/pgt_typecheck/src/lib.rs
+++ b/crates/pgt_typecheck/src/lib.rs
@@ -51,7 +51,7 @@ pub async fn check_sql(
     // each typecheck operation.
     conn.close_on_drop();
 
-    let res = params.conn.prepare(params.sql).await;
+    let res = conn.prepare(params.sql).await;
 
     match res {
         Ok(_) => Ok(None),