feat(router): implement ApiKeyInterface for MockDb (#1101)

Co-authored-by: Derek Leverenz <derel@dereleverenz.com>
Co-authored-by: Sanchith Hegde <22217505+SanchithHegde@users.noreply.github.com>

crates/router/src/db.rs

@@ -82,6 +82,7 @@ pub struct MockDb {
     processes: Arc<Mutex<Vec<storage::ProcessTracker>>>,
     connector_response: Arc<Mutex<Vec<storage::ConnectorResponse>>>,
     redis: Arc<redis_interface::RedisConnectionPool>,
+    api_keys: Arc<Mutex<Vec<storage::ApiKey>>>,
 }
 
 impl MockDb {
@@ -96,6 +97,7 @@ impl MockDb {
             processes: Default::default(),
             connector_response: Default::default(),
             redis: Arc::new(crate::connection::redis_connection(redis).await),
+            api_keys: Default::default(),
         }
     }
 }

crates/router/src/db/api_keys.rs

@@ -126,55 +126,252 @@ impl ApiKeyInterface for Store {
 impl ApiKeyInterface for MockDb {
     async fn insert_api_key(
         &self,
-        _api_key: storage::ApiKeyNew,
+        api_key: storage::ApiKeyNew,
     ) -> CustomResult<storage::ApiKey, errors::StorageError> {
-        // [#172]: Implement function for `MockDb`
+        let mut locked_api_keys = self.api_keys.lock().await;
-        Err(errors::StorageError::MockDbError)?
+        // don't allow duplicate key_ids, a those would be a unique constraint violation in the
+        // real db as it is used as the primary key
+        if locked_api_keys.iter().any(|k| k.key_id == api_key.key_id) {
+            Err(errors::StorageError::MockDbError)?;
+        }
+        let stored_key = storage::ApiKey {
+            key_id: api_key.key_id,
+            merchant_id: api_key.merchant_id,
+            name: api_key.name,
+            description: api_key.description,
+            hashed_api_key: api_key.hashed_api_key,
+            prefix: api_key.prefix,
+            created_at: api_key.created_at,
+            expires_at: api_key.expires_at,
+            last_used: api_key.last_used,
+        };
+        locked_api_keys.push(stored_key.clone());
+
+        Ok(stored_key)
     }
 
     async fn update_api_key(
         &self,
-        _merchant_id: String,
+        merchant_id: String,
-        _key_id: String,
+        key_id: String,
-        _api_key: storage::ApiKeyUpdate,
+        api_key: storage::ApiKeyUpdate,
     ) -> CustomResult<storage::ApiKey, errors::StorageError> {
-        // [#172]: Implement function for `MockDb`
+        let mut locked_api_keys = self.api_keys.lock().await;
-        Err(errors::StorageError::MockDbError)?
+        // find a key with the given merchant_id and key_id and update, otherwise return an error
+        let mut key_to_update = locked_api_keys
+            .iter_mut()
+            .find(|k| k.merchant_id == merchant_id && k.key_id == key_id)
+            .ok_or(errors::StorageError::MockDbError)?;
+
+        match api_key {
+            storage::ApiKeyUpdate::Update {
+                name,
+                description,
+                expires_at,
+                last_used,
+            } => {
+                if let Some(name) = name {
+                    key_to_update.name = name;
+                }
+                // only update these fields if the value was Some(_)
+                if description.is_some() {
+                    key_to_update.description = description;
+                }
+                if let Some(expires_at) = expires_at {
+                    key_to_update.expires_at = expires_at;
+                }
+                if last_used.is_some() {
+                    key_to_update.last_used = last_used
+                }
+            }
+            storage::ApiKeyUpdate::LastUsedUpdate { last_used } => {
+                key_to_update.last_used = Some(last_used);
+            }
+        }
+
+        Ok(key_to_update.clone())
     }
 
     async fn revoke_api_key(
         &self,
-        _merchant_id: &str,
+        merchant_id: &str,
-        _key_id: &str,
+        key_id: &str,
     ) -> CustomResult<bool, errors::StorageError> {
-        // [#172]: Implement function for `MockDb`
+        let mut locked_api_keys = self.api_keys.lock().await;
-        Err(errors::StorageError::MockDbError)?
+        // find the key to remove, if it exists
+        if let Some(pos) = locked_api_keys
+            .iter()
+            .position(|k| k.merchant_id == merchant_id && k.key_id == key_id)
+        {
+            // use `remove` instead of `swap_remove` so we have a consistent order, which might
+            // matter to someone using limit/offset in `list_api_keys_by_merchant_id`
+            locked_api_keys.remove(pos);
+            Ok(true)
+        } else {
+            Ok(false)
+        }
     }
 
     async fn find_api_key_by_merchant_id_key_id_optional(
         &self,
-        _merchant_id: &str,
+        merchant_id: &str,
-        _key_id: &str,
+        key_id: &str,
     ) -> CustomResult<Option<storage::ApiKey>, errors::StorageError> {
-        // [#172]: Implement function for `MockDb`
+        Ok(self
-        Err(errors::StorageError::MockDbError)?
+            .api_keys
+            .lock()
+            .await
+            .iter()
+            .find(|k| k.merchant_id == merchant_id && k.key_id == key_id)
+            .cloned())
     }
 
     async fn find_api_key_by_hash_optional(
         &self,
-        _hashed_api_key: storage::HashedApiKey,
+        hashed_api_key: storage::HashedApiKey,
     ) -> CustomResult<Option<storage::ApiKey>, errors::StorageError> {
-        // [#172]: Implement function for `MockDb`
+        Ok(self
-        Err(errors::StorageError::MockDbError)?
+            .api_keys
+            .lock()
+            .await
+            .iter()
+            .find(|k| k.hashed_api_key == hashed_api_key)
+            .cloned())
     }
 
     async fn list_api_keys_by_merchant_id(
         &self,
-        _merchant_id: &str,
+        merchant_id: &str,
-        _limit: Option<i64>,
+        limit: Option<i64>,
-        _offset: Option<i64>,
+        offset: Option<i64>,
     ) -> CustomResult<Vec<storage::ApiKey>, errors::StorageError> {
-        // [#172]: Implement function for `MockDb`
+        // mimic the SQL limit/offset behavior
-        Err(errors::StorageError::MockDbError)?
+        let offset: usize = if let Some(offset) = offset {
+            if offset < 0 {
+                Err(errors::StorageError::MockDbError)?;
+            }
+            offset
+                .try_into()
+                .map_err(|_| errors::StorageError::MockDbError)?
+        } else {
+            0
+        };
+
+        let limit: usize = if let Some(limit) = limit {
+            if limit < 0 {
+                Err(errors::StorageError::MockDbError)?;
+            }
+            limit
+                .try_into()
+                .map_err(|_| errors::StorageError::MockDbError)?
+        } else {
+            usize::MAX
+        };
+
+        let keys_for_merchant_id: Vec<storage::ApiKey> = self
+            .api_keys
+            .lock()
+            .await
+            .iter()
+            .filter(|k| k.merchant_id == merchant_id)
+            .skip(offset)
+            .take(limit)
+            .cloned()
+            .collect();
+
+        Ok(keys_for_merchant_id)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use time::macros::datetime;
+
+    use crate::{
+        db::{api_keys::ApiKeyInterface, MockDb},
+        types::storage,
+    };
+
+    #[allow(clippy::unwrap_used)]
+    #[tokio::test]
+    async fn test_mockdb_api_key_interface() {
+        let mockdb = MockDb::new(&Default::default()).await;
+
+        let key1 = mockdb
+            .insert_api_key(storage::ApiKeyNew {
+                key_id: "key_id1".into(),
+                merchant_id: "merchant1".into(),
+                name: "Key 1".into(),
+                description: None,
+                hashed_api_key: "hashed_key1".to_string().into(),
+                prefix: "abc".into(),
+                created_at: datetime!(2023-02-01 0:00),
+                expires_at: Some(datetime!(2023-03-01 0:00)),
+                last_used: None,
+            })
+            .await
+            .unwrap();
+
+        mockdb
+            .insert_api_key(storage::ApiKeyNew {
+                key_id: "key_id2".into(),
+                merchant_id: "merchant1".into(),
+                name: "Key 2".into(),
+                description: None,
+                hashed_api_key: "hashed_key2".to_string().into(),
+                prefix: "abc".into(),
+                created_at: datetime!(2023-03-01 0:00),
+                expires_at: None,
+                last_used: None,
+            })
+            .await
+            .unwrap();
+
+        let found_key1 = mockdb
+            .find_api_key_by_merchant_id_key_id_optional("merchant1", "key_id1")
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(found_key1.key_id, key1.key_id);
+        assert!(mockdb
+            .find_api_key_by_merchant_id_key_id_optional("merchant1", "does_not_exist")
+            .await
+            .unwrap()
+            .is_none());
+
+        mockdb
+            .update_api_key(
+                "merchant1".into(),
+                "key_id1".into(),
+                storage::ApiKeyUpdate::LastUsedUpdate {
+                    last_used: datetime!(2023-02-04 1:11),
+                },
+            )
+            .await
+            .unwrap();
+        let updated_key1 = mockdb
+            .find_api_key_by_merchant_id_key_id_optional("merchant1", "key_id1")
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(updated_key1.last_used, Some(datetime!(2023-02-04 1:11)));
+
+        assert_eq!(
+            mockdb
+                .list_api_keys_by_merchant_id("merchant1", None, None)
+                .await
+                .unwrap()
+                .len(),
+            2
+        );
+        mockdb.revoke_api_key("merchant1", "key_id1").await.unwrap();
+        assert_eq!(
+            mockdb
+                .list_api_keys_by_merchant_id("merchant1", None, None)
+                .await
+                .unwrap()
+                .len(),
+            1
+        );
     }
 }

crates/storage_models/src/api_keys.rs

@@ -3,7 +3,7 @@ use time::PrimitiveDateTime;
 
 use crate::schema::api_keys;
 
-#[derive(Debug, Identifiable, Queryable)]
+#[derive(Debug, Clone, Identifiable, Queryable)]
 #[diesel(table_name = api_keys, primary_key(key_id))]
 pub struct ApiKey {
     pub key_id: String,
@@ -77,7 +77,7 @@ impl From<ApiKeyUpdate> for ApiKeyUpdateInternal {
     }
 }
 
-#[derive(Debug, AsExpression)]
+#[derive(Debug, Clone, AsExpression, PartialEq)]
 #[diesel(sql_type = diesel::sql_types::Text)]
 pub struct HashedApiKey(String);