so far

Author: juleswritescode

crates/pg_lsp/src/b_server.rs

@@ -0,0 +1,159 @@
+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 tower_lsp::lsp_types::*;
+use tower_lsp::{Client, LanguageServer};
+
+use crate::client::client_flags::ClientFlags;
+use crate::db_connection::DbConnection;
+use crate::server::options::ClientConfigurationOptions;
+
+struct Server {
+    client: Client,
+    db: Mutex<Option<DbConnection>>,
+    ide: Arc<RwLock<Workspace>>,
+    client_capabilities: RwLock<Option<ClientFlags>>,
+}
+
+impl Server {
+    pub async fn new(client: Client) -> Self {
+        let ide = Arc::new(RwLock::new(Workspace::new()));
+        Self {
+            client,
+            db: Mutex::new(None),
+            ide,
+            client_capabilities: RwLock::new(None),
+        }
+    }
+
+    /// When the client sends a didChangeConfiguration notification, we need to parse the received JSON.
+    fn parse_options_from_client(
+        &self,
+        mut value: serde_json::Value,
+    ) -> Result<ClientConfigurationOptions> {
+        let options = match value.get_mut("pglsp") {
+            Some(section) => section.take(),
+            None => value,
+        };
+
+        let options = match serde_json::from_value::<ClientConfigurationOptions>(options) {
+            Ok(new_options) => Some(new_options),
+            Err(why) => {
+                let message = format!(
+                    "The texlab configuration is invalid; using the default settings instead.\nDetails: {why}"
+                );
+                let typ = MessageType::WARNING;
+                self.client
+                    .send_notification::<ShowMessage>(ShowMessageParams { message, typ });
+                None
+            }
+        };
+
+        Ok(options.unwrap_or_default())
+    }
+
+    async fn update_db_connection(
+        &self,
+        options: ClientConfigurationOptions,
+    ) -> anyhow::Result<()> {
+        if options.db_connection_string.is_none()
+            || self
+                .db
+                .lock()
+                .await
+                .as_ref()
+                // if the connection is already connected to the same database, do nothing
+                .is_some_and(|c| c.connected_to(options.db_connection_string.as_ref().unwrap()))
+        {
+            return Ok(());
+        }
+
+        let connection_string = options.db_connection_string.unwrap();
+
+        let mut db = DbConnection::new(connection_string).await?;
+
+        let ide = self.ide.clone();
+        db.listen_for_schema_updates(move |schema| {
+            let _guard = ide.blocking_write().set_schema_cache(schema);
+        });
+
+        let mut current_db = self.db.lock().await;
+        let old_db = current_db.replace(db);
+
+        if old_db.is_some() {
+            let old_db = old_db.unwrap();
+            old_db.close().await;
+        }
+
+        Ok(())
+    }
+}
+
+#[tower_lsp::async_trait]
+impl LanguageServer for Server {
+    async fn initialize(&self, params: InitializeParams) -> Result<InitializeResult> {
+        let flags = ClientFlags::from_initialize_request_params(&params);
+        self.client_capabilities.blocking_write().replace(flags);
+
+        Ok(InitializeResult {
+            server_info: None,
+            capabilities: ServerCapabilities {
+                text_document_sync: Some(TextDocumentSyncCapability::Options(
+                    TextDocumentSyncOptions {
+                        open_close: Some(true),
+                        change: Some(TextDocumentSyncKind::INCREMENTAL),
+                        will_save: None,
+                        will_save_wait_until: None,
+                        save: Some(TextDocumentSyncSaveOptions::SaveOptions(SaveOptions {
+                            include_text: Some(false),
+                        })),
+                    },
+                )),
+                hover_provider: Some(HoverProviderCapability::Simple(true)),
+                execute_command_provider: Some(ExecuteCommandOptions {
+                    commands: CommandType::ALL
+                        .iter()
+                        .map(|c| c.id().to_string())
+                        .collect(),
+                    ..Default::default()
+                }),
+                inlay_hint_provider: Some(OneOf::Left(true)),
+                code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
+                completion_provider: Some(CompletionOptions::default()),
+                ..ServerCapabilities::default()
+            },
+        })
+    }
+
+    async fn initialized(&self, _params: InitializedParams) {
+        self.client
+            .log_message(MessageType::INFO, "Postgres LSP Connected!")
+            .await;
+    }
+
+    async fn shutdown(&self) -> 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;
+            }
+            Err(e) => {
+                self.client
+                    .log_message(MessageType::ERROR, format!("Error parsing configuration: {}", e))
+                    .await;
+            }
+        };
+
+    }
+}

crates/pg_lsp/src/client.rs

@@ -1,52 +1,28 @@
 pub mod client_flags;
 
-use std::{
-    collections::HashMap,
-    sync::{
-        atomic::{AtomicI32, Ordering},
-        Arc, Mutex,
-    },
-};
+use anyhow::Result;
+use serde::Serialize;
+use tower_lsp::lsp_types::{notification::ShowMessage, MessageType, ShowMessageParams};
+use tower_lsp::Client;
 
-use anyhow::{bail, Result};
-use crossbeam_channel::Sender;
-use lsp_server::{ErrorCode, Message, Request, RequestId, Response};
-use lsp_types::{notification::ShowMessage, MessageType, ShowMessageParams};
-use serde::{de::DeserializeOwned, Serialize};
-
-use crate::server::options::Options;
-
-#[derive(Debug)]
-struct RawClient {
-    sender: Sender<Message>,
-    next_id: AtomicI32,
-    pending: Mutex<HashMap<RequestId, Sender<Response>>>,
-}
+use crate::server::options::ClientConfigurationOptions;
 
 #[derive(Debug, Clone)]
 pub struct LspClient {
-    raw: Arc<RawClient>,
+    client: Client,
 }
 
 impl LspClient {
-    pub fn new(sender: Sender<Message>) -> Self {
-        let raw = Arc::new(RawClient {
-            sender,
-            next_id: AtomicI32::new(1),
-            pending: Default::default(),
-        });
-
-        Self { raw }
+    pub fn new(client: Client) -> Self {
+        Self { client }
     }
 
     pub fn send_notification<N>(&self, params: N::Params) -> Result<()>
     where
-        N: lsp_types::notification::Notification,
+        N: tower_lsp::lsp_types::notification::Notification,
         N::Params: Serialize,
     {
-        self.raw
-            .sender
-            .send(lsp_server::Notification::new(N::METHOD.to_string(), params).into())?;
+        self.client.send_notification::<N>(params);
         Ok(())
     }
 
@@ -58,54 +34,19 @@ impl LspClient {
         });
     }
 
-    pub fn send_request<R>(&self, params: R::Params) -> Result<R::Result>
+    pub async fn send_request<R>(&self, params: R::Params) -> Result<R::Result>
     where
-        R: lsp_types::request::Request,
-        R::Params: Serialize,
-        R::Result: DeserializeOwned,
+        R: tower_lsp::lsp_types::request::Request,
     {
-        let id = RequestId::from(self.raw.next_id.fetch_add(1, Ordering::SeqCst));
+        let response = self.client.send_request::<R>(params).await?;
 
-        let (tx, rx) = crossbeam_channel::bounded(1);
-        self.raw.pending.lock().unwrap().insert(id.clone(), tx);
-
-        self.raw
-            .sender
-            .send(Request::new(id, R::METHOD.to_string(), params).into())?;
-
-        let response = rx.recv()?;
-        let result = match response.error {
-            Some(error) => bail!(error.message),
-            None => response.result.unwrap_or_default(),
-        };
-
-        Ok(serde_json::from_value(result)?)
-    }
-
-    pub fn send_response(&self, response: lsp_server::Response) -> Result<()> {
-        self.raw.sender.send(response.into())?;
-        Ok(())
-    }
-
-    pub fn send_error(&self, id: RequestId, code: ErrorCode, message: String) -> Result<()> {
-        self.send_response(lsp_server::Response::new_err(id, code as i32, message))?;
-        Ok(())
-    }
-
-    pub fn recv_response(&self, response: lsp_server::Response) -> Result<()> {
-        let tx = self
-            .raw
-            .pending
-            .lock()
-            .unwrap()
-            .remove(&response.id)
-            .expect("response with known request id received");
-
-        tx.send(response)?;
-        Ok(())
+        Ok(response)
     }
 
-    pub fn parse_options(&self, mut value: serde_json::Value) -> Result<Options> {
+    pub fn parse_options(
+        &self,
+        mut value: serde_json::Value,
+    ) -> Result<ClientConfigurationOptions> {
         // if there are multiple servers, we need to extract the options for pglsp first
         let options = match value.get_mut("pglsp") {
             Some(section) => section.take(),

crates/pg_lsp/src/client/client_flags.rs

@@ -1,4 +1,4 @@
-use lsp_types::InitializeParams;
+use tower_lsp::lsp_types::InitializeParams;
 
 /// Contains information about the client's capabilities.
 /// This is used to determine which features the server can use.

crates/pg_lsp/src/lib.rs

@@ -2,3 +2,5 @@ mod client;
 mod db_connection;
 pub mod server;
 mod utils;
+
+mod b_server;

crates/pg_lsp/src/server.rs

@@ -29,6 +29,7 @@ use pg_workspace::Workspace;
 use serde::{de::DeserializeOwned, Serialize};
 use std::{collections::HashSet, future::Future, sync::Arc, time::Duration};
 use text_size::TextSize;
+use tower_lsp::Client;
 
 use tokio::sync::mpsc;
 use tokio_util::sync::CancellationToken;
@@ -39,13 +40,13 @@ use crate::{
     utils::{file_path, from_proto, line_index_ext::LineIndexExt, normalize_uri, to_proto},
 };
 
-use self::{debouncer::EventDebouncer, options::Options};
+use self::{debouncer::EventDebouncer, options::ClientConfigurationOptions};
 use sqlx::{postgres::PgPool, Executor};
 
 #[derive(Debug)]
 enum InternalMessage {
     PublishDiagnostics(lsp_types::Url),
-    SetOptions(Options),
+    SetOptions(ClientConfigurationOptions),
     SetSchemaCache(SchemaCache),
     SetDatabaseConnection(DbConnection),
 }
@@ -61,7 +62,7 @@ fn get_client_receiver(
     tokio::task::spawn(async move {
         loop {
             let msg = match connection.receiver.recv() {
-                Ok(msg) => msg, 
+                Ok(msg) => msg,
                 Err(e) => {
                     eprint!("Connection was closed by LSP client: {}", e);
                     cancel_token.cancel();
@@ -97,12 +98,12 @@ pub struct Server {
 }
 
 impl Server {
-    pub fn init(connection: Connection) -> anyhow::Result<Self> {
-        let client = LspClient::new(connection.sender.clone());
+    pub fn init(client: Client) -> anyhow::Result<Self> {
+        let client = LspClient::new(client);
         let cancel_token = Arc::new(CancellationToken::new());
 
-        let (client_flags, client_rx) = Self::establish_client_connection(connection, &cancel_token)?;
-
+        let (client_flags, client_rx) =
+            Self::establish_client_connection(connection, &cancel_token)?;
 
         let ide = Arc::new(Workspace::new());
 
@@ -211,7 +212,7 @@ impl Server {
         });
     }
 
-    fn update_db_connection(&self, options: Options) -> anyhow::Result<()> {
+    fn update_db_connection(&self, options: ClientConfigurationOptions) -> anyhow::Result<()> {
         if options.db_connection_string.is_none()
             || self
                 .db_conn

crates/pg_lsp/src/server/options.rs

@@ -3,6 +3,6 @@ use serde::{Deserialize, Serialize};
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 #[serde(default)]
-pub struct Options {
+pub struct ClientConfigurationOptions {
     pub db_connection_string: Option<String>,
 }

crates/pg_schema_cache/src/lib.rs

@@ -4,11 +4,11 @@
 #![feature(future_join)]
 
 mod functions;
-mod versions;
 mod schema_cache;
 mod schemas;
 mod tables;
 mod types;
+mod versions;
 
 use sqlx::postgres::PgPool;