allow setting referral code if not already set (#141)

web3_proxy/src/frontend/users/mod.rs

@@ -8,7 +8,6 @@ pub mod subuser;
 
 use crate::app::Web3ProxyApp;
 use crate::errors::{Web3ProxyError, Web3ProxyErrorContext, Web3ProxyResponse};
-
 use axum::{
     headers::{authorization::Bearer, Authorization},
     response::IntoResponse,
@@ -16,9 +15,8 @@ use axum::{
 };
 use axum_macros::debug_handler;
 use check_if_email_exists::{check_email, CheckEmailInput, Reachable};
-use entities;
-use entities::user;
-use migration::sea_orm::{self, ActiveModelTrait};
+use entities::{self, referee, referrer, user};
+use migration::sea_orm::{self, ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter};
 use serde::Deserialize;
 use std::sync::Arc;
 
@@ -38,9 +36,11 @@ pub async fn user_get(
 }
 
 /// the JSON input to the `post_user` handler.
+/// TODO: what else can we update here? password hash? subscription to newsletter?
 #[derive(Debug, Deserialize)]
 pub struct UserPost {
     email: Option<String>,
+    referral_code: Option<String>,
 }
 
 /// `POST /user` -- modify the account connected to the bearer token in the `Authentication` header.
@@ -52,6 +52,8 @@ pub async fn user_post(
 ) -> Web3ProxyResponse {
     let (user, _semaphore) = app.bearer_is_authorized(bearer_token).await?;
 
+    let user_id = user.id;
+
     let mut user: user::ActiveModel = user.into();
 
     // update the email address
@@ -72,24 +74,56 @@ pub async fn user_post(
                 ));
             }
 
-            // TODO: do some basic validation
-            // TODO: don't set immediatly, send a confirmation email first
-            // TODO: compare first? or is sea orm smart enough to do that for us?
+            // TODO: send a confirmation email first before marking this email address as validated
             user.email = sea_orm::Set(Some(x));
         }
     }
 
-    // TODO: what else can we update here? password hash? subscription to newsletter?
+    let txn = app.db_transaction().await?;
+
+    // update the referral code IFF they do not already have one set
+    if let Some(x) = payload.referral_code {
+        let desired_referral_code = referrer::Entity::find()
+            .filter(referrer::Column::ReferralCode.like(&x))
+            .one(&txn)
+            .await?
+            .web3_context(format!("posted referral code does not exist"))?;
+
+        // make sure the user doesn't already have a referral code set
+        if let Some(existing_referee) = referee::Entity::find()
+            .filter(referee::Column::UserId.eq(user_id))
+            .one(&txn)
+            .await?
+        {
+            if existing_referee.used_referral_code == desired_referral_code.id {
+                // code was already set this code. nothing to do
+            } else {
+                // referral code cannot change! error!
+                return Err(Web3ProxyError::BadRequest(
+                    "referral code cannot be changed".into(),
+                ));
+            }
+        } else {
+            // no existing referral code. set one now
+            let new_referee = referee::ActiveModel {
+                used_referral_code: sea_orm::Set(desired_referral_code.id),
+                user_id: sea_orm::Set(user_id),
+                ..Default::default()
+            };
+
+            new_referee.save(&txn).await?;
+        }
+    }
 
     let user = if user.is_changed() {
-        let db_conn = app.db_conn().web3_context("Getting database connection")?;
-
-        user.save(&db_conn).await?
+        user.save(&txn).await?
     } else {
         // no changes. no need to touch the database
         user
     };
 
+    txn.commit().await?;
+
     let user: user::Model = user.try_into().web3_context("Returning updated user")?;
 
     Ok(Json(user).into_response())