Merge pull request #154 from juleswritescode/cache/result

psteinroe · web-flow · commit 97679c2d74a4 · 2024-12-14T21:15:52.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/pg_completions/Cargo.toml b/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
 
diff --git a/crates/pg_completions/src/complete.rs b/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
diff --git a/crates/pg_inlay_hints/Cargo.toml b/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
diff --git a/crates/pg_inlay_hints/src/functions_args.rs b/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),
diff --git a/crates/pg_lsp/src/db_connection.rs b/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);
                                 };
                             }
diff --git a/crates/pg_schema_cache/Cargo.toml b/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
diff --git a/crates/pg_schema_cache/src/functions.rs b/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()
     }
 }
diff --git a/crates/pg_schema_cache/src/lib.rs b/crates/pg_schema_cache/src/lib.rs
@@ -1,7 +1,6 @@
 //! The schema cache
 
 #![allow(dead_code)]
-#![feature(future_join)]
 
 mod functions;
 mod schema_cache;
@@ -10,25 +9,6 @@ mod tables;
 mod types;
 mod versions;
 
-use sqlx::postgres::PgPool;
-
 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;
-    }
-}
diff --git a/crates/pg_schema_cache/src/schema_cache.rs b/crates/pg_schema_cache/src/schema_cache.rs
@@ -1,5 +1,3 @@
-use std::future::join;
-
 use sqlx::postgres::PgPool;
 
 use crate::functions::Function;
@@ -22,23 +20,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, sqlx::Error> {
+        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 +69,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);
     }
diff --git a/crates/pg_schema_cache/src/schemas.rs b/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()
     }
 }
diff --git a/crates/pg_schema_cache/src/tables.rs b/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()
     }
 }
diff --git a/crates/pg_schema_cache/src/types.rs b/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()
     }
 }
diff --git a/crates/pg_schema_cache/src/versions.rs b/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()
     }
 
     /*
diff --git a/crates/pg_typecheck/Cargo.toml b/crates/pg_typecheck/Cargo.toml
@@ -24,6 +24,7 @@ async-std = "1.12.0"
 
 
 [dev-dependencies]
+pg_test_utils.workspace = true
 
 [lib]
 doctest = false
diff --git a/crates/pg_typecheck/src/lib.rs b/crates/pg_typecheck/src/lib.rs
@@ -85,23 +85,21 @@ pub async fn check_sql<'a>(params: TypecheckerParams<'a>) -> Vec<TypeError> {
 #[cfg(test)]
 mod tests {
     use async_std::task::block_on;
-    use sqlx::PgPool;
+    use pg_test_utils::test_database::get_new_test_db;
 
     use crate::{check_sql, TypecheckerParams};
 
     #[test]
-    fn test_check_sql() {
+    fn test_basic_type() {
         let input = "select id, unknown from contact;";
 
-        let conn_string = std::env::var("DATABASE_URL").unwrap();
-
-        let pool = block_on(PgPool::connect(conn_string.as_str())).unwrap();
+        let test_db = block_on(get_new_test_db());
 
         let root = pg_query_ext::parse(input).unwrap();
         let ast = pg_syntax::parse_syntax(input, &root).ast;
 
         let errs = block_on(check_sql(TypecheckerParams {
-            conn: &pool,
+            conn: &test_db,
             sql: input,
             ast: &root,
             enriched_ast: Some(&ast),
@@ -111,6 +109,6 @@ mod tests {
 
         let e = &errs[0];
 
-        assert_eq!(&input[e.range.unwrap()], "unknown");
+        assert_eq!(&input[e.range.unwrap()], "contact");
     }
 }
diff --git a/crates/pg_workspace_new/src/workspace/server.rs b/crates/pg_workspace_new/src/workspace/server.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;

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ impl DbConnection {`
`41`	`41`	`match res {`
`42`	`42`	`Ok(not) => {`
`43`	`43`	`if not.payload().to_string() == "reload schema" {`
`44`		`- let schema_cache = SchemaCache::load(&cloned_pool).await;`
	`44`	`+ let schema_cache = SchemaCache::load(&cloned_pool).await.unwrap();`
`45`	`45`	`ide.write().await.set_schema_cache(schema_cache);`
`46`	`46`	`};`
`47`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ pub struct Function {`
`69`	`69`	`impl SchemaCacheItem for Function {`
`70`	`70`	`type Item = Function;`
`71`	`71`
`72`		`- async fn load(pool: &PgPool) -> Vec<Function> {`
	`72`	`+ async fn load(pool: &PgPool) -> Result<Vec<Function>, sqlx::Error> {`
`73`	`73`	`sqlx::query_as!(`
`74`	`74`	`Function,`
`75`	`75`	`r#"`
`@@ -179,6 +179,5 @@ from`
`179`	`179`	`)`
`180`	`180`	`.fetch_all(pool)`
`181`	`181`	`.await`
`182`		`- .unwrap()`
`183`	`182`	`}`
`184`	`183`	`}`