feat: support an external identity store

src/environments.rs

@@ -5,6 +5,7 @@ use chrono::{DateTime, Utc};
 use parking_lot::RwLock;
 
 use crate::config::settings::EnvironmentKeyPair;
+use crate::identities::EnvironmentKeyType;
 
 /// A server-side (`ser.`) key together with the validity metadata the
 /// proxy config endpoint reports. Statically configured keys carry no
@@ -32,6 +33,14 @@ pub struct EnvironmentKeys {
 }
 
 impl EnvironmentKeys {
+    pub fn key_type(&self, environment_key: &str) -> EnvironmentKeyType {
+        if environment_key == self.client_key {
+            EnvironmentKeyType::Client
+        } else {
+            EnvironmentKeyType::Server
+        }
+    }
+
     /// The first server-side key still usable for upstream fetches.
     pub fn valid_server_key(&self) -> Option<&ServerKey> {
         self.server_keys.iter().find(|key| key.is_valid())

src/identities/mod.rs

@@ -0,0 +1,3 @@
+pub mod store;
+
+pub use store::{EnvironmentKeyType, IdentityStore, IdentityTraits};

src/identities/store.rs

@@ -0,0 +1,30 @@
+use crate::models::{IdentityWithTraits, TraitModel};
+use async_trait::async_trait;
+
+#[async_trait]
+pub trait IdentityStore: Send + Sync {
+    async fn load_traits(
+        &self,
+        environment: &serde_json::Value,
+        identity: &IdentityWithTraits,
+    ) -> IdentityTraits;
+
+    fn save_traits(
+        &self,
+        environment: &serde_json::Value,
+        identity: &IdentityWithTraits,
+        key_type: EnvironmentKeyType,
+    );
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct IdentityTraits {
+    pub traits: Vec<TraitModel>,
+    pub system_traits: Vec<TraitModel>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum EnvironmentKeyType {
+    Client,
+    Server,
+}

src/lib.rs

@@ -2,6 +2,7 @@ pub mod cache;
 pub mod config;
 pub mod environments;
 pub mod error;
+pub mod identities;
 pub mod models;
 pub mod routes;
 pub mod server;

src/main.rs

@@ -2,5 +2,5 @@ use edge_proxy::config::settings::get_settings;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    edge_proxy::run(get_settings()?).await
+    edge_proxy::run(get_settings()?, None).await
 }

src/routes/mod.rs

@@ -6,6 +6,7 @@ pub mod health;
 pub mod identities;
 
 use crate::config::AppSettings;
+use crate::identities::IdentityStore;
 use crate::services::EnvironmentService;
 use axum::{
     Router,
@@ -15,9 +16,16 @@ use axum::{
 use std::sync::Arc;
 use tower_http::{compression::CompressionLayer, normalize_path::NormalizePath, trace::TraceLayer};
 
-pub fn create_router(settings: AppSettings) -> (Router, Arc<EnvironmentService>) {
+pub fn create_router(
+    settings: AppSettings,
+    identity_store: Option<Arc<dyn IdentityStore>>,
+) -> (Router, Arc<EnvironmentService>) {
     let cors = cors::layer(&settings.allow_origins);
-    let environment_service = Arc::new(EnvironmentService::new(settings));
+    let environment_service = EnvironmentService::new(settings);
+    let environment_service = Arc::new(match identity_store {
+        Some(identity_store) => environment_service.with_identity_store(identity_store),
+        None => environment_service,
+    });
 
     let router = Router::new()
         // Health check routes

src/server.rs

@@ -1,12 +1,17 @@
 use crate::config::AppSettings;
 use crate::config::logging::setup_logging;
+use crate::identities::IdentityStore;
 use crate::routes::create_router;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::sync::Arc;
 use tracing::{info, warn};
 
 const DEFAULT_HOST: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
 
-pub async fn run(settings: AppSettings) -> anyhow::Result<()> {
+pub async fn run(
+    settings: AppSettings,
+    identity_store: Option<Arc<dyn IdentityStore>>,
+) -> anyhow::Result<()> {
     setup_logging(&settings.logging);
 
     info!("Starting Edge Proxy server...");
@@ -27,7 +32,7 @@ pub async fn run(settings: AppSettings) -> anyhow::Result<()> {
         );
     }
 
-    let (app, environment_service) = create_router(settings.clone());
+    let (app, environment_service) = create_router(settings.clone(), identity_store);
 
     // Refreshes must never overlap: a delayed older poll finishing after a
     // newer one could restore removed environments or rotated keys. The

src/services/environment.rs

@@ -2,8 +2,9 @@ use crate::cache::{EnvironmentsCache, LocalMemEnvironmentsCache};
 use crate::config::settings::AppSettings;
 use crate::environments::{EnvironmentIndex, EnvironmentKeys};
 use crate::error::{EdgeProxyError, Result};
+use crate::identities::IdentityStore;
 use crate::models::{
-    APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment,
+    APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, TraitModel,
 };
 use crate::services::feature_utils::filter_out_server_key_only_flag_results;
 use chrono::{DateTime, Utc};
@@ -23,6 +24,7 @@ pub struct EnvironmentService {
     pub settings: AppSettings,
     pub last_updated_at: Arc<RwLock<Option<DateTime<Utc>>>>,
     environments: EnvironmentIndex,
+    identity_store: Option<Arc<dyn IdentityStore>>,
 }
 
 impl EnvironmentService {
@@ -41,6 +43,7 @@ impl EnvironmentService {
             settings,
             last_updated_at: Arc::new(RwLock::new(None)),
             environments,
+            identity_store: None,
         }
     }
 
@@ -50,6 +53,11 @@ impl EnvironmentService {
         service
     }
 
+    pub fn with_identity_store(mut self, identity_store: Arc<dyn IdentityStore>) -> Self {
+        self.identity_store = Some(identity_store);
+        self
+    }
+
     pub async fn refresh_environment_caches(&self) -> bool {
         // Sync first so an environment added to the proxy config gets its
         // document fetched in the same pass.
@@ -341,7 +349,65 @@ impl EnvironmentService {
         environment_key: &str,
     ) -> Result<IdentityResponse> {
         let keys = self.resolve_key(environment_key)?;
+        match &self.identity_store {
+            None => {
+                self.evaluate_identity(
+                    &keys,
+                    environment_key,
+                    &identity.identifier,
+                    &identity.traits,
+                    &[],
+                )
+                .await
+            }
+            Some(identity_store) => {
+                self.evaluate_with_identity_store(
+                    identity_store.as_ref(),
+                    &keys,
+                    environment_key,
+                    identity,
+                )
+                .await
+            }
+        }
+    }
+
+    async fn evaluate_with_identity_store(
+        &self,
+        identity_store: &dyn IdentityStore,
+        keys: &EnvironmentKeys,
+        environment_key: &str,
+        identity: &IdentityWithTraits,
+    ) -> Result<IdentityResponse> {
+        let environment = self
+            .cache
+            .get_environment(&keys.client_key)
+            .await
+            .ok_or_else(|| {
+                EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string())
+            })?;
+        let traits = identity_store.load_traits(&environment, identity).await;
+        let response = self
+            .evaluate_identity(
+                keys,
+                environment_key,
+                &identity.identifier,
+                &traits.traits,
+                &traits.system_traits,
+            )
+            .await?;
+        identity_store.save_traits(&environment, identity, keys.key_type(environment_key));
+        Ok(response)
+    }
 
+    async fn evaluate_identity(
+        &self,
+        keys: &EnvironmentKeys,
+        environment_key: &str,
+        identifier: &str,
+        traits: &[TraitModel],
+        system_traits: &[TraitModel],
+    ) -> Result<IdentityResponse> {
         // Get pre-computed context from cache
         let context = self
             .cache
@@ -351,11 +417,15 @@ impl EnvironmentService {
                 EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string())
             })?;
 
-        let flagsmith_traits: Vec<FlagsmithTrait> =
-            identity.traits.iter().map(Into::into).collect();
+        let flagsmith_traits: Vec<FlagsmithTrait> = traits
+            .iter()
+            .filter(|t| !system_traits.iter().any(|s| s.trait_key == t.trait_key))
+            .chain(system_traits)
+            .map(Into::into)
+            .collect();
 
         let context_with_identity =
-            add_identity_to_context(&context, &identity.identifier, &flagsmith_traits);
+            add_identity_to_context(&context, identifier, &flagsmith_traits);
 
         let evaluation_result = get_evaluation_result(&context_with_identity);
 
@@ -373,11 +443,7 @@ impl EnvironmentService {
 
         Ok(IdentityResponse {
             flags: flag_results.iter().map(Into::into).collect(),
-            traits: identity
-                .traits
-                .iter()
-                .map(|t| t.to_response_json())
-                .collect(),
+            traits: traits.iter().map(|t| t.to_response_json()).collect(),
         })
     }
 

tests/fixtures.rs

@@ -89,7 +89,7 @@ pub fn environment_1() -> serde_json::Value {
                                         {
                                             "value": "test",
                                             "operator": "EQUAL",
-                                            "property": "first_name",
+                                            "property_": "first_name",
                                         }
                                     ],
                                     "type": "ANY",

tests/test_cors.rs

@@ -17,7 +17,7 @@ async fn server_allowing(allow_origins: &[&str]) -> TestServer {
         allow_origins: allow_origins.iter().map(|o| o.to_string()).collect(),
         ..AppSettings::new()
     };
-    let (app, service) = create_router(settings);
+    let (app, service) = create_router(settings, None);
     service
         .cache
         .put_environment(&environment_1_api_key(), environment_1())

tests/test_health.rs

@@ -7,7 +7,7 @@ use edge_proxy::routes::create_router;
 async fn test_liveness_check() {
     // Given
     let settings = AppSettings::new();
-    let (app, _) = create_router(settings);
+    let (app, _) = create_router(settings, None);
     let server = TestServer::new(app).unwrap();
 
     // When
@@ -21,7 +21,7 @@ async fn test_liveness_check() {
 async fn test_health_check_returns_200_if_cache_was_updated_recently() {
     // Given
     let settings = AppSettings::new();
-    let (app, service) = create_router(settings);
+    let (app, service) = create_router(settings, None);
     {
         let mut last_updated = service.last_updated_at.write().await;
         *last_updated = Some(Utc::now());
@@ -42,7 +42,7 @@ async fn test_health_check_returns_200_if_cache_was_updated_recently() {
 async fn test_health_check_returns_503_if_cache_was_not_updated() {
     // Given
     let settings = AppSettings::new();
-    let (app, _service) = create_router(settings);
+    let (app, _service) = create_router(settings, None);
     let server = TestServer::new(app).unwrap();
     // last_updated_at is None by default
 
@@ -62,7 +62,7 @@ async fn test_health_check_returns_503_if_cache_was_not_updated() {
 async fn test_health_check_returns_503_if_cache_is_stale() {
     // Given
     let settings = AppSettings::new();
-    let (app, service) = create_router(settings);
+    let (app, service) = create_router(settings, None);
     {
         let mut last_updated = service.last_updated_at.write().await;
         *last_updated = Some(Utc::now() - chrono::Duration::days(10));
@@ -88,7 +88,7 @@ async fn test_health_check_returns_200_if_cache_is_never_stale() {
     settings.health_check = HealthCheckSettings {
         environment_update_grace_period_seconds: None,
     };
-    let (app, service) = create_router(settings);
+    let (app, service) = create_router(settings, None);
     let last_update_time = Utc::now() - chrono::Duration::days(10);
     {
         let mut last_updated = service.last_updated_at.write().await;