fix(schema_cache): return results

crates/pg_completions/Cargo.toml

@@ -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/src/complete.rs

@@ -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_inlay_hints/Cargo.toml

@@ -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/src/functions_args.rs

@@ -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_lsp/src/db_connection.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_schema_cache/Cargo.toml

@@ -12,11 +12,17 @@ version              = "0.0.0"
 
 
 [dependencies]
+anyhow.workspace = true
 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
 
+[dev-dependencies]
+pg_test_utils.workspace = true
+
 [lib]
 doctest = false

crates/pg_schema_cache/src/diagnostics.rs

@@ -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/functions.rs

@@ -69,7 +69,7 @@ pub struct Function {
 impl SchemaCacheItem for Function {
     type Item = Function;
 
-    async fn load(pool: &PgPool) -> Vec<Function> {
+    async fn load(pool: &PgPool) -> Result<Vec<Function>, sqlx::Error> {
         sqlx::query_as!(
             Function,
             r#"
@@ -179,6 +179,5 @@ from
         )
         .fetch_all(pool)
         .await
-        .unwrap()
     }
 }

crates/pg_schema_cache/src/lib.rs

@@ -1,34 +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;
 
-use sqlx::postgres::PgPool;
-
+pub use diagnostics::SchemaCacheError;
 pub use functions::{Behavior, Function, FunctionArg, FunctionArgs};
 pub use schema_cache::SchemaCache;
 pub use tables::{ReplicaIdentity, Table};
-
-#[derive(Debug, Clone)]
-struct SchemaCacheManager {
-    pub cache: SchemaCache,
-}
-
-impl SchemaCacheManager {
-    pub async fn init(pool: &PgPool) -> Self {
-        SchemaCacheManager {
-            cache: SchemaCache::load(pool).await,
-        }
-    }
-
-    pub async fn reload_cache(&mut self, pool: &PgPool) {
-        self.cache = SchemaCache::load(pool).await;
-    }
-}

crates/pg_schema_cache/src/schema_cache.rs

@@ -1,7 +1,6 @@
-use std::future::join;
-
 use sqlx::postgres::PgPool;
 
+use crate::diagnostics::SchemaCacheError;
 use crate::functions::Function;
 use crate::schemas::Schema;
 use crate::tables::Table;
@@ -22,23 +21,22 @@ impl SchemaCache {
         SchemaCache::default()
     }
 
-    pub async fn load(pool: &PgPool) -> SchemaCache {
-        let (schemas, tables, functions, types, versions) = join!(
+    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),
-        )
-        .await;
+        )?;
 
-        SchemaCache {
+        Ok(SchemaCache {
             schemas,
             tables,
             functions,
             types,
             versions,
-        }
+        })
     }
 
     /// Applies an AST node to the repository
@@ -72,22 +70,21 @@ impl SchemaCache {
 pub trait SchemaCacheItem {
     type Item;
 
-    async fn load(pool: &PgPool) -> Vec<Self::Item>;
+    async fn load(pool: &PgPool) -> Result<Vec<Self::Item>, sqlx::Error>;
 }
 
 #[cfg(test)]
 mod tests {
-    use sqlx::PgPool;
+    use async_std::task::block_on;
+    use pg_test_utils::test_database::get_new_test_db;
 
     use crate::SchemaCache;
 
     #[test]
     fn test_schema_cache() {
-        let conn_string = std::env::var("DATABASE_URL").unwrap();
-
-        let pool = async_std::task::block_on(PgPool::connect(conn_string.as_str())).unwrap();
+        let test_db = block_on(get_new_test_db());
 
-        async_std::task::block_on(SchemaCache::load(&pool));
+        block_on(SchemaCache::load(&test_db)).expect("Couldn't load Schema Cache");
 
         assert!(true);
     }

crates/pg_schema_cache/src/schemas.rs

@@ -12,7 +12,7 @@ pub struct Schema {
 impl SchemaCacheItem for Schema {
     type Item = Schema;
 
-    async fn load(pool: &PgPool) -> Vec<Schema> {
+    async fn load(pool: &PgPool) -> Result<Vec<Schema>, sqlx::Error> {
         sqlx::query_as!(
             Schema,
             r#"select
@@ -33,6 +33,5 @@ where
         )
         .fetch_all(pool)
         .await
-        .unwrap()
     }
 }

crates/pg_schema_cache/src/tables.rs

@@ -41,7 +41,7 @@ pub struct Table {
 impl SchemaCacheItem for Table {
     type Item = Table;
 
-    async fn load(pool: &PgPool) -> Vec<Table> {
+    async fn load(pool: &PgPool) -> Result<Vec<Table>, sqlx::Error> {
         sqlx::query_as!(
             Table,
             r#"SELECT
@@ -87,6 +87,5 @@ group by
         )
         .fetch_all(pool)
         .await
-        .unwrap()
     }
 }

crates/pg_schema_cache/src/types.rs

@@ -50,7 +50,7 @@ pub struct PostgresType {
 impl SchemaCacheItem for PostgresType {
     type Item = PostgresType;
 
-    async fn load(pool: &PgPool) -> Vec<PostgresType> {
+    async fn load(pool: &PgPool) -> Result<Vec<PostgresType>, sqlx::Error> {
         sqlx::query_as!(
             PostgresType,
             r#"select
@@ -103,6 +103,5 @@ where
         )
         .fetch_all(pool)
         .await
-        .unwrap()
     }
 }

crates/pg_schema_cache/src/versions.rs

@@ -13,7 +13,7 @@ pub struct Version {
 impl SchemaCacheItem for Version {
     type Item = Version;
 
-    async fn load(pool: &PgPool) -> Vec<Version> {
+    async fn load(pool: &PgPool) -> Result<Vec<Version>, sqlx::Error> {
         sqlx::query_as!(
             Version,
             r#"select
@@ -29,7 +29,6 @@ impl SchemaCacheItem for Version {
         )
         .fetch_all(pool)
         .await
-        .unwrap()
     }
 
     /*

crates/pg_typecheck/Cargo.toml

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