refactor(api_keys): use merchant_id and key_id to query the table (#939)

Signed-off-by: Chethan <chethan.rao@juspay.in>

crates/api_models/src/api_keys.rs

@@ -134,10 +134,13 @@ pub struct UpdateApiKeyRequest {
 /// The response body for revoking an API Key.
 #[derive(Debug, Serialize, ToSchema)]
 pub struct RevokeApiKeyResponse {
+    /// The identifier for the Merchant Account.
+    #[schema(max_length = 64, example = "y3oqhf46pyzuxjbcn2giaqnb44")]
+    pub merchant_id: String,
+
     /// The identifier for the API Key.
     #[schema(max_length = 64, example = "5hEEqkgJUyuxgSKGArHA4mWSnX")]
     pub key_id: String,
-
     /// Indicates whether the API key was revoked or not.
     #[schema(example = "true")]
     pub revoked: bool,

crates/router/src/core/api_keys.rs

@@ -158,10 +158,11 @@ pub async fn create_api_key(
 #[instrument(skip_all)]
 pub async fn retrieve_api_key(
     store: &dyn StorageInterface,
+    merchant_id: &str,
     key_id: &str,
 ) -> RouterResponse<api::RetrieveApiKeyResponse> {
     let api_key = store
-        .find_api_key_by_key_id_optional(key_id)
+        .find_api_key_by_merchant_id_key_id_optional(merchant_id, key_id)
         .await
         .change_context(errors::ApiErrorResponse::InternalServerError) // If retrieve failed
         .attach_printable("Failed to retrieve new API key")?
@@ -173,11 +174,16 @@ pub async fn retrieve_api_key(
 #[instrument(skip_all)]
 pub async fn update_api_key(
     store: &dyn StorageInterface,
+    merchant_id: &str,
     key_id: &str,
     api_key: api::UpdateApiKeyRequest,
 ) -> RouterResponse<api::RetrieveApiKeyResponse> {
     let api_key = store
-        .update_api_key(key_id.to_owned(), api_key.foreign_into())
+        .update_api_key(
+            merchant_id.to_owned(),
+            key_id.to_owned(),
+            api_key.foreign_into(),
+        )
         .await
         .to_not_found_response(errors::ApiErrorResponse::ApiKeyNotFound)?;
 
@@ -187,16 +193,18 @@ pub async fn update_api_key(
 #[instrument(skip_all)]
 pub async fn revoke_api_key(
     store: &dyn StorageInterface,
+    merchant_id: &str,
     key_id: &str,
 ) -> RouterResponse<api::RevokeApiKeyResponse> {
     let revoked = store
-        .revoke_api_key(key_id)
+        .revoke_api_key(merchant_id, key_id)
         .await
         .to_not_found_response(errors::ApiErrorResponse::ApiKeyNotFound)?;
 
     metrics::API_KEY_REVOKED.add(&metrics::CONTEXT, 1, &[]);
 
     Ok(ApplicationResponse::Json(api::RevokeApiKeyResponse {
+        merchant_id: merchant_id.to_owned(),
         key_id: key_id.to_owned(),
         revoked,
     }))

crates/router/src/db/api_keys.rs

@@ -16,14 +16,20 @@ pub trait ApiKeyInterface {
 
     async fn update_api_key(
         &self,
+        merchant_id: String,
         key_id: String,
         api_key: storage::ApiKeyUpdate,
     ) -> CustomResult<storage::ApiKey, errors::StorageError>;
 
-    async fn revoke_api_key(&self, key_id: &str) -> CustomResult<bool, errors::StorageError>;
+    async fn revoke_api_key(
-
-    async fn find_api_key_by_key_id_optional(
         &self,
+        merchant_id: &str,
+        key_id: &str,
+    ) -> CustomResult<bool, errors::StorageError>;
+
+    async fn find_api_key_by_merchant_id_key_id_optional(
+        &self,
+        merchant_id: &str,
         key_id: &str,
     ) -> CustomResult<Option<storage::ApiKey>, errors::StorageError>;
 
@@ -56,30 +62,36 @@ impl ApiKeyInterface for Store {
 
     async fn update_api_key(
         &self,
+        merchant_id: String,
         key_id: String,
         api_key: storage::ApiKeyUpdate,
     ) -> CustomResult<storage::ApiKey, errors::StorageError> {
         let conn = connection::pg_connection_write(self).await?;
-        storage::ApiKey::update_by_key_id(&conn, key_id, api_key)
+        storage::ApiKey::update_by_merchant_id_key_id(&conn, merchant_id, key_id, api_key)
             .await
             .map_err(Into::into)
             .into_report()
     }
 
-    async fn revoke_api_key(&self, key_id: &str) -> CustomResult<bool, errors::StorageError> {
+    async fn revoke_api_key(
-        let conn = connection::pg_connection_write(self).await?;
-        storage::ApiKey::revoke_by_key_id(&conn, key_id)
-            .await
-            .map_err(Into::into)
-            .into_report()
-    }
-
-    async fn find_api_key_by_key_id_optional(
         &self,
+        merchant_id: &str,
+        key_id: &str,
+    ) -> CustomResult<bool, errors::StorageError> {
+        let conn = connection::pg_connection_write(self).await?;
+        storage::ApiKey::revoke_by_merchant_id_key_id(&conn, merchant_id, key_id)
+            .await
+            .map_err(Into::into)
+            .into_report()
+    }
+
+    async fn find_api_key_by_merchant_id_key_id_optional(
+        &self,
+        merchant_id: &str,
         key_id: &str,
     ) -> CustomResult<Option<storage::ApiKey>, errors::StorageError> {
         let conn = connection::pg_connection_read(self).await?;
-        storage::ApiKey::find_optional_by_key_id(&conn, key_id)
+        storage::ApiKey::find_optional_by_merchant_id_key_id(&conn, merchant_id, key_id)
             .await
             .map_err(Into::into)
             .into_report()
@@ -122,6 +134,7 @@ impl ApiKeyInterface for MockDb {
 
     async fn update_api_key(
         &self,
+        _merchant_id: String,
         _key_id: String,
         _api_key: storage::ApiKeyUpdate,
     ) -> CustomResult<storage::ApiKey, errors::StorageError> {
@@ -129,13 +142,18 @@ impl ApiKeyInterface for MockDb {
         Err(errors::StorageError::MockDbError)?
     }
 
-    async fn revoke_api_key(&self, _key_id: &str) -> CustomResult<bool, errors::StorageError> {
+    async fn revoke_api_key(
+        &self,
+        _merchant_id: &str,
+        _key_id: &str,
+    ) -> CustomResult<bool, errors::StorageError> {
         // [#172]: Implement function for `MockDb`
         Err(errors::StorageError::MockDbError)?
     }
 
-    async fn find_api_key_by_key_id_optional(
+    async fn find_api_key_by_merchant_id_key_id_optional(
         &self,
+        _merchant_id: &str,
         _key_id: &str,
     ) -> CustomResult<Option<storage::ApiKey>, errors::StorageError> {
         // [#172]: Implement function for `MockDb`

crates/router/src/routes/api_keys.rs

@@ -82,14 +82,16 @@ pub async fn api_key_retrieve(
     path: web::Path<(String, String)>,
 ) -> impl Responder {
     let flow = Flow::ApiKeyRetrieve;
-    let (_merchant_id, key_id) = path.into_inner();
+    let (merchant_id, key_id) = path.into_inner();
 
     api::server_wrap(
         flow,
         state.get_ref(),
         &req,
-        &key_id,
+        (&merchant_id, &key_id),
-        |state, _, key_id| api_keys::retrieve_api_key(&*state.store, key_id),
+        |state, _, (merchant_id, key_id)| {
+            api_keys::retrieve_api_key(&*state.store, merchant_id, key_id)
+        },
         &auth::AdminApiAuth,
     )
     .await
@@ -122,15 +124,17 @@ pub async fn api_key_update(
     json_payload: web::Json<api_types::UpdateApiKeyRequest>,
 ) -> impl Responder {
     let flow = Flow::ApiKeyUpdate;
-    let (_merchant_id, key_id) = path.into_inner();
+    let (merchant_id, key_id) = path.into_inner();
     let payload = json_payload.into_inner();
 
     api::server_wrap(
         flow,
         state.get_ref(),
         &req,
-        (&key_id, payload),
+        (&merchant_id, &key_id, payload),
-        |state, _, (key_id, payload)| api_keys::update_api_key(&*state.store, key_id, payload),
+        |state, _, (merchant_id, key_id, payload)| {
+            api_keys::update_api_key(&*state.store, merchant_id, key_id, payload)
+        },
         &auth::AdminApiAuth,
     )
     .await
@@ -162,14 +166,16 @@ pub async fn api_key_revoke(
     path: web::Path<(String, String)>,
 ) -> impl Responder {
     let flow = Flow::ApiKeyRevoke;
-    let (_merchant_id, key_id) = path.into_inner();
+    let (merchant_id, key_id) = path.into_inner();
 
     api::server_wrap(
         flow,
         state.get_ref(),
         &req,
-        &key_id,
+        (&merchant_id, &key_id),
-        |state, _, key_id| api_keys::revoke_api_key(&*state.store, key_id),
+        |state, _, (merchant_id, key_id)| {
+            api_keys::revoke_api_key(&*state.store, merchant_id, key_id)
+        },
         &auth::AdminApiAuth,
     )
     .await

crates/storage_models/src/query/api_keys.rs

@@ -1,4 +1,4 @@
-use diesel::{associations::HasTable, ExpressionMethods};
+use diesel::{associations::HasTable, BoolExpressionMethods, ExpressionMethods};
 use router_env::{instrument, tracing};
 
 use super::generics;
@@ -18,14 +18,22 @@ impl ApiKeyNew {
 
 impl ApiKey {
     #[instrument(skip(conn))]
-    pub async fn update_by_key_id(
+    pub async fn update_by_merchant_id_key_id(
         conn: &PgPooledConn,
+        merchant_id: String,
         key_id: String,
         api_key_update: ApiKeyUpdate,
     ) -> StorageResult<Self> {
-        match generics::generic_update_by_id::<<Self as HasTable>::Table, _, _, _>(
+        match generics::generic_update_with_unique_predicate_get_result::<
+            <Self as HasTable>::Table,
+            _,
+            _,
+            _,
+        >(
             conn,
-            key_id.clone(),
+            dsl::merchant_id
+                .eq(merchant_id.to_owned())
+                .and(dsl::key_id.eq(key_id.to_owned())),
             ApiKeyUpdateInternal::from(api_key_update),
         )
         .await
@@ -35,8 +43,13 @@ impl ApiKey {
                     Err(error.attach_printable("API key with the given key ID does not exist"))
                 }
                 errors::DatabaseError::NoFieldsToUpdate => {
-                    generics::generic_find_by_id::<<Self as HasTable>::Table, _, _>(conn, key_id)
+                    generics::generic_find_one::<<Self as HasTable>::Table, _, _>(
-                        .await
+                        conn,
+                        dsl::merchant_id
+                            .eq(merchant_id.to_owned())
+                            .and(dsl::key_id.eq(key_id.to_owned())),
+                    )
+                    .await
                 }
                 _ => Err(error),
             },
@@ -45,22 +58,31 @@ impl ApiKey {
     }
 
     #[instrument(skip(conn))]
-    pub async fn revoke_by_key_id(conn: &PgPooledConn, key_id: &str) -> StorageResult<bool> {
+    pub async fn revoke_by_merchant_id_key_id(
+        conn: &PgPooledConn,
+        merchant_id: &str,
+        key_id: &str,
+    ) -> StorageResult<bool> {
         generics::generic_delete::<<Self as HasTable>::Table, _>(
             conn,
-            dsl::key_id.eq(key_id.to_owned()),
+            dsl::merchant_id
+                .eq(merchant_id.to_owned())
+                .and(dsl::key_id.eq(key_id.to_owned())),
         )
         .await
     }
 
     #[instrument(skip(conn))]
-    pub async fn find_optional_by_key_id(
+    pub async fn find_optional_by_merchant_id_key_id(
         conn: &PgPooledConn,
+        merchant_id: &str,
         key_id: &str,
     ) -> StorageResult<Option<Self>> {
-        generics::generic_find_by_id_optional::<<Self as HasTable>::Table, _, _>(
+        generics::generic_find_one_optional::<<Self as HasTable>::Table, _, _>(
             conn,
-            key_id.to_owned(),
+            dsl::merchant_id
+                .eq(merchant_id.to_owned())
+                .and(dsl::key_id.eq(key_id.to_owned())),
         )
         .await
     }

migrations/2023-04-21-100150_create_index_for_api_keys/down.sql

@@ -0,0 +1 @@
+DROP INDEX api_keys_merchant_id_key_id_index;

migrations/2023-04-21-100150_create_index_for_api_keys/up.sql

@@ -0,0 +1 @@
+CREATE UNIQUE INDEX api_keys_merchant_id_key_id_index ON api_keys (merchant_id, key_id);