use proc_macro2::TokenStream;
use quote::quote;
use super::{context::Context, helpers::insert_bindings};
use crate::{entity::parse::ReturningMode, utils::tracing::instrument};
pub(super) fn tx_wrapping(streams: bool) -> (TokenStream, TokenStream, TokenStream) {
if streams {
(
quote! { let mut tx = self.begin().await?; },
quote! { tx.commit().await?; },
quote! { &mut *tx }
)
} else {
(TokenStream::new(), TokenStream::new(), quote! { self })
}
}
impl Context<'_> {
pub fn list_after_method(&self) -> TokenStream {
let Self {
entity_name,
row_name,
table,
columns_str,
id_name,
id_type,
soft_delete,
..
} = self;
let deleted_and = if *soft_delete {
" AND deleted_at IS NULL"
} else {
""
};
let deleted_where = if *soft_delete {
" WHERE deleted_at IS NULL"
} else {
""
};
let cursor_sql = format!(
"SELECT {columns_str} FROM {table} WHERE {id_name} < $1{deleted_and} \
ORDER BY {id_name} DESC LIMIT $2"
);
let head_sql = format!(
"SELECT {columns_str} FROM {table}{deleted_where} ORDER BY {id_name} DESC LIMIT $1"
);
let span = instrument(&entity_name.to_string(), "list_after");
quote! {
#span
async fn list_after(
&self,
cursor: Option<#id_type>,
limit: i64,
) -> Result<Vec<#entity_name>, Self::Error> {
let rows: Vec<#row_name> = match cursor {
Some(after) => {
sqlx::query_as(#cursor_sql)
.bind(&after)
.bind(limit)
.fetch_all(self).await?
}
None => {
sqlx::query_as(#head_sql)
.bind(limit)
.fetch_all(self).await?
}
};
Ok(rows.into_iter().map(#entity_name::from).collect())
}
}
}
pub fn create_method(&self) -> TokenStream {
if self.entity.create_fields().is_empty() {
return TokenStream::new();
}
let Self {
entity_name,
row_name,
insertable_name,
create_dto,
table,
columns_str,
placeholders_str,
entity,
returning,
streams,
..
} = self;
let bindings = insert_bindings(entity.all_fields());
let constraint_map_err = self.constraint_map_err();
let span = instrument(&entity_name.to_string(), "create");
let (tx_open, tx_close, executor) = tx_wrapping(*streams || self.outbox);
let outbox_created = self.outbox_created();
let notify = self.notify_created();
match returning {
ReturningMode::Full => {
quote! {
#span
async fn create(&self, dto: #create_dto) -> Result<#entity_name, Self::Error> {
#tx_open
let entity = #entity_name::from(dto);
let insertable = #insertable_name::from(&entity);
let row: #row_name = sqlx::query_as(
concat!("INSERT INTO ", #table, " (", #columns_str, ") VALUES (", #placeholders_str, ") RETURNING *")
)
#(#bindings)*
.fetch_one(#executor).await #constraint_map_err?;
let entity = #entity_name::from(row);
#outbox_created
#notify
#tx_close
Ok(entity)
}
}
}
ReturningMode::Id => {
let id_name = self.id_name;
quote! {
#span
async fn create(&self, dto: #create_dto) -> Result<#entity_name, Self::Error> {
#tx_open
let entity = #entity_name::from(dto);
let insertable = #insertable_name::from(&entity);
sqlx::query(concat!("INSERT INTO ", #table, " (", #columns_str, ") VALUES (", #placeholders_str, ") RETURNING ", stringify!(#id_name)))
#(#bindings)*
.execute(#executor).await #constraint_map_err?;
#outbox_created
#notify
#tx_close
Ok(entity)
}
}
}
ReturningMode::None => {
quote! {
#span
async fn create(&self, dto: #create_dto) -> Result<#entity_name, Self::Error> {
#tx_open
let entity = #entity_name::from(dto);
let insertable = #insertable_name::from(&entity);
sqlx::query(concat!("INSERT INTO ", #table, " (", #columns_str, ") VALUES (", #placeholders_str, ")"))
#(#bindings)*
.execute(#executor).await #constraint_map_err?;
#outbox_created
#notify
#tx_close
Ok(entity)
}
}
}
ReturningMode::Custom(columns) => {
let returning_cols = columns.join(", ");
quote! {
#span
async fn create(&self, dto: #create_dto) -> Result<#entity_name, Self::Error> {
#tx_open
let entity = #entity_name::from(dto);
let insertable = #insertable_name::from(&entity);
sqlx::query(::sqlx::AssertSqlSafe(format!("INSERT INTO {} ({}) VALUES ({}) RETURNING {}", #table, #columns_str, #placeholders_str, #returning_cols)))
#(#bindings)*
.execute(#executor).await #constraint_map_err?;
#outbox_created
#notify
#tx_close
Ok(entity)
}
}
}
}
}
pub fn find_by_id_method(&self) -> TokenStream {
let Self {
entity_name,
row_name,
table,
columns_str,
id_name,
id_type,
dialect,
soft_delete,
..
} = self;
let placeholder = dialect.placeholder(1);
let deleted_filter = if *soft_delete {
" AND deleted_at IS NULL"
} else {
""
};
let span = instrument(&entity_name.to_string(), "find_by_id");
quote! {
#span
async fn find_by_id(&self, id: #id_type) -> Result<Option<#entity_name>, Self::Error> {
let row: Option<#row_name> = sqlx::query_as(
::sqlx::AssertSqlSafe(format!("SELECT {} FROM {} WHERE {} = {}{}", #columns_str, #table, stringify!(#id_name), #placeholder, #deleted_filter))
).bind(&id).fetch_optional(self).await?;
Ok(row.map(#entity_name::from))
}
}
}
pub fn update_method(&self) -> TokenStream {
let update_fields = self.entity.update_fields();
if update_fields.is_empty() {
return TokenStream::new();
}
let Self {
entity_name,
row_name,
update_dto,
table,
id_name,
id_type,
dialect,
trait_name,
returning,
streams,
..
} = self;
let set_stmts = super::helpers::dynamic_set_stmts(&update_fields);
let set_binds = super::helpers::dynamic_set_binds(&update_fields);
let (version_stmts, version_where, version_bind) =
super::helpers::version_guard(self.entity, "e! { __idx + 1 });
let _ = dialect;
let span = instrument(&entity_name.to_string(), "update");
if *streams || self.outbox {
let fetch_old = self.fetch_old_for_update();
let outbox_updated = self.outbox_updated();
let notify = self.notify_updated();
return quote! {
#span
async fn update(&self, id: #id_type, dto: #update_dto) -> Result<#entity_name, Self::Error> {
#set_stmts
if __sets.is_empty() {
return <Self as #trait_name>::find_by_id(self, id).await?.ok_or_else(|| sqlx::Error::RowNotFound.into());
}
let mut tx = self.begin().await?;
#fetch_old
#version_stmts
let mut q = sqlx::query_as::<_, #row_name>(
::sqlx::AssertSqlSafe(format!("UPDATE {} SET {} WHERE {} = ${}{} RETURNING *", #table, __sets.join(", "), stringify!(#id_name), __idx, #version_where))
);
#set_binds
q = q.bind(&id);
#version_bind
let row: #row_name = q.fetch_optional(&mut *tx).await?
.ok_or_else(|| sqlx::Error::Protocol("row not found or version conflict".into()))?;
let entity = #entity_name::from(row);
#outbox_updated
#notify
tx.commit().await?;
Ok(entity)
}
};
}
match returning {
ReturningMode::Full => {
quote! {
#span
async fn update(&self, id: #id_type, dto: #update_dto) -> Result<#entity_name, Self::Error> {
#set_stmts
if __sets.is_empty() {
return <Self as #trait_name>::find_by_id(self, id).await?.ok_or_else(|| sqlx::Error::RowNotFound.into());
}
#version_stmts
let mut q = sqlx::query_as::<_, #row_name>(
::sqlx::AssertSqlSafe(format!("UPDATE {} SET {} WHERE {} = ${}{} RETURNING *", #table, __sets.join(", "), stringify!(#id_name), __idx, #version_where))
);
#set_binds
q = q.bind(&id);
#version_bind
let row: #row_name = q.fetch_optional(self).await?
.ok_or_else(|| sqlx::Error::Protocol("row not found or version conflict".into()))?;
Ok(#entity_name::from(row))
}
}
}
ReturningMode::Id | ReturningMode::None => {
quote! {
#span
async fn update(&self, id: #id_type, dto: #update_dto) -> Result<#entity_name, Self::Error> {
#set_stmts
if !__sets.is_empty() {
#version_stmts
let mut q = sqlx::query(::sqlx::AssertSqlSafe(format!("UPDATE {} SET {} WHERE {} = ${}{}", #table, __sets.join(", "), stringify!(#id_name), __idx, #version_where)));
#set_binds
q = q.bind(&id);
#version_bind
let __result = q.execute(self).await?;
if __result.rows_affected() == 0 {
return Err(sqlx::Error::Protocol("row not found or version conflict".into()).into());
}
}
<Self as #trait_name>::find_by_id(self, id).await?.ok_or_else(|| sqlx::Error::RowNotFound)
}
}
}
ReturningMode::Custom(columns) => {
let returning_cols = columns.join(", ");
quote! {
#span
async fn update(&self, id: #id_type, dto: #update_dto) -> Result<#entity_name, Self::Error> {
#set_stmts
if !__sets.is_empty() {
#version_stmts
let mut q = sqlx::query(::sqlx::AssertSqlSafe(format!("UPDATE {} SET {} WHERE {} = ${}{} RETURNING {}", #table, __sets.join(", "), stringify!(#id_name), __idx, #version_where, #returning_cols)));
#set_binds
q = q.bind(&id);
#version_bind
let __result = q.execute(self).await?;
if __result.rows_affected() == 0 {
return Err(sqlx::Error::Protocol("row not found or version conflict".into()).into());
}
}
<Self as #trait_name>::find_by_id(self, id).await?.ok_or_else(|| sqlx::Error::RowNotFound)
}
}
}
}
}
pub fn delete_method(&self) -> TokenStream {
let Self {
entity_name,
table,
id_name,
id_type,
dialect,
soft_delete,
streams,
..
} = self;
let placeholder = dialect.placeholder(1);
let (tx_open, tx_close, executor) = tx_wrapping(*streams || self.outbox);
let outbox_deleted = self.outbox_deleted();
let constraint_map_err = self.constraint_map_err();
if *soft_delete {
let notify = self.notify_soft_deleted();
let span = instrument(&entity_name.to_string(), "soft_delete");
quote! {
#span
async fn delete(&self, id: #id_type) -> Result<bool, Self::Error> {
#tx_open
let result = sqlx::query(::sqlx::AssertSqlSafe(format!(
"UPDATE {} SET deleted_at = NOW() WHERE {} = {} AND deleted_at IS NULL",
#table, stringify!(#id_name), #placeholder
))).bind(&id).execute(#executor).await #constraint_map_err?;
let deleted = result.rows_affected() > 0;
if deleted {
#outbox_deleted
#notify
}
#tx_close
Ok(deleted)
}
}
} else {
let notify = self.notify_hard_deleted();
let span = instrument(&entity_name.to_string(), "delete");
quote! {
#span
async fn delete(&self, id: #id_type) -> Result<bool, Self::Error> {
#tx_open
let result = sqlx::query(::sqlx::AssertSqlSafe(format!("DELETE FROM {} WHERE {} = {}", #table, stringify!(#id_name), #placeholder)))
.bind(&id).execute(#executor).await #constraint_map_err?;
let deleted = result.rows_affected() > 0;
if deleted {
#outbox_deleted
#notify
}
#tx_close
Ok(deleted)
}
}
}
}
pub fn list_method(&self) -> TokenStream {
let Self {
entity_name,
row_name,
table,
columns_str,
id_name,
dialect,
soft_delete,
..
} = self;
let limit_placeholder = dialect.placeholder(1);
let offset_placeholder = dialect.placeholder(2);
let where_clause = if *soft_delete {
"WHERE deleted_at IS NULL "
} else {
""
};
let span = instrument(&entity_name.to_string(), "list");
quote! {
#span
async fn list(&self, limit: i64, offset: i64) -> Result<Vec<#entity_name>, Self::Error> {
let rows: Vec<#row_name> = sqlx::query_as(
::sqlx::AssertSqlSafe(format!("SELECT {} FROM {} {}ORDER BY {} DESC LIMIT {} OFFSET {}",
#columns_str, #table, #where_clause, stringify!(#id_name), #limit_placeholder, #offset_placeholder))
).bind(limit).bind(offset).fetch_all(self).await?;
Ok(rows.into_iter().map(#entity_name::from).collect())
}
}
}
}
#[cfg(test)]
mod list_after_tests {
use quote::quote;
use syn::DeriveInput;
use super::super::context::Context;
use crate::entity::parse::EntityDef;
fn parse_entity(tokens: proc_macro2::TokenStream) -> EntityDef {
let input: DeriveInput = syn::parse2(tokens).expect("test entity must parse");
EntityDef::from_derive_input(&input).expect("test entity must be valid")
}
#[test]
fn list_after_generates_keyset_sql() {
let entity = parse_entity(quote! {
#[entity(table = "posts")]
pub struct Post {
#[id]
pub id: uuid::Uuid,
#[field(create, response)]
pub title: String,
}
});
let code = Context::new(&entity).list_after_method().to_string();
assert!(code.contains("id < $1"));
assert!(code.contains("ORDER BY id DESC LIMIT $2"));
assert!(code.contains("ORDER BY id DESC LIMIT $1"));
}
#[test]
fn list_after_respects_soft_delete() {
let entity = parse_entity(quote! {
#[entity(table = "posts", soft_delete)]
pub struct Post {
#[id]
pub id: uuid::Uuid,
#[field(create, response)]
pub title: String,
pub deleted_at: Option<chrono::DateTime<chrono::Utc>>,
}
});
let code = Context::new(&entity).list_after_method().to_string();
assert!(code.contains("AND deleted_at IS NULL"));
assert!(code.contains("WHERE deleted_at IS NULL"));
}
}
#[cfg(test)]
mod version_tests {
use quote::quote;
use syn::DeriveInput;
use super::super::context::Context;
use crate::entity::parse::EntityDef;
fn versioned_entity() -> EntityDef {
let input: DeriveInput = syn::parse_quote! {
#[entity(table = "orders")]
pub struct Order {
#[id]
pub id: uuid::Uuid,
#[field(create, update, response)]
pub note: String,
#[version]
#[field(response)]
#[auto]
pub version: i32,
}
};
EntityDef::from_derive_input(&input).unwrap()
}
#[test]
fn update_bumps_and_guards_version() {
let code = Context::new(&versioned_entity())
.update_method()
.to_string();
assert!(code.contains("version = version + 1"));
assert!(code.contains("AND version = ${}"));
assert!(code.contains("expected_version"));
assert!(code.contains("version conflict"));
let _ = quote!();
}
#[test]
fn update_without_version_has_no_guard() {
let input: DeriveInput = syn::parse_quote! {
#[entity(table = "orders")]
pub struct Order {
#[id]
pub id: uuid::Uuid,
#[field(create, update, response)]
pub note: String,
}
};
let entity = EntityDef::from_derive_input(&input).unwrap();
let code = Context::new(&entity).update_method().to_string();
assert!(!code.contains("expected_version"));
}
}