Add update() in TestDerbyConnectorRule

server/src/test/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinatorTest.java

@@ -383,17 +383,10 @@ public class IndexerSQLMetadataStorageCoordinatorTest
     for (final DataSegment segment : segments) {
       Assert.assertEquals(
           1,
-          (int) derbyConnector.getDBI().<Integer>withHandle(
+          (int) derbyConnectorRule.updateSegmentsTable(
-              handle -> {
+              "UPDATE %s SET used = false, used_status_last_updated = ? WHERE id = ?",
-                String request = StringUtils.format(
+              usedStatusLastUpdatedTime.toString(),
-                    "UPDATE %s SET used = false, used_status_last_updated = :used_status_last_updated WHERE id = :id",
+              segment.getId().toString()
-                    derbyConnectorRule.metadataTablesConfigSupplier().get().getSegmentsTable()
-                );
-                return handle.createStatement(request)
-                             .bind("id", segment.getId().toString())
-                             .bind("used_status_last_updated", usedStatusLastUpdatedTime.toString()
-                ).execute();
-              }
           )
       );
     }

server/src/test/java/org/apache/druid/metadata/SQLMetadataConnectorTest.java

@@ -30,13 +30,11 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
-import org.skife.jdbi.v2.Batch;
 import org.skife.jdbi.v2.DBI;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.exceptions.CallbackFailedException;
 import org.skife.jdbi.v2.exceptions.UnableToExecuteStatementException;
 import org.skife.jdbi.v2.exceptions.UnableToObtainConnectionException;
-import org.skife.jdbi.v2.tweak.HandleCallback;
 
 import java.sql.SQLException;
 import java.sql.SQLRecoverableException;
@@ -47,7 +45,6 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -175,29 +172,7 @@ public class SQLMetadataConnectorTest
   public void testAlterSegmentTableAddLastUsed()
   {
     connector.createSegmentTable();
-
+    derbyConnectorRule.updateSegmentsTable("ALTER TABLE %1$s DROP COLUMN USED_STATUS_LAST_UPDATED");
-    // Drop column used_status_last_updated to bring us in line with pre-upgrade state
-    derbyConnectorRule.getConnector().retryWithHandle(
-        new HandleCallback<Void>()
-        {
-          @Override
-          public Void withHandle(Handle handle)
-          {
-            final Batch batch = handle.createBatch();
-            batch.add(
-                StringUtils.format(
-                    "ALTER TABLE %1$s DROP COLUMN USED_STATUS_LAST_UPDATED",
-                    derbyConnectorRule.metadataTablesConfigSupplier()
-                                      .get()
-                                      .getSegmentsTable()
-                                      .toUpperCase(Locale.ENGLISH)
-                )
-            );
-            batch.execute();
-            return null;
-          }
-        }
-    );
 
     connector.alterSegmentTableAddUsedFlagLastUpdated();
     connector.tableHasColumn(

server/src/test/java/org/apache/druid/metadata/SqlSegmentsMetadataManagerTest.java

@@ -891,15 +891,6 @@ public class SqlSegmentsMetadataManagerTest
     Assert.assertEquals(0, getCountOfRowsWithLastUsedNull());
   }
 
-  private void updateSegmentPayload(DataSegment segment, byte[] payload)
-  {
-    executeUpdate(
-        "UPDATE %1$s SET PAYLOAD = ? WHERE ID = ?",
-        payload,
-        segment.getId().toString()
-    );
-  }
-
   private int getCountOfRowsWithLastUsedNull()
   {
     return derbyConnectorRule.getConnector().retryWithHandle(
@@ -912,9 +903,18 @@ public class SqlSegmentsMetadataManagerTest
     );
   }
 
+  private void updateSegmentPayload(DataSegment segment, byte[] payload)
+  {
+    derbyConnectorRule.updateSegmentsTable(
+        "UPDATE %1$s SET PAYLOAD = ? WHERE ID = ?",
+        payload,
+        segment.getId().toString()
+    );
+  }
+
   private void updateUsedStatusLastUpdated(DataSegment segment, DateTime newValue)
   {
-    executeUpdate(
+    derbyConnectorRule.updateSegmentsTable(
         "UPDATE %1$s SET USED_STATUS_LAST_UPDATED = ? WHERE ID = ?",
         newValue.toString(),
         segment.getId().toString()
@@ -923,22 +923,12 @@ public class SqlSegmentsMetadataManagerTest
 
   private void updateUsedStatusLastUpdatedToNull(DataSegment segment)
   {
-    executeUpdate(
+    derbyConnectorRule.updateSegmentsTable(
         "UPDATE %1$s SET USED_STATUS_LAST_UPDATED = NULL WHERE ID = ?",
         segment.getId().toString()
     );
   }
 
-  private void executeUpdate(String sqlFormat, Object... args)
-  {
-    derbyConnectorRule.getConnector().retryWithHandle(
-        handle -> handle.update(
-            StringUtils.format(sqlFormat, getSegmentsTable()),
-            args
-        )
-    );
-  }
-
   /**
    * Alters the column used_status_last_updated to be nullable. This is used to
    * test backward compatibility with versions of Druid without this column
@@ -946,7 +936,7 @@ public class SqlSegmentsMetadataManagerTest
    */
   private void allowUsedFlagLastUpdatedToBeNullable()
   {
-    executeUpdate("ALTER TABLE %1$s ALTER COLUMN USED_STATUS_LAST_UPDATED NULL");
+    derbyConnectorRule.updateSegmentsTable("ALTER TABLE %1$s ALTER COLUMN USED_STATUS_LAST_UPDATED NULL");
   }
 
   private String getSegmentsTable()

server/src/test/java/org/apache/druid/metadata/TestDerbyConnector.java

@@ -29,6 +29,7 @@ import org.skife.jdbi.v2.DBI;
 import org.skife.jdbi.v2.exceptions.UnableToObtainConnectionException;
 
 import java.sql.SQLException;
+import java.util.Locale;
 import java.util.UUID;
 
 public class TestDerbyConnector extends DerbyConnector
@@ -135,5 +136,24 @@ public class TestDerbyConnector extends DerbyConnector
     {
       return dbTables;
     }
+
+
+    public Integer updateSegmentsTable(String sqlFormat, Object... args)
+    {
+      return this.getConnector().retryWithHandle(
+          handle -> handle.update(
+              StringUtils.format(sqlFormat, getSegmentsTable()),
+              args
+          )
+      );
+    }
+
+    private String getSegmentsTable()
+    {
+      return this.metadataTablesConfigSupplier()
+                 .get()
+                 .getSegmentsTable()
+                 .toUpperCase(Locale.ENGLISH);
+    }
   }
 }

server/src/test/java/org/apache/druid/server/coordinator/duty/KillUnusedSegmentsTest.java

@@ -926,6 +926,15 @@ public class KillUnusedSegmentsTest
   }
 
   private void updateUsedStatusLastUpdated(DataSegment segment, DateTime lastUpdatedTime)
+  {
+    derbyConnectorRule.updateSegmentsTable(
+        "UPDATE %1$s SET USED_STATUS_LAST_UPDATED = ? WHERE ID = ?",
+        lastUpdatedTime.toString(),
+        segment.getId().toString()
+    );
+  }
+
+  private void updateUsedStatusLastUpdatedBoilerPlate(DataSegment segment, DateTime lastUpdatedTime)
   {
     derbyConnectorRule.getConnector().retryWithHandle(
         handle -> handle.update(