udpate database via messages

Author: juleswritescode

crates/pg_lsp/src/client.rs

@@ -50,6 +50,14 @@ impl LspClient {
         Ok(())
     }
 
+    /// This will ignore any errors that occur while sending the notification.
+    pub fn send_info_notification(&self, message: &str) {
+        let _ = self.send_notification::<ShowMessage>(ShowMessageParams {
+            message: message.into(),
+            typ: MessageType::INFO,
+        });
+    }
+
     pub fn send_request<R>(&self, params: R::Params) -> Result<R::Result>
     where
         R: lsp_types::request::Request,

crates/pg_lsp/src/db_connection.rs

@@ -1,10 +1,12 @@
 use pg_schema_cache::SchemaCache;
 use sqlx::{postgres::PgListener, PgPool};
+use tokio::task::JoinHandle;
 
 #[derive(Debug)]
 pub(crate) struct DbConnection {
     pub pool: PgPool,
     connection_string: String,
+    schema_update_handle: Option<JoinHandle<()>>,
 }
 
 impl DbConnection {
@@ -13,49 +15,52 @@ impl DbConnection {
         Ok(Self {
             pool,
             connection_string: connection_string,
+            schema_update_handle: None,
         })
     }
 
-    pub(crate) async fn refresh_db_connection(
-        self,
-        connection_string: Option<String>,
-    ) -> anyhow::Result<Self> {
-        if connection_string.is_none()
-            || connection_string.as_ref() == Some(&self.connection_string)
-        {
-            return Ok(self);
-        }
+    pub(crate) fn connected_to(&self, connection_string: &str) -> bool {
+        connection_string == self.connection_string
+    }
 
+    pub(crate) async fn close(self) {
+        if self.schema_update_handle.is_some() {
+            self.schema_update_handle.unwrap().abort();
+        }
         self.pool.close().await;
-
-        let conn = DbConnection::new(connection_string.unwrap()).await?;
-
-        Ok(conn)
     }
 
-    pub(crate) async fn start_listening<F>(&self, on_schema_update: F) -> anyhow::Result<()>
+    pub(crate) async fn listen_for_schema_updates<F>(
+        &mut self,
+        on_schema_update: F,
+    ) -> anyhow::Result<()>
     where
-        F: Fn() -> () + Send + 'static,
+        F: Fn(SchemaCache) -> () + Send + 'static,
     {
         let mut listener = PgListener::connect_with(&self.pool).await?;
         listener.listen_all(["postgres_lsp", "pgrst"]).await?;
 
-        loop {
-            match listener.recv().await {
-                Ok(notification) => {
-                    if notification.payload().to_string() == "reload schema" {
-                        on_schema_update();
+        let pool = self.pool.clone();
+
+        let handle: JoinHandle<()> = tokio::spawn(async move {
+            loop {
+                match listener.recv().await {
+                    Ok(not) => {
+                        if not.payload().to_string() == "reload schema" {
+                            let schema_cache = SchemaCache::load(&pool).await;
+                            on_schema_update(schema_cache);
+                        };
+                    }
+                    Err(why) => {
+                        eprintln!("Error receiving notification: {:?}", why);
+                        break;
                     }
-                }
-                Err(e) => {
-                    eprintln!("Listener error: {}", e);
-                    return Err(e.into());
                 }
             }
-        }
-    }
+        });
+
+        self.schema_update_handle = Some(handle);
 
-    pub(crate) async fn get_schema_cache(&self) -> SchemaCache {
-        SchemaCache::load(&self.pool).await
+        Ok(())
     }
 }

crates/pg_lsp/src/server.rs

@@ -40,17 +40,14 @@ use crate::{
 };
 
 use self::{debouncer::EventDebouncer, options::Options};
-use sqlx::{
-    postgres::{PgListener, PgPool},
-    Executor,
-};
+use sqlx::{postgres::PgPool, Executor};
 
 #[derive(Debug)]
 enum InternalMessage {
     PublishDiagnostics(lsp_types::Url),
     SetOptions(Options),
-    RefreshSchemaCache,
     SetSchemaCache(SchemaCache),
+    SetDatabaseConnection(DbConnection),
 }
 
 /// `lsp-servers` `Connection` type uses a crossbeam channel, which is not compatible with tokio's async runtime.
@@ -210,29 +207,54 @@ impl Server {
         });
     }
 
-    async fn update_options(&mut self, options: Options) -> anyhow::Result<()> {
-        if options.db_connection_string.is_none() {
+    fn update_db_connection(&self, options: Options) -> anyhow::Result<()> {
+        if options.db_connection_string.is_none()
+            || self
+                .db_conn
+                .as_ref()
+                .is_some_and(|c| c.connected_to(options.db_connection_string.as_ref().unwrap()))
+        {
             return Ok(());
         }
 
-        let new_conn = if self.db_conn.is_none() {
-            DbConnection::new(options.db_connection_string.clone().unwrap()).await?
-        } else {
-            let current_conn = self.db_conn.take().unwrap();
-            current_conn
-                .refresh_db_connection(options.db_connection_string.clone())
-                .await?
-        };
+        let connection_string = options.db_connection_string.unwrap();
 
         let internal_tx = self.internal_tx.clone();
+        let client = self.client.clone();
         self.spawn_with_cancel(async move {
-            new_conn.start_listening(move || {
+            match DbConnection::new(connection_string.into()).await {
+                Ok(conn) => {
+                    internal_tx
+                        .send(InternalMessage::SetDatabaseConnection(conn))
+                        .unwrap();
+                }
+                Err(why) => {
+                    client.send_info_notification(&format!("Unable to update database connection: {}", why));
+                    
+                }
+            }
+        });
+
+        Ok(())
+    }
+
+    async fn listen_for_schema_updates(&mut self) -> anyhow::Result<()> {
+        if self.db_conn.is_none() {
+            eprintln!("Error trying to listen for schema updates: No database connection");
+            return Ok(());
+        }
+
+        let internal_tx = self.internal_tx.clone();
+        self.db_conn
+            .as_mut()
+            .unwrap()
+            .listen_for_schema_updates(move |schema_cache| {
                 internal_tx
-                    .send(InternalMessage::RefreshSchemaCache)
+                    .send(InternalMessage::SetSchemaCache(schema_cache))
                     .unwrap();
-            // TODO: handle result
-            }).await.unwrap()
-        });
+                // TODO: handle result
+            })
+            .await?;
 
         Ok(())
     }
@@ -692,26 +714,6 @@ impl Server {
         });
     }
 
-    async fn refresh_schema_cache(&self) {
-        if self.db_conn.is_none() {
-            return;
-        }
-
-        let tx = self.internal_tx.clone();
-        let conn = self.db_conn.as_ref().unwrap().pool.clone();
-        let client = self.client.clone();
-
-        client
-            .send_notification::<ShowMessage>(ShowMessageParams {
-                typ: lsp_types::MessageType::INFO,
-                message: "Refreshing schema cache...".to_string(),
-            })
-            .unwrap();
-        let schema_cache = SchemaCache::load(&conn).await;
-        tx.send(InternalMessage::SetSchemaCache(schema_cache))
-            .unwrap();
-    }
-
     fn did_change_configuration(
         &mut self,
         params: DidChangeConfigurationParams,
@@ -720,7 +722,7 @@ impl Server {
             self.pull_options();
         } else {
             let options = self.client.parse_options(params.settings)?;
-            self.update_options(options);
+            self.update_db_connection(options);
         }
 
         Ok(())
@@ -744,14 +746,14 @@ impl Server {
                 msg = self.client_rx.recv() => {
                     match msg {
                         None => panic!("The LSP's client closed, but not via an 'exit' method. This should never happen."),
-                        Some(m) => self.handle_message(m)
+                        Some(m) => self.handle_message(m).await
                     }
                 },
             }?;
         }
     }
 
-    fn handle_message(&mut self, msg: Message) -> anyhow::Result<()> {
+    async fn handle_message(&mut self, msg: Message) -> anyhow::Result<()> {
         match msg {
             Message::Request(request) => {
                 if let Some(response) = dispatch::RequestDispatcher::new(request)
@@ -768,7 +770,8 @@ impl Server {
             Message::Notification(notification) => {
                 dispatch::NotificationDispatcher::new(notification)
                     .on::<DidChangeConfiguration, _>(|params| {
-                        self.did_change_configuration(params)
+                        self.did_change_configuration(params);
+                        Ok(())
                     })?
                     .on::<DidCloseTextDocument, _>(|params| self.did_close(params))?
                     .on::<DidOpenTextDocument, _>(|params| self.did_open(params))?
@@ -788,17 +791,24 @@ impl Server {
     async fn handle_internal_message(&mut self, msg: InternalMessage) -> anyhow::Result<()> {
         match msg {
             InternalMessage::SetSchemaCache(c) => {
+                self.client
+                    .send_info_notification("Refreshing Schema Cache...");
                 self.ide.set_schema_cache(c);
+                self.client.send_info_notification("Updated Schema Cache.");
                 self.compute_now();
             }
-            InternalMessage::RefreshSchemaCache => {
-                self.refresh_schema_cache().await;
-            }
             InternalMessage::PublishDiagnostics(uri) => {
                 self.publish_diagnostics(uri)?;
             }
             InternalMessage::SetOptions(options) => {
-                self.update_options(options);
+                self.update_db_connection(options);
+            }
+            InternalMessage::SetDatabaseConnection(conn) => {
+                let current = self.db_conn.replace(conn);
+                if current.is_some() {
+                    current.unwrap().close().await
+                }
+                self.listen_for_schema_updates();
             }
         }