feat: Cache for Storage Secrets (lakekeeper#1485)

Author: c-thiel

crates/lakekeeper/src/api/management/v1/warehouse/mod.rs

@@ -378,7 +378,7 @@ pub trait Service<C: CatalogStore, A: Authorizer, S: SecretStore> {
                 context
                     .v1_state
                     .secrets
-                    .create_secret(storage_credential)
+                    .create_storage_secret(storage_credential)
                     .await?,
             )
         } else {
@@ -832,7 +832,7 @@ pub trait Service<C: CatalogStore, A: Authorizer, S: SecretStore> {
                 context
                     .v1_state
                     .secrets
-                    .create_secret(storage_credential)
+                    .create_storage_secret(storage_credential)
                     .await?,
             )
         } else {
@@ -920,7 +920,7 @@ pub trait Service<C: CatalogStore, A: Authorizer, S: SecretStore> {
                 context
                     .v1_state
                     .secrets
-                    .create_secret(new_storage_credential)
+                    .create_storage_secret(new_storage_credential)
                     .await?,
             )
         } else {

crates/lakekeeper/src/config.rs

@@ -436,6 +436,8 @@ pub(crate) struct Cache {
     pub(crate) warehouse: WarehouseCache,
     /// Namespace cache configuration.
     pub(crate) namespace: NamespaceCache,
+    /// Secrets cache configuration.
+    pub(crate) secrets: SecretsCache,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -492,6 +494,25 @@ impl std::default::Default for NamespaceCache {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(default)]
+pub(crate) struct SecretsCache {
+    pub(crate) enabled: bool,
+    pub(crate) capacity: u64,
+    /// Time-to-live for cache entries in seconds. Defaults to 60 seconds.
+    pub(crate) time_to_live_secs: u64,
+}
+
+impl std::default::Default for SecretsCache {
+    fn default() -> Self {
+        Self {
+            enabled: true,
+            capacity: 500,
+            time_to_live_secs: 600,
+        }
+    }
+}
+
 impl Default for DynAppConfig {
     fn default() -> Self {
         Self {

crates/lakekeeper/src/implementations/kv2.rs

@@ -6,7 +6,10 @@ use iceberg_ext::catalog::rest::IcebergErrorResponse;
 use serde::{de::DeserializeOwned, Serialize};
 use tokio::{sync::RwLock, time::Sleep};
 use uuid::Uuid;
-use vaultrs::client::{Client, VaultClient};
+use vaultrs::{
+    client::{Client, VaultClient},
+    error::ClientError,
+};
 use vaultrs_login::{engines::userpass::UserpassLogin, LoginMethod};
 
 use crate::{
@@ -21,27 +24,35 @@ use crate::{
 #[async_trait::async_trait]
 impl SecretStore for SecretsState {
     /// Get the secret for a given warehouse.
-    async fn get_secret_by_id<S: DeserializeOwned>(
+    async fn get_secret_by_id_impl<S: DeserializeOwned>(
         &self,
         secret_id: SecretId,
-    ) -> Result<Secret<S>> {
+    ) -> Result<Option<Secret<S>>> {
         // it seems there is no atomic get for metadata and secret so we read_metadata, and then
         // read the secret with the current version defined in the previously read metadata
         let metadata = vaultrs::kv2::read_metadata(
             &*self.vault_client.read().await,
             self.secret_mount.as_str(),
             &secret_ident_to_key(secret_id),
         )
-        .await
-        .map_err(|err| {
-            IcebergErrorResponse::from(ErrorModel::internal(
-                "secret metadata read failure",
-                "SecretReadFailed",
-                Some(Box::new(err)),
-            ))
-        })?;
+        .await;
+
+        let metadata = match metadata {
+            Ok(meta) => meta,
+            Err(err) => {
+                if matches!(&err, ClientError::APIError { code: 404, .. }) {
+                    return Ok(None);
+                }
 
-        Ok(Secret {
+                return Err(IcebergErrorResponse::from(ErrorModel::internal(
+                    "secret metadata read failure",
+                    "SecretReadFailed",
+                    Some(Box::new(err)),
+                )));
+            }
+        };
+
+        Ok(Some(Secret {
             secret_id,
             secret: vaultrs::kv2::read_version::<S>(
                 &*self.vault_client.read().await,
@@ -71,11 +82,11 @@ impl SecretStore for SecretsState {
                     Some(Box::new(err)),
                 ))
             })?),
-        })
+        }))
     }
 
     /// Create a new secret
-    async fn create_secret<S: Send + Sync + Serialize + std::fmt::Debug>(
+    async fn create_secret_impl<S: Send + Sync + Serialize + std::fmt::Debug>(
         &self,
         secret: S,
     ) -> Result<SecretId> {
@@ -98,7 +109,7 @@ impl SecretStore for SecretsState {
     }
 
     /// Delete a secret
-    async fn delete_secret(&self, secret_id: &SecretId) -> Result<()> {
+    async fn delete_secret_impl(&self, secret_id: &SecretId) -> Result<()> {
         Ok(vaultrs::kv2::delete_metadata(
             &*self.vault_client.read().await,
             self.secret_mount.as_str(),
@@ -274,14 +285,26 @@ mod tests {
             })
             .into();
 
-            let secret_id = state.create_secret(secret.clone()).await.unwrap();
+            let secret_id = state.create_storage_secret(secret.clone()).await.unwrap();
 
-            let read_secret = state
-                .get_secret_by_id::<StorageCredential>(secret_id)
+            let read_secret = state.require_storage_secret_by_id(secret_id).await.unwrap();
+
+            assert_eq!(&*read_secret.secret, &secret);
+        }
+
+        #[tokio::test]
+        async fn test_read_missing_secret() {
+            let state = SecretsState::from_config(CONFIG.kv2.as_ref().unwrap())
                 .await
                 .unwrap();
 
-            assert_eq!(read_secret.secret, secret);
+            let secret_id = SecretId::from(Uuid::new_v4());
+
+            let read_secret = state
+                .get_secret_by_id_impl::<StorageCredential>(secret_id)
+                .await;
+
+            assert!(read_secret.unwrap().is_none());
         }
 
         #[tokio::test]
@@ -298,13 +321,13 @@ mod tests {
             .into();
 
             let secret_id = state
-                .create_secret(secret.clone())
+                .create_storage_secret(secret.clone())
                 .await
                 .expect("create secret failed");
 
             state.delete_secret(&secret_id).await.unwrap();
 
-            let read_secret = state.get_secret_by_id::<StorageCredential>(secret_id).await;
+            let read_secret = state.require_storage_secret_by_id(secret_id).await;
 
             assert!(read_secret.is_err());
         }

crates/lakekeeper/src/implementations/mod.rs

@@ -74,35 +74,35 @@ pub enum Secrets {
 
 #[async_trait]
 impl SecretStore for Secrets {
-    async fn get_secret_by_id<S: SecretInStorage + serde::de::DeserializeOwned>(
+    async fn get_secret_by_id_impl<S: SecretInStorage + serde::de::DeserializeOwned>(
         &self,
         secret_id: SecretId,
-    ) -> crate::api::Result<Secret<S>> {
+    ) -> crate::api::Result<Option<Secret<S>>> {
         match self {
             #[cfg(feature = "sqlx-postgres")]
-            Self::Postgres(state) => state.get_secret_by_id(secret_id).await,
-            Self::KV2(state) => state.get_secret_by_id(secret_id).await,
+            Self::Postgres(state) => state.get_secret_by_id_impl(secret_id).await,
+            Self::KV2(state) => state.get_secret_by_id_impl(secret_id).await,
         }
     }
 
-    async fn create_secret<
+    async fn create_secret_impl<
         S: SecretInStorage + Send + Sync + serde::Serialize + std::fmt::Debug,
     >(
         &self,
         secret: S,
     ) -> crate::api::Result<SecretId> {
         match self {
             #[cfg(feature = "sqlx-postgres")]
-            Self::Postgres(state) => state.create_secret(secret).await,
-            Self::KV2(state) => state.create_secret(secret).await,
+            Self::Postgres(state) => state.create_secret_impl(secret).await,
+            Self::KV2(state) => state.create_secret_impl(secret).await,
         }
     }
 
-    async fn delete_secret(&self, secret_id: &SecretId) -> crate::api::Result<()> {
+    async fn delete_secret_impl(&self, secret_id: &SecretId) -> crate::api::Result<()> {
         match self {
             #[cfg(feature = "sqlx-postgres")]
-            Self::Postgres(state) => state.delete_secret(secret_id).await,
-            Self::KV2(state) => state.delete_secret(secret_id).await,
+            Self::Postgres(state) => state.delete_secret_impl(secret_id).await,
+            Self::KV2(state) => state.delete_secret_impl(secret_id).await,
         }
     }
 }

crates/lakekeeper/src/implementations/postgres/secrets.rs

@@ -51,17 +51,17 @@ impl SecretsState {
 #[async_trait::async_trait]
 impl SecretStore for SecretsState {
     /// Get the secret for a given warehouse.
-    async fn get_secret_by_id<S: for<'de> Deserialize<'de>>(
+    async fn get_secret_by_id_impl<S: for<'de> Deserialize<'de>>(
         &self,
         secret_id: SecretId,
-    ) -> Result<Secret<S>> {
+    ) -> Result<Option<Secret<S>>> {
         struct SecretRow {
             secret: Option<String>,
             created_at: chrono::DateTime<chrono::Utc>,
             updated_at: Option<chrono::DateTime<chrono::Utc>>,
         }
 
-        let secret: SecretRow = sqlx::query_as!(
+        let secret = sqlx::query_as!(
             SecretRow,
             r#"
             SELECT 
@@ -74,25 +74,24 @@ impl SecretStore for SecretsState {
             secret_id.as_uuid(),
             CONFIG.pg_encryption_key
         )
-        .fetch_one(&self.read_write.read_pool)
+        .fetch_optional(&self.read_write.read_pool)
         .await
-        .map_err(|e| match e {
-            sqlx::Error::RowNotFound => ErrorModel::builder()
-                .code(StatusCode::NOT_FOUND.into())
-                .message("Secret not found".to_string())
-                .r#type("SecretNotFound".to_string())
-                .stack(vec![format!("secret_id: {}", secret_id), e.to_string()])
-                .build(),
-            _ => ErrorModel::builder()
+        .map_err(|e| {
+            ErrorModel::builder()
                 .code(StatusCode::INTERNAL_SERVER_ERROR.into())
                 .message("Error fetching secret".to_string())
                 .r#type("SecretFetchError".to_string())
                 .stack(vec![format!("secret_id: {}", secret_id), e.to_string()])
-                .build(),
+                .build()
         })?;
 
+        let Some(secret) = secret else {
+            return Ok(None);
+        };
+
         let inner =
-            serde_json::from_str(&secret.secret.unwrap_or("{}".to_string())).map_err(|_e| {
+            serde_json::from_str(&secret.secret.unwrap_or("{}".to_string())).map_err(|e| {
+                tracing::error!("Error parsing secret id {}: {}", secret_id, e);
                 ErrorModel::builder()
                     .code(StatusCode::INTERNAL_SERVER_ERROR.into())
                     .message("Error parsing secret".to_string())
@@ -102,16 +101,16 @@ impl SecretStore for SecretsState {
                     .build()
             })?;
 
-        Ok(Secret {
+        Ok(Some(Secret {
             secret_id,
             secret: inner,
             created_at: secret.created_at,
             updated_at: secret.updated_at,
-        })
+        }))
     }
 
     /// Create a new secret
-    async fn create_secret<S: Send + Sync + Serialize + std::fmt::Debug>(
+    async fn create_secret_impl<S: Send + Sync + Serialize + std::fmt::Debug>(
         &self,
         secret: S,
     ) -> Result<SecretId> {
@@ -149,7 +148,7 @@ impl SecretStore for SecretsState {
     }
 
     /// Delete a secret
-    async fn delete_secret(&self, secret_id: &SecretId) -> Result<()> {
+    async fn delete_secret_impl(&self, secret_id: &SecretId) -> Result<()> {
         sqlx::query!(
             r#"
             DELETE FROM secret
@@ -189,14 +188,21 @@ mod tests {
         })
         .into();
 
-        let secret_id = state.create_secret(secret.clone()).await.unwrap();
+        let secret_id = state.create_storage_secret(secret.clone()).await.unwrap();
 
-        let read_secret = state
-            .get_secret_by_id::<StorageCredential>(secret_id)
-            .await
-            .unwrap();
+        let read_secret = state.require_storage_secret_by_id(secret_id).await.unwrap();
+
+        assert_eq!(&*read_secret.secret, &secret);
+    }
 
-        assert_eq!(read_secret.secret, secret);
+    #[sqlx::test]
+    async fn test_read_missing_secret(pool: sqlx::PgPool) {
+        let state = SecretsState::from_pools(pool.clone(), pool);
+        let missing_secret_id = SecretId::from(uuid::Uuid::new_v4());
+        let read_secret = state
+            .get_secret_by_id_impl::<StorageCredential>(missing_secret_id)
+            .await;
+        assert!(read_secret.unwrap().is_none());
     }
 
     #[sqlx::test]
@@ -210,11 +216,11 @@ mod tests {
         })
         .into();
 
-        let secret_id = state.create_secret(secret.clone()).await.unwrap();
+        let secret_id = state.create_storage_secret(secret.clone()).await.unwrap();
 
         state.delete_secret(&secret_id).await.unwrap();
 
-        let read_secret = state.get_secret_by_id::<StorageCredential>(secret_id).await;
+        let read_secret = state.require_storage_secret_by_id(secret_id).await;
 
         assert!(read_secret.is_err());
     }

crates/lakekeeper/src/server/mod.rs

@@ -108,9 +108,11 @@ fn require_warehouse_id(prefix: Option<&Prefix>) -> std::result::Result<Warehous
 pub(crate) async fn maybe_get_secret<S: SecretStore>(
     secret: Option<crate::SecretId>,
     state: &S,
-) -> Result<Option<StorageCredential>, IcebergErrorResponse> {
+) -> Result<Option<Arc<StorageCredential>>, IcebergErrorResponse> {
     if let Some(secret_id) = secret {
-        Ok(Some(state.get_secret_by_id(secret_id).await?.secret))
+        Ok(Some(
+            state.require_storage_secret_by_id(secret_id).await?.secret,
+        ))
     } else {
         Ok(None)
     }

crates/lakekeeper/src/server/s3_signer/sign.rs

@@ -22,8 +22,7 @@ use crate::{
         },
         secrets::SecretStore,
         storage::{
-            s3::S3UrlStyleDetectionMode, S3Credential, S3Profile, StorageCredential,
-            StorageProfile, ValidationError,
+            s3::S3UrlStyleDetectionMode, S3Credential, S3Profile, StorageProfile, ValidationError,
         },
         AuthZTableInfo, CatalogNamespaceOps, CatalogStore, CatalogTabularOps, CatalogWarehouseOps,
         GetTabularInfoByLocationError, ResolvedWarehouse, State, TableId, TableInfo,
@@ -214,7 +213,7 @@ impl<C: CatalogStore, A: Authorizer + Clone, S: SecretStore>
                 state
                     .v1_state
                     .secrets
-                    .get_secret_by_id::<StorageCredential>(storage_secret_id)
+                    .require_storage_secret_by_id(storage_secret_id)
                     .await?
                     .secret,
             )