From 173a206829e5fff7580fae068928a79f10726124 Mon Sep 17 00:00:00 2001 From: Vishesh Garg Date: Mon, 22 Apr 2024 15:18:49 +0530 Subject: [PATCH] Fix incorrect check of InvalidFieldException to InvalidFieldFault while generating MSQ Error Report (#16273) InvalidFieldFault is incorrectly checked as InvalidFieldException in mapQueryColumnNameToOutputColumnName. This fixes the bug. --- .../apache/druid/msq/exec/ControllerImpl.java | 13 ++-- .../frame/write/InvalidFieldException.java | 23 ++++++ .../frame/write/RowBasedFrameWriterTest.java | 74 +++++++++++++++++++ 3 files changed, 104 insertions(+), 6 deletions(-) create mode 100644 processing/src/test/java/org/apache/druid/frame/write/RowBasedFrameWriterTest.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java index 0d1eb5004a71..81e6ddd88caa 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java @@ -124,6 +124,7 @@ import org.apache.druid.msq.indexing.error.InsertCannotBeEmptyFault; import org.apache.druid.msq.indexing.error.InsertLockPreemptedFault; import org.apache.druid.msq.indexing.error.InsertTimeOutOfBoundsFault; +import org.apache.druid.msq.indexing.error.InvalidFieldFault; import org.apache.druid.msq.indexing.error.InvalidNullByteFault; import org.apache.druid.msq.indexing.error.MSQErrorReport; import org.apache.druid.msq.indexing.error.MSQException; @@ -3159,17 +3160,17 @@ private MSQErrorReport mapQueryColumnNameToOutputColumnName( .build(), task.getQuerySpec().getColumnMappings() ); - } else if (workerErrorReport.getFault() instanceof InvalidFieldException) { - InvalidFieldException ife = (InvalidFieldException) workerErrorReport.getFault(); + } else if (workerErrorReport.getFault() instanceof InvalidFieldFault) { + InvalidFieldFault iff = (InvalidFieldFault) workerErrorReport.getFault(); return MSQErrorReport.fromException( workerErrorReport.getTaskId(), workerErrorReport.getHost(), workerErrorReport.getStageNumber(), InvalidFieldException.builder() - .source(ife.getSource()) - .rowNumber(ife.getRowNumber()) - .column(ife.getColumn()) - .errorMsg(ife.getErrorMsg()) + .source(iff.getSource()) + .rowNumber(iff.getRowNumber()) + .column(iff.getColumn()) + .errorMsg(iff.getErrorMsg()) .build(), task.getQuerySpec().getColumnMappings() ); diff --git a/processing/src/main/java/org/apache/druid/frame/write/InvalidFieldException.java b/processing/src/main/java/org/apache/druid/frame/write/InvalidFieldException.java index 3a8bd133c796..a39b70517deb 100644 --- a/processing/src/main/java/org/apache/druid/frame/write/InvalidFieldException.java +++ b/processing/src/main/java/org/apache/druid/frame/write/InvalidFieldException.java @@ -23,6 +23,7 @@ import org.apache.druid.java.util.common.StringUtils; import javax.annotation.Nullable; +import java.util.Objects; public class InvalidFieldException extends RuntimeException { @@ -79,6 +80,28 @@ public String getErrorMsg() return errorMsg; } + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + InvalidFieldException that = (InvalidFieldException) o; + return Objects.equals(source, that.source) + && Objects.equals(column, that.column) + && Objects.equals(rowNumber, that.rowNumber) + && Objects.equals(errorMsg, that.errorMsg); + } + + @Override + public int hashCode() + { + return Objects.hash(source, column, rowNumber, errorMsg); + } + public static Builder builder() { return new Builder(); diff --git a/processing/src/test/java/org/apache/druid/frame/write/RowBasedFrameWriterTest.java b/processing/src/test/java/org/apache/druid/frame/write/RowBasedFrameWriterTest.java new file mode 100644 index 000000000000..bb7f9e558788 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/frame/write/RowBasedFrameWriterTest.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.frame.write; + +import com.google.common.collect.ImmutableList; +import org.apache.druid.frame.allocation.AppendableMemory; +import org.apache.druid.frame.allocation.HeapMemoryAllocator; +import org.apache.druid.frame.field.LongFieldWriter; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.easymock.EasyMock; +import org.junit.Assert; +import org.junit.Test; + +import java.util.Collections; + +public class RowBasedFrameWriterTest +{ + @Test + public void testAddSelectionWithException() + { + String colName = "colName"; + String errorMsg = "Frame writer exception"; + + final RowSignature signature = RowSignature.builder().add(colName, ColumnType.LONG).build(); + + LongFieldWriter fieldWriter = EasyMock.mock(LongFieldWriter.class); + EasyMock.expect(fieldWriter.writeTo( + EasyMock.anyObject(), + EasyMock.anyLong(), + EasyMock.anyLong() + )).andThrow(new RuntimeException(errorMsg)); + + EasyMock.replay(fieldWriter); + + RowBasedFrameWriter rowBasedFrameWriter = new RowBasedFrameWriter( + signature, + Collections.emptyList(), + ImmutableList.of(fieldWriter), + null, + null, + AppendableMemory.create(HeapMemoryAllocator.unlimited()), + AppendableMemory.create(HeapMemoryAllocator.unlimited()) + ); + + InvalidFieldException expectedException = new InvalidFieldException.Builder() + .column(colName) + .errorMsg(errorMsg) + .build(); + + InvalidFieldException actualException = Assert.assertThrows( + InvalidFieldException.class, + rowBasedFrameWriter::addSelection + ); + Assert.assertEquals(expectedException, actualException); + } +}