diff --git a/server/src/test/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinatorTest.java b/server/src/test/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinatorTest.java index da7bf422692..53e9aa28953 100644 --- a/server/src/test/java/org/apache/druid/metadata/IndexerSQLMetadataStorageCoordinatorTest.java +++ b/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().withHandle( - handle -> { - String request = StringUtils.format( - "UPDATE %s SET used = false, used_status_last_updated = :used_status_last_updated WHERE id = :id", - derbyConnectorRule.metadataTablesConfigSupplier().get().getSegmentsTable() - ); - return handle.createStatement(request) - .bind("id", segment.getId().toString()) - .bind("used_status_last_updated", usedStatusLastUpdatedTime.toString() - ).execute(); - } + (int) derbyConnectorRule.updateSegmentsTable( + "UPDATE %s SET used = false, used_status_last_updated = ? WHERE id = ?", + usedStatusLastUpdatedTime.toString(), + segment.getId().toString() ) ); } diff --git a/server/src/test/java/org/apache/druid/metadata/SQLMetadataConnectorTest.java b/server/src/test/java/org/apache/druid/metadata/SQLMetadataConnectorTest.java index 1c8f6493e22..fa8d75e4a49 100644 --- a/server/src/test/java/org/apache/druid/metadata/SQLMetadataConnectorTest.java +++ b/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(); - - // Drop column used_status_last_updated to bring us in line with pre-upgrade state - derbyConnectorRule.getConnector().retryWithHandle( - new HandleCallback() - { - @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; - } - } - ); + derbyConnectorRule.updateSegmentsTable("ALTER TABLE %1$s DROP COLUMN USED_STATUS_LAST_UPDATED"); connector.alterSegmentTableAddUsedFlagLastUpdated(); connector.tableHasColumn( diff --git a/server/src/test/java/org/apache/druid/metadata/SqlSegmentsMetadataManagerTest.java b/server/src/test/java/org/apache/druid/metadata/SqlSegmentsMetadataManagerTest.java index 85e5021f6c6..7871f2f18b8 100644 --- a/server/src/test/java/org/apache/druid/metadata/SqlSegmentsMetadataManagerTest.java +++ b/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() diff --git a/server/src/test/java/org/apache/druid/metadata/TestDerbyConnector.java b/server/src/test/java/org/apache/druid/metadata/TestDerbyConnector.java index e5460ce402b..f8983c8a000 100644 --- a/server/src/test/java/org/apache/druid/metadata/TestDerbyConnector.java +++ b/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); + } } } diff --git a/server/src/test/java/org/apache/druid/server/coordinator/duty/KillUnusedSegmentsTest.java b/server/src/test/java/org/apache/druid/server/coordinator/duty/KillUnusedSegmentsTest.java index 649900c841f..6144dbc360a 100644 --- a/server/src/test/java/org/apache/druid/server/coordinator/duty/KillUnusedSegmentsTest.java +++ b/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(