fix: separate macros for tx and raw

Author: ayushjain17

crates/service_utils/src/db.rs

@@ -8,58 +8,51 @@ pub mod utils;
 pub type PgSchemaConnectionPool = Pool<ConnectionManager<PgConnection>>;
 
 /// Helper macro to run a database query with connection management and error handling.
-/// Supports both raw connection mode and transaction mode.
 /// Example usage:
 /// ```rust,ignore
-/// // Using raw connection mode
 /// run_query!(db_pool, conn, {
 ///     // Your query logic here, using `conn` as the database connection
 /// });
-/// // Using transaction mode
-/// run_query!(db_pool, tx conn, {
-///     // Your query logic here, using `conn` as the database connection within a transaction
-/// });
 /// ```
 #[macro_export]
 macro_rules! run_query {
-    // Public API — raw connection mode
-    ($db_pool:expr, $conn:ident, $body:expr) => {
-        run_query!(@execute $db_pool, raw, $conn, $body)
-    };
-
-    // Public API — transaction mode
-    ($db_pool:expr, tx $conn:ident, $body:expr) => {
-        run_query!(@execute $db_pool, tx, $conn, $body)
-    };
-
-    // Shared connection acquisition and execution logic
-    (@execute $db_pool:expr, $mode:ident, $conn:ident, $body:expr) => {{
-        match $db_pool.get() {
-            Ok(mut $conn) => {
-                diesel::Connection::set_prepared_statement_cache_size(
-                    &mut $conn,
-                    diesel::connection::CacheSize::Disabled,
-                );
-
-                run_query!(@dispatch $mode, $conn, $body)
-            }
-            Err(e) => Err(
-                superposition_macros::unexpected_error!(
-                    "Unable to get db connection from pool, error: {}", e
-                ),
-            ),
-        }
+    ($db_pool:expr, $conn:ident, $body:expr) => {{
+        let mut $conn = $db_pool.get().map_err(|e| {
+            superposition_macros::unexpected_error!(
+                "Unable to get db connection from pool, error: {}",
+                e
+            )
+        })?;
+        diesel::Connection::set_prepared_statement_cache_size(
+            &mut $conn,
+            diesel::connection::CacheSize::Disabled,
+        );
+
+        $body
     }};
+}
 
-    // Dispatch logic based on execution mode
-    (@dispatch raw, $conn:ident, $body:expr) => {{
-        let $conn = &mut $conn;
-        $body.map_err(Into::into)
+/// Helper macro to run a database query within a transaction, with connection management and error handling.
+/// Example usage:
+/// ```rust,ignore
+/// run_tx_query!(db_pool, |conn| {
+///     // Your transactional query logic here, using `conn` as the database connection
+/// });
+/// ```
+#[macro_export]
+macro_rules! run_tx_query {
+    ($db_pool:expr, $query_fn:expr) => {{
+        let mut conn = $db_pool.get().map_err(|e| {
+            superposition_macros::unexpected_error!(
+                "Unable to get db connection from pool, error: {}",
+                e
+            )
+        })?;
+        diesel::Connection::set_prepared_statement_cache_size(
+            &mut conn,
+            diesel::connection::CacheSize::Disabled,
+        );
+
+        diesel::Connection::transaction(&mut conn, $query_fn)
     }};
-
-
-    // For transaction mode, we wrap the body in a transaction block
-    (@dispatch tx, $conn:ident, $body:expr) => {
-        diesel::Connection::transaction(&mut $conn, |$conn| $body).map_err(Into::into)
-    };
 }

crates/service_utils/src/helpers.rs

@@ -212,15 +212,17 @@ pub fn get_workspace(
     workspace_schema_name: &SchemaName,
     db_pool: &PgSchemaConnectionPool,
 ) -> result::Result<Workspace> {
-    run_query!(
+    let workspace = run_query!(
         db_pool,
         conn,
         workspaces::dsl::workspaces
             .filter(
                 workspaces::workspace_schema_name.eq(workspace_schema_name.to_string()),
             )
-            .get_result::<Workspace>(conn)
-    )
+            .get_result::<Workspace>(&mut conn)
+    )?;
+
+    Ok(workspace)
 }
 
 fn has_pattern_in_headers(headers: &CustomHeaders) -> (bool, bool) {

crates/superposition/src/organisation/handlers.rs

@@ -59,7 +59,7 @@ pub async fn create_handler(
     let new_org = run_query!(state.db_pool, conn, {
         diesel::insert_into(organisations::table)
             .values(&new_org)
-            .get_result(conn)
+            .get_result(&mut conn)
     })?;
 
     Ok(Json(new_org))
@@ -83,7 +83,7 @@ pub async fn update_handler(
         diesel::update(organisations::table)
             .filter(organisations::id.eq(org_id))
             .set((req, updated_at.eq(now), updated_by.eq(user.get_email())))
-            .get_result(conn)
+            .get_result(&mut conn)
     })?;
 
     Ok(Json(updated_org))
@@ -95,11 +95,13 @@ pub async fn get_handler(
     org_id: Path<String>,
     state: Data<AppState>,
 ) -> superposition::Result<Json<Organisation>> {
-    let org = run_query!(state.db_pool, conn, {
+    let org = run_query!(
+        state.db_pool,
+        conn,
         organisations::table
             .find(org_id.as_str())
-            .first::<Organisation>(conn)
-    })?;
+            .first::<Organisation>(&mut conn)
+    )?;
 
     Ok(Json(org))
 }
@@ -111,11 +113,13 @@ pub async fn list_handler(
     filters: Query<PaginationParams>,
 ) -> superposition::Result<Json<PaginatedResponse<Organisation>>> {
     if let Some(true) = filters.all {
-        let result = run_query!(state.db_pool, conn, {
+        let result = run_query!(
+            state.db_pool,
+            conn,
             organisations::table
                 .order(organisations::created_at.desc())
-                .get_results(conn)
-        })?;
+                .get_results(&mut conn)
+        )?;
 
         return Ok(Json(PaginatedResponse::all(result)));
     }
@@ -124,7 +128,7 @@ pub async fn list_handler(
     let total_items = run_query!(
         state.db_pool,
         conn,
-        organisations::table.count().get_result(conn)
+        organisations::table.count().get_result(&mut conn)
     )?;
 
     // Set up pagination
@@ -141,7 +145,7 @@ pub async fn list_handler(
     }
 
     // Get paginated results
-    let data = run_query!(state.db_pool, conn, builder.load(conn))?;
+    let data = run_query!(state.db_pool, conn, builder.load(&mut conn))?;
 
     let total_pages = (total_items as f64 / limit as f64).ceil() as i64;
 

crates/superposition/src/resolve/handlers.rs

@@ -44,9 +44,11 @@ async fn resolve_with_exp_handler(
     let query_filters = query_filters.into_inner();
     let identifier_query = identifier_query.into_inner();
     // TODO: Granularise the connection usage in this function once all crates are migrated
-    let max_created_at = run_query!(state.db_pool, conn, {
-        get_max_created_at(conn, &workspace_context.schema_name)
-    })
+    let max_created_at = run_query!(
+        state.db_pool,
+        conn,
+        get_max_created_at(&mut conn, &workspace_context.schema_name)
+    )
     .map_err(|e| log::error!("failed to fetch max timestamp from event_log : {e}"))
     .ok();
 
@@ -64,49 +66,56 @@ async fn resolve_with_exp_handler(
     let config_ver = config_version.to_owned();
 
     // TODO: Granularise the connection usage in this function once all crates are migrated
-    let mut config = run_query!(state.db_pool, conn, {
+    let mut config = run_query!(
+        state.db_pool,
+        conn,
         generate_config_from_version(
             &mut config_version,
-            conn,
+            &mut conn,
             &workspace_context.schema_name,
         )
-    })?;
+    )?;
 
     if let (None, Some(identifier)) = (config_ver, identifier_query.identifier) {
         let context_map: &Map<String, Value> = &query_data;
         // TODO: Granularise the connection usage in this function once all crates are migrated
-        let (applicable_variants, _) = run_query!(state.db_pool, conn, {
+        let (applicable_variants, _) = run_query!(
+            state.db_pool,
+            conn,
             get_applicable_variants_helper(
-                conn,
+                &mut conn,
                 context_map.clone(),
                 &config.dimensions,
                 identifier,
                 &workspace_context,
             )
-        })?;
+        )?;
         query_data.insert("variantIds".to_string(), applicable_variants.into());
     }
 
     // TODO: Granularise the connection usage in this function once all crates are migrated
-    let resolved_config = run_query!(state.db_pool, conn, {
+    let resolved_config = run_query!(
+        state.db_pool,
+        conn,
         resolve(
             &mut config,
             query_data,
             merge_strategy,
-            conn,
+            &mut conn,
             &query_filters,
             &workspace_context,
             &state.master_encryption_key,
         )
-    })?;
+    )?;
 
     let mut resp = HttpResponse::Ok();
     add_last_modified_to_header(max_created_at, is_smithy, &mut resp);
     // TODO: Granularise the connection usage in this function once all crates are migrated
-    run_query!(state.db_pool, conn, {
-        add_audit_id_to_header(conn, &mut resp, &workspace_context.schema_name);
-        Ok::<(), superposition::AppError>(())
-    })?;
+    run_query!(
+        state.db_pool,
+        conn,
+        add_audit_id_to_header(&mut conn, &mut resp, &workspace_context.schema_name)
+    );
     add_config_version_to_header(&config_version, &mut resp);
     Ok(resp.json(resolved_config))
 }

crates/superposition/src/webhooks/handlers.rs

@@ -7,12 +7,12 @@ use chrono::Utc;
 use context_aware_config::helpers::validate_change_reason;
 use diesel::{ExpressionMethods, PgArrayExpressionMethods, QueryDsl, RunQueryDsl};
 use service_utils::{
-    run_query,
+    run_query, run_tx_query,
     service::types::{AppState, WorkspaceContext},
 };
 use superposition_derives::authorized;
 use superposition_types::{
-    PaginatedResponse, User,
+    DBConnection, PaginatedResponse, User,
     api::webhook::{CreateWebhookRequest, UpdateWebhookRequest, WebhookName},
     custom_query::PaginationParams,
     database::{
@@ -48,7 +48,7 @@ async fn create_handler(
         validate_change_reason(
             &workspace_context,
             &req.change_reason,
-            conn,
+            &mut conn,
             &state.master_encryption_key,
         )
     )?;
@@ -84,7 +84,7 @@ async fn create_handler(
         diesel::insert_into(webhooks::table)
             .values(&webhook_data)
             .schema_name(&workspace_context.schema_name)
-            .get_result::<Webhook>(conn)
+            .get_result::<Webhook>(&mut conn)
     )?;
 
     Ok(Json(created))
@@ -109,7 +109,7 @@ async fn update_handler(
         validate_change_reason(
             &workspace_context,
             &req.change_reason,
-            conn,
+            &mut conn,
             &state.master_encryption_key,
         )
     )?;
@@ -134,7 +134,7 @@ async fn update_handler(
                 last_modified_by.eq(user.get_email()),
             ))
             .schema_name(&workspace_context.schema_name)
-            .get_result::<Webhook>(conn)
+            .get_result::<Webhook>(&mut conn)
     )?;
 
     Ok(Json(update))
@@ -163,11 +163,13 @@ async fn list_handler(
     pagination: Query<PaginationParams>,
 ) -> superposition::Result<Json<PaginatedResponse<Webhook>>> {
     if let Some(true) = pagination.all {
-        let result: Vec<Webhook> = run_query!(state.db_pool, conn, {
+        let result: Vec<Webhook> = run_query!(
+            state.db_pool,
+            conn,
             webhooks
                 .schema_name(&workspace_context.schema_name)
-                .get_results(conn)
-        })?;
+                .get_results(&mut conn)
+        )?;
         return Ok(Json(PaginatedResponse::all(result)));
     }
 
@@ -177,7 +179,7 @@ async fn list_handler(
         webhooks
             .count()
             .schema_name(&workspace_context.schema_name)
-            .get_result(conn)
+            .get_result(&mut conn)
     )?;
     let limit = pagination.count.unwrap_or(10);
     let mut builder = webhooks
@@ -189,7 +191,7 @@ async fn list_handler(
         let offset = (page - 1) * limit;
         builder = builder.offset(offset);
     }
-    let data: Vec<Webhook> = run_query!(state.db_pool, conn, builder.load(conn))?;
+    let data: Vec<Webhook> = run_query!(state.db_pool, conn, builder.load(&mut conn))?;
     let total_pages = (total_items as f64 / limit as f64).ceil() as i64;
 
     Ok(Json(PaginatedResponse {
@@ -209,7 +211,7 @@ async fn delete_handler(
 ) -> superposition::Result<HttpResponse> {
     let w_name: String = params.into_inner().into();
 
-    run_query!(state.db_pool, tx conn, {
+    run_tx_query!(state.db_pool, |conn: &mut DBConnection| {
         diesel::update(webhooks::table)
             .filter(webhooks::name.eq(&w_name))
             .set((
@@ -233,11 +235,13 @@ async fn get_by_event_handler(
     state: Data<AppState>,
 ) -> superposition::Result<Json<Webhook>> {
     let event = params.into_inner();
-    let webhook_row = run_query!(state.db_pool, conn, {
+    let webhook_row = run_query!(
+        state.db_pool,
+        conn,
         webhooks
             .filter(webhooks::events.contains(vec![event]))
             .schema_name(&workspace_context.schema_name)
-            .first(conn)
-    })?;
+            .first(&mut conn)
+    )?;
     Ok(Json(webhook_row))
 }