From 2cb4ddbd90a10e43c1f28707a5feb69e2223a7a4 Mon Sep 17 00:00:00 2001
From: Sage Griffin <sage@sagetheprogrammer.com>
Date: Sun, 5 Jul 2026 15:24:40 -0600
Subject: [PATCH] Remove redundant same-shard check for cross-shard updates

We don't need to check this in insert_row, we already checked if we're
on the same shard in UpdateMulti::execute_internal before calling this
---
 .../client/query_engine/multi_step/update.rs  | 74 ++++++++-----------
 1 file changed, 29 insertions(+), 45 deletions(-)

diff --git a/pgdog/src/frontend/client/query_engine/multi_step/update.rs b/pgdog/src/frontend/client/query_engine/multi_step/update.rs
index 9570d2b24..a589a958b 100644
--- a/pgdog/src/frontend/client/query_engine/multi_step/update.rs
+++ b/pgdog/src/frontend/client/query_engine/multi_step/update.rs
@@ -105,60 +105,44 @@ impl<'a> UpdateMulti<'a> {
             &row.data_row,
         )?;
         self.route(&mut request, context)?;
+        debug!("[update] executing multi-shard insert/delete");
 
-        let original_shard = context.client_request.route().shard();
-        let new_shard = request.route().shard();
+        // Check if we are allowed to do this operation by the config.
+        if self.engine.backend.cluster()?.rewrite().shard_key == RewriteMode::Error {
+            self.engine
+                .error_response(context, ErrorResponse::from_err(&UpdateError::Disabled))
+                .await?;
+            return Ok(());
+        }
 
-        // The new row maps to the same shard as the old row.
-        // We don't need to do the multi-step UPDATE anymore.
-        // Forward the original request as-is.
-        if original_shard.is_direct() && new_shard == original_shard {
-            debug!("[update] selected row is on the same shard");
-            self.execute_original(context).await
-        } else {
-            debug!("[update] executing multi-shard insert/delete");
-
-            // Check if we are allowed to do this operation by the config.
-            if self.engine.backend.cluster()?.rewrite().shard_key == RewriteMode::Error {
-                self.engine
-                    .error_response(context, ErrorResponse::from_err(&UpdateError::Disabled))
-                    .await?;
-                return Ok(());
-            }
+        if !context.in_transaction() || !self.engine.backend.is_multishard()
+        // Do this check at the last possible moment.
+        // Just in case we change how transactions are
+        // routed in the future.
+        {
+            self.engine.cleanup_backend(context)?;
+            return Err(UpdateError::TransactionRequired.into());
+        }
 
-            if !context.in_transaction() || !self.engine.backend.is_multishard()
-            // Do this check at the last possible moment.
-            // Just in case we change how transactions are
-            // routed in the future.
-            {
-                self.engine.cleanup_backend(context)?;
-                return Err(UpdateError::TransactionRequired.into());
-            }
+        if self.has_destructive_on_delete_reference(context)? {
+            return Err(UpdateError::ForeignKeyOnDelete.into());
+        }
 
-            if self.has_destructive_on_delete_reference(context)? {
-                return Err(UpdateError::ForeignKeyOnDelete.into());
-            }
+        self.delete_row(context).await?;
+        self.execute_request_internal(context, &mut request, self.rewrite.insert.is_returning())
+            .await?;
 
-            self.delete_row(context).await?;
-            self.execute_request_internal(
+        self.engine
+            .process_server_message(context, CommandComplete::new("UPDATE 1").message()?) // We only allow to update one row at a time.
+            .await?;
+        self.engine
+            .process_server_message(
                 context,
-                &mut request,
-                self.rewrite.insert.is_returning(),
+                ReadyForQuery::in_transaction(context.in_transaction()).message()?,
             )
             .await?;
 
-            self.engine
-                .process_server_message(context, CommandComplete::new("UPDATE 1").message()?) // We only allow to update one row at a time.
-                .await?;
-            self.engine
-                .process_server_message(
-                    context,
-                    ReadyForQuery::in_transaction(context.in_transaction()).message()?,
-                )
-                .await?;
-
-            Ok(())
-        }
+        Ok(())
     }
 
     fn has_destructive_on_delete_reference(