feat(users): add support for caching and resolving last used lineage context (#7871)

Co-authored-by: hyperswitch-bot[bot] <148525504+hyperswitch-bot[bot]@users.noreply.github.com>

crates/router/src/consts/user.rs

@@ -41,3 +41,6 @@ pub const REDIS_SSO_TTL: i64 = 5 * 60; // 5 minutes
 pub const DEFAULT_PROFILE_NAME: &str = "default";
 pub const DEFAULT_PRODUCT_TYPE: common_enums::MerchantProductType =
     common_enums::MerchantProductType::Orchestration;
+
+pub const LINEAGE_CONTEXT_TIME_EXPIRY_IN_SECS: i64 = 60 * 60 * 24 * 7; // 7 days
+pub const LINEAGE_CONTEXT_PREFIX: &str = "LINEAGE_CONTEXT_";

crates/router/src/core/user.rs

@@ -3169,6 +3169,23 @@ pub async fn switch_org_for_user(
         }
     };
 
+    let lineage_context = domain::LineageContext {
+        user_id: user_from_token.user_id.clone(),
+        merchant_id: merchant_id.clone(),
+        role_id: role_id.clone(),
+        org_id: request.org_id.clone(),
+        profile_id: profile_id.clone(),
+        tenant_id: user_from_token
+            .tenant_id
+            .as_ref()
+            .unwrap_or(&state.tenant.tenant_id)
+            .clone(),
+    };
+
+    lineage_context
+        .try_set_lineage_context_in_cache(&state, user_from_token.user_id.as_str())
+        .await;
+
     let token = utils::user::generate_jwt_auth_token_with_attributes(
         &state,
         user_from_token.user_id,
@@ -3364,6 +3381,23 @@ pub async fn switch_merchant_for_user_in_org(
         }
     };
 
+    let lineage_context = domain::LineageContext {
+        user_id: user_from_token.user_id.clone(),
+        merchant_id: merchant_id.clone(),
+        role_id: role_id.clone(),
+        org_id: org_id.clone(),
+        profile_id: profile_id.clone(),
+        tenant_id: user_from_token
+            .tenant_id
+            .as_ref()
+            .unwrap_or(&state.tenant.tenant_id)
+            .clone(),
+    };
+
+    lineage_context
+        .try_set_lineage_context_in_cache(&state, user_from_token.user_id.as_str())
+        .await;
+
     let token = utils::user::generate_jwt_auth_token_with_attributes(
         &state,
         user_from_token.user_id,
@@ -3480,6 +3514,23 @@ pub async fn switch_profile_for_user_in_org_and_merchant(
         }
     };
 
+    let lineage_context = domain::LineageContext {
+        user_id: user_from_token.user_id.clone(),
+        merchant_id: user_from_token.merchant_id.clone(),
+        role_id: role_id.clone(),
+        org_id: user_from_token.org_id.clone(),
+        profile_id: profile_id.clone(),
+        tenant_id: user_from_token
+            .tenant_id
+            .as_ref()
+            .unwrap_or(&state.tenant.tenant_id)
+            .clone(),
+    };
+
+    lineage_context
+        .try_set_lineage_context_in_cache(&state, user_from_token.user_id.as_str())
+        .await;
+
     let token = utils::user::generate_jwt_auth_token_with_attributes(
         &state,
         user_from_token.user_id,

crates/router/src/types/domain/user.rs

@@ -23,7 +23,7 @@ use hyperswitch_domain_models::api::ApplicationResponse;
 use masking::{ExposeInterface, PeekInterface, Secret};
 use once_cell::sync::Lazy;
 use rand::distributions::{Alphanumeric, DistString};
-use router_env::env;
+use router_env::{env, logger};
 use time::PrimitiveDateTime;
 use unicode_segmentation::UnicodeSegmentation;
 #[cfg(feature = "keymanager_create")]
@@ -39,7 +39,7 @@ use crate::{
     routes::SessionState,
     services::{self, authentication::UserFromToken},
     types::{domain, transformers::ForeignFrom},
-    utils::user::password,
+    utils::{self, user::password},
 };
 
 pub mod dashboard_metadata;
@@ -1458,3 +1458,54 @@ where
             .change_context(UserErrors::InternalServerError)
     }
 }
+
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct LineageContext {
+    pub user_id: String,
+    pub merchant_id: id_type::MerchantId,
+    pub role_id: String,
+    pub org_id: id_type::OrganizationId,
+    pub profile_id: id_type::ProfileId,
+    pub tenant_id: id_type::TenantId,
+}
+
+impl LineageContext {
+    pub async fn try_get_lineage_context_from_cache(
+        state: &SessionState,
+        user_id: &str,
+    ) -> Option<Self> {
+        // The errors are not handled here because we don't want to fail the request if the cache operation fails.
+        // The errors are logged for debugging purposes.
+        match utils::user::get_lineage_context_from_cache(state, user_id).await {
+            Ok(Some(ctx)) => Some(ctx),
+            Ok(None) => {
+                logger::debug!("Lineage context not found in Redis for user {}", user_id);
+                None
+            }
+            Err(e) => {
+                logger::error!(
+                    "Failed to retrieve lineage context from Redis for user {}: {:?}",
+                    user_id,
+                    e
+                );
+                None
+            }
+        }
+    }
+
+    pub async fn try_set_lineage_context_in_cache(&self, state: &SessionState, user_id: &str) {
+        // The errors are not handled here because we don't want to fail the request if the cache operation fails.
+        // The errors are logged for debugging purposes.
+        if let Err(e) =
+            utils::user::set_lineage_context_in_cache(state, user_id, self.clone()).await
+        {
+            logger::error!(
+                "Failed to set lineage context in Redis for user {}: {:?}",
+                user_id,
+                e
+            );
+        } else {
+            logger::debug!("Lineage context cached for user {}", user_id);
+        }
+    }
+}

crates/router/src/types/domain/user/decision_manager.rs

@@ -1,6 +1,9 @@
 use common_enums::TokenPurpose;
 use common_utils::id_type;
-use diesel_models::{enums::UserStatus, user_role::UserRole};
+use diesel_models::{
+    enums::{UserRoleVersion, UserStatus},
+    user_role::UserRole,
+};
 use error_stack::ResultExt;
 use masking::Secret;
 
@@ -10,6 +13,7 @@ use crate::{
     db::user_role::ListUserRolesByUserIdPayload,
     routes::SessionState,
     services::authentication as auth,
+    types::domain::LineageContext,
     utils,
 };
 
@@ -124,23 +128,93 @@ impl JWTFlow {
         next_flow: &NextFlow,
         user_role: &UserRole,
     ) -> UserResult<Secret<String>> {
+        let user_id = next_flow.user.get_user_id();
+        let cached_lineage_context =
+            LineageContext::try_get_lineage_context_from_cache(state, user_id).await;
+
+        let new_lineage_context = match cached_lineage_context {
+            Some(ctx) => {
+                let tenant_id = ctx.tenant_id.clone();
+                let user_role_match_v1 = state
+                    .global_store
+                    .find_user_role_by_user_id_and_lineage(
+                        &ctx.user_id,
+                        &tenant_id,
+                        &ctx.org_id,
+                        &ctx.merchant_id,
+                        &ctx.profile_id,
+                        UserRoleVersion::V1,
+                    )
+                    .await
+                    .is_ok();
+
+                if user_role_match_v1 {
+                    ctx
+                } else {
+                    let user_role_match_v2 = state
+                        .global_store
+                        .find_user_role_by_user_id_and_lineage(
+                            &ctx.user_id,
+                            &tenant_id,
+                            &ctx.org_id,
+                            &ctx.merchant_id,
+                            &ctx.profile_id,
+                            UserRoleVersion::V2,
+                        )
+                        .await
+                        .is_ok();
+
+                    if user_role_match_v2 {
+                        ctx
+                    } else {
+                        // fallback to default lineage if cached context is invalid
+                        Self::resolve_lineage_from_user_role(state, user_role, user_id).await?
+                    }
+                }
+            }
+            None =>
+            // no cached context found
+            {
+                Self::resolve_lineage_from_user_role(state, user_role, user_id).await?
+            }
+        };
+
+        new_lineage_context
+            .try_set_lineage_context_in_cache(state, user_id)
+            .await;
+
+        auth::AuthToken::new_token(
+            new_lineage_context.user_id,
+            new_lineage_context.merchant_id,
+            new_lineage_context.role_id,
+            &state.conf,
+            new_lineage_context.org_id,
+            new_lineage_context.profile_id,
+            Some(new_lineage_context.tenant_id),
+        )
+        .await
+        .map(|token| token.into())
+    }
+
+    pub async fn resolve_lineage_from_user_role(
+        state: &SessionState,
+        user_role: &UserRole,
+        user_id: &str,
+    ) -> UserResult<LineageContext> {
         let org_id = utils::user_role::get_single_org_id(state, user_role).await?;
         let merchant_id =
             utils::user_role::get_single_merchant_id(state, user_role, &org_id).await?;
         let profile_id =
             utils::user_role::get_single_profile_id(state, user_role, &merchant_id).await?;
 
-        auth::AuthToken::new_token(
+        Ok(LineageContext {
-            next_flow.user.get_user_id().to_string(),
+            user_id: user_id.to_string(),
-            merchant_id,
-            user_role.role_id.clone(),
-            &state.conf,
             org_id,
+            merchant_id,
             profile_id,
-            Some(user_role.tenant_id.clone()),
+            role_id: user_role.role_id.clone(),
-        )
+            tenant_id: user_role.tenant_id.clone(),
-        .await
+        })
-        .map(|token| token.into())
     }
 }
 

crates/router/src/utils/user.rs

@@ -12,7 +12,10 @@ use redis_interface::RedisConnectionPool;
 use router_env::env;
 
 use crate::{
-    consts::user::{REDIS_SSO_PREFIX, REDIS_SSO_TTL},
+    consts::user::{
+        LINEAGE_CONTEXT_PREFIX, LINEAGE_CONTEXT_TIME_EXPIRY_IN_SECS, REDIS_SSO_PREFIX,
+        REDIS_SSO_TTL,
+    },
     core::errors::{StorageError, UserErrors, UserResult},
     routes::SessionState,
     services::{
@@ -20,7 +23,7 @@ use crate::{
         authorization::roles::RoleInfo,
     },
     types::{
-        domain::{self, MerchantAccount, UserFromStorage},
+        domain::{self, LineageContext, MerchantAccount, UserFromStorage},
         transformers::ForeignFrom,
     },
 };
@@ -339,3 +342,49 @@ pub async fn validate_email_domain_auth_type_using_db(
     .then_some(())
     .ok_or(UserErrors::InvalidUserAuthMethodOperation.into())
 }
+
+pub async fn get_lineage_context_from_cache(
+    state: &SessionState,
+    user_id: &str,
+) -> UserResult<Option<LineageContext>> {
+    let connection = get_redis_connection(state)?;
+    let key = format!("{}{}", LINEAGE_CONTEXT_PREFIX, user_id);
+    let lineage_context = connection
+        .get_key::<Option<String>>(&key.into())
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .attach_printable("Failed to get lineage context from redis")?;
+
+    match lineage_context {
+        Some(json_str) => {
+            let ctx = serde_json::from_str::<LineageContext>(&json_str)
+                .change_context(UserErrors::InternalServerError)
+                .attach_printable("Failed to deserialize LineageContext from JSON")?;
+            Ok(Some(ctx))
+        }
+        None => Ok(None),
+    }
+}
+
+pub async fn set_lineage_context_in_cache(
+    state: &SessionState,
+    user_id: &str,
+    lineage_context: LineageContext,
+) -> UserResult<()> {
+    let connection = get_redis_connection(state)?;
+    let key = format!("{}{}", LINEAGE_CONTEXT_PREFIX, user_id);
+    let serialized_lineage_context: String = serde_json::to_string(&lineage_context)
+        .change_context(UserErrors::InternalServerError)
+        .attach_printable("Failed to serialize LineageContext")?;
+    connection
+        .set_key_with_expiry(
+            &key.into(),
+            serialized_lineage_context,
+            LINEAGE_CONTEXT_TIME_EXPIRY_IN_SECS,
+        )
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .attach_printable("Failed to set lineage context in redis")?;
+
+    Ok(())
+}