Add update() in TestDerbyConnectorRule

This commit is contained in:
Abhishek Balaji Radhakrishnan 2024-03-12 10:04:18 +05:30
parent 0a615f16de
commit c916cd6363
5 changed files with 46 additions and 59 deletions

View File

@ -383,17 +383,10 @@ public class IndexerSQLMetadataStorageCoordinatorTest
for (final DataSegment segment : segments) {
Assert.assertEquals(
1,
(int) derbyConnector.getDBI().<Integer>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()
)
);
}

View File

@ -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<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;
}
}
);
derbyConnectorRule.updateSegmentsTable("ALTER TABLE %1$s DROP COLUMN USED_STATUS_LAST_UPDATED");
connector.alterSegmentTableAddUsedFlagLastUpdated();
connector.tableHasColumn(

View File

@ -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()

View File

@ -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);
}
}
}

View File

@ -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(