todo: adjust result types to use tower_lsp::jsonrpc

Author: juleswritescode

crates/pg_lsp/src/b_server.rs

@@ -3,8 +3,8 @@ use std::sync::Arc;
 use notification::ShowMessage;
 use pg_commands::CommandType;
 use pg_workspace::Workspace;
-use tokio::sync::{Mutex, RwLock};
-use tower_lsp::jsonrpc::Result;
+use tokio::sync::RwLock;
+use tower_lsp::jsonrpc::Error;
 use tower_lsp::lsp_types::*;
 use tower_lsp::{Client, LanguageServer};
 
@@ -14,7 +14,7 @@ use crate::server::options::ClientConfigurationOptions;
 
 struct Server {
     client: Client,
-    db: Mutex<Option<DbConnection>>,
+    db: RwLock<Option<DbConnection>>,
     ide: Arc<RwLock<Workspace>>,
     client_capabilities: RwLock<Option<ClientFlags>>,
 }
@@ -24,7 +24,7 @@ impl Server {
         let ide = Arc::new(RwLock::new(Workspace::new()));
         Self {
             client,
-            db: Mutex::new(None),
+            db: RwLock::new(None),
             ide,
             client_capabilities: RwLock::new(None),
         }
@@ -34,13 +34,13 @@ impl Server {
     fn parse_options_from_client(
         &self,
         mut value: serde_json::Value,
-    ) -> Result<ClientConfigurationOptions> {
+    ) -> Option<ClientConfigurationOptions> {
         let options = match value.get_mut("pglsp") {
             Some(section) => section.take(),
             None => value,
         };
 
-        let options = match serde_json::from_value::<ClientConfigurationOptions>(options) {
+        match serde_json::from_value::<ClientConfigurationOptions>(options) {
             Ok(new_options) => Some(new_options),
             Err(why) => {
                 let message = format!(
@@ -51,19 +51,21 @@ impl Server {
                     .send_notification::<ShowMessage>(ShowMessageParams { message, typ });
                 None
             }
-        };
-
-        Ok(options.unwrap_or_default())
+        }
     }
 
+    /// `update_db_connection` will update `Self`'s database connection.
+    /// 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.
     async fn update_db_connection(
         &self,
         options: ClientConfigurationOptions,
     ) -> anyhow::Result<()> {
         if options.db_connection_string.is_none()
             || self
                 .db
-                .lock()
+                .read()
                 .await
                 .as_ref()
                 // if the connection is already connected to the same database, do nothing
@@ -81,7 +83,7 @@ impl Server {
             let _guard = ide.blocking_write().set_schema_cache(schema);
         });
 
-        let mut current_db = self.db.lock().await;
+        let mut current_db = self.db.blocking_write();
         let old_db = current_db.replace(db);
 
         if old_db.is_some() {
@@ -91,11 +93,48 @@ impl Server {
 
         Ok(())
     }
+
+    async fn request_opts_from_client(&self) -> Option<ClientConfigurationOptions> {
+        let params = ConfigurationParams {
+            items: vec![ConfigurationItem {
+                section: Some("pglsp".to_string()),
+                scope_uri: None,
+            }],
+        };
+
+        match self
+            .client
+            .send_request::<request::WorkspaceConfiguration>(params)
+            .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()
+                    .next()
+                    .expect("workspace/configuration request did not yield expected response.");
+
+                let opts = self.parse_options_from_client(relevant);
+
+                opts
+            }
+            Err(why) => {
+                let message = format!(
+                    "Unable to pull client options via workspace/configuration request: {}",
+                    why
+                );
+                println!("{}", message);
+                self.client.log_message(MessageType::ERROR, message);
+                None
+            }
+        }
+    }
 }
 
 #[tower_lsp::async_trait]
 impl LanguageServer for Server {
-    async fn initialize(&self, params: InitializeParams) -> Result<InitializeResult> {
+    async fn initialize(&self, params: InitializeParams) -> tower_lsp::jsonrpc::Result<InitializeResult> {
         let flags = ClientFlags::from_initialize_request_params(&params);
         self.client_capabilities.blocking_write().replace(flags);
 
@@ -135,25 +174,86 @@ impl LanguageServer for Server {
             .await;
     }
 
-    async fn shutdown(&self) -> Result<()> {
+    async fn shutdown(&self) -> anyhow::Result<()> {
         self.client
             .log_message(MessageType::INFO, "Postgres LSP terminated.")
             .await;
         Ok(())
     }
 
-
     async fn did_change_configuration(&self, params: DidChangeConfigurationParams) {
-        match self.parse_options_from_client(params.settings) {
-            Ok(opts) => {
-                self.update_db_connection(opts).await;
+        let capabilities = self.client_capabilities.read().await;
+
+        if capabilities.as_ref().unwrap().supports_pull_opts {
+            let opts = self.request_opts_from_client().await;
+            if opts.is_some() {
+                self.update_db_connection(opts.unwrap()).await;
+                return;
+            }
+        }
+
+        let opts = self.parse_options_from_client(params.settings);
+
+        if opts.is_some() {
+            self.update_db_connection(opts.unwrap()).await;
+        }
+    }
+
+    async fn execute_command(
+        &self,
+        params: ExecuteCommandParams,
+    ) -> tower_lsp::jsonrpc::Result<Option<serde_json::Value>> {
+        match CommandType::from_id(params.command.replace("pglsp.", "").as_str()) {
+            Some(CommandType::ExecuteStatement) => {
+                if params.arguments.is_empty() {
+                    return tower_lsp::jsonrpc::Result::Err(Error::new("No arguments provided!"));
+                }
+
+                let stmt = params
+                    .arguments
+                    .into_iter()
+                    .next()
+                    .map(|v| serde_json::from_value(v))
+                    .unwrap()?;
+
+                let conn = self.db.read().await;
+                match conn
+                    .as_ref()
+                    .expect("No connection to the database.")
+                    .run_stmt(stmt)
+                    .await
+                {
+                    Ok(pg_result) => {
+                        self.client
+                            .send_notification::<ShowMessage>(ShowMessageParams {
+                                typ: MessageType::INFO,
+                                message: format!(
+                                    "Success! Affected rows: {}",
+                                    pg_result.rows_affected()
+                                ),
+                            })
+                            .await;
+                    }
+                    Err(why) => {
+                        self.client
+                            .send_notification::<ShowMessage>(ShowMessageParams {
+                                typ: MessageType::ERROR,
+                                message: format!("Error! Statement exectuion failed: {}", why),
+                            })
+                            .await;
+                    }
+                };
             }
-            Err(e) => {
+            None => {
                 self.client
-                    .log_message(MessageType::ERROR, format!("Error parsing configuration: {}", e))
+                    .show_message(
+                        MessageType::ERROR,
+                        format!("Unknown command: {}", params.command),
+                    )
                     .await;
             }
         };
 
+        Ok(None)
     }
 }

crates/pg_lsp/src/client/client_flags.rs

@@ -5,22 +5,22 @@ use tower_lsp::lsp_types::InitializeParams;
 #[derive(Debug, Clone)]
 pub struct ClientFlags {
     /// If `true`, the server can pull configuration from the client.
-    pub has_configuration: bool,
+    pub supports_pull_opts: bool,
 
     /// If `true`, the client notifies the server when its configuration changes.
-    pub will_push_configuration: bool,
+    pub supports_dynamic_registration: bool,
 }
 
 impl ClientFlags {
     pub(crate) fn from_initialize_request_params(params: &InitializeParams) -> Self {
-        let has_configuration = params
+        let supports_pull_opts = params
             .capabilities
             .workspace
             .as_ref()
             .and_then(|w| w.configuration)
             .unwrap_or(false);
 
-        let will_push_configuration = params
+        let supports_dynamic_registration = params
             .capabilities
             .workspace
             .as_ref()
@@ -29,8 +29,8 @@ impl ClientFlags {
             .unwrap_or(false);
 
         Self {
-            has_configuration,
-            will_push_configuration,
+            supports_pull_opts,
+            supports_dynamic_registration,
         }
     }
 }

crates/pg_lsp/src/db_connection.rs

@@ -1,5 +1,9 @@
+use pg_commands::ExecuteStatementCommand;
 use pg_schema_cache::SchemaCache;
-use sqlx::{postgres::PgListener, PgPool};
+use sqlx::{
+    postgres::{PgListener, PgQueryResult},
+    PgPool,
+};
 use tokio::task::JoinHandle;
 
 #[derive(Debug)]
@@ -19,6 +23,14 @@ impl DbConnection {
         })
     }
 
+    /// TODO: this should simply take a `Command` type, and the individual
+    /// enums should have their deps included (i.e. `ExecuteStatement(String)`)
+    pub async fn run_stmt(&self, stmt: String) -> anyhow::Result<PgQueryResult> {
+        let command = ExecuteStatementCommand::new(stmt);
+        let pool = self.pool.clone();
+        command.run(Some(pool)).await
+    }
+
     pub(crate) fn connected_to(&self, connection_string: &str) -> bool {
         connection_string == self.connection_string
     }

crates/pg_lsp/src/server.rs

@@ -725,7 +725,7 @@ impl Server {
         &mut self,
         params: DidChangeConfigurationParams,
     ) -> anyhow::Result<()> {
-        if self.client_flags.has_configuration {
+        if self.client_flags.supports_pull_opts {
             self.pull_options();
         } else {
             let options = self.client.parse_options(params.settings)?;
@@ -913,11 +913,11 @@ impl Server {
     }
 
     pub async fn run(mut self) -> anyhow::Result<()> {
-        if self.client_flags.will_push_configuration {
+        if self.client_flags.supports_dynamic_registration {
             self.register_configuration();
         }
 
-        if self.client_flags.has_configuration {
+        if self.client_flags.supports_pull_opts {
             self.pull_options();
         }