Skip to content

Fix Cannot mark an unqueryable datasource's segments used / unused #16127

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
abhishekrb19 merged 4 commits into from
Mar 15, 2024
Merged

Fix Cannot mark an unqueryable datasource's segments used / unused #16127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c627cbe
* fix
zachjsh Mar 14, 2024
f451788
* address review comments
zachjsh Mar 15, 2024
67e2c7b
* all remove the short-circuit for markUnused api
zachjsh Mar 15, 2024
4a965c9
* add test
zachjsh Mar 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 72 additions & 76 deletions server/src/main/java/org/apache/druid/server/http/DataSourcesResource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,12 +166,12 @@ public Response getQueryableDataSources(
@Path("/{dataSourceName}")
@Produces(MediaType.APPLICATION_JSON)
@ResourceFilters(DatasourceResourceFilter.class)
public Response getDataSource(
public Response getQueryableDataSource(
@PathParam("dataSourceName") final String dataSourceName,
@QueryParam("full") final String full
)
{
final ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
final ImmutableDruidDataSource dataSource = getQueryableDataSource(dataSourceName);

if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
Expand Down Expand Up @@ -209,31 +209,41 @@ public Response markAsUsedNonOvershadowedSegments(
MarkDataSourceSegmentsPayload payload
)
{
SegmentUpdateOperation operation = () -> {
final Interval interval = payload.getInterval();
if (interval != null) {
return segmentsMetadataManager.markAsUsedNonOvershadowedSegmentsInInterval(dataSourceName, interval);
} else {
final Set<String> segmentIds = payload.getSegmentIds();
if (segmentIds == null || segmentIds.isEmpty()) {
return 0;
}
if (payload == null || !payload.isValid()) {
log.warn("Invalid request payload: [%s]", payload);

Check notice

Code scanning / CodeQL

Use of default toString()

Default toString(): MarkDataSourceSegmentsPayload inherits toString() from Object, and so is not suitable for printing.
return Response
.status(Response.Status.BAD_REQUEST)
.entity("Invalid request payload, either interval or segmentIds array must be specified")
.build();
} else {
SegmentUpdateOperation operation = () -> {

final Interval interval = payload.getInterval();
if (interval != null) {
return segmentsMetadataManager.markAsUsedNonOvershadowedSegmentsInInterval(dataSourceName, interval);
} else {
final Set<String> segmentIds = payload.getSegmentIds();
if (segmentIds == null || segmentIds.isEmpty()) {
return 0;
}

// Validate segmentIds
final List<String> invalidSegmentIds = new ArrayList<>();
for (String segmentId : segmentIds) {
if (SegmentId.iteratePossibleParsingsWithDataSource(dataSourceName, segmentId).isEmpty()) {
invalidSegmentIds.add(segmentId);
// Validate segmentIds
final List<String> invalidSegmentIds = new ArrayList<>();
for (String segmentId : segmentIds) {
if (SegmentId.iteratePossibleParsingsWithDataSource(dataSourceName, segmentId).isEmpty()) {
invalidSegmentIds.add(segmentId);
}
}
if (!invalidSegmentIds.isEmpty()) {
throw InvalidInput.exception("Could not parse invalid segment IDs[%s]", invalidSegmentIds);
}

return segmentsMetadataManager.markAsUsedNonOvershadowedSegments(dataSourceName, segmentIds);
}
if (!invalidSegmentIds.isEmpty()) {
throw InvalidInput.exception("Could not parse invalid segment IDs[%s]", invalidSegmentIds);
}
};

return segmentsMetadataManager.markAsUsedNonOvershadowedSegments(dataSourceName, segmentIds);
}
};
return performSegmentUpdate(dataSourceName, payload, operation);
return performSegmentUpdate(dataSourceName, operation);
}
}

@POST
Expand All @@ -246,61 +256,47 @@ public Response markSegmentsAsUnused(
final MarkDataSourceSegmentsPayload payload,
@Context final HttpServletRequest req
)
{
SegmentUpdateOperation operation = () -> {
final Interval interval = payload.getInterval();
final int numUpdatedSegments;
if (interval != null) {
numUpdatedSegments = segmentsMetadataManager.markAsUnusedSegmentsInInterval(dataSourceName, interval);
} else {
final Set<SegmentId> segmentIds =
payload.getSegmentIds()
.stream()
.map(idStr -> SegmentId.tryParse(dataSourceName, idStr))
.filter(Objects::nonNull)
.collect(Collectors.toSet());

// Filter out segmentIds that do not belong to this datasource
numUpdatedSegments = segmentsMetadataManager.markSegmentsAsUnused(
segmentIds.stream()
.filter(segmentId -> segmentId.getDataSource().equals(dataSourceName))
.collect(Collectors.toSet())
);
}
auditManager.doAudit(
AuditEntry.builder()
.key(dataSourceName)
.type("segment.markUnused")
.payload(payload)
.auditInfo(AuthorizationUtils.buildAuditInfo(req))
.request(AuthorizationUtils.buildRequestInfo("coordinator", req))
.build()
);
return numUpdatedSegments;
};
return performSegmentUpdate(dataSourceName, payload, operation);
}

private Response performSegmentUpdate(
String dataSourceName,
MarkDataSourceSegmentsPayload payload,
SegmentUpdateOperation operation
)
{
if (payload == null || !payload.isValid()) {
log.warn("Invalid request payload: [%s]", payload);
return Response
.status(Response.Status.BAD_REQUEST)
.entity("Invalid request payload, either interval or segmentIds array must be specified")
.build();
} else {
SegmentUpdateOperation operation = () -> {
final Interval interval = payload.getInterval();
final int numUpdatedSegments;
if (interval != null) {
numUpdatedSegments = segmentsMetadataManager.markAsUnusedSegmentsInInterval(dataSourceName, interval);
} else {
final Set<SegmentId> segmentIds =
payload.getSegmentIds()
.stream()
.map(idStr -> SegmentId.tryParse(dataSourceName, idStr))
.filter(Objects::nonNull)
.collect(Collectors.toSet());

// Filter out segmentIds that do not belong to this datasource
numUpdatedSegments = segmentsMetadataManager.markSegmentsAsUnused(
segmentIds.stream()
.filter(segmentId -> segmentId.getDataSource().equals(dataSourceName))
.collect(Collectors.toSet())
);
}
auditManager.doAudit(
AuditEntry.builder()
.key(dataSourceName)
.type("segment.markUnused")
.payload(payload)
.auditInfo(AuthorizationUtils.buildAuditInfo(req))
.request(AuthorizationUtils.buildRequestInfo("coordinator", req))
.build()
);
return numUpdatedSegments;
};
return performSegmentUpdate(dataSourceName, operation);
}

final ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
}

return performSegmentUpdate(dataSourceName, operation);
}

private static Response logAndCreateDataSourceNotFoundResponse(String dataSourceName)
Expand Down Expand Up @@ -434,7 +430,7 @@ public Response getIntervalsWithServedSegmentsOrAllServedSegmentsPerIntervals(
)
{
if (simple == null && full == null) {
final ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
final ImmutableDruidDataSource dataSource = getQueryableDataSource(dataSourceName);
if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
}
Expand All @@ -460,7 +456,7 @@ public Response getServedSegmentsInInterval(
{
final Interval theInterval = Intervals.of(interval.replace('_', '/'));
if (simple == null && full == null) {
final ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
final ImmutableDruidDataSource dataSource = getQueryableDataSource(dataSourceName);
if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
}
Expand Down Expand Up @@ -617,7 +613,7 @@ private Response getServedSegmentsInInterval(
Predicate<Interval> intervalFilter
)
{
final ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
final ImmutableDruidDataSource dataSource = getQueryableDataSource(dataSourceName);

if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
Expand Down Expand Up @@ -667,7 +663,7 @@ public Response getAllServedSegments(
@QueryParam("full") String full
)
{
ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
ImmutableDruidDataSource dataSource = getQueryableDataSource(dataSourceName);
if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
}
Expand All @@ -689,7 +685,7 @@ public Response getServedSegment(
@PathParam("segmentId") String segmentId
)
{
ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
ImmutableDruidDataSource dataSource = getQueryableDataSource(dataSourceName);
if (dataSource == null) {
return logAndCreateDataSourceNotFoundResponse(dataSourceName);
}
Expand Down Expand Up @@ -747,7 +743,7 @@ public Response getTiersWhereSegmentsAreServed(@PathParam("dataSourceName") Stri
}

@Nullable
private ImmutableDruidDataSource getDataSource(final String dataSourceName)
private ImmutableDruidDataSource getQueryableDataSource(final String dataSourceName)
{
List<DruidDataSource> dataSources = serverInventoryView
.getInventory()
Expand Down
Loading