three to go

Author: juleswritescode

crates/pg_lsp/src/b_server.rs

@@ -7,22 +7,22 @@ use tower_lsp::{Client, LanguageServer};
 
 use crate::client::client_flags::ClientFlags;
 use crate::server::options::ClientConfigurationOptions;
+use crate::session::Session;
 use crate::utils::file_path;
 use crate::utils::normalize_uri;
 use crate::utils::to_proto;
-use crate::workspace_handler::WorkspaceHandler;
 
 struct Server {
     client: Client,
-    workspace_handler: WorkspaceHandler,
+    session: Session,
     client_capabilities: RwLock<Option<ClientFlags>>,
 }
 
 impl Server {
     pub async fn new(client: Client) -> Self {
         Self {
             client,
-            workspace_handler: WorkspaceHandler::new(),
+            session: Session::new(),
             client_capabilities: RwLock::new(None),
         }
     }
@@ -91,10 +91,8 @@ impl Server {
     async fn publish_diagnostics(&self, mut uri: Url) {
         normalize_uri(&mut uri);
 
-        let diagnostics = self
-            .workspace_handler
-            .get_diagnostics(file_path(&uri))
-            .await;
+        let url = file_path(&uri);
+        let diagnostics = self.session.get_diagnostics(url).await;
 
         let diagnostics: Vec<Diagnostic> = diagnostics
             .into_iter()
@@ -177,22 +175,79 @@ impl LanguageServer for Server {
                 .is_some_and(|o| o.db_connection_string.is_some())
             {
                 let conn_str = opts.unwrap().db_connection_string.unwrap();
-                self.workspace_handler.change_db(conn_str).await;
+                self.session.change_db(conn_str).await;
                 return;
             }
         }
 
+        // if we couldn't pull settings from the client,
+        // we'll try parsing the passed in params.
         let opts = self.parse_options_from_client(params.settings);
 
         if opts
             .as_ref()
             .is_some_and(|o| o.db_connection_string.is_some())
         {
             let conn_str = opts.unwrap().db_connection_string.unwrap();
-            self.workspace_handler.change_db(conn_str).await;
+            self.session.change_db(conn_str).await;
+        }
+    }
+
+    async fn did_open(&self, params: DidOpenTextDocumentParams) {
+        let mut uri = params.text_document.uri;
+        normalize_uri(&mut uri);
+
+        let changed_urls = self
+            .session
+            .apply_doc_changes(
+                file_path(url),
+                params.text_document.version,
+                params.text_document.text,
+            )
+            .await;
+
+        for url in changed_urls {
+            self.publish_diagnostics(url).await;
+        }
+    }
+
+    async fn did_save(&self, params: DidSaveTextDocumentParams) {
+        let mut uri = params.text_document.uri;
+        normalize_uri(&mut uri);
+
+        self.publish_diagnostics(uri).await;
+
+        // TODO: "Compute Now"
+        let changed_urls = self.session.recompute_and_get_changed_files();
+        for url in changed_urls {
+            self.publish_diagnostics(url).await;
         }
     }
 
+    async fn did_change(&self, params: DidChangeTextDocumentParams) {
+        todo!()
+    }
+
+    async fn did_close(&self, params: DidSaveTextDocumentParams) {
+        let mut uri = params.text_document.uri;
+        normalize_uri(&mut uri);
+        let path = file_path(&uri);
+
+        self.session.on_file_closed(path);
+    }
+
+    async fn code_action(&self, params: CodeActionParams) -> Result<Option<CodeActionResponse>> {
+        let mut uri = params.text_document.uri;
+        normalize_uri(&mut uri);
+
+        let path = file_path(&uri);
+        let range = params.range;
+
+        let actions = self.session.get_available_code_actions(path, range);
+
+        Ok(actions)
+    }
+
     async fn execute_command(
         &self,
         params: ExecuteCommandParams,
@@ -207,7 +262,7 @@ impl LanguageServer for Server {
                 let stmt = serde_json::from_value(params)
                     .map_err(|_| jsonrpc::Error::invalid_request())?;
 
-                match self.workspace_handler.run_stmt(stmt).await {
+                match self.session.run_stmt(stmt).await {
                     Ok(rows_affected) => {
                         self.client
                             .send_notification::<ShowMessage>(ShowMessageParams {

crates/pg_lsp/src/db_connection.rs

@@ -8,7 +8,7 @@ use tokio::task::JoinHandle;
 
 #[derive(Debug)]
 pub(crate) struct DbConnection {
-    pub pool: PgPool,
+    pool: PgPool,
     connection_string: String,
     schema_update_handle: Option<JoinHandle<()>>,
 }
@@ -75,4 +75,8 @@ impl DbConnection {
 
         Ok(())
     }
+
+    pub(crate) fn get_pool(&self) -> PgPool {
+        self.pool.clone()
+    }
 }

crates/pg_lsp/src/lib.rs

@@ -4,4 +4,4 @@ pub mod server;
 mod utils;
 
 mod b_server;
-mod workspace_handler;
+mod session;

crates/pg_lsp/src/server.rs

@@ -197,7 +197,9 @@ impl Server {
                     })
                     .unwrap();
             }
+
             let changed = cloned_ide.compute(conn);
+
             let urls = HashSet::<&str>::from_iter(
                 changed.iter().map(|f| f.document_url.to_str().unwrap()),
             );
@@ -375,7 +377,7 @@ impl Server {
             DocumentChange::new(params.text_document.version, changes),
         );
 
-        let conn = self.db_conn.as_ref().map(|p| p.pool.clone());
+        let conn = self.db_conn.as_ref().map(|p| p.get_pool());
         self.compute_debouncer.put(conn);
 
         Ok(())

crates/pg_lsp/src/session.rs

@@ -0,0 +1,180 @@
+use std::{collections::HashSet, sync::Arc};
+
+use pg_base_db::{Change, DocumentChange, PgLspPath};
+use pg_diagnostics::Diagnostic;
+use pg_workspace::Workspace;
+use tokio::sync::RwLock;
+use tower_lsp::lsp_types::{CodeAction, Range};
+
+use crate::{db_connection::DbConnection, utils::line_index_ext::LineIndexExt};
+
+pub struct Session {
+    db: RwLock<Option<DbConnection>>,
+    ide: Arc<RwLock<Workspace>>,
+}
+
+impl Session {
+    pub fn new() -> Self {
+        let ide = Arc::new(RwLock::new(Workspace::new()));
+        Self {
+            db: RwLock::new(None),
+            ide,
+        }
+    }
+
+    /// `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.
+    pub async fn change_db(&self, connection_string: String) -> anyhow::Result<()> {
+        if self
+            .db
+            .read()
+            .await
+            .as_ref()
+            // if the connection is already connected to the same database, do nothing
+            .is_some_and(|c| c.connected_to(&connection_string))
+        {
+            return Ok(());
+        }
+
+        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.blocking_write();
+        let old_db = current_db.replace(db);
+
+        if old_db.is_some() {
+            let old_db = old_db.unwrap();
+            old_db.close().await;
+        }
+
+        Ok(())
+    }
+
+    pub async fn run_stmt(&self, stmt: String) -> anyhow::Result<u64> {
+        let db = self.db.read().await;
+        db.as_ref()
+            .expect("No Db Connection")
+            .run_stmt(stmt)
+            .await
+            .map(|pg_query_result| pg_query_result.rows_affected())
+    }
+
+    pub async fn on_file_closed(&self, path: PgLspPath) {
+        let ide = self.ide.read().await;
+        ide.remove_document(path);
+    }
+
+    pub async fn get_diagnostics(&self, path: PgLspPath) -> Vec<(Diagnostic, Range)> {
+        let ide = self.ide.read().await;
+
+        // make sure there are documents at the provided path before
+        // trying to collect diagnostics.
+        let doc = ide.documents.get(&path);
+        if doc.is_none() {
+            return vec![];
+        }
+
+        ide.diagnostics(&path)
+            .into_iter()
+            .map(|d| {
+                let range = doc
+                    .as_ref()
+                    .unwrap()
+                    .line_index
+                    .line_col_lsp_range(d.range)
+                    .unwrap();
+                (d, range)
+            })
+            .collect()
+    }
+
+    pub async fn apply_doc_changes(
+        &self,
+        path: PgLspPath,
+        version: i32,
+        text: String,
+    ) -> HashSet<String> {
+        {
+            let ide = self.ide.read().await;
+
+            let doc = ide.documents.get(&path);
+            if doc.is_none() {
+                return HashSet::new();
+            }
+
+            ide.apply_change(
+                path,
+                DocumentChange::new(version, vec![Change { range: None, text }]),
+            );
+        }
+
+        self.recompute_and_get_changed_files()
+    }
+
+    pub async fn recompute_and_get_changed_files(&self) -> HashSet<String> {
+        let ide = self.ide.read().await;
+
+        let db = self.db.read().await;
+        let pool = db.as_ref().map(|d| d.get_pool());
+
+        let changed_files = ide.compute(pool);
+
+        changed_files
+            .into_iter()
+            .map(|f| f.document_url.to_string_lossy().to_string())
+            .collect()
+    }
+
+    pub async fn get_available_code_actions(
+        &self,
+        path: PgLspPath,
+        range: Range,
+    ) -> Option<Vec<CodeAction>> {
+        let ide = self.ide.read().await;
+        let db = self.db.read().await;
+
+        let doc = ide.documents.get(&path);
+        if doc.is_none() || db.is_none() {
+            return None;
+        }
+
+        let doc = doc.unwrap();
+
+        let range = doc.line_index.offset_lsp_range(range).unwrap();
+
+        // for now, we only provide `ExcecuteStatementCommand`s.
+        let actions = doc
+            .statements_at_range(&range)
+            .into_iter()
+            .map(|stmt| {
+                let cmd = ExecuteStatementCommand::command_type();
+                let title = format!(
+                    "Execute '{}'",
+                    ExecuteStatementCommand::trim_statement(stmt.text.clone(), 50)
+                );
+                CodeAction {
+                    title: title.clone(),
+                    kind: None,
+                    edit: None,
+                    command: Some(Command {
+                        title,
+                        command: format!("pglsp.{}", cmd.id()),
+                        arguments: Some(vec![serde_json::to_value(stmt.text.clone()).unwrap()]),
+                    }),
+                    diagnostics: None,
+                    is_preferred: None,
+                    disabled: None,
+                    data: None,
+                }
+            })
+            .collect();
+
+        Some(actions)
+    }
+}