lift cache key

Author: juleswritescode

crates/pg_schema_cache/src/schema_cache.rs

@@ -9,8 +9,6 @@ use crate::versions::Version;
 
 #[derive(Debug, Clone, Default)]
 pub struct SchemaCache {
-    cached_conn_str: String,
-
     pub schemas: Vec<Schema>,
     pub tables: Vec<Table>,
     pub functions: Vec<Function>,
@@ -31,7 +29,6 @@ impl SchemaCache {
         )?;
 
         Ok(SchemaCache {
-            cached_conn_str: SchemaCache::pool_to_conn_str(pool),
             schemas,
             tables,
             functions,
@@ -41,30 +38,6 @@ impl SchemaCache {
         })
     }
 
-    fn pool_to_conn_str(pool: &PgPool) -> String {
-        let conn = pool.connect_options();
-
-        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
-            ),
-        }
-    }
-
-    pub fn has_already_cached_connection(&self, pool: &PgPool) -> bool {
-        self.cached_conn_str == SchemaCache::pool_to_conn_str(pool)
-    }
-
     /// Applies an AST node to the repository
     ///
     /// For example, alter table add column will add the column to the table if it does not exist

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

@@ -7,48 +7,77 @@ use crate::WorkspaceError;
 
 use super::async_helper::run_async;
 
-#[derive(Debug)]
 pub(crate) struct SchemaCacheHandle<'a> {
-    inner: RwLockReadGuard<'a, SchemaCache>,
+    inner: RwLockReadGuard<'a, SchemaCacheManagerInner>,
 }
 
 impl<'a> SchemaCacheHandle<'a> {
-    pub(crate) fn new(cache: &'a RwLock<SchemaCache>) -> Self {
+    pub(crate) fn new(cache: &'a RwLock<SchemaCacheManagerInner>) -> Self {
         Self {
             inner: cache.read().unwrap(),
         }
     }
 
-    pub(crate) fn wrap(inner: RwLockReadGuard<'a, SchemaCache>) -> Self {
+    pub(crate) fn wrap(inner: RwLockReadGuard<'a, SchemaCacheManagerInner>) -> Self {
         Self { inner }
     }
 }
 
 impl AsRef<SchemaCache> for SchemaCacheHandle<'_> {
     fn as_ref(&self) -> &SchemaCache {
-        &self.inner
+        &self.inner.cache
     }
 }
 
+#[derive(Default)]
+pub(crate) struct SchemaCacheManagerInner {
+    cache: SchemaCache,
+    conn_str: String,
+}
+
 #[derive(Default)]
 pub struct SchemaCacheManager {
-    cache: RwLock<SchemaCache>,
+    inner: RwLock<SchemaCacheManagerInner>,
 }
 
 impl SchemaCacheManager {
     pub fn load(&self, pool: PgPool) -> Result<SchemaCacheHandle, WorkspaceError> {
-        let cache_lock = self.cache.read().unwrap();
+        let inner = self.inner.read().unwrap();
 
-        if cache_lock.has_already_cached_connection(&pool) {
-            Ok(SchemaCacheHandle::wrap(cache_lock))
+        if pool_to_conn_str(&pool) == inner.conn_str {
+            Ok(SchemaCacheHandle::wrap(inner))
         } else {
+            let new_conn_str = pool_to_conn_str(&pool);
+
             let maybe_refreshed = run_async(async move { SchemaCache::load(&pool).await })?;
             let refreshed = maybe_refreshed?;
 
-            let mut cache = self.cache.write().unwrap();
-            *cache = refreshed;
+            let mut inner = self.inner.write().unwrap();
+
+            inner.cache = refreshed;
+            inner.conn_str = new_conn_str;
 
-            Ok(SchemaCacheHandle::new(&self.cache))
+            Ok(SchemaCacheHandle::new(&self.inner))
         }
     }
 }
+
+fn pool_to_conn_str(pool: &PgPool) -> String {
+    let conn = pool.connect_options();
+
+    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
+        ),
+    }
+}