From 36edbce03667bdcb58911455b74b8e64a169c14c Mon Sep 17 00:00:00 2001 From: Laksh Singla Date: Mon, 9 Oct 2023 20:05:48 +0530 Subject: [PATCH 1/3] Fix compilation failure in master (#15111) Merging since it's a dev blocker. --- .../org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java index 8ee9e78c8388..6ec17687c45e 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java @@ -95,7 +95,8 @@ public SqlEngine createEngine( queryJsonMapper, injector, new MSQTestTaskActionClient(queryJsonMapper), - workerMemoryParameters + workerMemoryParameters, + ImmutableList.of() ); return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper); } From b0edbc3d912628e936ec2af06549e1b5b8f11898 Mon Sep 17 00:00:00 2001 From: Laksh Singla Date: Mon, 9 Oct 2023 20:31:07 +0530 Subject: [PATCH 2/3] MSQ writes out string arrays instead of MVDs by default (#15093) MSQ uses the string dimension schema for ARRAY typed columns, which creates MVDs instead of string arrays as required. Therefore someone trying to ingest columns of type ARRAY from an external data source or another data source would get STRING columns in the newly generated segments. This patch changes the following: - Use auto dimension schema to ingest the ARRAY columns, which will create columns with the desired type. - Add an undocumented flag ingestStringArraysAsMVDs to preserve the legacy behavior. Legacy behaviour is turned on by default. - Create MSQArraysInsertTest and refactor some of the tests in MSQInsertTest. --- .../apache/druid/msq/exec/ControllerImpl.java | 18 +- .../external/ExternalInputSliceReader.java | 5 +- .../druid/msq/util/ArrayIngestMode.java | 45 ++ .../druid/msq/util/DimensionSchemaUtils.java | 93 ++- .../msq/util/MultiStageQueryContext.java | 16 +- .../apache/druid/msq/exec/MSQArraysTest.java | 727 ++++++++++++++++++ .../apache/druid/msq/exec/MSQInsertTest.java | 289 ------- .../msq/util/MultiStageQueryContextTest.java | 159 ++-- 8 files changed, 965 insertions(+), 387 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/ArrayIngestMode.java create mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.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 58768644bf69..6f46007d93c0 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 @@ -167,6 +167,7 @@ import org.apache.druid.msq.shuffle.input.DurableStorageInputChannelFactory; import org.apache.druid.msq.shuffle.input.WorkerInputChannelFactory; import org.apache.druid.msq.statistics.PartialKeyStatisticsInformation; +import org.apache.druid.msq.util.ArrayIngestMode; import org.apache.druid.msq.util.DimensionSchemaUtils; import org.apache.druid.msq.util.IntervalUtils; import org.apache.druid.msq.util.MSQFutureUtils; @@ -1999,6 +2000,17 @@ private static Pair, List> makeDimensio final Query query ) { + // Log a warning unconditionally if arrayIngestMode is MVD, since the behaviour is incorrect, and is subject to + // deprecation and removal in future + if (MultiStageQueryContext.getArrayIngestMode(query.context()) == ArrayIngestMode.MVD) { + log.warn( + "'%s' is set to 'mvd' in the query's context. This ingests the string arrays as multi-value " + + "strings instead of arrays, and is preserved for legacy reasons when MVDs were the only way to ingest string " + + "arrays in Druid. It is incorrect behaviour and will likely be removed in the future releases of Druid", + MultiStageQueryContext.CTX_ARRAY_INGEST_MODE + ); + } + final List dimensions = new ArrayList<>(); final List aggregators = new ArrayList<>(); @@ -2076,7 +2088,8 @@ private static Pair, List> makeDimensio DimensionSchemaUtils.createDimensionSchema( outputColumnName, type, - MultiStageQueryContext.useAutoColumnSchemas(query.context()) + MultiStageQueryContext.useAutoColumnSchemas(query.context()), + MultiStageQueryContext.getArrayIngestMode(query.context()) ) ); } else if (!isRollupQuery) { @@ -2125,7 +2138,8 @@ private static void populateDimensionsAndAggregators( DimensionSchemaUtils.createDimensionSchema( outputColumn, type, - MultiStageQueryContext.useAutoColumnSchemas(context) + MultiStageQueryContext.useAutoColumnSchemas(context), + MultiStageQueryContext.getArrayIngestMode(context) ) ); } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java index 084d58e217d6..714e8dc3a639 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java @@ -119,10 +119,9 @@ private static Iterator inputSourceSegmentIterator( new DimensionsSpec( signature.getColumnNames().stream().map( column -> - DimensionSchemaUtils.createDimensionSchema( + DimensionSchemaUtils.createDimensionSchemaForExtern( column, - signature.getColumnType(column).orElse(null), - false + signature.getColumnType(column).orElse(null) ) ).collect(Collectors.toList()) ), diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/ArrayIngestMode.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/ArrayIngestMode.java new file mode 100644 index 000000000000..ff6b4718ad85 --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/ArrayIngestMode.java @@ -0,0 +1,45 @@ +/* + * 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.msq.util; + +/** + * Values that the query context flag 'arrayIngestMode' can take to specify the behaviour of ingestion of arrays via + * MSQ's INSERT queries + */ +public enum ArrayIngestMode +{ + /** + * Disables the ingestion of arrays via MSQ's INSERT queries. + */ + NONE, + + /** + * String arrays are ingested as MVDs. This is to preserve the legacy behaviour of Druid and will be removed in the + * future, since MVDs are not true array types and the behaviour is incorrect. + * This also disables the ingestion of numeric arrays + */ + MVD, + + /** + * Allows numeric and string arrays to be ingested as arrays. This should be the preferred method of ingestion, + * unless bound by compatibility reasons to use 'mvd' + */ + ARRAY +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java index 2efc94740ac7..98d94518bde8 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java @@ -24,7 +24,9 @@ import org.apache.druid.data.input.impl.FloatDimensionSchema; import org.apache.druid.data.input.impl.LongDimensionSchema; import org.apache.druid.data.input.impl.StringDimensionSchema; +import org.apache.druid.error.InvalidInput; import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.segment.AutoTypeColumnSchema; import org.apache.druid.segment.DimensionHandlerUtils; import org.apache.druid.segment.column.ColumnCapabilities; @@ -40,15 +42,31 @@ */ public class DimensionSchemaUtils { + + /** + * Creates a dimension schema for creating {@link org.apache.druid.data.input.InputSourceReader}. + */ + public static DimensionSchema createDimensionSchemaForExtern(final String column, @Nullable final ColumnType type) + { + return createDimensionSchema( + column, + type, + false, + // Least restrictive mode since we do not have any type restrictions while reading the extern files. + ArrayIngestMode.ARRAY + ); + } + public static DimensionSchema createDimensionSchema( final String column, @Nullable final ColumnType type, - boolean useAutoType + boolean useAutoType, + ArrayIngestMode arrayIngestMode ) { if (useAutoType) { // for complex types that are not COMPLEX, we still want to use the handler since 'auto' typing - // only works for the 'standard' built-in typesg + // only works for the 'standard' built-in types if (type != null && type.is(ValueType.COMPLEX) && !ColumnType.NESTED_DATA.equals(type)) { final ColumnCapabilities capabilities = ColumnCapabilitiesImpl.createDefault().setType(type); return DimensionHandlerUtils.getHandlerFromCapabilities(column, capabilities, null) @@ -57,35 +75,54 @@ public static DimensionSchema createDimensionSchema( return new AutoTypeColumnSchema(column); } else { - // if schema information not available, create a string dimension + // if schema information is not available, create a string dimension if (type == null) { return new StringDimensionSchema(column); - } - - switch (type.getType()) { - case STRING: - return new StringDimensionSchema(column); - case LONG: - return new LongDimensionSchema(column); - case FLOAT: - return new FloatDimensionSchema(column); - case DOUBLE: - return new DoubleDimensionSchema(column); - case ARRAY: - switch (type.getElementType().getType()) { - case STRING: - return new StringDimensionSchema(column, DimensionSchema.MultiValueHandling.ARRAY, null); - case LONG: - case FLOAT: - case DOUBLE: - return new AutoTypeColumnSchema(column); - default: - throw new ISE("Cannot create dimension for type [%s]", type.toString()); + } else if (type.getType() == ValueType.STRING) { + return new StringDimensionSchema(column); + } else if (type.getType() == ValueType.LONG) { + return new LongDimensionSchema(column); + } else if (type.getType() == ValueType.FLOAT) { + return new FloatDimensionSchema(column); + } else if (type.getType() == ValueType.DOUBLE) { + return new DoubleDimensionSchema(column); + } else if (type.getType() == ValueType.ARRAY) { + ValueType elementType = type.getElementType().getType(); + if (elementType == ValueType.STRING) { + if (arrayIngestMode == ArrayIngestMode.NONE) { + throw InvalidInput.exception( + "String arrays can not be ingested when '%s' is set to '%s'. Either set '%s' in query context " + + "to 'array' to ingest the string array as an array, or ingest it as an MVD by explicitly casting the " + + "array to an MVD with ARRAY_TO_MV function.", + MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, + StringUtils.toLowerCase(arrayIngestMode.name()), + MultiStageQueryContext.CTX_ARRAY_INGEST_MODE + ); + } else if (arrayIngestMode == ArrayIngestMode.MVD) { + return new StringDimensionSchema(column, DimensionSchema.MultiValueHandling.ARRAY, null); + } else { + // arrayIngestMode == ArrayIngestMode.ARRAY would be true + return new AutoTypeColumnSchema(column); + } + } else if (elementType.isNumeric()) { + // ValueType == LONG || ValueType == FLOAT || ValueType == DOUBLE + if (arrayIngestMode == ArrayIngestMode.ARRAY) { + return new AutoTypeColumnSchema(column); + } else { + throw InvalidInput.exception( + "Numeric arrays can only be ingested when '%s' is set to 'array' in the MSQ query's context. " + + "Current value of the parameter [%s]", + MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, + StringUtils.toLowerCase(arrayIngestMode.name()) + ); } - default: - final ColumnCapabilities capabilities = ColumnCapabilitiesImpl.createDefault().setType(type); - return DimensionHandlerUtils.getHandlerFromCapabilities(column, capabilities, null) - .getDimensionSchema(capabilities); + } else { + throw new ISE("Cannot create dimension for type [%s]", type.toString()); + } + } else { + final ColumnCapabilities capabilities = ColumnCapabilitiesImpl.createDefault().setType(type); + return DimensionHandlerUtils.getHandlerFromCapabilities(column, capabilities, null) + .getDimensionSchema(capabilities); } } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java index 6e477d0c364b..613fac6203c2 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/MultiStageQueryContext.java @@ -74,6 +74,10 @@ * {@link org.apache.druid.segment.AutoTypeColumnSchema} for all 'standard' type columns during segment generation, * see {@link DimensionSchemaUtils#createDimensionSchema} for more details. * + *
  • arrayIngestMode: Tri-state query context that controls the behaviour and support of arrays that are + * ingested via MSQ. If set to 'none', arrays are not allowed to be ingested in MSQ. If set to 'array', array types + * can be ingested as expected. If set to 'mvd', numeric arrays can not be ingested, and string arrays will be + * ingested as MVDs (this is kept for legacy purpose). * **/ public class MultiStageQueryContext @@ -127,6 +131,11 @@ public class MultiStageQueryContext public static final String CTX_INDEX_SPEC = "indexSpec"; public static final String CTX_USE_AUTO_SCHEMAS = "useAutoColumnSchemas"; + public static final boolean DEFAULT_USE_AUTO_SCHEMAS = false; + + public static final String CTX_ARRAY_INGEST_MODE = "arrayIngestMode"; + public static final ArrayIngestMode DEFAULT_ARRAY_INGEST_MODE = ArrayIngestMode.MVD; + private static final Pattern LOOKS_LIKE_JSON_ARRAY = Pattern.compile("^\\s*\\[.*", Pattern.DOTALL); @@ -266,7 +275,12 @@ public static IndexSpec getIndexSpec(final QueryContext queryContext, final Obje public static boolean useAutoColumnSchemas(final QueryContext queryContext) { - return queryContext.getBoolean(CTX_USE_AUTO_SCHEMAS, false); + return queryContext.getBoolean(CTX_USE_AUTO_SCHEMAS, DEFAULT_USE_AUTO_SCHEMAS); + } + + public static ArrayIngestMode getArrayIngestMode(final QueryContext queryContext) + { + return queryContext.getEnum(CTX_ARRAY_INGEST_MODE, ArrayIngestMode.class, DEFAULT_ARRAY_INGEST_MODE); } /** diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java new file mode 100644 index 000000000000..d2696f232820 --- /dev/null +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java @@ -0,0 +1,727 @@ +/* + * 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.msq.exec; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import org.apache.druid.data.input.impl.JsonInputFormat; +import org.apache.druid.data.input.impl.LocalInputSource; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.msq.indexing.MSQSpec; +import org.apache.druid.msq.indexing.MSQTuningConfig; +import org.apache.druid.msq.indexing.destination.TaskReportMSQDestination; +import org.apache.druid.msq.test.MSQTestBase; +import org.apache.druid.msq.util.MultiStageQueryContext; +import org.apache.druid.query.NestedDataTestUtils; +import org.apache.druid.query.Query; +import org.apache.druid.query.expression.TestExprMacroTable; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.virtual.ExpressionVirtualColumn; +import org.apache.druid.sql.calcite.external.ExternalDataSource; +import org.apache.druid.sql.calcite.filtration.Filtration; +import org.apache.druid.sql.calcite.planner.ColumnMapping; +import org.apache.druid.sql.calcite.planner.ColumnMappings; +import org.apache.druid.timeline.SegmentId; +import org.apache.druid.utils.CompressionUtils; +import org.hamcrest.CoreMatchers; +import org.junit.Test; +import org.junit.internal.matchers.ThrowableMessageMatcher; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Tests INSERT and SELECT behaviour of MSQ with arrays and MVDs + */ +@RunWith(Parameterized.class) +public class MSQArraysTest extends MSQTestBase +{ + + @Parameterized.Parameters(name = "{index}:with context {0}") + public static Collection data() + { + Object[][] data = new Object[][]{ + {DEFAULT, DEFAULT_MSQ_CONTEXT}, + {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT}, + {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT}, + {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT} + }; + return Arrays.asList(data); + } + + @Parameterized.Parameter(0) + public String contextName; + + @Parameterized.Parameter(1) + public Map context; + + /** + * Tests the behaviour of INSERT query when arrayIngestMode is set to none (default) and the user tries to ingest + * string arrays + */ + @Test + public void testInsertStringArrayWithArrayIngestModeNone() + { + + final Map adjustedContext = new HashMap<>(context); + adjustedContext.put(MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, "none"); + + testIngestQuery().setSql( + "INSERT INTO foo1 SELECT MV_TO_ARRAY(dim3) AS dim3 FROM foo GROUP BY 1 PARTITIONED BY ALL TIME") + .setQueryContext(adjustedContext) + .setExpectedExecutionErrorMatcher(CoreMatchers.allOf( + CoreMatchers.instanceOf(ISE.class), + ThrowableMessageMatcher.hasMessage(CoreMatchers.containsString( + "String arrays can not be ingested when 'arrayIngestMode' is set to 'none'")) + )) + .verifyExecutionError(); + } + + + /** + * Tests the behaviour of INSERT query when arrayIngestMode is set to mvd (default) and the only array type to be + * ingested is string array + */ + @Test + public void testInsertOnFoo1WithMultiValueToArrayGroupByWithDefaultContext() + { + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .add("dim3", ColumnType.STRING) + .build(); + + testIngestQuery().setSql( + "INSERT INTO foo1 SELECT MV_TO_ARRAY(dim3) AS dim3 FROM foo GROUP BY 1 PARTITIONED BY ALL TIME") + .setExpectedDataSource("foo1") + .setExpectedRowSignature(rowSignature) + .setQueryContext(context) + .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0))) + .setExpectedResultRows(expectedMultiValueFooRowsToArray()) + .verifyResults(); + } + + /** + * Tests the INSERT query when 'auto' type is set + */ + @Test + public void testInsertArraysAutoType() throws IOException + { + List expectedRows = Arrays.asList( + new Object[]{1672531200000L, null, null, null}, + new Object[]{1672531200000L, null, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, + new Object[]{1672531200000L, new Object[]{"d", "e"}, new Object[]{1L, 4L}, new Object[]{2.2, 3.3, 4.0}}, + new Object[]{1672531200000L, new Object[]{"a", "b"}, null, null}, + new Object[]{1672531200000L, new Object[]{"a", "b"}, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, + new Object[]{1672531200000L, new Object[]{"b", "c"}, new Object[]{1L, 2L, 3L, 4L}, new Object[]{1.1, 3.3}}, + new Object[]{1672531200000L, new Object[]{"a", "b", "c"}, new Object[]{2L, 3L}, new Object[]{3.3, 4.4, 5.5}}, + new Object[]{1672617600000L, null, null, null}, + new Object[]{1672617600000L, null, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, + new Object[]{1672617600000L, new Object[]{"d", "e"}, new Object[]{1L, 4L}, new Object[]{2.2, 3.3, 4.0}}, + new Object[]{1672617600000L, new Object[]{"a", "b"}, null, null}, + new Object[]{1672617600000L, new Object[]{"a", "b"}, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, + new Object[]{1672617600000L, new Object[]{"b", "c"}, new Object[]{1L, 2L, 3L, 4L}, new Object[]{1.1, 3.3}}, + new Object[]{1672617600000L, new Object[]{"a", "b", "c"}, new Object[]{2L, 3L}, new Object[]{3.3, 4.4, 5.5}} + ); + + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .add("arrayString", ColumnType.STRING_ARRAY) + .add("arrayLong", ColumnType.LONG_ARRAY) + .add("arrayDouble", ColumnType.DOUBLE_ARRAY) + .build(); + + final Map adjustedContext = new HashMap<>(context); + adjustedContext.put(MultiStageQueryContext.CTX_USE_AUTO_SCHEMAS, true); + + final File tmpFile = temporaryFolder.newFile(); + final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader() + .getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); + final InputStream decompressing = CompressionUtils.decompress( + resourceStream, + NestedDataTestUtils.ARRAY_TYPES_DATA_FILE + ); + Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING); + decompressing.close(); + + final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile); + + testIngestQuery().setSql(" INSERT INTO foo1 SELECT\n" + + " TIME_PARSE(\"timestamp\") as __time,\n" + + " arrayString,\n" + + " arrayLong,\n" + + " arrayDouble\n" + + "FROM TABLE(\n" + + " EXTERN(\n" + + " '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n" + + " '{\"type\": \"json\"}',\n" + + " '[{\"name\": \"timestamp\", \"type\": \"STRING\"}, {\"name\": \"arrayString\", \"type\": \"COMPLEX\"}, {\"name\": \"arrayLong\", \"type\": \"COMPLEX\"}, {\"name\": \"arrayDouble\", \"type\": \"COMPLEX\"}]'\n" + + " )\n" + + ") PARTITIONED BY ALL") + .setQueryContext(adjustedContext) + .setExpectedResultRows(expectedRows) + .setExpectedDataSource("foo1") + .setExpectedRowSignature(rowSignature) + .verifyResults(); + } + + /** + * Tests the behaviour of INSERT query when arrayIngestMode is set to mvd and the user tries to ingest numeric array + * types as well + */ + @Test + public void testInsertArraysWithStringArraysAsMVDs() throws IOException + { + RowSignature rowSignatureWithoutTimeAndStringColumns = + RowSignature.builder() + .add("arrayLong", ColumnType.LONG_ARRAY) + .add("arrayLongNulls", ColumnType.LONG_ARRAY) + .add("arrayDouble", ColumnType.DOUBLE_ARRAY) + .add("arrayDoubleNulls", ColumnType.DOUBLE_ARRAY) + .build(); + + + RowSignature fileSignature = RowSignature.builder() + .add("timestamp", ColumnType.STRING) + .add("arrayString", ColumnType.STRING_ARRAY) + .add("arrayStringNulls", ColumnType.STRING_ARRAY) + .addAll(rowSignatureWithoutTimeAndStringColumns) + .build(); + + final Map adjustedContext = new HashMap<>(context); + adjustedContext.put(MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, "mvd"); + + final File tmpFile = temporaryFolder.newFile(); + final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader() + .getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); + final InputStream decompressing = CompressionUtils.decompress( + resourceStream, + NestedDataTestUtils.ARRAY_TYPES_DATA_FILE + ); + Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING); + decompressing.close(); + + final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile); + + testIngestQuery().setSql(" INSERT INTO foo1 SELECT\n" + + " TIME_PARSE(\"timestamp\") as __time,\n" + + " arrayString,\n" + + " arrayStringNulls,\n" + + " arrayLong,\n" + + " arrayLongNulls,\n" + + " arrayDouble,\n" + + " arrayDoubleNulls\n" + + "FROM TABLE(\n" + + " EXTERN(\n" + + " '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n" + + " '{\"type\": \"json\"}',\n" + + " '" + queryFramework().queryJsonMapper().writeValueAsString(fileSignature) + "'\n" + + " )\n" + + ") PARTITIONED BY ALL") + .setQueryContext(adjustedContext) + .setExpectedExecutionErrorMatcher(CoreMatchers.allOf( + CoreMatchers.instanceOf(ISE.class), + ThrowableMessageMatcher.hasMessage(CoreMatchers.containsString( + "Numeric arrays can only be ingested when")) + )) + .verifyExecutionError(); + } + + /** + * Tests the behaviour of INSERT query when arrayIngestMode is set to array and the user tries to ingest all + * array types + */ + @Test + public void testInsertArraysAsArrays() throws IOException + { + final List expectedRows = Arrays.asList( + new Object[]{ + 1672531200000L, + null, + null, + new Object[]{1L, 2L, 3L}, + new Object[]{}, + new Object[]{1.1d, 2.2d, 3.3d}, + null + }, + new Object[]{ + 1672531200000L, + null, + new Object[]{"a", "b"}, + null, + new Object[]{2L, 3L}, + null, + new Object[]{null} + }, + new Object[]{ + 1672531200000L, + new Object[]{"d", "e"}, + new Object[]{"b", "b"}, + new Object[]{1L, 4L}, + new Object[]{1L}, + new Object[]{2.2d, 3.3d, 4.0d}, + null + }, + new Object[]{ + 1672531200000L, + new Object[]{"a", "b"}, + null, + null, + new Object[]{null, 2L, 9L}, + null, + new Object[]{999.0d, 5.5d, null} + }, + new Object[]{ + 1672531200000L, + new Object[]{"a", "b"}, + new Object[]{"a", "b"}, + new Object[]{1L, 2L, 3L}, + new Object[]{1L, null, 3L}, + new Object[]{1.1d, 2.2d, 3.3d}, + new Object[]{1.1d, 2.2d, null} + }, + new Object[]{ + 1672531200000L, + new Object[]{"b", "c"}, + new Object[]{"d", null, "b"}, + new Object[]{1L, 2L, 3L, 4L}, + new Object[]{1L, 2L, 3L}, + new Object[]{1.1d, 3.3d}, + new Object[]{null, 2.2d, null} + }, + new Object[]{ + 1672531200000L, + new Object[]{"a", "b", "c"}, + new Object[]{null, "b"}, + new Object[]{2L, 3L}, + null, + new Object[]{3.3d, 4.4d, 5.5d}, + new Object[]{999.0d, null, 5.5d} + }, + new Object[]{ + 1672617600000L, + null, + null, + new Object[]{1L, 2L, 3L}, + null, + new Object[]{1.1d, 2.2d, 3.3d}, + new Object[]{} + }, + new Object[]{ + 1672617600000L, + null, + new Object[]{"a", "b"}, + null, + new Object[]{2L, 3L}, + null, + new Object[]{null, 1.1d} + }, + new Object[]{ + 1672617600000L, + new Object[]{"d", "e"}, + new Object[]{"b", "b"}, + new Object[]{1L, 4L}, + new Object[]{null}, + new Object[]{2.2d, 3.3d, 4.0}, + null + }, + new Object[]{ + 1672617600000L, + new Object[]{"a", "b"}, + new Object[]{null}, + null, + new Object[]{null, 2L, 9L}, + null, + new Object[]{999.0d, 5.5d, null} + }, + new Object[]{ + 1672617600000L, + new Object[]{"a", "b"}, + new Object[]{}, + new Object[]{1L, 2L, 3L}, + new Object[]{1L, null, 3L}, + new Object[]{1.1d, 2.2d, 3.3d}, + new Object[]{1.1d, 2.2d, null} + }, + new Object[]{ + 1672617600000L, + new Object[]{"b", "c"}, + new Object[]{"d", null, "b"}, + new Object[]{1L, 2L, 3L, 4L}, + new Object[]{1L, 2L, 3L}, + new Object[]{1.1d, 3.3d}, + new Object[]{null, 2.2d, null} + }, + new Object[]{ + 1672617600000L, + new Object[]{"a", "b", "c"}, + new Object[]{null, "b"}, + new Object[]{2L, 3L}, + null, + new Object[]{3.3d, 4.4d, 5.5d}, + new Object[]{999.0d, null, 5.5d} + } + ); + + RowSignature rowSignatureWithoutTimeColumn = + RowSignature.builder() + .add("arrayString", ColumnType.STRING_ARRAY) + .add("arrayStringNulls", ColumnType.STRING_ARRAY) + .add("arrayLong", ColumnType.LONG_ARRAY) + .add("arrayLongNulls", ColumnType.LONG_ARRAY) + .add("arrayDouble", ColumnType.DOUBLE_ARRAY) + .add("arrayDoubleNulls", ColumnType.DOUBLE_ARRAY) + .build(); + + RowSignature fileSignature = RowSignature.builder() + .add("timestamp", ColumnType.STRING) + .addAll(rowSignatureWithoutTimeColumn) + .build(); + + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .addAll(rowSignatureWithoutTimeColumn) + .build(); + + final Map adjustedContext = new HashMap<>(context); + adjustedContext.put(MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, "array"); + + final File tmpFile = temporaryFolder.newFile(); + final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader() + .getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); + final InputStream decompressing = CompressionUtils.decompress( + resourceStream, + NestedDataTestUtils.ARRAY_TYPES_DATA_FILE + ); + Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING); + decompressing.close(); + + final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile); + + testIngestQuery().setSql(" INSERT INTO foo1 SELECT\n" + + " TIME_PARSE(\"timestamp\") as __time,\n" + + " arrayString,\n" + + " arrayStringNulls,\n" + + " arrayLong,\n" + + " arrayLongNulls,\n" + + " arrayDouble,\n" + + " arrayDoubleNulls\n" + + "FROM TABLE(\n" + + " EXTERN(\n" + + " '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n" + + " '{\"type\": \"json\"}',\n" + + " '" + queryFramework().queryJsonMapper().writeValueAsString(fileSignature) + "'\n" + + " )\n" + + ") PARTITIONED BY ALL") + .setQueryContext(adjustedContext) + .setExpectedResultRows(expectedRows) + .setExpectedDataSource("foo1") + .setExpectedRowSignature(rowSignature) + .verifyResults(); + } + + @Test + public void testSelectOnArraysWithArrayIngestModeAsNone() throws IOException + { + testSelectOnArrays("none"); + } + + @Test + public void testSelectOnArraysWithArrayIngestModeAsMVD() throws IOException + { + testSelectOnArrays("mvd"); + } + + @Test + public void testSelectOnArraysWithArrayIngestModeAsArray() throws IOException + { + testSelectOnArrays("array"); + } + + // Tests the behaviour of the select with the given arrayIngestMode. The expectation should be the same, since the + // arrayIngestMode should only determine how the array gets ingested at the end. + public void testSelectOnArrays(String arrayIngestMode) throws IOException + { + final List expectedRows = Arrays.asList( + new Object[]{ + 1672531200000L, + Arrays.asList("a", "b"), + Arrays.asList("a", "b"), + Arrays.asList(1L, 2L, 3L), + Arrays.asList(1L, null, 3L), + Arrays.asList(1.1d, 2.2d, 3.3d), + Arrays.asList(1.1d, 2.2d, null) + }, + new Object[]{ + 1672531200000L, + Arrays.asList("a", "b", "c"), + Arrays.asList(null, "b"), + Arrays.asList(2L, 3L), + null, + Arrays.asList(3.3d, 4.4d, 5.5d), + Arrays.asList(999.0d, null, 5.5d), + }, + new Object[]{ + 1672531200000L, + Arrays.asList("b", "c"), + Arrays.asList("d", null, "b"), + Arrays.asList(1L, 2L, 3L, 4L), + Arrays.asList(1L, 2L, 3L), + Arrays.asList(1.1d, 3.3d), + Arrays.asList(null, 2.2d, null) + }, + new Object[]{ + 1672531200000L, + Arrays.asList("d", "e"), + Arrays.asList("b", "b"), + Arrays.asList(1L, 4L), + Collections.singletonList(1L), + Arrays.asList(2.2d, 3.3d, 4.0d), + null + }, + new Object[]{ + 1672531200000L, + null, + null, + Arrays.asList(1L, 2L, 3L), + Collections.emptyList(), + Arrays.asList(1.1d, 2.2d, 3.3d), + null + }, + new Object[]{ + 1672531200000L, + Arrays.asList("a", "b"), + null, + null, + Arrays.asList(null, 2L, 9L), + null, + Arrays.asList(999.0d, 5.5d, null) + }, + new Object[]{ + 1672531200000L, + null, + Arrays.asList("a", "b"), + null, + Arrays.asList(2L, 3L), + null, + Collections.singletonList(null) + }, + new Object[]{ + 1672617600000L, + Arrays.asList("a", "b"), + Collections.emptyList(), + Arrays.asList(1L, 2L, 3L), + Arrays.asList(1L, null, 3L), + Arrays.asList(1.1d, 2.2d, 3.3d), + Arrays.asList(1.1d, 2.2d, null) + }, + new Object[]{ + 1672617600000L, + Arrays.asList("a", "b", "c"), + Arrays.asList(null, "b"), + Arrays.asList(2L, 3L), + null, + Arrays.asList(3.3d, 4.4d, 5.5d), + Arrays.asList(999.0d, null, 5.5d) + }, + new Object[]{ + 1672617600000L, + Arrays.asList("b", "c"), + Arrays.asList("d", null, "b"), + Arrays.asList(1L, 2L, 3L, 4L), + Arrays.asList(1L, 2L, 3L), + Arrays.asList(1.1d, 3.3d), + Arrays.asList(null, 2.2d, null) + }, + new Object[]{ + 1672617600000L, + Arrays.asList("d", "e"), + Arrays.asList("b", "b"), + Arrays.asList(1L, 4L), + Collections.singletonList(null), + Arrays.asList(2.2d, 3.3d, 4.0), + null + }, + new Object[]{ + 1672617600000L, + null, + null, + Arrays.asList(1L, 2L, 3L), + null, + Arrays.asList(1.1d, 2.2d, 3.3d), + Collections.emptyList() + }, + new Object[]{ + 1672617600000L, + Arrays.asList("a", "b"), + Collections.singletonList(null), + null, + Arrays.asList(null, 2L, 9L), + null, + Arrays.asList(999.0d, 5.5d, null) + }, + new Object[]{ + 1672617600000L, + null, + Arrays.asList("a", "b"), + null, + Arrays.asList(2L, 3L), + null, + Arrays.asList(null, 1.1d), + } + ); + + RowSignature rowSignatureWithoutTimeColumn = + RowSignature.builder() + .add("arrayString", ColumnType.STRING_ARRAY) + .add("arrayStringNulls", ColumnType.STRING_ARRAY) + .add("arrayLong", ColumnType.LONG_ARRAY) + .add("arrayLongNulls", ColumnType.LONG_ARRAY) + .add("arrayDouble", ColumnType.DOUBLE_ARRAY) + .add("arrayDoubleNulls", ColumnType.DOUBLE_ARRAY) + .build(); + + RowSignature fileSignature = RowSignature.builder() + .add("timestamp", ColumnType.STRING) + .addAll(rowSignatureWithoutTimeColumn) + .build(); + + RowSignature rowSignature = RowSignature.builder() + .add("__time", ColumnType.LONG) + .addAll(rowSignatureWithoutTimeColumn) + .build(); + + RowSignature scanSignature = RowSignature.builder() + .add("arrayDouble", ColumnType.DOUBLE_ARRAY) + .add("arrayDoubleNulls", ColumnType.DOUBLE_ARRAY) + .add("arrayLong", ColumnType.LONG_ARRAY) + .add("arrayLongNulls", ColumnType.LONG_ARRAY) + .add("arrayString", ColumnType.STRING_ARRAY) + .add("arrayStringNulls", ColumnType.STRING_ARRAY) + .add("v0", ColumnType.LONG) + .build(); + + final Map adjustedContext = new HashMap<>(context); + adjustedContext.put(MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, arrayIngestMode); + + final File tmpFile = temporaryFolder.newFile(); + final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader() + .getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); + final InputStream decompressing = CompressionUtils.decompress( + resourceStream, + NestedDataTestUtils.ARRAY_TYPES_DATA_FILE + ); + Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING); + decompressing.close(); + + final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile); + + Query expectedQuery = newScanQueryBuilder() + .dataSource(new ExternalDataSource( + new LocalInputSource(null, null, ImmutableList.of(tmpFile)), + new JsonInputFormat(null, null, null, null, null), + fileSignature + )) + .intervals(querySegmentSpec(Filtration.eternity())) + .columns( + "arrayDouble", + "arrayDoubleNulls", + "arrayLong", + "arrayLongNulls", + "arrayString", + "arrayStringNulls", + "v0" + ) + .virtualColumns(new ExpressionVirtualColumn( + "v0", + "timestamp_parse(\"timestamp\",null,'UTC')", + ColumnType.LONG, + TestExprMacroTable.INSTANCE + )) + .context(defaultScanQueryContext(adjustedContext, scanSignature)) + .build(); + + testSelectQuery().setSql("SELECT\n" + + " TIME_PARSE(\"timestamp\") as __time,\n" + + " arrayString,\n" + + " arrayStringNulls,\n" + + " arrayLong,\n" + + " arrayLongNulls,\n" + + " arrayDouble,\n" + + " arrayDoubleNulls\n" + + "FROM TABLE(\n" + + " EXTERN(\n" + + " '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n" + + " '{\"type\": \"json\"}',\n" + + " '" + queryFramework().queryJsonMapper().writeValueAsString(fileSignature) + "'\n" + + " )\n" + + ")") + .setQueryContext(adjustedContext) + .setExpectedMSQSpec(MSQSpec + .builder() + .query(expectedQuery) + .columnMappings(new ColumnMappings(ImmutableList.of( + new ColumnMapping("v0", "__time"), + new ColumnMapping("arrayString", "arrayString"), + new ColumnMapping("arrayStringNulls", "arrayStringNulls"), + new ColumnMapping("arrayLong", "arrayLong"), + new ColumnMapping("arrayLongNulls", "arrayLongNulls"), + new ColumnMapping("arrayDouble", "arrayDouble"), + new ColumnMapping("arrayDoubleNulls", "arrayDoubleNulls") + ))) + .tuningConfig(MSQTuningConfig.defaultConfig()) + .destination(TaskReportMSQDestination.INSTANCE) + .build() + ) + .setExpectedRowSignature(rowSignature) + .setExpectedResultRows(expectedRows) + .verifyResults(); + } + + + private List expectedMultiValueFooRowsToArray() + { + List expectedRows = new ArrayList<>(); + expectedRows.add(new Object[]{0L, null}); + if (!useDefault) { + expectedRows.add(new Object[]{0L, ""}); + } + + expectedRows.addAll(ImmutableList.of( + new Object[]{0L, ImmutableList.of("a", "b")}, + new Object[]{0L, ImmutableList.of("b", "c")}, + new Object[]{0L, "d"} + )); + return expectedRows; + } +} diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java index e54027c2449b..b43dd72e88c8 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java @@ -38,7 +38,6 @@ import org.apache.druid.msq.test.MSQTestBase; import org.apache.druid.msq.test.MSQTestFileUtils; import org.apache.druid.msq.util.MultiStageQueryContext; -import org.apache.druid.query.NestedDataTestUtils; import org.apache.druid.query.QueryContexts; import org.apache.druid.query.aggregation.LongSumAggregatorFactory; import org.apache.druid.query.aggregation.hyperloglog.HyperUniquesAggregatorFactory; @@ -46,7 +45,6 @@ import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.column.ValueType; import org.apache.druid.timeline.SegmentId; -import org.apache.druid.utils.CompressionUtils; import org.hamcrest.CoreMatchers; import org.junit.Test; import org.junit.internal.matchers.ThrowableMessageMatcher; @@ -54,16 +52,11 @@ import org.junit.runners.Parameterized; import org.mockito.Mockito; -import javax.annotation.Nonnull; import java.io.File; import java.io.IOException; -import java.io.InputStream; -import java.nio.file.Files; -import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; -import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -735,22 +728,6 @@ public void testInsertOnFoo1WithMultiValueMeasureGroupBy() } - @Test - public void testInsertOnFoo1WithMultiValueToArrayGroupBy() - { - RowSignature rowSignature = RowSignature.builder() - .add("__time", ColumnType.LONG) - .add("dim3", ColumnType.STRING).build(); - - testIngestQuery().setSql( - "INSERT INTO foo1 SELECT MV_TO_ARRAY(dim3) AS dim3 FROM foo GROUP BY 1 PARTITIONED BY ALL TIME") - .setExpectedDataSource("foo1") - .setExpectedRowSignature(rowSignature) - .setQueryContext(context) - .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0))) - .setExpectedResultRows(expectedMultiValueFooRowsToArray()) - .verifyResults(); - } @Test public void testInsertOnFoo1WithAutoTypeArrayGroupBy() @@ -1407,251 +1384,6 @@ public void testCorrectNumberOfWorkersUsedAutoModeWithBytesLimit() throws IOExce .verifyResults(); } - @Test - public void testInsertArraysAutoType() throws IOException - { - List expectedRows = Arrays.asList( - new Object[]{1672531200000L, null, null, null}, - new Object[]{1672531200000L, null, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, - new Object[]{1672531200000L, new Object[]{"d", "e"}, new Object[]{1L, 4L}, new Object[]{2.2, 3.3, 4.0}}, - new Object[]{1672531200000L, new Object[]{"a", "b"}, null, null}, - new Object[]{1672531200000L, new Object[]{"a", "b"}, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, - new Object[]{1672531200000L, new Object[]{"b", "c"}, new Object[]{1L, 2L, 3L, 4L}, new Object[]{1.1, 3.3}}, - new Object[]{1672531200000L, new Object[]{"a", "b", "c"}, new Object[]{2L, 3L}, new Object[]{3.3, 4.4, 5.5}}, - new Object[]{1672617600000L, null, null, null}, - new Object[]{1672617600000L, null, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, - new Object[]{1672617600000L, new Object[]{"d", "e"}, new Object[]{1L, 4L}, new Object[]{2.2, 3.3, 4.0}}, - new Object[]{1672617600000L, new Object[]{"a", "b"}, null, null}, - new Object[]{1672617600000L, new Object[]{"a", "b"}, new Object[]{1L, 2L, 3L}, new Object[]{1.1, 2.2, 3.3}}, - new Object[]{1672617600000L, new Object[]{"b", "c"}, new Object[]{1L, 2L, 3L, 4L}, new Object[]{1.1, 3.3}}, - new Object[]{1672617600000L, new Object[]{"a", "b", "c"}, new Object[]{2L, 3L}, new Object[]{3.3, 4.4, 5.5}} - ); - - RowSignature rowSignature = RowSignature.builder() - .add("__time", ColumnType.LONG) - .add("arrayString", ColumnType.STRING_ARRAY) - .add("arrayLong", ColumnType.LONG_ARRAY) - .add("arrayDouble", ColumnType.DOUBLE_ARRAY) - .build(); - - final Map adjustedContext = new HashMap<>(context); - adjustedContext.put(MultiStageQueryContext.CTX_USE_AUTO_SCHEMAS, true); - - final File tmpFile = temporaryFolder.newFile(); - final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader().getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); - final InputStream decompressing = CompressionUtils.decompress(resourceStream, NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); - Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING); - decompressing.close(); - - final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile); - - testIngestQuery().setSql(" INSERT INTO foo1 SELECT\n" - + " TIME_PARSE(\"timestamp\") as __time,\n" - + " arrayString,\n" - + " arrayLong,\n" - + " arrayDouble\n" - + "FROM TABLE(\n" - + " EXTERN(\n" - + " '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n" - + " '{\"type\": \"json\"}',\n" - + " '[{\"name\": \"timestamp\", \"type\": \"STRING\"}, {\"name\": \"arrayString\", \"type\": \"COMPLEX\"}, {\"name\": \"arrayLong\", \"type\": \"COMPLEX\"}, {\"name\": \"arrayDouble\", \"type\": \"COMPLEX\"}]'\n" - + " )\n" - + ") PARTITIONED BY day") - .setQueryContext(adjustedContext) - .setExpectedResultRows(expectedRows) - .setExpectedDataSource("foo1") - .setExpectedRowSignature(rowSignature) - .verifyResults(); - } - - @Test - public void testInsertArrays() throws IOException - { - List expectedRows = Arrays.asList( - new Object[]{ - 1672531200000L, - null, - null, - new Object[]{1L, 2L, 3L}, - new Object[]{}, - new Object[]{1.1d, 2.2d, 3.3d}, - null - }, - new Object[]{ - 1672531200000L, - null, - Arrays.asList("a", "b"), - null, - new Object[]{2L, 3L}, - null, - new Object[]{null} - }, - new Object[]{ - 1672531200000L, - Arrays.asList("a", "b"), - null, - null, - new Object[]{null, 2L, 9L}, - null, - new Object[]{999.0d, 5.5d, null} - }, - new Object[]{ - 1672531200000L, - Arrays.asList("a", "b"), - Arrays.asList("a", "b"), - new Object[]{1L, 2L, 3L}, - new Object[]{1L, null, 3L}, - new Object[]{1.1d, 2.2d, 3.3d}, - new Object[]{1.1d, 2.2d, null} - }, - new Object[]{ - 1672531200000L, - Arrays.asList("a", "b", "c"), - Arrays.asList(null, "b"), - new Object[]{2L, 3L}, - null, - new Object[]{3.3d, 4.4d, 5.5d}, - new Object[]{999.0d, null, 5.5d} - }, - new Object[]{ - 1672531200000L, - Arrays.asList("b", "c"), - Arrays.asList("d", null, "b"), - new Object[]{1L, 2L, 3L, 4L}, - new Object[]{1L, 2L, 3L}, - new Object[]{1.1d, 3.3d}, - new Object[]{null, 2.2d, null} - }, - new Object[]{ - 1672531200000L, - Arrays.asList("d", "e"), - Arrays.asList("b", "b"), - new Object[]{1L, 4L}, - new Object[]{1L}, - new Object[]{2.2d, 3.3d, 4.0d}, - null - }, - new Object[]{ - 1672617600000L, - null, - null, - new Object[]{1L, 2L, 3L}, - null, - new Object[]{1.1d, 2.2d, 3.3d}, - new Object[]{} - }, - new Object[]{ - 1672617600000L, - null, - Arrays.asList("a", "b"), - null, - new Object[]{2L, 3L}, - null, - new Object[]{null, 1.1d} - }, - new Object[]{ - 1672617600000L, - Arrays.asList("a", "b"), - null, - null, - new Object[]{null, 2L, 9L}, - null, - new Object[]{999.0d, 5.5d, null} - }, - new Object[]{ - 1672617600000L, - Arrays.asList("a", "b"), - Collections.emptyList(), - new Object[]{1L, 2L, 3L}, - new Object[]{1L, null, 3L}, - new Object[]{1.1d, 2.2d, 3.3d}, - new Object[]{1.1d, 2.2d, null} - }, - new Object[]{ - 1672617600000L, - Arrays.asList("a", "b", "c"), - Arrays.asList(null, "b"), - new Object[]{2L, 3L}, - null, - new Object[]{3.3d, 4.4d, 5.5d}, - new Object[]{999.0d, null, 5.5d} - }, - new Object[]{ - 1672617600000L, - Arrays.asList("b", "c"), - Arrays.asList("d", null, "b"), - new Object[]{1L, 2L, 3L, 4L}, - new Object[]{1L, 2L, 3L}, - new Object[]{1.1d, 3.3d}, - new Object[]{null, 2.2d, null} - }, - new Object[]{ - 1672617600000L, - Arrays.asList("d", "e"), - Arrays.asList("b", "b"), - new Object[]{1L, 4L}, - new Object[]{null}, - new Object[]{2.2d, 3.3d, 4.0}, - null - } - ); - - RowSignature rowSignatureWithoutTimeAndStringColumns = - RowSignature.builder() - .add("arrayLong", ColumnType.LONG_ARRAY) - .add("arrayLongNulls", ColumnType.LONG_ARRAY) - .add("arrayDouble", ColumnType.DOUBLE_ARRAY) - .add("arrayDoubleNulls", ColumnType.DOUBLE_ARRAY) - .build(); - - - RowSignature fileSignature = RowSignature.builder() - .add("timestamp", ColumnType.STRING) - .add("arrayString", ColumnType.STRING_ARRAY) - .add("arrayStringNulls", ColumnType.STRING_ARRAY) - .addAll(rowSignatureWithoutTimeAndStringColumns) - .build(); - - // MSQ writes strings instead of string arrays - RowSignature rowSignature = RowSignature.builder() - .add("__time", ColumnType.LONG) - .add("arrayString", ColumnType.STRING) - .add("arrayStringNulls", ColumnType.STRING) - .addAll(rowSignatureWithoutTimeAndStringColumns) - .build(); - - final Map adjustedContext = new HashMap<>(context); - final File tmpFile = temporaryFolder.newFile(); - final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader().getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); - final InputStream decompressing = CompressionUtils.decompress(resourceStream, NestedDataTestUtils.ARRAY_TYPES_DATA_FILE); - Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING); - decompressing.close(); - - final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile); - - testIngestQuery().setSql(" INSERT INTO foo1 SELECT\n" - + " TIME_PARSE(\"timestamp\") as __time,\n" - + " arrayString,\n" - + " arrayStringNulls,\n" - + " arrayLong,\n" - + " arrayLongNulls,\n" - + " arrayDouble,\n" - + " arrayDoubleNulls\n" - + "FROM TABLE(\n" - + " EXTERN(\n" - + " '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n" - + " '{\"type\": \"json\"}',\n" - + " '" + queryFramework().queryJsonMapper().writeValueAsString(fileSignature) + "'\n" - + " )\n" - + ") PARTITIONED BY day") - .setQueryContext(adjustedContext) - .setExpectedResultRows(expectedRows) - .setExpectedDataSource("foo1") - .setExpectedRowSignature(rowSignature) - .verifyResults(); - } - - @Nonnull private List expectedFooRows() { List expectedRows = new ArrayList<>(); @@ -1668,7 +1400,6 @@ private List expectedFooRows() return expectedRows; } - @Nonnull private List expectedFooRowsWithAggregatedComplexColumn() { List expectedRows = new ArrayList<>(); @@ -1687,7 +1418,6 @@ private List expectedFooRowsWithAggregatedComplexColumn() return expectedRows; } - @Nonnull private List expectedMultiValueFooRows() { List expectedRows = new ArrayList<>(); @@ -1704,24 +1434,6 @@ private List expectedMultiValueFooRows() return expectedRows; } - @Nonnull - private List expectedMultiValueFooRowsToArray() - { - List expectedRows = new ArrayList<>(); - expectedRows.add(new Object[]{0L, null}); - if (!useDefault) { - expectedRows.add(new Object[]{0L, ""}); - } - - expectedRows.addAll(ImmutableList.of( - new Object[]{0L, ImmutableList.of("a", "b")}, - new Object[]{0L, ImmutableList.of("b", "c")}, - new Object[]{0L, "d"} - )); - return expectedRows; - } - - @Nonnull private List expectedMultiValueFooRowsGroupBy() { List expectedRows = new ArrayList<>(); @@ -1737,7 +1449,6 @@ private List expectedMultiValueFooRowsGroupBy() return expectedRows; } - @Nonnull private Set expectedFooSegments() { Set expectedSegments = new TreeSet<>(); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/MultiStageQueryContextTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/MultiStageQueryContextTest.java index 830b414daedb..5bfb4d2eb279 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/MultiStageQueryContextTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/util/MultiStageQueryContextTest.java @@ -39,6 +39,7 @@ import java.util.List; import java.util.Map; +import static org.apache.druid.msq.util.MultiStageQueryContext.CTX_ARRAY_INGEST_MODE; import static org.apache.druid.msq.util.MultiStageQueryContext.CTX_DURABLE_SHUFFLE_STORAGE; import static org.apache.druid.msq.util.MultiStageQueryContext.CTX_FAULT_TOLERANCE; import static org.apache.druid.msq.util.MultiStageQueryContext.CTX_FINALIZE_AGGREGATIONS; @@ -54,46 +55,46 @@ public class MultiStageQueryContextTest { @Test - public void isDurableShuffleStorageEnabled_noParameterSetReturnsDefaultValue() + public void isDurableShuffleStorageEnabled_unset_returnsDefaultValue() { Assert.assertFalse(MultiStageQueryContext.isDurableStorageEnabled(QueryContext.empty())); } @Test - public void isDurableShuffleStorageEnabled_parameterSetReturnsCorrectValue() + public void isDurableShuffleStorageEnabled_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_DURABLE_SHUFFLE_STORAGE, "true"); Assert.assertTrue(MultiStageQueryContext.isDurableStorageEnabled(QueryContext.of(propertyMap))); } @Test - public void isFaultToleranceEnabled_noParameterSetReturnsDefaultValue() + public void isFaultToleranceEnabled_unset_returnsDefaultValue() { Assert.assertFalse(MultiStageQueryContext.isFaultToleranceEnabled(QueryContext.empty())); } @Test - public void isFaultToleranceEnabled_parameterSetReturnsCorrectValue() + public void isFaultToleranceEnabled_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_FAULT_TOLERANCE, "true"); Assert.assertTrue(MultiStageQueryContext.isFaultToleranceEnabled(QueryContext.of(propertyMap))); } @Test - public void isFinalizeAggregations_noParameterSetReturnsDefaultValue() + public void isFinalizeAggregations_unset_returnsDefaultValue() { Assert.assertTrue(MultiStageQueryContext.isFinalizeAggregations(QueryContext.empty())); } @Test - public void isFinalizeAggregations_parameterSetReturnsCorrectValue() + public void isFinalizeAggregations_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_FINALIZE_AGGREGATIONS, "false"); Assert.assertFalse(MultiStageQueryContext.isFinalizeAggregations(QueryContext.of(propertyMap))); } @Test - public void getAssignmentStrategy_noParameterSetReturnsDefaultValue() + public void getAssignmentStrategy_unset_returnsDefaultValue() { Assert.assertEquals( WorkerAssignmentStrategy.MAX, @@ -102,7 +103,7 @@ public void getAssignmentStrategy_noParameterSetReturnsDefaultValue() } @Test - public void testGetMaxInputBytesPerWorker() + public void getMaxInputBytesPerWorker_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(MultiStageQueryContext.CTX_MAX_INPUT_BYTES_PER_WORKER, 1024); @@ -112,7 +113,7 @@ public void testGetMaxInputBytesPerWorker() } @Test - public void getAssignmentStrategy_parameterSetReturnsCorrectValue() + public void getAssignmentStrategy_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_TASK_ASSIGNMENT_STRATEGY, "AUTO"); Assert.assertEquals( @@ -122,27 +123,20 @@ public void getAssignmentStrategy_parameterSetReturnsCorrectValue() } @Test - public void getMaxNumTasks_noParameterSetReturnsDefaultValue() + public void getMaxNumTasks_unset_returnsDefaultValue() { Assert.assertEquals(DEFAULT_MAX_NUM_TASKS, MultiStageQueryContext.getMaxNumTasks(QueryContext.empty())); } @Test - public void getMaxNumTasks_parameterSetReturnsCorrectValue() + public void getMaxNumTasks_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_MAX_NUM_TASKS, 101); Assert.assertEquals(101, MultiStageQueryContext.getMaxNumTasks(QueryContext.of(propertyMap))); } @Test - public void getMaxNumTasks_legacyParameterSetReturnsCorrectValue() - { - Map propertyMap = ImmutableMap.of(CTX_MAX_NUM_TASKS, 101); - Assert.assertEquals(101, MultiStageQueryContext.getMaxNumTasks(QueryContext.of(propertyMap))); - } - - @Test - public void getRowsPerSegment_noParameterSetReturnsDefaultValue() + public void getRowsPerSegment_unset_returnsDefaultValue() { Assert.assertEquals( MultiStageQueryContext.DEFAULT_ROWS_PER_SEGMENT, @@ -151,14 +145,14 @@ public void getRowsPerSegment_noParameterSetReturnsDefaultValue() } @Test - public void getRowsPerSegment_parameterSetReturnsCorrectValue() + public void getRowsPerSegment_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_ROWS_PER_SEGMENT, 10); Assert.assertEquals(10, MultiStageQueryContext.getRowsPerSegment(QueryContext.of(propertyMap))); } @Test - public void getRowsInMemory_noParameterSetReturnsDefaultValue() + public void getRowsInMemory_unset_returnsDefaultValue() { Assert.assertEquals( MultiStageQueryContext.DEFAULT_ROWS_IN_MEMORY, @@ -167,12 +161,91 @@ public void getRowsInMemory_noParameterSetReturnsDefaultValue() } @Test - public void getRowsInMemory_parameterSetReturnsCorrectValue() + public void getRowsInMemory_set_returnsCorrectValue() { Map propertyMap = ImmutableMap.of(CTX_ROWS_IN_MEMORY, 10); Assert.assertEquals(10, MultiStageQueryContext.getRowsInMemory(QueryContext.of(propertyMap))); } + @Test + public void getSortOrder_unset_returnsDefaultValue() + { + Assert.assertEquals(Collections.emptyList(), MultiStageQueryContext.getSortOrder(QueryContext.empty())); + } + + @Test + public void getSortOrder_set_returnsCorrectValue() + { + Map propertyMap = ImmutableMap.of(CTX_SORT_ORDER, "a, b,\"c,d\""); + Assert.assertEquals( + ImmutableList.of("a", "b", "c,d"), + MultiStageQueryContext.getSortOrder(QueryContext.of(propertyMap)) + ); + } + + @Test + public void getMSQMode_unset_returnsDefaultValue() + { + Assert.assertEquals("strict", MultiStageQueryContext.getMSQMode(QueryContext.empty())); + } + + @Test + public void getMSQMode_set_returnsCorrectValue() + { + Map propertyMap = ImmutableMap.of(CTX_MSQ_MODE, "nonStrict"); + Assert.assertEquals("nonStrict", MultiStageQueryContext.getMSQMode(QueryContext.of(propertyMap))); + } + + @Test + public void getSelectDestination_unset_returnsDefaultValue() + { + Assert.assertEquals(MSQSelectDestination.TASKREPORT, MultiStageQueryContext.getSelectDestination(QueryContext.empty())); + } + + @Test + public void useAutoColumnSchemes_unset_returnsDefaultValue() + { + Assert.assertFalse(MultiStageQueryContext.useAutoColumnSchemas(QueryContext.empty())); + } + + @Test + public void useAutoColumnSchemes_set_returnsCorrectValue() + { + Map propertyMap = ImmutableMap.of(CTX_USE_AUTO_SCHEMAS, true); + Assert.assertTrue(MultiStageQueryContext.useAutoColumnSchemas(QueryContext.of(propertyMap))); + } + + @Test + public void arrayIngestMode_unset_returnsDefaultValue() + { + Assert.assertEquals(ArrayIngestMode.MVD, MultiStageQueryContext.getArrayIngestMode(QueryContext.empty())); + } + + @Test + public void arrayIngestMode_set_returnsCorrectValue() + { + Assert.assertEquals( + ArrayIngestMode.NONE, + MultiStageQueryContext.getArrayIngestMode(QueryContext.of(ImmutableMap.of(CTX_ARRAY_INGEST_MODE, "none"))) + ); + + Assert.assertEquals( + ArrayIngestMode.MVD, + MultiStageQueryContext.getArrayIngestMode(QueryContext.of(ImmutableMap.of(CTX_ARRAY_INGEST_MODE, "mvd"))) + ); + + Assert.assertEquals( + ArrayIngestMode.ARRAY, + MultiStageQueryContext.getArrayIngestMode(QueryContext.of(ImmutableMap.of(CTX_ARRAY_INGEST_MODE, "array"))) + ); + + Assert.assertThrows( + BadQueryContextException.class, + () -> + MultiStageQueryContext.getArrayIngestMode(QueryContext.of(ImmutableMap.of(CTX_ARRAY_INGEST_MODE, "dummy"))) + ); + } + @Test public void testDecodeSortOrder() { @@ -221,48 +294,6 @@ public void testGetIndexSpec() ); } - @Test - public void getSortOrderNoParameterSetReturnsDefaultValue() - { - Assert.assertEquals(Collections.emptyList(), MultiStageQueryContext.getSortOrder(QueryContext.empty())); - } - - @Test - public void getSortOrderParameterSetReturnsCorrectValue() - { - Map propertyMap = ImmutableMap.of(CTX_SORT_ORDER, "a, b,\"c,d\""); - Assert.assertEquals( - ImmutableList.of("a", "b", "c,d"), - MultiStageQueryContext.getSortOrder(QueryContext.of(propertyMap)) - ); - } - - @Test - public void getMSQModeNoParameterSetReturnsDefaultValue() - { - Assert.assertEquals("strict", MultiStageQueryContext.getMSQMode(QueryContext.empty())); - } - - @Test - public void getMSQModeParameterSetReturnsCorrectValue() - { - Map propertyMap = ImmutableMap.of(CTX_MSQ_MODE, "nonStrict"); - Assert.assertEquals("nonStrict", MultiStageQueryContext.getMSQMode(QueryContext.of(propertyMap))); - } - - @Test - public void limitSelectResultReturnsDefaultValue() - { - Assert.assertEquals(MSQSelectDestination.TASKREPORT, MultiStageQueryContext.getSelectDestination(QueryContext.empty())); - } - - @Test - public void testUseAutoSchemas() - { - Map propertyMap = ImmutableMap.of(CTX_USE_AUTO_SCHEMAS, true); - Assert.assertTrue(MultiStageQueryContext.useAutoColumnSchemas(QueryContext.of(propertyMap))); - } - private static List decodeSortOrder(@Nullable final String input) { return MultiStageQueryContext.decodeSortOrder(input); From 90a1458ac9b81bd3bac443790242353bb701aadf Mon Sep 17 00:00:00 2001 From: Abhishek Agarwal <1477457+abhishekagarwal87@users.noreply.github.com> Date: Mon, 9 Oct 2023 20:45:10 +0530 Subject: [PATCH 3/3] Parse passwords containing colon correctly (#15109) --- .../BasicHTTPAuthenticator.java | 20 +++- .../BasicHTTPAuthenticatorTest.java | 109 ++++++++++-------- 2 files changed, 79 insertions(+), 50 deletions(-) diff --git a/extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authentication/BasicHTTPAuthenticator.java b/extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authentication/BasicHTTPAuthenticator.java index 600af931f031..85cc60d2e76b 100644 --- a/extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authentication/BasicHTTPAuthenticator.java +++ b/extensions-core/druid-basic-security/src/main/java/org/apache/druid/security/basic/authentication/BasicHTTPAuthenticator.java @@ -182,15 +182,27 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo return; } - String[] splits = decodedUserSecret.split(":"); - if (splits.length != 2) { + /* From https://www.rfc-editor.org/rfc/rfc7617.html, we can assume that userid won't include a colon but password + can. + + The user-id and password MUST NOT contain any control characters (see + "CTL" in Appendix B.1 of [RFC5234]). + + Furthermore, a user-id containing a colon character is invalid, as + the first colon in a user-pass string separates user-id and password + from one another; text after the first colon is part of the password. + User-ids containing colons cannot be encoded in user-pass strings. + + */ + int split = decodedUserSecret.indexOf(':'); + if (split < 0) { // The decoded user secret is not of the right format httpResp.sendError(HttpServletResponse.SC_UNAUTHORIZED); return; } - String user = splits[0]; - char[] password = splits[1].toCharArray(); + String user = decodedUserSecret.substring(0, split); + char[] password = decodedUserSecret.substring(split + 1).toCharArray(); // If any authentication error occurs we send a 401 response immediately and do not proceed further down the filter chain. // If the authentication result is null and skipOnFailure property is false, we send a 401 response and do not proceed diff --git a/extensions-core/druid-basic-security/src/test/java/org/apache/druid/security/authentication/BasicHTTPAuthenticatorTest.java b/extensions-core/druid-basic-security/src/test/java/org/apache/druid/security/authentication/BasicHTTPAuthenticatorTest.java index 84bfdcf56b1c..bf0cf1778a14 100644 --- a/extensions-core/druid-basic-security/src/test/java/org/apache/druid/security/authentication/BasicHTTPAuthenticatorTest.java +++ b/extensions-core/druid-basic-security/src/test/java/org/apache/druid/security/authentication/BasicHTTPAuthenticatorTest.java @@ -112,55 +112,21 @@ public void testGoodPassword() throws IOException, ServletException } @Test - public void testGoodPasswordWithValidator() throws IOException, ServletException + public void testGoodNonEmptyPasswordWithValidator() throws IOException, ServletException { - CredentialsValidator validator = EasyMock.createMock(CredentialsValidator.class); - BasicHTTPAuthenticator authenticatorWithValidator = new BasicHTTPAuthenticator( - CACHE_MANAGER_PROVIDER, - "basic", - "basic", - null, - null, - false, - null, null, - false, - validator - ); - - String header = StringUtils.utf8Base64("userA:helloworld"); - header = StringUtils.format("Basic %s", header); - - EasyMock - .expect( - validator.validateCredentials(EasyMock.eq("basic"), EasyMock.eq("basic"), EasyMock.eq("userA"), EasyMock.aryEq("helloworld".toCharArray())) - ) - .andReturn( - new AuthenticationResult("userA", "basic", "basic", null) - ) - .times(1); - EasyMock.replay(validator); - - HttpServletRequest req = EasyMock.createMock(HttpServletRequest.class); - EasyMock.expect(req.getHeader("Authorization")).andReturn(header); - req.setAttribute( - AuthConfig.DRUID_AUTHENTICATION_RESULT, - new AuthenticationResult("userA", "basic", "basic", null) - ); - EasyMock.expectLastCall().times(1); - EasyMock.replay(req); - - HttpServletResponse resp = EasyMock.createMock(HttpServletResponse.class); - EasyMock.replay(resp); - - FilterChain filterChain = EasyMock.createMock(FilterChain.class); - filterChain.doFilter(req, resp); - EasyMock.expectLastCall().times(1); - EasyMock.replay(filterChain); + testGoodPasswordWithValidator("userA", "helloworld"); + } - Filter authenticatorFilter = authenticatorWithValidator.getFilter(); - authenticatorFilter.doFilter(req, resp, filterChain); + @Test + public void testGoodEmptyPasswordWithValidator() throws IOException, ServletException + { + testGoodPasswordWithValidator("userA", ""); + } - EasyMock.verify(req, resp, validator, filterChain); + @Test + public void testGoodColonInPasswordWithValidator() throws IOException, ServletException + { + testGoodPasswordWithValidator("userA", "hello:hello"); } @Test @@ -396,4 +362,55 @@ public void testMissingHeader() throws IOException, ServletException EasyMock.verify(req, resp, filterChain); } + + private void testGoodPasswordWithValidator(String username, String password) throws IOException, ServletException + { + CredentialsValidator validator = EasyMock.createMock(CredentialsValidator.class); + BasicHTTPAuthenticator authenticatorWithValidator = new BasicHTTPAuthenticator( + CACHE_MANAGER_PROVIDER, + "basic", + "basic", + null, + null, + false, + null, null, + false, + validator + ); + + String header = StringUtils.utf8Base64(username + ":" + password); + header = StringUtils.format("Basic %s", header); + + EasyMock + .expect( + validator.validateCredentials(EasyMock.eq("basic"), EasyMock.eq("basic"), EasyMock.eq(username), EasyMock.aryEq(password.toCharArray())) + ) + .andReturn( + new AuthenticationResult(username, "basic", "basic", null) + ) + .times(1); + EasyMock.replay(validator); + + HttpServletRequest req = EasyMock.createMock(HttpServletRequest.class); + EasyMock.expect(req.getHeader("Authorization")).andReturn(header); + req.setAttribute( + AuthConfig.DRUID_AUTHENTICATION_RESULT, + new AuthenticationResult(username, "basic", "basic", null) + ); + EasyMock.expectLastCall().times(1); + EasyMock.replay(req); + + HttpServletResponse resp = EasyMock.createMock(HttpServletResponse.class); + EasyMock.replay(resp); + + FilterChain filterChain = EasyMock.createMock(FilterChain.class); + filterChain.doFilter(req, resp); + EasyMock.expectLastCall().times(1); + EasyMock.replay(filterChain); + + Filter authenticatorFilter = authenticatorWithValidator.getFilter(); + authenticatorFilter.doFilter(req, resp, filterChain); + + EasyMock.verify(req, resp, validator, filterChain); + } }