fix(workspace): gracefully handle failing db connections

crates/pg_completions/src/context.rs

@@ -266,7 +266,7 @@ mod tests {
                 position: (position as u32).into(),
                 text,
                 tree: Some(&tree),
-                schema: &pg_schema_cache::SchemaCache::new(),
+                schema: &pg_schema_cache::SchemaCache::default(),
             };
 
             let ctx = CompletionContext::new(&params);
@@ -298,7 +298,7 @@ mod tests {
                 position: (position as u32).into(),
                 text,
                 tree: Some(&tree),
-                schema: &pg_schema_cache::SchemaCache::new(),
+                schema: &pg_schema_cache::SchemaCache::default(),
             };
 
             let ctx = CompletionContext::new(&params);
@@ -332,7 +332,7 @@ mod tests {
                 position: (position as u32).into(),
                 text,
                 tree: Some(&tree),
-                schema: &pg_schema_cache::SchemaCache::new(),
+                schema: &pg_schema_cache::SchemaCache::default(),
             };
 
             let ctx = CompletionContext::new(&params);
@@ -357,7 +357,7 @@ mod tests {
                 position: (position as u32).into(),
                 text,
                 tree: Some(&tree),
-                schema: &pg_schema_cache::SchemaCache::new(),
+                schema: &pg_schema_cache::SchemaCache::default(),
             };
 
             let ctx = CompletionContext::new(&params);
@@ -385,7 +385,7 @@ mod tests {
             position: (position as u32).into(),
             text,
             tree: Some(&tree),
-            schema: &pg_schema_cache::SchemaCache::new(),
+            schema: &pg_schema_cache::SchemaCache::default(),
         };
 
         let ctx = CompletionContext::new(&params);
@@ -411,7 +411,7 @@ mod tests {
             position: (position as u32).into(),
             text,
             tree: Some(&tree),
-            schema: &pg_schema_cache::SchemaCache::new(),
+            schema: &pg_schema_cache::SchemaCache::default(),
         };
 
         let ctx = CompletionContext::new(&params);
@@ -436,7 +436,7 @@ mod tests {
             position: (position as u32).into(),
             text,
             tree: Some(&tree),
-            schema: &pg_schema_cache::SchemaCache::new(),
+            schema: &pg_schema_cache::SchemaCache::default(),
         };
 
         let ctx = CompletionContext::new(&params);

crates/pg_configuration/src/database.rs

@@ -26,6 +26,10 @@ pub struct DatabaseConfiguration {
     /// The name of the database.
     #[partial(bpaf(long("database")))]
     pub database: String,
+
+    /// The connection timeout in seconds.
+    #[partial(bpaf(long("conn_timeout")))]
+    pub conn_timeout_secs: Option<u16>,
 }
 
 impl Default for DatabaseConfiguration {
@@ -36,15 +40,7 @@ impl Default for DatabaseConfiguration {
             username: "postgres".to_string(),
             password: "postgres".to_string(),
             database: "postgres".to_string(),
+            conn_timeout_secs: Some(10),
         }
     }
 }
-
-impl DatabaseConfiguration {
-    pub fn to_connection_string(&self) -> String {
-        format!(
-            "postgres://{}:{}@{}:{}/{}",
-            self.username, self.password, self.host, self.port, self.database
-        )
-    }
-}

crates/pg_configuration/src/lib.rs

@@ -104,6 +104,7 @@ impl PartialConfiguration {
                 username: Some("postgres".to_string()),
                 password: Some("postgres".to_string()),
                 database: Some("postgres".to_string()),
+                conn_timeout_secs: Some(10),
             }),
         }
     }

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>,
@@ -18,10 +20,6 @@ pub struct SchemaCache {
 }
 
 impl SchemaCache {
-    pub fn new() -> SchemaCache {
-        SchemaCache::default()
-    }
-
     pub async fn load(pool: &PgPool) -> Result<SchemaCache, sqlx::Error> {
         let (schemas, tables, functions, types, versions, columns) = futures_util::try_join!(
             Schema::load(pool),
@@ -33,6 +31,7 @@ impl SchemaCache {
         )?;
 
         Ok(SchemaCache {
+            cached_conn_str: SchemaCache::pool_to_conn_str(pool),
             schemas,
             tables,
             functions,
@@ -42,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/settings.rs

@@ -5,6 +5,7 @@ use std::{
     num::NonZeroU64,
     path::{Path, PathBuf},
     sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
+    time::Duration,
 };
 
 use ignore::gitignore::{Gitignore, GitignoreBuilder};
@@ -266,6 +267,7 @@ pub struct DatabaseSettings {
     pub username: String,
     pub password: String,
     pub database: String,
+    pub conn_timeout_secs: Duration,
 }
 
 impl Default for DatabaseSettings {
@@ -276,19 +278,11 @@ impl Default for DatabaseSettings {
             username: "postgres".to_string(),
             password: "postgres".to_string(),
             database: "postgres".to_string(),
+            conn_timeout_secs: Duration::from_secs(10),
         }
     }
 }
 
-impl DatabaseSettings {
-    pub fn to_connection_string(&self) -> String {
-        format!(
-            "postgres://{}:{}@{}:{}/{}",
-            self.username, self.password, self.host, self.port, self.database
-        )
-    }
-}
-
 impl From<PartialDatabaseConfiguration> for DatabaseSettings {
     fn from(value: PartialDatabaseConfiguration) -> Self {
         let d = DatabaseSettings::default();
@@ -298,6 +292,10 @@ impl From<PartialDatabaseConfiguration> for DatabaseSettings {
             username: value.username.unwrap_or(d.username),
             password: value.password.unwrap_or(d.password),
             database: value.database.unwrap_or(d.database),
+            conn_timeout_secs: value
+                .conn_timeout_secs
+                .map(|s| Duration::from_secs(s.into()))
+                .unwrap_or(d.conn_timeout_secs),
         }
     }
 }

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,