diff --git a/query-service-impl/build.gradle.kts b/query-service-impl/build.gradle.kts index 7e73a015..2f5aa7ce 100644 --- a/query-service-impl/build.gradle.kts +++ b/query-service-impl/build.gradle.kts @@ -22,8 +22,8 @@ dependencies { implementation("org.apache.calcite:calcite-babel:1.34.0") { because("CVE-2022-39135") } - implementation("org.apache.avro:avro:1.11.3") { - because("CVE-2023-39410") + implementation("org.apache.avro:avro:1.11.4") { + because("CVE-2024-47561") } implementation("org.apache.commons:commons-compress:1.26.0") { because("CVE-2024-25710") diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/HandlerScopedMaskingConfig.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/HandlerScopedMaskingConfig.java new file mode 100644 index 00000000..bb4e3bfc --- /dev/null +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/HandlerScopedMaskingConfig.java @@ -0,0 +1,119 @@ +package org.hypertrace.core.query.service; + +import com.typesafe.config.Config; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; +import lombok.Value; +import lombok.experimental.NonFinal; +import lombok.extern.slf4j.Slf4j; + +@Slf4j +public class HandlerScopedMaskingConfig { + private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingConfig"; + private Map> tenantToTimeRangeMaskedAttributes = + Collections.emptyMap(); + + public HandlerScopedMaskingConfig(Config config) { + if (config.hasPath(TENANT_SCOPED_MASKS_CONFIG_KEY)) { + this.tenantToTimeRangeMaskedAttributes = + config.getConfigList(TENANT_SCOPED_MASKS_CONFIG_KEY).stream() + .map(TenantMaskingConfig::new) + .collect( + Collectors.toMap( + TenantMaskingConfig::getTenantId, + TenantMaskingConfig::getTimeRangeToMaskedAttributes)); + } + } + + public Set getMaskedAttributes(ExecutionContext executionContext) { + String tenantId = executionContext.getTenantId(); + HashSet maskedAttributes = new HashSet<>(); + if (!tenantToTimeRangeMaskedAttributes.containsKey(tenantId)) { + return maskedAttributes; + } + + Optional queryTimeRange = executionContext.getQueryTimeRange(); + Instant queryStartTime = Instant.MIN, queryEndTime = Instant.MAX; + if (queryTimeRange.isPresent()) { + queryStartTime = queryTimeRange.get().getStartTime(); + queryEndTime = queryTimeRange.get().getEndTime(); + } + for (TimeRangeToMaskedAttributes timeRangeAndMasks : + tenantToTimeRangeMaskedAttributes.get(tenantId)) { + if (isTimeRangeOverlap(timeRangeAndMasks, queryStartTime, queryEndTime)) { + maskedAttributes.addAll(timeRangeAndMasks.maskedAttributes); + } + } + return maskedAttributes; + } + + private static boolean isTimeRangeOverlap( + TimeRangeToMaskedAttributes timeRangeAndMasks, Instant queryStartTime, Instant queryEndTime) { + return !(timeRangeAndMasks.startTimeMillis.isAfter(queryEndTime) + || timeRangeAndMasks.endTimeMillis.isBefore(queryStartTime)); + } + + @Value + @NonFinal + static class TenantMaskingConfig { + private static final String TENANT_ID_CONFIG_KEY = "tenantId"; + private static final String TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY = + "timeRangeToMaskedAttributes"; + String tenantId; + List timeRangeToMaskedAttributes; + + private TenantMaskingConfig(Config config) { + this.tenantId = config.getString(TENANT_ID_CONFIG_KEY); + this.timeRangeToMaskedAttributes = + config.getConfigList(TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY).stream() + .map(TimeRangeToMaskedAttributes::new) + .filter( + timeRangeToMaskedAttributes -> { + if (!timeRangeToMaskedAttributes.isValid()) { + log.warn( + "Invalid masking configuration for tenant: {}. Either the time range is missing or the mask list is empty.", + this.tenantId); + return false; + } + return true; + }) + .collect(Collectors.toList()); + } + } + + @NonFinal + static class TimeRangeToMaskedAttributes { + private static final String START_TIME_CONFIG_PATH = "startTimeMillis"; + private static final String END_TIME_CONFIG_PATH = "endTimeMillis"; + private static final String MASK_ATTRIBUTES_CONFIG_PATH = "maskedAttributes"; + Instant startTimeMillis = null; + Instant endTimeMillis = null; + ArrayList maskedAttributes = new ArrayList<>(); + + private TimeRangeToMaskedAttributes(Config config) { + if (config.hasPath(START_TIME_CONFIG_PATH) && config.hasPath(END_TIME_CONFIG_PATH)) { + Instant startTimeMillis = Instant.ofEpochMilli(config.getLong(START_TIME_CONFIG_PATH)); + Instant endTimeMillis = Instant.ofEpochMilli(config.getLong(END_TIME_CONFIG_PATH)); + + if (startTimeMillis.isBefore(endTimeMillis)) { + this.startTimeMillis = startTimeMillis; + this.endTimeMillis = endTimeMillis; + if (config.hasPath(MASK_ATTRIBUTES_CONFIG_PATH)) { + maskedAttributes = new ArrayList<>(config.getStringList(MASK_ATTRIBUTES_CONFIG_PATH)); + } + } + } + } + + boolean isValid() { + return startTimeMillis != null && endTimeMillis != null && !maskedAttributes.isEmpty(); + } + } +} diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java index a9f55115..8d07b725 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandler.java @@ -28,6 +28,7 @@ import org.apache.pinot.client.ResultSetGroup; import org.hypertrace.core.query.service.ExecutionContext; import org.hypertrace.core.query.service.HandlerScopedFiltersConfig; +import org.hypertrace.core.query.service.HandlerScopedMaskingConfig; import org.hypertrace.core.query.service.QueryCost; import org.hypertrace.core.query.service.RequestHandler; import org.hypertrace.core.query.service.api.Expression; @@ -58,6 +59,10 @@ public class PinotBasedRequestHandler implements RequestHandler { private static final String START_TIME_ATTRIBUTE_NAME_CONFIG_KEY = "startTimeAttributeName"; private static final String SLOW_QUERY_THRESHOLD_MS_CONFIG = "slowQueryThresholdMs"; + private static final String DEFAULT_MASKED_VALUE = "*"; + // This is how empty list is represented in Pinot + private static final String ARRAY_TYPE_MASKED_VALUE = "[\"\"]"; + private static final int DEFAULT_SLOW_QUERY_THRESHOLD_MS = 3000; private static final Set GTE_OPERATORS = Set.of(Operator.GE, Operator.GT, Operator.EQ); @@ -67,6 +72,7 @@ public class PinotBasedRequestHandler implements RequestHandler { private QueryRequestToPinotSQLConverter request2PinotSqlConverter; private final PinotMapConverter pinotMapConverter; private HandlerScopedFiltersConfig handlerScopedFiltersConfig; + private HandlerScopedMaskingConfig handlerScopedMaskingConfig; // The implementations of ResultSet are package private and hence there's no way to determine the // shape of the results // other than to do string comparison on the simple class names. In order to be able to unit test @@ -143,6 +149,7 @@ private void processConfig(Config config) { this.handlerScopedFiltersConfig = new HandlerScopedFiltersConfig(config, this.startTimeAttributeName); + this.handlerScopedMaskingConfig = new HandlerScopedMaskingConfig(config); LOG.info( "Using {}ms as the threshold for logging slow queries of handler: {}", slowQueryThreshold, @@ -424,7 +431,7 @@ public Observable handleRequest( LOG.debug("Query results: [ {} ]", resultSetGroup.toString()); } // need to merge data especially for Pinot. That's why we need to track the map columns - return this.convert(resultSetGroup, executionContext.getSelectedColumns()) + return this.convert(resultSetGroup, executionContext) .doOnComplete( () -> { long requestTimeMs = stopwatch.stop().elapsed(TimeUnit.MILLISECONDS); @@ -493,17 +500,21 @@ private Filter rewriteLeafFilter( return queryFilter; } - Observable convert(ResultSetGroup resultSetGroup, LinkedHashSet selectedAttributes) { + Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executionContext) { List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { + LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); + Set maskedAttributes = + handlerScopedMaskingConfig.getMaskedAttributes(executionContext); ResultSet resultSet = resultSetGroup.getResultSet(0); // Pinot has different Response format for selection and aggregation/group by query. if (resultSetTypePredicateProvider.isSelectionResultSetType(resultSet)) { // map merging is only supported in the selection. Filtering and Group by has its own // syntax in Pinot - handleSelection(resultSetGroup, rowBuilderList, selectedAttributes); + handleSelection(resultSetGroup, rowBuilderList, selectedAttributes, maskedAttributes); } else if (resultSetTypePredicateProvider.isResultTableResultSetType(resultSet)) { - handleTableFormatResultSet(resultSetGroup, rowBuilderList); + handleTableFormatResultSet( + resultSetGroup, rowBuilderList, selectedAttributes, maskedAttributes); } else { handleAggregationAndGroupBy(resultSetGroup, rowBuilderList); } @@ -516,7 +527,8 @@ Observable convert(ResultSetGroup resultSetGroup, LinkedHashSet sel private void handleSelection( ResultSetGroup resultSetGroup, List rowBuilderList, - LinkedHashSet selectedAttributes) { + LinkedHashSet selectedAttributes, + Set maskedAttributes) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); @@ -536,7 +548,11 @@ private void handleSelection( for (String logicalName : selectedAttributes) { // colVal will never be null. But getDataRow can throw a runtime exception if it failed // to retrieve data - String colVal = resultAnalyzer.getDataFromRow(rowId, logicalName); + String colVal = + maskedAttributes.contains(logicalName) + ? DEFAULT_MASKED_VALUE + : resultAnalyzer.getDataFromRow(rowId, logicalName); + builder.addColumn(Value.newBuilder().setString(colVal).build()); } } @@ -588,10 +604,15 @@ private void handleAggregationAndGroupBy( } private void handleTableFormatResultSet( - ResultSetGroup resultSetGroup, List rowBuilderList) { + ResultSetGroup resultSetGroup, + List rowBuilderList, + LinkedHashSet selectedAttributes, + Set maskedAttributes) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); + PinotResultAnalyzer resultAnalyzer = + PinotResultAnalyzer.create(resultSet, selectedAttributes, viewDefinition); for (int rowIdx = 0; rowIdx < resultSet.getRowCount(); rowIdx++) { Builder builder; builder = Row.newBuilder(); @@ -602,8 +623,13 @@ private void handleTableFormatResultSet( // Read the key and value column values. The columns should be side by side. That's how // the Pinot query // is structured + String logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); String mapKeys = resultSet.getString(rowIdx, colIdx); - String mapVals = resultSet.getString(rowIdx, colIdx + 1); + String mapVals = + maskedAttributes.contains(logicalName) + ? ARRAY_TYPE_MASKED_VALUE + : resultSet.getString(rowIdx, colIdx + 1); + try { builder.addColumn( Value.newBuilder().setString(pinotMapConverter.merge(mapKeys, mapVals)).build()); @@ -615,7 +641,11 @@ private void handleTableFormatResultSet( // advance colIdx by 1 since we have read 2 columns colIdx++; } else { - String val = resultSet.getString(rowIdx, colIdx); + String logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + String val = + maskedAttributes.contains(logicalName) + ? DEFAULT_MASKED_VALUE + : resultSet.getString(rowIdx, colIdx); builder.addColumn(Value.newBuilder().setString(val).build()); } } diff --git a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotResultAnalyzer.java b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotResultAnalyzer.java index 62c91131..6225562f 100644 --- a/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotResultAnalyzer.java +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/pinot/PinotResultAnalyzer.java @@ -27,6 +27,7 @@ class PinotResultAnalyzer { private final ViewDefinition viewDefinition; private final Map attributeLogRateLimitter; private final PinotMapConverter pinotMapConverter; + private final Map indexToLogicalName; PinotResultAnalyzer( ResultSet resultSet, @@ -34,10 +35,12 @@ class PinotResultAnalyzer { ViewDefinition viewDefinition, Map mapLogicalNameToKeyIndex, Map mapLogicalNameToValueIndex, - Map logicalNameToPhysicalNameIndex) { + Map logicalNameToPhysicalNameIndex, + Map indexToLogicalName) { this.mapLogicalNameToKeyIndex = mapLogicalNameToKeyIndex; this.mapLogicalNameToValueIndex = mapLogicalNameToValueIndex; this.logicalNameToPhysicalNameIndex = logicalNameToPhysicalNameIndex; + this.indexToLogicalName = indexToLogicalName; this.resultSet = resultSet; this.viewDefinition = viewDefinition; this.attributeLogRateLimitter = new HashMap<>(); @@ -53,6 +56,7 @@ static PinotResultAnalyzer create( Map mapLogicalNameToKeyIndex = new HashMap<>(); Map mapLogicalNameToValueIndex = new HashMap<>(); Map logicalNameToPhysicalNameIndex = new HashMap<>(); + Map indexToLogicalName = new HashMap<>(); for (String logicalName : selectedAttributes) { if (viewDefinition.isMap(logicalName)) { @@ -62,8 +66,10 @@ static PinotResultAnalyzer create( String physName = resultSet.getColumnName(colIndex); if (physName.equalsIgnoreCase(keyPhysicalName)) { mapLogicalNameToKeyIndex.put(logicalName, colIndex); + indexToLogicalName.put(colIndex, logicalName); } else if (physName.equalsIgnoreCase(valuePhysicalName)) { mapLogicalNameToValueIndex.put(logicalName, colIndex); + indexToLogicalName.put(colIndex, logicalName); } } } else { @@ -73,21 +79,24 @@ static PinotResultAnalyzer create( String physName = resultSet.getColumnName(colIndex); if (physName.equalsIgnoreCase(names.get(0))) { logicalNameToPhysicalNameIndex.put(logicalName, colIndex); + indexToLogicalName.put(colIndex, logicalName); break; } } } } - LOG.info("Map LogicalName to Key Index: {} ", mapLogicalNameToKeyIndex); - LOG.info("Map LogicalName to Value Index: {}", mapLogicalNameToValueIndex); - LOG.info("Attributes to Index: {}", logicalNameToPhysicalNameIndex); + LOG.debug("Map LogicalName to Key Index: {} ", mapLogicalNameToKeyIndex); + LOG.debug("Map LogicalName to Value Index: {}", mapLogicalNameToValueIndex); + LOG.debug("Attributes to Index: {}", logicalNameToPhysicalNameIndex); + LOG.debug("Index to LogicalName: {}", indexToLogicalName); return new PinotResultAnalyzer( resultSet, selectedAttributes, viewDefinition, mapLogicalNameToKeyIndex, mapLogicalNameToValueIndex, - logicalNameToPhysicalNameIndex); + logicalNameToPhysicalNameIndex, + indexToLogicalName); } @VisibleForTesting @@ -149,4 +158,8 @@ String getDataFromRow(int rowIndex, String logicalName) { } return result; } + + String getLogicalNameFromColIdx(Integer colIdx) { + return indexToLogicalName.get(colIdx); + } } diff --git a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java index af1296f7..09997626 100644 --- a/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java +++ b/query-service-impl/src/test/java/org/hypertrace/core/query/service/pinot/PinotBasedRequestHandlerTest.java @@ -18,7 +18,8 @@ import com.typesafe.config.ConfigFactory; import io.reactivex.rxjava3.core.Observable; import java.io.IOException; -import java.util.LinkedHashSet; +import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Objects; @@ -1353,7 +1354,9 @@ public void testConvertSimpleSelectionsQueryResultSet() throws IOException { ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), resultTable); + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), + resultTable); } @Test @@ -1371,7 +1374,9 @@ public void testConvertAggregationColumnsQueryResultSet() throws IOException { ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), resultTable); + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), + resultTable); } @Test @@ -1432,7 +1437,9 @@ public void testConvertSelectionsWithMapKeysAndValuesQueryResultSet() throws IOE }; verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), expectedRows); + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), + expectedRows); } @Test @@ -1467,7 +1474,9 @@ public void testConvertMultipleResultSetsInFResultSetGroup() throws IOException }; verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), expectedRows); + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), + expectedRows); } @Test @@ -1756,6 +1765,426 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { } } + @Test + public void testMaskColumnValueSelection() throws IOException { + for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { + if (!isPinotConfig(config)) { + continue; + } + + if (!config.getString("name").equals("span-event-view-handler")) { + continue; + } + + // Mock the PinotClient + PinotClient pinotClient = mock(PinotClient.class); + PinotClientFactory factory = mock(PinotClientFactory.class); + when(factory.getPinotClient(any())).thenReturn(pinotClient); + + String[][] resultTable = + new String[][] { + {"test-span-id-1", "trace-id-1"}, + {"test-span-id-2", "trace-id-1"}, + {"test-span-id-3", "trace-id-1"}, + {"test-span-id-4", "trace-id-2"} + }; + List columnNames = List.of("span_id", "trace_id"); + ResultSet resultSet = mockResultSet(4, 2, columnNames, resultTable); + ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); + + PinotBasedRequestHandler handler = + new PinotBasedRequestHandler( + config.getString("name"), + config.getConfig("requestHandlerInfo"), + new ResultSetTypePredicateProvider() { + @Override + public boolean isSelectionResultSetType(ResultSet resultSet) { + return true; + } + + @Override + public boolean isResultTableResultSetType(ResultSet resultSet) { + return false; + } + }, + factory); + + QueryRequest request = + QueryRequest.newBuilder() + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.id")) + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(99))) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.LT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(300)))) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + String expectedQuery = + "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND ( start_time_millis > ? AND start_time_millis < ? )"; + Params params = + Params.newBuilder() + .addStringParam("maskTenant") + .addLongParam(99) + .addLongParam(300) + .build(); + when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); + + String[][] expectedTable = + new String[][] { + {"*", "trace-id-1"}, + {"*", "trace-id-1"}, + {"*", "trace-id-1"}, + {"*", "trace-id-2"} + }; + + verifyResponseRows(handler.handleRequest(request, context), expectedTable); + } + } + + @Test + public void testMaskColumnValueTableFormat() throws IOException { + for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { + if (!isPinotConfig(config)) { + continue; + } + + if (!config.getString("name").equals("span-event-view-handler")) { + continue; + } + + // Mock the PinotClient + PinotClient pinotClient = mock(PinotClient.class); + PinotClientFactory factory = mock(PinotClientFactory.class); + when(factory.getPinotClient(any())).thenReturn(pinotClient); + + String[][] resultTable = + new String[][] { + {"test-span-id-1", "trace-id-1"}, + {"test-span-id-2", "trace-id-1"}, + {"test-span-id-3", "trace-id-1"}, + {"test-span-id-4", "trace-id-2"} + }; + List columnNames = List.of("span_id", "trace_id"); + ResultSet resultSet = mockResultSet(4, 2, columnNames, resultTable); + ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); + + PinotBasedRequestHandler handler = + new PinotBasedRequestHandler( + config.getString("name"), + config.getConfig("requestHandlerInfo"), + new ResultSetTypePredicateProvider() { + @Override + public boolean isSelectionResultSetType(ResultSet resultSet) { + return false; + } + + @Override + public boolean isResultTableResultSetType(ResultSet resultSet) { + return true; + } + }, + factory); + + QueryRequest request = + QueryRequest.newBuilder() + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.id")) + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.traceId")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(99))) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.LT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(300)))) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + String expectedQuery = + "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND ( start_time_millis > ? AND start_time_millis < ? )"; + Params params = + Params.newBuilder() + .addStringParam("maskTenant") + .addLongParam(99) + .addLongParam(300) + .build(); + when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); + + String[][] expectedTable = + new String[][] { + {"*", "trace-id-1"}, + {"*", "trace-id-1"}, + {"*", "trace-id-1"}, + {"*", "trace-id-2"} + }; + + verifyResponseRows(handler.handleRequest(request, context), expectedTable); + } + } + + @Test + public void testMaskMapColumnValue() throws IOException { + for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { + if (!isPinotConfig(config)) { + continue; + } + + if (!config.getString("name").equals("span-event-view-handler")) { + continue; + } + + // Mock the PinotClient + PinotClient pinotClient = mock(PinotClient.class); + PinotClientFactory factory = mock(PinotClientFactory.class); + when(factory.getPinotClient(any())).thenReturn(pinotClient); + + String[][] resultTable = + new String[][] { + { + stringify(new ArrayList<>(Arrays.asList("k1", "k2", "k3"))), + stringify(new ArrayList<>(Arrays.asList("v1", "v2", "v3"))) + }, + }; + List columnNames = + List.of( + "tags" + ViewDefinition.MAP_KEYS_SUFFIX, "tags" + ViewDefinition.MAP_VALUES_SUFFIX); + ResultSet resultSet = mockResultSet(1, 2, columnNames, resultTable); + ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); + + PinotBasedRequestHandler handler = + new PinotBasedRequestHandler( + config.getString("name"), + config.getConfig("requestHandlerInfo"), + new ResultSetTypePredicateProvider() { + @Override + public boolean isSelectionResultSetType(ResultSet resultSet) { + return false; + } + + @Override + public boolean isResultTableResultSetType(ResultSet resultSet) { + return true; + } + }, + factory); + + QueryRequest request = + QueryRequest.newBuilder() + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.spanTags")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(99))) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.LT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(300)))) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + String expectedQuery = + "Select tags__KEYS, tags__VALUES FROM spanEventView WHERE tenant_id = ? AND ( start_time_millis > ? AND start_time_millis < ? )"; + Params params = + Params.newBuilder() + .addStringParam("maskTenant") + .addLongParam(99) + .addLongParam(300) + .build(); + when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); + + String[][] expectedTable = + new String[][] { + { + "{\"k1\":\"\",\"k2\":\"\",\"k3\":\"\"}", + }, + }; + + verifyResponseRows(handler.handleRequest(request, context), expectedTable); + } + } + + @Test + public void testNoMaskColumnValueTenantMismatch() throws IOException { + for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { + if (!isPinotConfig(config)) { + continue; + } + + if (!config.getString("name").equals("span-event-view-handler")) { + continue; + } + + // Mock the PinotClient + PinotClient pinotClient = mock(PinotClient.class); + PinotClientFactory factory = mock(PinotClientFactory.class); + when(factory.getPinotClient(any())).thenReturn(pinotClient); + + String[][] resultTable = + new String[][] { + { + stringify(new ArrayList<>(Arrays.asList("k1", "k2", "k3"))), + stringify(new ArrayList<>(Arrays.asList("v1", "v2", "v3"))) + }, + }; + List columnNames = + List.of( + "tags" + ViewDefinition.MAP_KEYS_SUFFIX, "tags" + ViewDefinition.MAP_VALUES_SUFFIX); + ResultSet resultSet = mockResultSet(1, 2, columnNames, resultTable); + ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); + + PinotBasedRequestHandler handler = + new PinotBasedRequestHandler( + config.getString("name"), + config.getConfig("requestHandlerInfo"), + new ResultSetTypePredicateProvider() { + @Override + public boolean isSelectionResultSetType(ResultSet resultSet) { + return false; + } + + @Override + public boolean isResultTableResultSetType(ResultSet resultSet) { + return true; + } + }, + factory); + + QueryRequest request = + QueryRequest.newBuilder() + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.spanTags")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(99)))) + .build(); + ExecutionContext context = new ExecutionContext("noMaskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + String expectedQuery = + "Select tags__KEYS, tags__VALUES FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; + Params params = Params.newBuilder().addStringParam("noMaskTenant").addLongParam(99).build(); + when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); + + String[][] expectedTable = + new String[][] { + { + "{\"k1\":\"v1\",\"k2\":\"v2\",\"k3\":\"v3\"}", + }, + }; + + verifyResponseRows(handler.handleRequest(request, context), expectedTable); + } + } + + @Test + public void testNoTimerangeOverlapForMasking() throws IOException { + for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { + if (!isPinotConfig(config)) { + continue; + } + + if (!config.getString("name").equals("span-event-view-handler")) { + continue; + } + + // Mock the PinotClient + PinotClient pinotClient = mock(PinotClient.class); + PinotClientFactory factory = mock(PinotClientFactory.class); + when(factory.getPinotClient(any())).thenReturn(pinotClient); + + String[][] resultTable = + new String[][] { + { + stringify(new ArrayList<>(Arrays.asList("k1", "k2", "k3"))), + stringify(new ArrayList<>(Arrays.asList("v1", "v2", "v3"))) + }, + }; + List columnNames = + List.of( + "tags" + ViewDefinition.MAP_KEYS_SUFFIX, "tags" + ViewDefinition.MAP_VALUES_SUFFIX); + ResultSet resultSet = mockResultSet(1, 2, columnNames, resultTable); + ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); + + PinotBasedRequestHandler handler = + new PinotBasedRequestHandler( + config.getString("name"), + config.getConfig("requestHandlerInfo"), + new ResultSetTypePredicateProvider() { + @Override + public boolean isSelectionResultSetType(ResultSet resultSet) { + return false; + } + + @Override + public boolean isResultTableResultSetType(ResultSet resultSet) { + return true; + } + }, + factory); + + QueryRequest request = + QueryRequest.newBuilder() + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.spanTags")) + .setFilter( + Filter.newBuilder() + .setOperator(Operator.AND) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.GT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(1000))) + .addChildFilter( + QueryRequestBuilderUtils.createFilter( + "EVENT.startTime", + Operator.LT, + QueryRequestBuilderUtils.createLongLiteralValueExpression(2000)))) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + String expectedQuery = + "Select tags__KEYS, tags__VALUES FROM spanEventView WHERE tenant_id = ? AND ( start_time_millis > ? AND start_time_millis < ? )"; + Params params = + Params.newBuilder() + .addStringParam("maskTenant") + .addLongParam(1000) + .addLongParam(2000) + .build(); + when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); + + String[][] expectedTable = + new String[][] { + { + "{\"k1\":\"v1\",\"k2\":\"v2\",\"k3\":\"v3\"}", + }, + }; + + verifyResponseRows(handler.handleRequest(request, context), expectedTable); + } + } + @Test public void testViewColumnFilterRemovalComplexCase() throws IOException { for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index 5b954ec1..832108c4 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -91,6 +91,26 @@ service.config = { ] } ] + tenantScopedMaskingConfig = [ + { + tenantId = "maskTenant" + timeRangeToMaskedAttributes = [ + { + startTimeMillis = 0 + endTimeMillis = 500 + maskedAttributes = [ + "EVENT.id", + "EVENT.spanTags" + ] + }, + # Invalid mask added to test warning + { + maskedAttributes = [ + ] + } + ] + } + ] viewDefinition = { viewName = spanEventView mapFields = ["tags"] diff --git a/query-service/build.gradle.kts b/query-service/build.gradle.kts index ebf991c4..dce842b1 100644 --- a/query-service/build.gradle.kts +++ b/query-service/build.gradle.kts @@ -27,7 +27,7 @@ dependencies { integrationTestImplementation("org.apache.kafka:kafka-clients:7.2.1-ccs") integrationTestImplementation("org.apache.kafka:kafka-streams:7.2.1-ccs") - integrationTestImplementation("org.apache.avro:avro:1.11.3") + integrationTestImplementation("org.apache.avro:avro:1.11.4") integrationTestImplementation("com.google.guava:guava:32.1.2-jre") integrationTestImplementation("org.hypertrace.core.datamodel:data-model:0.1.12") integrationTestImplementation("org.hypertrace.core.kafkastreams.framework:kafka-streams-serdes:0.3.2") diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index 35befa57..e797d039 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -93,6 +93,23 @@ service.config = { # } # ] # } +# ] +# tenantScopedMaskingConfig = [ +# { +# "tenantId": "testTenant", +# "timeRangeToMaskedAttributes": [ +# { +# #Start and end time filter must always be provided +# "startTimeMillis": 0, +# "endTimeMillis": 1000, +# "maskAttributes": [ +# # Masking is only supported for attributes which have string data +# "attribute_1", +# "attribute_2", +# ] +# }, +# ] +# } # ] viewDefinition = { viewName = spanEventView