progress

Author: psteinroe

crates/pgt_workspace/src/settings.rs

@@ -12,18 +12,18 @@ use tracing::trace;
 
 use ignore::gitignore::{Gitignore, GitignoreBuilder};
 use pgt_configuration::{
+    ConfigurationDiagnostic, LinterConfiguration, PartialConfiguration,
     database::PartialDatabaseConfiguration,
     diagnostics::InvalidIgnorePattern,
     files::FilesConfiguration,
     migrations::{MigrationsConfiguration, PartialMigrationsConfiguration},
-    ConfigurationDiagnostic, LinterConfiguration, PartialConfiguration,
 };
 use pgt_fs::PgTPath;
 
 use crate::{
+    WorkspaceError,
     matcher::Matcher,
     workspace::{ProjectKey, WorkspaceData},
-    WorkspaceError,
 };
 
 #[derive(Debug, Default)]
@@ -133,8 +133,7 @@ impl WorkspaceSettings {
         for (key, path_to_settings) in iter {
             trace!(
                 "Workspace path {:?}, file path {:?}",
-                path_to_settings.path,
-                path
+                path_to_settings.path, path
             );
             trace!("Iter key: {:?}", key);
             if key == self.current_project {

crates/pgt_workspace/src/workspace.rs

@@ -6,24 +6,24 @@ use pgt_configuration::{PartialConfiguration, RuleSelector};
 use pgt_fs::PgTPath;
 use pgt_text_size::TextRange;
 use serde::{Deserialize, Serialize};
-use slotmap::{new_key_type, DenseSlotMap};
+use slotmap::{DenseSlotMap, new_key_type};
 
 use crate::{
+    WorkspaceError,
     features::{
         code_actions::{
             CodeActionsParams, CodeActionsResult, ExecuteStatementParams, ExecuteStatementResult,
         },
         completions::{CompletionsResult, GetCompletionsParams},
         diagnostics::{PullDiagnosticsParams, PullDiagnosticsResult},
     },
-    WorkspaceError,
 };
 
 mod client;
 mod server;
 
-pub(crate) use server::parsed_document::*;
 pub use server::StatementId;
+pub(crate) use server::parsed_document::*;
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]

crates/pgt_workspace/src/workspace/client.rs

@@ -1,13 +1,16 @@
 use crate::workspace::ServerInfo;
 use crate::{TransportError, Workspace, WorkspaceError};
-use serde::{de::DeserializeOwned, Deserialize, Serialize};
+use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use serde_json::json;
 use std::{
     panic::RefUnwindSafe,
     sync::atomic::{AtomicU64, Ordering},
 };
 
-use super::{CloseFileParams, GetFileContentParams, IsPathIgnoredParams, OpenFileParams};
+use super::{
+    CloseFileParams, GetFileContentParams, IsPathIgnoredParams, OpenFileParams, ProjectKey,
+    RegisterProjectFolderParams, UnregisterProjectFolderParams,
+};
 
 pub struct WorkspaceClient<T> {
     transport: T,

crates/pgt_workspace/src/workspace/server.rs

@@ -10,15 +10,15 @@ use async_helper::run_async;
 use connection_manager::ConnectionManager;
 use dashmap::DashMap;
 use document::Document;
-use futures::{stream, StreamExt};
+use futures::{StreamExt, stream};
 use parsed_document::{
     AsyncDiagnosticsMapper, CursorPositionFilter, DefaultMapper, ExecuteStatementMapper,
     ParsedDocument, SyncDiagnosticsMapper,
 };
 use pgt_analyse::{AnalyserOptions, AnalysisFilter};
 use pgt_analyser::{Analyser, AnalyserConfig, AnalyserContext};
 use pgt_diagnostics::{
-    serde::Diagnostic as SDiagnostic, Diagnostic, DiagnosticExt, Error, Severity,
+    Diagnostic, DiagnosticExt, Error, Severity, serde::Diagnostic as SDiagnostic,
 };
 use pgt_fs::{ConfigName, PgTPath};
 use pgt_typecheck::{IdentifierType, TypecheckParams, TypedIdentifier};
@@ -27,20 +27,17 @@ use sqlx::{Executor, PgPool};
 use tracing::{debug, info};
 
 use crate::{
+    WorkspaceError,
     configuration::to_analyser_rules,
     features::{
         code_actions::{
             self, CodeAction, CodeActionKind, CodeActionsResult, CommandAction,
             CommandActionCategory, ExecuteStatementParams, ExecuteStatementResult,
         },
-        completions::{get_statement_for_completions, CompletionsResult, GetCompletionsParams},
+        completions::{CompletionsResult, GetCompletionsParams, get_statement_for_completions},
         diagnostics::{PullDiagnosticsParams, PullDiagnosticsResult},
     },
-    settings::{
-        Settings, SettingsHandle, SettingsHandleMut, WorkspaceSettings, WorkspaceSettingsHandle,
-        WorkspaceSettingsHandleMut,
-    },
-    WorkspaceError,
+    settings::{WorkspaceSettings, WorkspaceSettingsHandle, WorkspaceSettingsHandleMut},
 };
 
 use super::{
@@ -55,6 +52,7 @@ mod analyser;
 mod annotation;
 mod async_helper;
 mod change;
+mod connection_key;
 mod connection_manager;
 pub(crate) mod document;
 mod migration;
@@ -444,15 +442,14 @@ impl Workspace for WorkspaceServer {
         if let Some(pool) = self.get_current_connection() {
             let path_clone = params.path.clone();
             let schema_cache = self.schema_cache.load(pool.clone())?;
-            let schema_cache_arc = schema_cache.get_arc();
             let input = parser.iter(AsyncDiagnosticsMapper).collect::<Vec<_>>();
             // sorry for the ugly code :(
             let async_results = run_async(async move {
                 stream::iter(input)
                     .map(|(_id, range, content, ast, cst, sign)| {
                         let pool = pool.clone();
                         let path = path_clone.clone();
-                        let schema_cache = Arc::clone(&schema_cache_arc);
+                        let schema_cache = Arc::clone(&schema_cache);
                         async move {
                             if let Some(ast) = ast {
                                 pgt_typecheck::check_sql(TypecheckParams {

crates/pgt_workspace/src/workspace/server/connection_key.rs

@@ -0,0 +1,44 @@
+use sqlx::PgPool;
+
+use crate::settings::DatabaseSettings;
+
+/// A unique identifier for database connection settings
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub(crate) struct ConnectionKey {
+    pub host: String,
+    pub port: u16,
+    pub username: String,
+    pub database: String,
+}
+
+impl From<&DatabaseSettings> for ConnectionKey {
+    fn from(settings: &DatabaseSettings) -> Self {
+        Self {
+            host: settings.host.clone(),
+            port: settings.port,
+            username: settings.username.clone(),
+            database: settings.database.clone(),
+        }
+    }
+}
+
+impl From<&PgPool> for ConnectionKey {
+    fn from(pool: &PgPool) -> Self {
+        let conn = pool.connect_options();
+
+        match conn.get_database() {
+            None => Self {
+                host: conn.get_host().to_string(),
+                port: conn.get_port(),
+                username: conn.get_username().to_string(),
+                database: String::new(),
+            },
+            Some(db) => Self {
+                host: conn.get_host().to_string(),
+                port: conn.get_port(),
+                username: conn.get_username().to_string(),
+                database: db.to_string(),
+            },
+        }
+    }
+}

crates/pgt_workspace/src/workspace/server/connection_manager.rs

@@ -1,31 +1,11 @@
 use std::time::{Duration, Instant};
 
 use dashmap::DashMap;
-use sqlx::{pool::PoolOptions, postgres::PgConnectOptions, PgPool, Postgres};
+use sqlx::{PgPool, Postgres, pool::PoolOptions, postgres::PgConnectOptions};
 
 use crate::settings::DatabaseSettings;
 
-/// A unique identifier for database connection settings
-#[derive(Clone, PartialEq, Eq, Hash)]
-struct ConnectionKey {
-    host: String,
-    port: u16,
-    username: String,
-    password: String,
-    database: String,
-}
-
-impl From<&DatabaseSettings> for ConnectionKey {
-    fn from(settings: &DatabaseSettings) -> Self {
-        Self {
-            host: settings.host.clone(),
-            port: settings.port,
-            username: settings.username.clone(),
-            password: settings.password.clone(),
-            database: settings.database.clone(),
-        }
-    }
-}
+use super::connection_key::ConnectionKey;
 
 /// Cached connection pool with last access time
 struct CachedPool {
@@ -49,6 +29,7 @@ impl ConnectionManager {
     /// Get a connection pool for the given database settings.
     /// If a pool already exists for these settings, it will be returned.
     /// If not, a new pool will be created if connections are enabled.
+    /// Will also clean up idle connections that haven't been accessed for a while.
     pub(crate) fn get_pool(&self, settings: &DatabaseSettings) -> Option<PgPool> {
         let key = ConnectionKey::from(settings);
 
@@ -84,8 +65,8 @@ impl ConnectionManager {
         let cached_pool = CachedPool {
             pool: pool.clone(),
             last_accessed: Instant::now(),
-            // TODO: add this to the db settings, for now default to one minute
-            idle_timeout: Duration::from_secs(60),
+            // TODO: add this to the db settings, for now default to five minutes
+            idle_timeout: Duration::from_secs(60 * 5),
         };
 
         self.pools.insert(key, cached_pool);

crates/pgt_workspace/src/workspace/server/schema_cache_manager.rs

@@ -1,97 +1,30 @@
-use std::sync::{Arc, RwLock, RwLockReadGuard};
+use std::sync::Arc;
 
+use dashmap::DashMap;
 use pgt_schema_cache::SchemaCache;
 use sqlx::PgPool;
 
 use crate::WorkspaceError;
 
-use super::async_helper::run_async;
-
-pub(crate) struct SchemaCacheHandle<'a> {
-    inner: RwLockReadGuard<'a, SchemaCacheManagerInner>,
-}
-
-impl<'a> SchemaCacheHandle<'a> {
-    pub(crate) fn new(cache: &'a RwLock<SchemaCacheManagerInner>) -> Self {
-        Self {
-            inner: cache.read().unwrap(),
-        }
-    }
-
-    pub(crate) fn wrap(inner: RwLockReadGuard<'a, SchemaCacheManagerInner>) -> Self {
-        Self { inner }
-    }
-
-    pub fn get_arc(&self) -> Arc<SchemaCache> {
-        Arc::clone(&self.inner.cache)
-    }
-}
-
-impl AsRef<SchemaCache> for SchemaCacheHandle<'_> {
-    fn as_ref(&self) -> &SchemaCache {
-        &self.inner.cache
-    }
-}
-
-#[derive(Default)]
-pub(crate) struct SchemaCacheManagerInner {
-    cache: Arc<SchemaCache>,
-    conn_str: String,
-}
+use super::{async_helper::run_async, connection_key::ConnectionKey};
 
 #[derive(Default)]
 pub struct SchemaCacheManager {
-    inner: RwLock<SchemaCacheManagerInner>,
+    schemas: DashMap<ConnectionKey, Arc<SchemaCache>>,
 }
 
 impl SchemaCacheManager {
-    pub fn load(&self, pool: PgPool) -> Result<SchemaCacheHandle, WorkspaceError> {
-        let new_conn_str = pool_to_conn_str(&pool);
+    pub fn load(&self, pool: PgPool) -> Result<Arc<SchemaCache>, WorkspaceError> {
+        let key: ConnectionKey = (&pool).into();
 
-        {
-            // return early if the connection string is the same
-            let inner = self.inner.read().unwrap();
-            if new_conn_str == inner.conn_str {
-                tracing::info!("Same connection string, no updates.");
-                return Ok(SchemaCacheHandle::wrap(inner));
-            }
+        if let Some(cache) = self.schemas.get(&key) {
+            return Ok(Arc::clone(&*cache));
         }
 
-        let maybe_refreshed = run_async(async move { SchemaCache::load(&pool).await })?;
-        let refreshed = maybe_refreshed?;
-
-        {
-            // write lock must be dropped before we return the reference below, hence the block
-            let mut inner = self.inner.write().unwrap();
-
-            // Double-check that we still need to refresh (another thread might have done it)
-            if new_conn_str != inner.conn_str {
-                inner.cache = Arc::new(refreshed);
-                inner.conn_str = new_conn_str;
-                tracing::info!("Refreshed connection.");
-            }
-        }
-
-        Ok(SchemaCacheHandle::new(&self.inner))
-    }
-}
+        let schema_cache = Arc::new(run_async(async move { SchemaCache::load(&pool).await })??);
 
-fn pool_to_conn_str(pool: &PgPool) -> String {
-    let conn = pool.connect_options();
+        self.schemas.insert(key, Arc::clone(&schema_cache));
 
-    match conn.get_database() {
-        None => format!(
-            "postgres://{}:<redacted_pw>@{}:{}",
-            conn.get_username(),
-            conn.get_host(),
-            conn.get_port()
-        ),
-        Some(db) => format!(
-            "postgres://{}:<redacted_pw>@{}:{}/{}",
-            conn.get_username(),
-            conn.get_host(),
-            conn.get_port(),
-            db
-        ),
+        Ok(schema_cache)
     }
 }