fix(pglsp): properly set up workspace/configuration flow

crates/pg_lsp/src/client/client_config_opts.rs

@@ -1,7 +1,25 @@
 use serde::Deserialize;
 
-// TODO: Check that the Opts are correct (existed in server.rs)
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug)]
 pub struct ClientConfigurationOptions {
-    pub db_connection_string: Option<String>,
+    #[serde(rename(deserialize = "databaseUrl"))]
+    pub(crate) db_connection_string: Option<String>,
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use crate::client::client_config_opts::ClientConfigurationOptions;
+
+    #[test]
+    fn test_json_parsing() {
+        let config = json!({
+            "databaseUrl": "cool-shit"
+        });
+
+        let parsed: ClientConfigurationOptions = serde_json::from_value(config).unwrap();
+
+        assert_eq!(parsed.db_connection_string, Some("cool-shit".into()));
+    }
 }

crates/pg_lsp/src/db_connection.rs

@@ -13,14 +13,20 @@ pub(crate) struct DbConnection {
 }
 
 impl DbConnection {
+    #[tracing::instrument(name = "Setting up new Database Connection…", skip(ide))]
     pub(crate) async fn new(
         connection_string: String,
         ide: Arc<RwLock<Workspace>>,
     ) -> Result<Self, sqlx::Error> {
+        tracing::info!("Trying to connect to pool…");
         let pool = PgPool::connect(&connection_string).await?;
+        tracing::info!("Connected to Pool.");
 
         let mut listener = PgListener::connect_with(&pool).await?;
+        tracing::info!("Connected to Listener.");
+
         listener.listen_all(["postgres_lsp", "pgrst"]).await?;
+        tracing::info!("Listening!");
 
         let (close_tx, close_rx) = tokio::sync::oneshot::channel::<()>();
 
@@ -52,6 +58,7 @@ impl DbConnection {
                 }
             }
         });
+        tracing::info!("Set up schema update handle.");
 
         Ok(Self {
             pool,

crates/pg_lsp/src/server.rs

@@ -66,40 +66,72 @@ impl LspServer {
         }
     }
 
-    #[tracing::instrument(name = "Requesting Configuration from Client", skip(self))]
-    async fn request_config_from_client(&self) -> Option<ClientConfigurationOptions> {
-        let params = ConfigurationParams {
-            items: vec![ConfigurationItem {
-                section: Some("pglsp".to_string()),
-                scope_uri: None,
-            }],
+    #[tracing::instrument(name = "Processing Config", skip(self))]
+    async fn process_config(&self, opts: Option<ClientConfigurationOptions>) -> anyhow::Result<()> {
+        if opts
+            .as_ref()
+            .is_some_and(|o| o.db_connection_string.is_some())
+        {
+            let conn_str = opts.unwrap().db_connection_string.unwrap();
+            self.session.change_db(conn_str).await
+        } else {
+            Ok(())
+        }
+    }
+
+    async fn parse_and_handle_config_from_client(&self, value: serde_json::Value) {
+        let parsed = self.parse_config_from_client(value).await;
+        match self.process_config(parsed).await {
+            Ok(_) => {}
+            Err(e) => {
+                self.client
+                    .show_message(
+                        MessageType::ERROR,
+                        format!("Unable to parse received config: {e:?}"),
+                    )
+                    .await;
+            }
         };
+    }
+
+    #[tracing::instrument(name = "Requesting & Handling Configuration from Client", skip(self))]
+    async fn request_and_handle_config_from_client(&self) {
+        let config_items = vec![ConfigurationItem {
+            section: Some("pglsp".to_string()),
+            scope_uri: None,
+        }];
 
         tracing::info!("sending workspace/configuration request");
-        match self
-            .client
-            .send_request::<request::WorkspaceConfiguration>(params)
-            .await
-        {
+        let config = match self.client.configuration(config_items).await {
             Ok(json) => {
                 // The client reponse fits the requested `ConfigurationParams.items`,
                 // so the first value is what we're looking for.
-                let relevant = json
-                    .into_iter()
+                json.into_iter()
                     .next()
-                    .expect("workspace/configuration request did not yield expected response.");
-
-                self.parse_config_from_client(relevant).await
+                    .expect("workspace/configuration request did not yield expected response.")
             }
             Err(why) => {
                 let message = format!(
                     "Unable to pull client options via workspace/configuration request: {}",
                     why
                 );
                 self.client.log_message(MessageType::ERROR, message).await;
-                None
+                return;
             }
-        }
+        };
+
+        let parsed = self.parse_config_from_client(config).await;
+        match self.process_config(parsed).await {
+            Ok(()) => {}
+            Err(e) => {
+                self.client
+                    .log_message(
+                        MessageType::ERROR,
+                        format!("Unable to process config from client: {e:?}"),
+                    )
+                    .await
+            }
+        };
     }
 
     #[tracing::instrument(
@@ -185,7 +217,11 @@ impl LanguageServer for LspServer {
         self.client
             .show_message(MessageType::INFO, "Initialize Request received")
             .await;
+
         let flags = ClientFlags::from_initialize_request_params(&params);
+
+        tracing::info!("flags: {:?}", flags);
+
         self.client_capabilities.write().await.replace(flags);
 
         Ok(InitializeResult {
@@ -220,6 +256,12 @@ impl LanguageServer for LspServer {
 
     #[tracing::instrument(name = "initialized", skip(self, _params))]
     async fn initialized(&self, _params: InitializedParams) {
+        let capabilities = self.client_capabilities.read().await;
+
+        if capabilities.as_ref().unwrap().supports_pull_opts {
+            self.request_and_handle_config_from_client().await;
+        }
+
         self.client
             .log_message(MessageType::INFO, "Postgres LSP Connected!")
             .await;
@@ -245,51 +287,11 @@ impl LanguageServer for LspServer {
         let capabilities = self.client_capabilities.read().await;
 
         if capabilities.as_ref().unwrap().supports_pull_opts {
-            let opts = self.request_config_from_client().await;
-            if opts
-                .as_ref()
-                .is_some_and(|o| o.db_connection_string.is_some())
-            {
-                let conn_str = opts.unwrap().db_connection_string.unwrap();
-                match self.session.change_db(conn_str).await {
-                    Ok(_) => {}
-                    Err(err) => {
-                        self.client
-                            .show_message(
-                                MessageType::ERROR,
-                                format!("Pulled Client Options but failed to set them: {}", err),
-                            )
-                            .await
-                    }
-                }
-                return;
-            }
-        }
-
-        // if we couldn't pull settings from the client,
-        // we'll try parsing the passed in params.
-        let opts = self.parse_config_from_client(params.settings).await;
-
-        if opts
-            .as_ref()
-            .is_some_and(|o| o.db_connection_string.is_some())
-        {
-            let conn_str = opts.unwrap().db_connection_string.unwrap();
-            match self.session.change_db(conn_str).await {
-                Ok(_) => {}
-                Err(err) => {
-                    self.client
-                        .show_message(
-                            MessageType::ERROR,
-                            format!(
-                                "Used Client Options from params but failed to set them: {}",
-                                err
-                            ),
-                        )
-                        .await
-                }
-            }
-        }
+            self.request_and_handle_config_from_client().await
+        } else {
+            self.parse_and_handle_config_from_client(params.settings)
+                .await
+        };
     }
 
     #[tracing::instrument(
@@ -332,13 +334,9 @@ impl LanguageServer for LspServer {
 
         self.publish_diagnostics(uri).await;
 
-        // TODO: "Compute Now"
         let changed_urls = self.session.recompute_and_get_changed_files().await;
         for url in changed_urls {
             let url = Url::from_file_path(url.as_path()).expect("Expected absolute File Path");
-
-            tracing::info!("publishing diagnostics: {}", url);
-
             self.publish_diagnostics(url).await;
         }
     }

crates/pg_lsp/src/session.rs

@@ -43,6 +43,7 @@ impl Session {
     /// If the passed-in connection string is the same that we're already connected to, it's a noop.
     /// Otherwise, it'll first open a new connection, replace `Self`'s connection, and then close
     /// the old one.
+    #[tracing::instrument(name = "Updating DB Connection", skip(self))]
     pub async fn change_db(&self, connection_string: String) -> anyhow::Result<()> {
         if self
             .db
@@ -55,17 +56,20 @@ impl Session {
             return Ok(());
         }
 
+        tracing::info!("Setting up new Database connection");
         let new_db = DbConnection::new(connection_string, Arc::clone(&self.ide)).await?;
+        tracing::info!("Set up new connection, trying to acquire write lock…");
 
         let mut current_db = self.db.write().await;
         let old_db = current_db.replace(new_db);
-        drop(current_db);
 
         if old_db.is_some() {
+            tracing::info!("Dropping previous Database Connection.");
             let old_db = old_db.unwrap();
             old_db.close().await;
         }
 
+        tracing::info!("Successfully set up new connection.");
         Ok(())
     }
 

docker-compose.yml

@@ -1,6 +1,7 @@
 version: "3.8"
 services:
   db:
+    # postgres://postgres:[email protected]:5432/postgres
     image: postgres
     restart: always
     environment:

editors/code/package.json

@@ -84,8 +84,13 @@
                     ],
                     "default": "off",
                     "description": "Traces the communication between VS Code and the language server."
+                },
+                "pglsp.databaseUrl": {
+                    "type": "string",
+                    "default": "",
+                    "description": "Your Postgres Database URL"
                 }
             }
         }
     }
-}
+}