fix(workspace): gracefully handle failing db connections (#194)

Author: juleswritescode

crates/pg_cli/src/cli_options.rs

@@ -18,6 +18,10 @@ pub struct CliOptions {
     #[bpaf(long("use-server"), switch, fallback(false))]
     pub use_server: bool,
 
+    /// Skip connecting to the database and only run checks that don't require a database connection.
+    #[bpaf(long("skip-db"), switch, fallback(false))]
+    pub skip_db: bool,
+
     /// Print additional diagnostics, and some diagnostics show more information. Also, print out what files were processed and which ones were modified.
     #[bpaf(long("verbose"), switch, fallback(false))]
     pub verbose: bool,

crates/pg_cli/src/commands/mod.rs

@@ -35,8 +35,10 @@ pub enum PgLspCommand {
     Check {
         #[bpaf(external(partial_configuration), hide_usage, optional)]
         configuration: Option<PartialConfiguration>,
+
         #[bpaf(external, hide_usage)]
         cli_options: CliOptions,
+
         /// Use this option when you want to format code piped from `stdin`, and print the output to `stdout`.
         ///
         /// The file doesn't need to exist on disk, what matters is the extension of the file. Based on the extension, we know how to check the code.
@@ -286,6 +288,7 @@ pub(crate) trait CommandRunner: Sized {
             configuration,
             vcs_base_path,
             gitignore_matches,
+            skip_db: cli_options.skip_db,
         })?;
 
         let execution = self.get_execution(cli_options, console, workspace)?;

crates/pg_cli/src/execute/process_file/check.rs

@@ -28,6 +28,7 @@ pub(crate) fn check_with_guard<'ctx>(
             let (only, skip) = (Vec::new(), Vec::new());
 
             let max_diagnostics = ctx.remaining_diagnostics.load(Ordering::Relaxed);
+
             let pull_diagnostics_result = workspace_file
                 .guard()
                 .pull_diagnostics(

crates/pg_cli/src/execute/traverse.rs

@@ -123,6 +123,7 @@ pub(crate) fn traverse(
     let skipped = skipped.load(Ordering::Relaxed);
     let suggested_fixes_skipped = printer.skipped_fixes();
     let diagnostics_not_printed = printer.not_printed_diagnostics();
+
     Ok(TraverseResult {
         summary: TraversalSummary {
             changed,
@@ -381,6 +382,7 @@ impl<'ctx> DiagnosticsPrinter<'ctx> {
                 }
             }
         }
+
         diagnostics_to_print
     }
 }

crates/pg_cli/src/reporter/gitlab.rs

@@ -13,8 +13,8 @@ use std::{
 };
 
 pub struct GitLabReporter {
-    pub execution: Execution,
-    pub diagnostics: DiagnosticsPayload,
+    pub(crate) execution: Execution,
+    pub(crate) diagnostics: DiagnosticsPayload,
 }
 
 impl Reporter for GitLabReporter {

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_secs"), fallback(Some(10)), debug_fallback))]
+    pub conn_timeout_secs: 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: 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_lsp/src/session.rs

@@ -464,6 +464,7 @@ impl Session {
                             configuration: fs_configuration,
                             vcs_base_path,
                             gitignore_matches,
+                            skip_db: false,
                         });
 
                         if let Err(error) = result {