add update! macro

Author: trevyn

turbosql-impl/src/lib.rs

@@ -26,9 +26,9 @@ const MIGRATIONS_FILENAME: &str = "migrations.toml";
 #[cfg(feature = "test")]
 const MIGRATIONS_FILENAME: &str = "test.migrations.toml";
 
+mod delete;
 mod insert;
 mod update;
-mod delete;
 
 #[derive(Debug, Clone)]
 struct Table {
@@ -80,6 +80,11 @@ struct ExecuteTokens {
 	tokens: proc_macro2::TokenStream,
 }
 
+#[derive(Debug)]
+struct UpdateTokens {
+	tokens: proc_macro2::TokenStream,
+}
+
 #[derive(Clone, Debug)]
 struct SingleColumn {
 	table: Ident,
@@ -230,8 +235,9 @@ fn _extract_stmt_members(stmt: &Statement, span: &Span) -> MembersAndCasters {
 enum ParseStatementType {
 	Execute,
 	Select,
+	Update,
 }
-use ParseStatementType::{Execute, Select};
+use ParseStatementType::{Execute, Select, Update};
 
 #[derive(Debug)]
 struct StatementInfo {
@@ -414,22 +420,25 @@ fn do_parse_tokens(
 		None
 	};
 
-	let (sql, params, sql_and_parameters_tokens) = parse_interpolated_sql(input)?;
+	let (mut sql, params, sql_and_parameters_tokens) = parse_interpolated_sql(input)?;
 
 	// Try validating SQL as-is
 
-	let stmt_info = sql.as_ref().and_then(|s| validate_sql(s).ok());
+	let mut stmt_info = sql.as_ref().and_then(|s| validate_sql(s).ok());
 
-	// Try adding SELECT if it didn't validate
+	// Try adding SELECT or UPDATE if it didn't validate
 
-	let (sql, stmt_info) = match (sql, stmt_info) {
-		(Some(sql), None) => {
-			let sql_with_select = format!("SELECT {}", sql);
-			let stmt_info = validate_sql(&sql_with_select).ok();
-			(Some(if stmt_info.is_some() { sql_with_select } else { sql }), stmt_info)
+	if let (ty @ (Select | Update), Some(orig_sql), None) = (&statement_type, &sql, &stmt_info) {
+		let sql_modified = match ty {
+			Select => format!("SELECT {}", orig_sql),
+			Update => format!("UPDATE {}", orig_sql),
+			_ => unreachable!(),
+		};
+		if let Ok(stmt_info_modified) = validate_sql(&sql_modified) {
+			sql = Some(sql_modified);
+			stmt_info = Some(stmt_info_modified);
 		}
-		t => t,
-	};
+	}
 
 	if is_rust_analyzer() {
 		return Ok(if let Some(ty) = result_type {
@@ -547,7 +556,7 @@ fn do_parse_tokens(
 	// if we return no columns, this should be an execute
 
 	if stmt_info.column_names.is_empty() {
-		if !matches!(statement_type, Execute) {
+		if matches!(statement_type, Select) {
 			abort_call_site!("No rows returned from SQL, use execute! instead.");
 		}
 
@@ -660,6 +669,12 @@ impl Parse for ExecuteTokens {
 	}
 }
 
+impl Parse for UpdateTokens {
+	fn parse(input: ParseStream) -> syn::Result<Self> {
+		Ok(UpdateTokens { tokens: do_parse_tokens(input, Update)? })
+	}
+}
+
 /// Executes a SQL statement. On success, returns the number of rows that were changed or inserted or deleted.
 #[proc_macro]
 #[proc_macro_error]
@@ -676,6 +691,14 @@ pub fn select(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
 	proc_macro::TokenStream::from(tokens)
 }
 
+/// Executes a SQL statement with optionally automatic `UPDATE` clause. On success, returns the number of rows that were changed.
+#[proc_macro]
+#[proc_macro_error]
+pub fn update(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+	let UpdateTokens { tokens } = parse_macro_input!(input);
+	proc_macro::TokenStream::from(tokens)
+}
+
 /// Derive this on a `struct` to create a corresponding SQLite table and `Turbosql` trait methods.
 #[proc_macro_derive(Turbosql, attributes(turbosql))]
 #[proc_macro_error]

turbosql/src/lib.rs

@@ -8,7 +8,7 @@ compile_error!("turbosql must be tested with '--features test -- --test-threads=
 include!("lib_inner.rs");
 
 #[cfg(target_arch = "wasm32")]
-pub use turbosql_impl::{execute, select, Turbosql};
+pub use turbosql_impl::{execute, select, update, Turbosql};
 
 #[cfg(target_arch = "wasm32")]
 pub fn now_ms() -> i64 {

turbosql/src/lib_inner.rs

@@ -21,7 +21,7 @@ pub use rusqlite::{
 pub use serde::Serialize;
 #[doc(hidden)]
 pub use serde_json;
-pub use turbosql_impl::{execute, select, Turbosql};
+pub use turbosql_impl::{execute, select, update, Turbosql};
 
 /// Wrapper for `Vec<u8>` that may one day impl `Read`, `Write` and `Seek` traits.
 pub type Blob = Vec<u8>;

turbosql/tests/integration_test.rs

@@ -7,7 +7,7 @@ compile_error!("turbosql must be tested with '--features test -- --test-threads=
 #[cfg(not(test))]
 compile_error!("integration_tests.rs must be run in test mode");
 
-use turbosql::{execute, select, Blob, Turbosql};
+use turbosql::{execute, select, update, Blob, Turbosql};
 
 #[derive(Turbosql, Default, Debug, PartialEq, Clone)]
 struct PersonIntegrationTest {
@@ -301,6 +301,8 @@ fn integration_test() {
 
 	execute!("INSERT INTO personintegrationtest(field_u8, field_i8) VALUES (" 1, 2 ")").unwrap();
 
+	update!("personintegrationtest SET field_u8 = " 0).unwrap();
+
 	// DELETE
 
 	assert!(execute!("DELETE FROM personintegrationtest").is_ok());