feat(upsert): TTL

example/src/entities/person/model.rs

@@ -22,6 +22,7 @@ pub enum PersonKind {
 /// Represents a person in the database
 #[entity]
 #[upsert_query(table = "person", name = UpsertPerson)]
+#[upsert_query(table = "person", name = UpsertPersonWithTTL, ttl)]
 pub struct PersonEntity {
     /// The id of the person
     #[entity(primary_key)]
@@ -109,4 +110,54 @@ mod test {
 
         assert_eq!(values, result_values);
     }
+
+    #[test]
+    fn test_upsert_ttl() {
+        let upsert = UpsertPersonWithTTL {
+            id: v1_uuid(),
+            email: MaybeUnset::Set("[email protected]".to_string()),
+            age: MaybeUnset::Unset,
+            kind: MaybeUnset::Set(PersonKind::Parent),
+            data: MaybeUnset::Set(Some(PersonData {
+                stripe_id: Some("stripe_id".to_string()),
+            })),
+            created_at: MaybeUnset::Unset,
+
+            set_ttl: 300,
+        };
+
+        let query = <UpsertPersonWithTTL as Query>::query();
+        let values = <UpsertPersonWithTTL as Query>::bind(&upsert).unwrap();
+
+        assert_eq!(
+            query,
+            r#"update person using ttl :set_ttl set "email" = :email, "age" = :age, "data" = :data, "kind" = :kind, "createdAt" = :created_at where "id" = :id;"#
+        );
+
+        let mut result_values = SerializedValues::new();
+        result_values
+            .add_named_value("email", &upsert.email)
+            .expect("failed to add value");
+        result_values
+            .add_named_value("age", &upsert.age)
+            .expect("failed to add value");
+        result_values
+            .add_named_value("data", &upsert.data)
+            .expect("failed to add value");
+        result_values
+            .add_named_value("kind", &upsert.kind)
+            .expect("failed to add value");
+        result_values
+            .add_named_value("created_at", &upsert.created_at)
+            .expect("failed to add value");
+        result_values
+            .add_named_value("id", &upsert.id)
+            .expect("failed to add value");
+
+        result_values
+            .add_named_value("set_ttl", &upsert.set_ttl)
+            .expect("failed to add value");
+
+        assert_eq!(values, result_values);
+    }
 }

example/src/entities/person/queries.rs

@@ -1,4 +1,4 @@
-use super::model::UpsertPerson;
+use super::model::{UpsertPerson, UpsertPersonWithTTL};
 use scyllax::prelude::*;
 use uuid::Uuid;
 
@@ -10,6 +10,7 @@ create_query_collection!(
         GetPersonByEmail,
         DeletePersonById,
         UpsertPerson,
+        UpsertPersonWithTTL,
     ]
 );
 

example/src/main.rs

@@ -1,8 +1,8 @@
 //! Example
-use example::entities::person::{
-    model::{PersonData, PersonKind, UpsertPerson},
+use example::entities::{person::{
+    model::{PersonData, PersonKind, UpsertPerson, UpsertPersonWithTTL},
     queries::{DeletePersonById, GetPeopleByIds, GetPersonByEmail, GetPersonById, PersonQueries},
-};
+}, PersonEntity};
 use scyllax::prelude::*;
 use scyllax::{executor::create_session, util::v1_uuid};
 use tracing_subscriber::prelude::*;
@@ -22,33 +22,18 @@ async fn main() -> anyhow::Result<()> {
     let session = create_session(known_nodes, default_keyspace).await?;
     let executor = Executor::<PersonQueries>::new(session).await?;
 
-    let query = GetPersonByEmail {
-        email: "[email protected]".to_string(),
-    };
-    let res_one = executor
-        .execute_read(&query)
-        .await?
-        .expect("person not found");
-    tracing::info!("GetPersonByEmail returned: {:?}", res_one);
-
-    let query = GetPersonById { id: res_one.id };
-    let res_two = executor
-        .execute_read(&query)
-        .await?
-        .expect("person not found");
-    tracing::info!("GetPersonById returned: {:?}", res_two);
-    assert_eq!(res_one, res_two);
+    let by_email_res = by_email(&executor, "[email protected]".to_string()).await?;
+    let by_id_res = by_id(&executor, by_email_res.id).await?;
+    assert_eq!(by_email_res, by_id_res);
 
     let ids = [
         "e01e84d6-414c-11ee-be56-0242ac120002",
         "e01e880a-414c-11ee-be56-0242ac120002",
     ]
-    .iter()
-    .map(|s| Uuid::parse_str(s).unwrap())
-    .collect::<Vec<_>>();
-    let query = GetPeopleByIds { ids, rowlimit: 10 };
-    let res = executor.execute_read(&query).await?;
-    tracing::info!("GetPeopleByIds returned: {:?}", res);
+        .iter()
+        .map(|s| Uuid::parse_str(s).unwrap())
+        .collect::<Vec<_>>();
+    by_ids(&executor, ids).await?;
 
     let upsert_id = v1_uuid();
     let query = UpsertPerson {
@@ -68,5 +53,54 @@ async fn main() -> anyhow::Result<()> {
     let res = executor.execute_write(&delete).await?;
     tracing::info!("DeletePersonById returned: {:?}", res);
 
+    let upsert_ttl_id = v1_uuid();
+    let query = UpsertPersonWithTTL {
+        id: upsert_ttl_id,
+        email: "[email protected]".to_string().into(),
+        age: MaybeUnset::Set(Some(42)),
+        data: MaybeUnset::Set(Some(PersonData {
+            stripe_id: Some("stripe_id".to_string()),
+        })),
+        kind: MaybeUnset::Set(PersonKind::Parent),
+        created_at: MaybeUnset::Unset,
+
+        // 5 minutes
+        set_ttl: 300,
+    };
+    let res = executor.execute_write(&query).await?;
+    tracing::info!("UpsertPersonWithTTL returned: {:?}", res);
+
     Ok(())
 }
+
+async fn by_email(executor: &Executor<PersonQueries>, email: String) -> anyhow::Result<PersonEntity> {
+    let res = executor
+        .execute_read(&GetPersonByEmail {
+            email
+        })
+        .await?
+        .expect("person not found");
+
+    tracing::info!("GetPersonByEmail returned: {:?}", res);
+
+    Ok(res)
+}
+
+async fn by_id(executor: &Executor<PersonQueries>, id: Uuid) -> anyhow::Result<PersonEntity> {
+    let res = executor
+        .execute_read(&GetPersonById { id })
+        .await?
+        .expect("person not found");
+
+    tracing::info!("GetPersonById returned: {:?}", res);
+
+    Ok(res)
+}
+
+async fn by_ids(executor: &Executor<PersonQueries>, ids: Vec<Uuid>) -> anyhow::Result<Vec<PersonEntity>> {
+    let res = executor.execute_read(&GetPeopleByIds { ids, rowlimit: 10 }).await?;
+
+    tracing::info!("GetPeopleByIds returned: {:?}", res);
+
+    Ok(res)
+}

scyllax-macros-core/src/queries/upsert.rs

@@ -9,6 +9,7 @@ use crate::entity::{EntityDerive, EntityDeriveColumn};
 pub(crate) struct UpsertQueryOptions {
     pub name: syn::Ident,
     pub table: String,
+    pub ttl: Option<bool>,
 }
 
 /// Attribute expand
@@ -89,6 +90,16 @@ pub(crate) fn upsert_impl(
         })
         .collect::<Vec<_>>();
 
+
+    let ttl = if opt.ttl.unwrap_or(false) {
+        quote! {
+            #[doc = "The ttl of the row in seconds"]
+            pub set_ttl: i32,
+        }
+    } else {
+        quote! {}
+    };
+
     let docs = format!(
         "Upserts a {} into the `{}` table",
         struct_ident, upsert_table
@@ -99,6 +110,7 @@ pub(crate) fn upsert_impl(
         pub struct #upsert_struct {
             #(#expanded_pks,)*
             #(#maybe_unset_fields,)*
+            #ttl
         }
     };
 
@@ -149,7 +161,14 @@ pub(crate) fn upsert_impl(
 
     // if there are no set clauses, then we need to do an insert
     // because we can't do an update with no set clauses
-    let query = build_query(upsert_table, set_clauses, where_clauses);
+    let query = build_query(opt, upsert_table, set_clauses, where_clauses);
+    let ttl_sv_push = if opt.ttl.unwrap_or(false) {
+        quote! {
+            values.add_named_value("set_ttl", &self.set_ttl)?;
+        }
+    } else {
+        quote! {}
+    };
 
     quote! {
         #input
@@ -167,6 +186,7 @@ pub(crate) fn upsert_impl(
 
                 #(#set_sv_push)*
                 #(#where_sv_push)*
+                #ttl_sv_push
 
                 Ok(values)
             }
@@ -177,12 +197,18 @@ pub(crate) fn upsert_impl(
 }
 
 fn build_query(
+    args: &UpsertQueryOptions,
     table: &String,
     set_clauses: Vec<String>,
     where_clauses: Vec<(String, String)>,
 ) -> String {
+    let ttl = match args.ttl.unwrap_or(false) {
+        true => " using ttl :set_ttl",
+        _ => "",
+    };
+
     if set_clauses.is_empty() {
-        let mut query = format!("insert into {table}");
+        let mut query = format!("insert into {table}{ttl}");
         let (cols, named_var) = where_clauses.into_iter().unzip::<_, _, Vec<_>, Vec<_>>();
         let cols = cols.join(", ");
         let named_var = named_var
@@ -195,7 +221,7 @@ fn build_query(
 
         query
     } else {
-        let mut query = format!("update {table} set ");
+        let mut query = format!("update {table}{ttl} set ");
         let query_set = set_clauses.join(", ");
         query.push_str(&query_set);
 
@@ -215,7 +241,9 @@ fn build_query(
 
 #[cfg(test)]
 mod tests {
-    use super::build_query;
+    use syn::{parse::Parser, parse_str};
+
+    use super::*;
 
     fn get_set_clauses() -> Vec<String> {
         vec![
@@ -238,6 +266,11 @@ mod tests {
     #[test]
     fn test_update() {
         let query = build_query(
+            &UpsertQueryOptions {
+                name: syn::parse_str::<syn::Ident>("UpdatePerson").unwrap(),
+                table: "person".to_string(),
+                ttl: None,
+            },
             &"person".to_string(),
             get_set_clauses(),
             get_where_clauses(),
@@ -249,13 +282,50 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_update_ttl() {
+        let query = build_query(
+            &UpsertQueryOptions {
+                name: syn::parse_str::<syn::Ident>("UpdatePerson").unwrap(),
+                table: "person".to_string(),
+                ttl: Some(true),
+            },
+            &"person".to_string(),
+            get_set_clauses(),
+            get_where_clauses(),
+        );
+
+        assert_eq!(
+            query,
+            "update person using ttl :set_ttl set name = :name, email = :email, \"createdAt\" = :created_at where id = :id and \"orgId\" = :org_id;",
+        );
+    }
+
     #[test]
     fn test_insert() {
-        let query = build_query(&"person".to_string(), vec![], get_where_clauses());
+        let query = build_query(&UpsertQueryOptions {
+            name: syn::parse_str::<syn::Ident>("UpdatePerson").unwrap(),
+            table: "person".to_string(),
+            ttl: Default::default(),
+        }, &"person".to_string(), vec![], get_where_clauses());
 
         assert_eq!(
             query,
             "insert into person (id, \"orgId\") values (:id, :org_id);",
         );
     }
+
+    #[test]
+    fn test_insert_ttl() {
+        let query = build_query(&UpsertQueryOptions {
+            name: syn::parse_str::<syn::Ident>("UpdatePerson").unwrap(),
+            table: "person".to_string(),
+            ttl: Some(true),
+        }, &"person".to_string(), vec![], get_where_clauses());
+
+        assert_eq!(
+            query,
+            "insert into person using ttl :set_ttl (id, \"orgId\") values (:id, :org_id);",
+        );
+    }
 }