diff --git a/example/src/entities/person/model.rs b/example/src/entities/person/model.rs index 7908686..76ef453 100644 --- a/example/src/entities/person/model.rs +++ b/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("foo21@scyllax.local".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 = ::query(); + let values = ::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); + } } diff --git a/example/src/entities/person/queries.rs b/example/src/entities/person/queries.rs index aa1c363..60bc478 100644 --- a/example/src/entities/person/queries.rs +++ b/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, ] ); diff --git a/example/src/main.rs b/example/src/main.rs index b201de6..aa67111 100644 --- a/example/src/main.rs +++ b/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::::new(session).await?; - let query = GetPersonByEmail { - email: "foo1@scyllax.local".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, "foo1@scyllax.local".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::>(); - 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::>(); + 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: "foo42@scyllax.local".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, email: String) -> anyhow::Result { + 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, id: Uuid) -> anyhow::Result { + 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, ids: Vec) -> anyhow::Result> { + let res = executor.execute_read(&GetPeopleByIds { ids, rowlimit: 10 }).await?; + + tracing::info!("GetPeopleByIds returned: {:?}", res); + + Ok(res) +} diff --git a/scyllax-macros-core/src/queries/upsert.rs b/scyllax-macros-core/src/queries/upsert.rs index e5426b1..8352958 100644 --- a/scyllax-macros-core/src/queries/upsert.rs +++ b/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, } /// Attribute expand @@ -89,6 +90,16 @@ pub(crate) fn upsert_impl( }) .collect::>(); + + 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, 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 { vec![ @@ -238,6 +266,11 @@ mod tests { #[test] fn test_update() { let query = build_query( + &UpsertQueryOptions { + name: syn::parse_str::("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::("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::("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::("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);", + ); + } }