feat(schema_cache): implement diagnostics

Author: juleswritescode

Cargo.lock

@@ -17,14 +17,16 @@ async-std = "1.12.0"
 text-size.workspace = true
 
 pg_schema_cache.workspace = true
-pg_test_utils.workspace   = true
 tree-sitter.workspace     = true
 tree_sitter_sql.workspace = true
 
 sqlx.workspace = true
 
 tokio = { version = "1.41.1", features = ["full"] }
 
+[dev-dependencies]
+pg_test_utils.workspace   = true
+
 [lib]
 doctest = false
 

crates/pg_completions/Cargo.toml

@@ -72,7 +72,9 @@ mod tests {
             .expect("Error loading sql language");
 
         let tree = parser.parse(input, None).unwrap();
-        let schema_cache = SchemaCache::load(&test_db).await;
+        let schema_cache = SchemaCache::load(&test_db)
+            .await
+            .expect("Couldn't load Schema Cache");
 
         let p = CompletionParams {
             position: ((input.len() - 1) as u32).into(),
@@ -117,7 +119,9 @@ mod tests {
             .await
             .expect("Failed to execute setup query");
 
-        let schema_cache = SchemaCache::load(&test_db).await;
+        let schema_cache = SchemaCache::load(&test_db)
+            .await
+            .expect("Couldn't load Schema Cache");
 
         let mut parser = tree_sitter::Parser::new();
         parser
@@ -180,7 +184,9 @@ mod tests {
             .await
             .expect("Failed to execute setup query");
 
-        let schema_cache = SchemaCache::load(&test_db).await;
+        let schema_cache = SchemaCache::load(&test_db)
+            .await
+            .expect("Couldn't load SchemaCache");
 
         let mut parser = tree_sitter::Parser::new();
         parser

crates/pg_completions/src/complete.rs

@@ -23,6 +23,8 @@ tree_sitter_sql.workspace  = true
 
 [dev-dependencies]
 async-std = "1.12.0"
+pg_test_utils.workspace   = true
+
 
 [lib]
 doctest = false

crates/pg_inlay_hints/Cargo.toml

@@ -80,7 +80,7 @@ fn resolve_func_arg_hint(
 mod tests {
     use async_std::task::block_on;
     use pg_schema_cache::SchemaCache;
-    use sqlx::PgPool;
+    use pg_test_utils::test_database::get_new_test_db;
 
     use crate::{
         functions_args::FunctionArgHint,
@@ -89,17 +89,14 @@ mod tests {
 
     #[test]
     fn test_function_args() {
+        let test_db = block_on(get_new_test_db());
         let input = "select lower('TEST')";
 
-        let conn_string = std::env::var("DATABASE_URL").unwrap();
-
-        let pool = block_on(PgPool::connect(conn_string.as_str())).unwrap();
-
         let root = pg_query_ext::parse(input).unwrap();
-
         let res = pg_syntax::parse_syntax(input, &root);
 
-        let schema_cache = block_on(SchemaCache::load(&pool));
+        let schema_cache =
+            block_on(SchemaCache::load(&test_db)).expect("Couldn't load Schema Cache");
 
         let hints = FunctionArgHint::find_all(InlayHintsParams {
             ast: Some(&root),

crates/pg_inlay_hints/src/functions_args.rs

@@ -41,7 +41,7 @@ impl DbConnection {
                         match res {
                             Ok(not) => {
                                 if not.payload().to_string() == "reload schema" {
-                                    let schema_cache = SchemaCache::load(&cloned_pool).await;
+                                    let schema_cache = SchemaCache::load(&cloned_pool).await.unwrap();
                                     ide.write().await.set_schema_cache(schema_cache);
                                 };
                             }

crates/pg_lsp/src/db_connection.rs

@@ -17,6 +17,8 @@ async-std            = { version = "1.12.0" }
 futures-util = "0.3.31"
 serde.workspace      = true
 serde_json.workspace = true
+pg_diagnostics.workspace = true
+pg_console.workspace = true
 
 sqlx.workspace = true
 

crates/pg_schema_cache/Cargo.toml

@@ -0,0 +1,38 @@
+use pg_diagnostics::Diagnostic;
+use std::fmt::Debug;
+
+#[derive(Debug, Diagnostic)]
+pub enum SchemaCacheError {
+    DatabaseConnectionError(DatabaseConnectionError),
+}
+
+#[derive(Debug, Diagnostic)]
+#[diagnostic(
+    category = "database/connection",
+    severity = Error,
+    message(
+        description = "Unable to Load Database Schema",
+        message("Database Error Message:: "{self.message})
+    )
+)]
+pub struct DatabaseConnectionError {
+    pub message: String,
+    pub code: Option<String>,
+}
+
+impl From<sqlx::Error> for SchemaCacheError {
+    fn from(err: sqlx::Error) -> Self {
+        let db_err = err.as_database_error();
+        if let Some(db_err) = db_err {
+            Self::DatabaseConnectionError(DatabaseConnectionError {
+                message: db_err.message().to_string(),
+                code: db_err.code().map(|c| c.to_string()),
+            })
+        } else {
+            Self::DatabaseConnectionError(DatabaseConnectionError {
+                message: err.to_string(),
+                code: None,
+            })
+        }
+    }
+}

crates/pg_schema_cache/src/diagnostics.rs

@@ -1,15 +1,16 @@
 //! The schema cache
 
 #![allow(dead_code)]
-#![feature(future_join)]
 
+mod diagnostics;
 mod functions;
 mod schema_cache;
 mod schemas;
 mod tables;
 mod types;
 mod versions;
 
+pub use diagnostics::SchemaCacheError;
 pub use functions::{Behavior, Function, FunctionArg, FunctionArgs};
 pub use schema_cache::SchemaCache;
 pub use tables::{ReplicaIdentity, Table};

crates/pg_schema_cache/src/lib.rs

@@ -1,6 +1,6 @@
-use anyhow::Context;
 use sqlx::postgres::PgPool;
 
+use crate::diagnostics::SchemaCacheError;
 use crate::functions::Function;
 use crate::schemas::Schema;
 use crate::tables::Table;
@@ -21,15 +21,14 @@ impl SchemaCache {
         SchemaCache::default()
     }
 
-    pub async fn load(pool: &PgPool) -> anyhow::Result<SchemaCache> {
+    pub async fn load(pool: &PgPool) -> Result<SchemaCache, SchemaCacheError> {
         let (schemas, tables, functions, types, versions) = futures_util::try_join!(
             Schema::load(pool),
             Table::load(pool),
             Function::load(pool),
             PostgresType::load(pool),
             Version::load(pool),
-        )
-        .with_context(|| format!("Unable to load Schema Cache"))?;
+        )?;
 
         Ok(SchemaCache {
             schemas,
@@ -86,7 +85,7 @@ mod tests {
 
         let pool = async_std::task::block_on(PgPool::connect(conn_string.as_str())).unwrap();
 
-        async_std::task::block_on(SchemaCache::load(&pool));
+        async_std::task::block_on(SchemaCache::load(&pool)).expect("Couldn't load Schema Cache");
 
         assert!(true);
     }

crates/pg_schema_cache/src/schema_cache.rs

@@ -5,6 +5,7 @@ use pg_diagnostics::{
     category, Advices, Category, Diagnostic, DiagnosticTags, LogCategory, Severity, Visit,
 };
 use pg_fs::FileSystemDiagnostic;
+use pg_schema_cache::SchemaCacheError;
 use serde::{Deserialize, Serialize};
 use std::error::Error;
 use std::fmt;
@@ -239,6 +240,19 @@ impl From<sqlx::Error> for WorkspaceError {
     }
 }
 
+impl From<SchemaCacheError> for WorkspaceError {
+    fn from(err: SchemaCacheError) -> Self {
+        match err {
+            SchemaCacheError::DatabaseConnectionError(db_err) => {
+                WorkspaceError::DatabaseConnectionError(DatabaseConnectionError {
+                    message: db_err.message,
+                    code: db_err.code,
+                })
+            }
+        }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize, Diagnostic)]
 #[diagnostic(
     category = "internalError/fs",

crates/pg_workspace_new/src/diagnostics.rs

@@ -128,10 +128,8 @@ impl WorkspaceServer {
         tracing::info!("Reloading schema cache");
         // TODO return error if db connection is not available
         if let Some(c) = self.connection.read().unwrap().get_pool() {
-            let schema_cache = run_async(async move {
-                // TODO load should return a Result
-                SchemaCache::load(&c).await
-            })?;
+            let maybe_schema_cache = run_async(async move { SchemaCache::load(&c).await })?;
+            let schema_cache = maybe_schema_cache?;
 
             let mut cache = self.schema_cache.write().unwrap();
             *cache = schema_cache;