fix(workspace)?: dont hang on db connection errors

Author: juleswritescode

crates/pg_lsp/src/handlers/completions.rs

@@ -1,6 +1,6 @@
 use crate::session::Session;
 use anyhow::Result;
-use pg_workspace::workspace;
+use pg_workspace::{workspace, WorkspaceError};
 use tower_lsp::lsp_types::{self, CompletionItem, CompletionItemLabelDetails};
 
 #[tracing::instrument(level = "trace", skip_all)]
@@ -26,12 +26,22 @@ pub fn get_completions(
         pg_lsp_converters::negotiated_encoding(client_capabilities),
     )?;
 
-    let completion_result = session
+    let completion_result = match session
         .workspace
         .get_completions(workspace::CompletionParams {
             path,
             position: offset,
-        })?;
+        }) {
+        Ok(result) => result,
+        Err(e) => match e {
+            WorkspaceError::DatabaseConnectionError(_) => {
+                return Ok(lsp_types::CompletionResponse::Array(vec![]));
+            }
+            _ => {
+                return Err(e.into());
+            }
+        },
+    };
 
     let items: Vec<CompletionItem> = completion_result
         .into_iter()

crates/pg_schema_cache/src/schema_cache.rs

@@ -9,6 +9,8 @@ use crate::versions::Version;
 
 #[derive(Debug, Clone, Default)]
 pub struct SchemaCache {
+    cached_conn_str: String,
+
     pub schemas: Vec<Schema>,
     pub tables: Vec<Table>,
     pub functions: Vec<Function>,
@@ -29,6 +31,7 @@ impl SchemaCache {
         )?;
 
         Ok(SchemaCache {
+            cached_conn_str: SchemaCache::pool_to_conn_str(pool),
             schemas,
             tables,
             functions,
@@ -38,6 +41,22 @@ impl SchemaCache {
         })
     }
 
+    fn pool_to_conn_str(pool: &PgPool) -> String {
+        let conn = pool.connect_options();
+
+        format!(
+            "postgres://{}:<redacted_pw>@{}:{}/{}",
+            conn.get_username(),
+            conn.get_host(),
+            conn.get_port(),
+            conn.get_database().unwrap_or("default")
+        )
+    }
+
+    pub fn has_already_cached_connection(&self, pool: &PgPool) -> bool {
+        self.cached_conn_str == SchemaCache::pool_to_conn_str(pool)
+    }
+
     /// Applies an AST node to the repository
     ///
     /// For example, alter table add column will add the column to the table if it does not exist

crates/pg_workspace/src/workspace.rs

@@ -119,9 +119,6 @@ pub trait Workspace: Send + Sync + RefUnwindSafe {
         params: CompletionParams,
     ) -> Result<pg_completions::CompletionResult, WorkspaceError>;
 
-    /// Refresh the schema cache for this workspace
-    fn refresh_schema_cache(&self) -> Result<(), WorkspaceError>;
-
     /// Update the global settings for this workspace
     fn update_settings(&self, params: UpdateSettingsParams) -> Result<(), WorkspaceError>;
 

crates/pg_workspace/src/workspace/client.rs

@@ -117,10 +117,6 @@ where
         self.request("pglsp/get_file_content", params)
     }
 
-    fn refresh_schema_cache(&self) -> Result<(), WorkspaceError> {
-        self.request("pglsp/refresh_schema_cache", ())
-    }
-
     fn pull_diagnostics(
         &self,
         params: super::PullDiagnosticsParams,