feat: Use Variant type in model

juspay · Sep 16, 2024 · 470881c · 470881c
1 parent 0667967
commit 470881c
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 64 deletions.
diff --git a/crates/experimentation_platform/src/api/experiments/handlers.rs b/crates/experimentation_platform/src/api/experiments/handlers.rs
@@ -34,13 +34,13 @@ use super::{
         AuditQueryFilters, ConcludeExperimentRequest, ContextAction, ContextBulkResponse,
         ContextMoveReq, ContextPutReq, ExperimentCreateRequest, ExperimentCreateResponse,
         ExperimentResponse, ExperimentsResponse, ListFilters, OverrideKeysUpdateRequest,
-        RampRequest, Variant,
+        RampRequest,
     },
 };
 
-use crate::{
-    db::models::{EventLog, Experiment, ExperimentStatusType},
-    db::schema::{event_log::dsl as event_log, experiments::dsl as experiments},
+use crate::db::{
+    models::{EventLog, Experiment, ExperimentStatusType, Variant, Variants},
+    schema::{event_log::dsl as event_log, experiments::dsl as experiments},
 };
 
 use serde_json::{json, Map, Value};
@@ -271,7 +271,7 @@ async fn create(
         traffic_percentage: 0,
         status: ExperimentStatusType::CREATED,
         context: Value::Object(req.context.clone().into_inner().into()),
-        variants: serde_json::to_value(variants).unwrap(),
+        variants: Variants::new(variants),
         last_modified_by: user.get_email(),
         chosen_variant: None,
     };
@@ -344,17 +344,9 @@ pub async fn conclude(
     })?;
 
     let mut operations: Vec<ContextAction> = vec![];
-    let experiment_variants: Vec<Variant> = serde_json::from_value(experiment.variants)
-        .map_err(|err| {
-        log::error!(
-            "failed parse eixisting experiment variant while concluding with error: {}",
-            err
-        );
-        unexpected_error!("Something went wrong, failed to conclude experiment")
-    })?;
 
     let mut is_valid_winner_variant = false;
-    for variant in experiment_variants {
+    for variant in experiment.variants.into_inner() {
         let context_id = variant.context_id.ok_or_else(|| {
             log::error!("context id not available for variant {:?}", variant.id);
             unexpected_error!("Something went wrong, failed to conclude experiment")
@@ -547,15 +539,7 @@ async fn ramp(
 
     let old_traffic_percentage = experiment.traffic_percentage as u8;
     let new_traffic_percentage = req.traffic_percentage as u8;
-    let experiment_variants: Vec<Variant> = serde_json::from_value(experiment.variants)
-        .map_err(|e| {
-        log::error!(
-            "failed to parse existing experiment variants while ramping {}",
-            e
-        );
-        unexpected_error!("Something went wrong, failed to ramp traffic percentage")
-    })?;
-    let variants_count = experiment_variants.len() as u8;
+    let variants_count = experiment.variants.into_inner().len() as u8;
     let max = 100 / variants_count;
 
     if matches!(experiment.status, ExperimentStatusType::CONCLUDED) {
@@ -619,11 +603,7 @@ async fn update_overrides(
         ));
     }
 
-    let experiment_variants: Vec<Variant> = serde_json::from_value(experiment.variants)
-        .map_err(|err| {
-        log::error!("failed to parse exisiting variants with error {}", err);
-        unexpected_error!("Something went wrong, failed to update experiment")
-    })?;
+    let experiment_variants: Vec<Variant> = experiment.variants.into_inner();
 
     let id_to_existing_variant: HashMap<String, &Variant> = HashMap::from_iter(
         experiment_variants

diff --git a/crates/experimentation_platform/src/api/experiments/helpers.rs b/crates/experimentation_platform/src/api/experiments/helpers.rs
@@ -1,5 +1,4 @@
-use super::types::{Variant, VariantType};
-use crate::db::models::{Experiment, ExperimentStatusType};
+use crate::db::models::{Experiment, ExperimentStatusType, Variant, VariantType};
 use diesel::pg::PgConnection;
 use diesel::{BoolExpressionMethods, ExpressionMethods, QueryDsl, RunQueryDsl};
 use serde_json::{Map, Value};

diff --git a/crates/experimentation_platform/src/api/experiments/types.rs b/crates/experimentation_platform/src/api/experiments/types.rs
@@ -4,22 +4,7 @@ use serde_json::{Map, Value};
 use service_utils::helpers::deserialize_stringified_list;
 use superposition_types::{Condition, Exp, Overrides};
 
-use crate::db::models::{self, ExperimentStatusType};
-
-#[derive(Deserialize, Serialize, Clone, PartialEq, Debug)]
-pub enum VariantType {
-    CONTROL,
-    EXPERIMENTAL,
-}
-
-#[derive(Deserialize, Serialize, Clone)]
-pub struct Variant {
-    pub id: String,
-    pub variant_type: VariantType,
-    pub context_id: Option<String>,
-    pub override_id: Option<String>,
-    pub overrides: Exp<Overrides>,
-}
+use crate::db::models::{self, ExperimentStatusType, Variant};
 
 /********** Experiment Create Req Types ************/
 
@@ -60,7 +45,7 @@ pub struct ExperimentResponse {
     pub traffic_percentage: i32,
 
     pub context: Value,
-    pub variants: Value,
+    pub variants: Vec<Variant>,
     pub last_modified_by: String,
     pub chosen_variant: Option<String>,
 }
@@ -79,7 +64,7 @@ impl From<models::Experiment> for ExperimentResponse {
             traffic_percentage: experiment.traffic_percentage,
 
             context: experiment.context,
-            variants: experiment.variants,
+            variants: experiment.variants.into_inner(),
             last_modified_by: experiment.last_modified_by,
             chosen_variant: experiment.chosen_variant,
         }

diff --git a/crates/experimentation_platform/src/db/models.rs b/crates/experimentation_platform/src/db/models.rs
@@ -2,10 +2,16 @@ use crate::db::schema::*;
 use chrono::{DateTime, NaiveDateTime, Utc};
 
 use diesel::{
-    query_builder::QueryId, Insertable, Queryable, QueryableByName, Selectable,
+    deserialize::{self, FromSql, FromSqlRow},
+    expression::AsExpression,
+    query_builder::QueryId,
+    serialize::{self, Output, ToSql},
+    sql_types::Json,
+    Insertable, Queryable, QueryableByName, Selectable,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use superposition_types::{Exp, Overrides};
 
 #[derive(
     Debug,
@@ -25,6 +31,58 @@ pub enum ExperimentStatusType {
     INPROGRESS,
 }
 
+#[derive(Deserialize, Serialize, Clone, PartialEq, Debug)]
+pub enum VariantType {
+    CONTROL,
+    EXPERIMENTAL,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Variant {
+    pub id: String,
+    pub variant_type: VariantType,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub context_id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub override_id: Option<String>,
+    pub overrides: Exp<Overrides>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, AsExpression, FromSqlRow)]
+#[diesel(sql_type = Json)]
+pub struct Variants(Vec<Variant>);
+
+impl Variants {
+    pub fn new(data: Vec<Variant>) -> Self {
+        Self(data)
+    }
+
+    pub fn into_inner(self) -> Vec<Variant> {
+        self.0
+    }
+}
+
+impl FromSql<Json, diesel::pg::Pg> for Variants {
+    fn from_sql(bytes: diesel::pg::PgValue<'_>) -> deserialize::Result<Self> {
+        let value =
+            <serde_json::Value as FromSql<Json, diesel::pg::Pg>>::from_sql(bytes)?;
+        Ok(serde_json::from_value(value)?)
+    }
+}
+
+impl ToSql<Json, diesel::pg::Pg> for Variants {
+    fn to_sql<'b>(
+        &'b self,
+        out: &mut Output<'b, '_, diesel::pg::Pg>,
+    ) -> serialize::Result {
+        let value = serde_json::to_value(self)?;
+        <serde_json::Value as ToSql<Json, diesel::pg::Pg>>::to_sql(
+            &value,
+            &mut out.reborrow(),
+        )
+    }
+}
+
 #[derive(QueryableByName, Queryable, Selectable, Insertable, Serialize, Clone, Debug)]
 #[diesel(check_for_backend(diesel::pg::Pg))]
 #[diesel(primary_key(id))]
@@ -40,7 +98,7 @@ pub struct Experiment {
     pub traffic_percentage: i32,
 
     pub context: Value,
-    pub variants: Value,
+    pub variants: Variants,
     pub last_modified_by: String,
     pub chosen_variant: Option<String>,
 }

diff --git a/crates/experimentation_platform/tests/experimentation_tests.rs b/crates/experimentation_platform/tests/experimentation_tests.rs
@@ -1,6 +1,8 @@
 use chrono::Utc;
 use experimentation_platform::api::experiments::helpers;
-use experimentation_platform::db::models::{Experiment, ExperimentStatusType};
+use experimentation_platform::db::models::{
+    Experiment, ExperimentStatusType, Variant, Variants,
+};
 use serde_json::{json, Map, Value};
 use service_utils::helpers::extract_dimensions;
 use service_utils::service::types::ExperimentationFlags;
@@ -52,7 +54,7 @@ fn experiment_gen(
     override_keys: &[String],
     context: &Map<String, Value>,
     status: ExperimentStatusType,
-    variants: &Value,
+    variants: &Vec<Variant>,
 ) -> Experiment {
     Experiment {
         id: 123456789,
@@ -66,7 +68,7 @@ fn experiment_gen(
         override_keys: override_keys.to_vec(),
         status,
         context: json!(context.clone()),
-        variants: variants.clone(),
+        variants: Variants::new(variants.clone()),
         chosen_variant: None,
     }
 }
@@ -231,7 +233,7 @@ fn test_is_valid_experiment_no_restrictions_overlapping_experiment(
         &["key1".to_string(), "key2".to_string()],
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -271,7 +273,7 @@ fn test_is_valid_experiment_no_restrictions_non_overlapping_experiment(
             Dimensions::Client("testclient2".to_string()),
         ]),
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -310,7 +312,7 @@ fn test_is_valid_experiment_restrict_same_keys_overlapping_ctx_overlapping_exper
         &experiment_override_keys,
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -347,7 +349,7 @@ fn test_is_valid_experiment_restrict_same_keys_overlapping_ctx_overlapping_exper
         &["key1".to_string(), "key3".to_string()],
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -384,7 +386,7 @@ fn test_is_valid_experiment_restrict_same_keys_overlapping_ctx_overlapping_exper
         &["key3".to_string(), "key4".to_string()],
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -423,7 +425,7 @@ fn test_is_valid_experiment_restrict_diff_keys_overlapping_ctx_overlapping_exper
         &experiment_override_keys,
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -460,7 +462,7 @@ fn test_is_valid_experiment_restrict_diff_keys_overlapping_ctx_overlapping_exper
         &["key1".to_string(), "key3".to_string()],
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -497,7 +499,7 @@ fn test_is_valid_experiment_restrict_diff_keys_overlapping_ctx_overlapping_exper
         &["key3".to_string(), "key4".to_string()],
         &experiment_context,
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -539,7 +541,7 @@ fn test_is_valid_experiment_restrict_same_keys_non_overlapping_ctx_non_overlappi
             Dimensions::Client("testclient2".to_string()),
         ]),
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -579,7 +581,7 @@ fn test_is_valid_experiment_restrict_same_keys_non_overlapping_ctx_non_overlappi
             Dimensions::Client("testclient2".to_string()),
         ]),
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(
@@ -619,7 +621,7 @@ fn test_is_valid_experiment_restrict_same_keys_non_overlapping_ctx_non_overlappi
             Dimensions::Client("testclient2".to_string()),
         ]),
         ExperimentStatusType::CREATED,
-        &json!(""),
+        &vec![],
     )];
 
     assert_eq!(