From 46b1921ef58e598e8e5c27fc92e8e3e05b4b54cb Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Fri, 27 Sep 2024 13:23:27 +0530 Subject: [PATCH 01/26] create masking config --- .../service/HandlerScopedMaskingConfig.java | 170 ++++++++++++++++++ .../pinot/PinotBasedRequestHandler.java | 11 ++ .../src/test/resources/application.conf | 21 +++ .../resources/configs/common/application.conf | 21 +++ 4 files changed, 223 insertions(+) create mode 100644 query-service-impl/src/main/java/org/hypertrace/core/query/service/HandlerScopedMaskingConfig.java 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..9ad9fabd --- /dev/null +++ b/query-service-impl/src/main/java/org/hypertrace/core/query/service/HandlerScopedMaskingConfig.java @@ -0,0 +1,170 @@ +package org.hypertrace.core.query.service; + +import com.typesafe.config.Config; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.stream.Collectors; +import lombok.Value; +import lombok.experimental.NonFinal; +import lombok.extern.slf4j.Slf4j; +import org.apache.pinot.client.ResultSet; +import org.hypertrace.core.query.service.api.Row; + +@Slf4j +public class HandlerScopedMaskingConfig { + private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; + private final Map> tenantToMaskValuesMap; + private final String tenantColumnName; + private final String timeFilterColumnName; + private int tenantColumnIndex; + private int timeFilterColumnIndex; + private final HashMap columnNameToIndexMap = new HashMap<>(); + + List getMaskValues(String tenantId) { + if (tenantToMaskValuesMap.containsKey(tenantId)) return tenantToMaskValuesMap.get(tenantId); + + return new ArrayList<>(); + } + + public HandlerScopedMaskingConfig( + Config config, Optional timeFilterColumnName, String tenantColumnName) { + if (config.hasPath(TENANT_SCOPED_MASKS_CONFIG_KEY)) { + this.tenantToMaskValuesMap = + config.getConfigList(TENANT_SCOPED_MASKS_CONFIG_KEY).stream() + .map(maskConfig -> new TenantMasks(maskConfig, timeFilterColumnName)) + .collect( + Collectors.toMap( + tenantFilters -> tenantFilters.tenantId, + tenantFilters -> tenantFilters.maskValues)); + } else { + this.tenantToMaskValuesMap = Collections.emptyMap(); + } + + this.tenantColumnName = tenantColumnName; + this.timeFilterColumnName = timeFilterColumnName.orElse(null); + } + + public void parseColumns(ResultSet resultSet) { + timeFilterColumnIndex = -1; + tenantColumnIndex = -1; + columnNameToIndexMap.clear(); + + for (int colIdx = 0; colIdx < resultSet.getColumnCount(); colIdx++) { + if (Objects.equals(this.tenantColumnName, resultSet.getColumnName(colIdx))) { + this.tenantColumnIndex = colIdx; + } else if (Objects.equals(this.timeFilterColumnName, resultSet.getColumnName(colIdx))) { + this.timeFilterColumnIndex = colIdx; + } + + columnNameToIndexMap.put(resultSet.getColumnName(colIdx), colIdx); + } + } + + public Row mask(Row row) { + if (this.tenantColumnIndex == -1 || this.timeFilterColumnIndex == -1) { + return row; + } + List masks = + getMaskValues(row.getColumn(this.tenantColumnIndex).getString()); + + Row.Builder maskedRowBuilder = Row.newBuilder(row); + + for (MaskValuesForTimeRange mask : masks) { + boolean toBeMasked = true; + if (mask.getEndTimeMillis().isPresent()) { + if (mask.getEndTimeMillis().get() < row.getColumn(this.timeFilterColumnIndex).getLong()) { + toBeMasked = false; + } + } + if (mask.getStartTimeMillis().isPresent()) { + if (mask.getStartTimeMillis().get() > row.getColumn(this.timeFilterColumnIndex).getLong()) { + toBeMasked = false; + } + } + + if (toBeMasked) { + for (String columnName : mask.maskValues.getColumnToMaskedValue().keySet()) { + int colIdx = columnNameToIndexMap.get(columnName); + org.hypertrace.core.query.service.api.Value value = + org.hypertrace.core.query.service.api.Value.newBuilder() + .setString(mask.maskValues.getColumnToMaskedValue().get(columnName)) + .build(); + maskedRowBuilder.setColumn(colIdx, value); + } + } + } + + return maskedRowBuilder.build(); + } + + @Value + @NonFinal + private class TenantMasks { + private static final String TENANT_ID_CONFIG_KEY = "tenantId"; + private static final String TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY = "timeRangeAndMaskValues"; + String tenantId; + String startTimeAttributeName; + List maskValues; + + private TenantMasks(Config config, Optional startTimeAttributeName) { + this.tenantId = config.getString(TENANT_ID_CONFIG_KEY); + this.startTimeAttributeName = startTimeAttributeName.orElse(null); + this.maskValues = + config.getConfigList(TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY).stream() + .map(MaskValuesForTimeRange::new) + .collect(Collectors.toList()); + } + } + + @Value + private class MaskValues { + Map columnToMaskedValue; + + MaskValues(Map columnToMaskedValue) { + this.columnToMaskedValue = columnToMaskedValue; + } + } + + @Value + @NonFinal + class MaskValuesForTimeRange { + private static final String START_TIME_CONFIG_PATH = "startTimeMillis"; + private static final String END_TIME_CONFIG_PATH = "endTimeMillis"; + private static final String MASK_VALUE_CONFIG_PATH = "maskValues"; + private static final String ATTRIBUTE_ID_CONFIG_PATH = "attributeId"; + private static final String MASKED_VALUE_CONFIG_PATH = "maskedValue"; + Optional startTimeMillis; + Optional endTimeMillis; + MaskValues maskValues; + + private MaskValuesForTimeRange(Config config) { + if (config.hasPath(START_TIME_CONFIG_PATH) && config.hasPath(END_TIME_CONFIG_PATH)) { + this.startTimeMillis = Optional.of(config.getLong(START_TIME_CONFIG_PATH)); + this.endTimeMillis = Optional.of(config.getLong(END_TIME_CONFIG_PATH)); + } else { + startTimeMillis = Optional.empty(); + endTimeMillis = Optional.empty(); + } + if (config.hasPath(MASK_VALUE_CONFIG_PATH)) { + List maskedValuesList = + new ArrayList<>(config.getConfigList(MASK_VALUE_CONFIG_PATH)); + HashMap maskedValuesMap = new HashMap<>(); + maskedValuesList.forEach( + maskedValue -> { + maskedValuesMap.put( + maskedValue.getString(ATTRIBUTE_ID_CONFIG_PATH), + maskedValue.getString(MASKED_VALUE_CONFIG_PATH)); + }); + + maskValues = new MaskValues(maskedValuesMap); + } else { + maskValues = new MaskValues(new HashMap<>()); + } + } + } +} 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..5d8f241a 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; @@ -67,6 +68,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 +145,8 @@ private void processConfig(Config config) { this.handlerScopedFiltersConfig = new HandlerScopedFiltersConfig(config, this.startTimeAttributeName); + this.handlerScopedMaskingConfig = + new HandlerScopedMaskingConfig(config, this.startTimeAttributeName, tenantColumnName); LOG.info( "Using {}ms as the threshold for logging slow queries of handler: {}", slowQueryThreshold, @@ -497,6 +501,7 @@ Observable convert(ResultSetGroup resultSetGroup, LinkedHashSet sel List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { ResultSet resultSet = resultSetGroup.getResultSet(0); + handlerScopedMaskingConfig.parseColumns(resultSet); // 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 @@ -508,8 +513,10 @@ Observable convert(ResultSetGroup resultSetGroup, LinkedHashSet sel handleAggregationAndGroupBy(resultSetGroup, rowBuilderList); } } + return Observable.fromIterable(rowBuilderList) .map(Builder::build) + .map(row -> handlerScopedMaskingConfig.mask(row)) .doOnNext(row -> LOG.debug("collect a row: {}", row)); } @@ -678,4 +685,8 @@ private boolean isInvalidExpression(Expression expression) { && viewDefinition.getColumnType(expression.getAttributeExpression().getAttributeId()) != ValueType.STRING_MAP; } + + HandlerScopedMaskingConfig getHandlerScopedMaskingConfig() { + return handlerScopedMaskingConfig; + } } diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index 5b954ec1..f05d40c5 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -91,6 +91,27 @@ service.config = { ] } ] + tenantScopedMaskingCriteria = [ + { + "tenantId": "testTenant", + "timeRangeAndMaskValues": [ + { + "startTimeMillis": 0, + "endTimeMillis": -1, + "maskValues": [ + { + "attributeId": "column_name1", + "maskedValue": "*" + }, + { + "attributeId": "column_name2", + "maskedValue": "*" + } + ] + }, + ] + } + ] viewDefinition = { viewName = spanEventView mapFields = ["tags"] diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index 35befa57..5233b3f5 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -94,6 +94,27 @@ service.config = { # ] # } # ] + tenantScopedMaskingCriteria = [ + { + "tenantId": "testTenant", + "timeRangeAndMaskValues": [ + { + "startTimeMillis": 0, + "endTimeMillis": -1, + "maskValues": [ + { + "attributeId": "column_name1", + "maskedValue": "*" + }, + { + "attributeId": "column_name2", + "maskedValue": "*" + } + ] + }, + ] + } + ] viewDefinition = { viewName = spanEventView mapFields = ["tags"] From 5e424ae8cfcc28a23a41417837a116a160c93b0c Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Mon, 30 Sep 2024 17:26:01 +0530 Subject: [PATCH 02/26] WIP: time range filter --- .../service/HandlerScopedMaskingConfig.java | 30 +++---- .../pinot/PinotBasedRequestHandler.java | 10 ++- .../pinot/PinotBasedRequestHandlerTest.java | 88 ++++++++++++++++++- .../src/test/resources/application.conf | 12 +-- 4 files changed, 109 insertions(+), 31 deletions(-) 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 index 9ad9fabd..76cd5076 100644 --- 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 @@ -19,9 +19,8 @@ public class HandlerScopedMaskingConfig { private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; private final Map> tenantToMaskValuesMap; - private final String tenantColumnName; + private String tenantId; private final String timeFilterColumnName; - private int tenantColumnIndex; private int timeFilterColumnIndex; private final HashMap columnNameToIndexMap = new HashMap<>(); @@ -32,7 +31,7 @@ List getMaskValues(String tenantId) { } public HandlerScopedMaskingConfig( - Config config, Optional timeFilterColumnName, String tenantColumnName) { + Config config, String timeFilterColumnName, String tenantColumnName) { if (config.hasPath(TENANT_SCOPED_MASKS_CONFIG_KEY)) { this.tenantToMaskValuesMap = config.getConfigList(TENANT_SCOPED_MASKS_CONFIG_KEY).stream() @@ -45,19 +44,17 @@ public HandlerScopedMaskingConfig( this.tenantToMaskValuesMap = Collections.emptyMap(); } - this.tenantColumnName = tenantColumnName; - this.timeFilterColumnName = timeFilterColumnName.orElse(null); + this.timeFilterColumnName = timeFilterColumnName; } - public void parseColumns(ResultSet resultSet) { + public void parseColumns(ResultSet resultSet, String tenantId) { timeFilterColumnIndex = -1; - tenantColumnIndex = -1; columnNameToIndexMap.clear(); + this.tenantId = tenantId; for (int colIdx = 0; colIdx < resultSet.getColumnCount(); colIdx++) { - if (Objects.equals(this.tenantColumnName, resultSet.getColumnName(colIdx))) { - this.tenantColumnIndex = colIdx; - } else if (Objects.equals(this.timeFilterColumnName, resultSet.getColumnName(colIdx))) { + String temp = resultSet.getColumnName(colIdx); + if (Objects.equals(this.timeFilterColumnName, resultSet.getColumnName(colIdx))) { this.timeFilterColumnIndex = colIdx; } @@ -66,11 +63,14 @@ public void parseColumns(ResultSet resultSet) { } public Row mask(Row row) { - if (this.tenantColumnIndex == -1 || this.timeFilterColumnIndex == -1) { + if (this.timeFilterColumnIndex == -1) { + return row; + } + + List masks = getMaskValues(this.tenantId); + if (masks.isEmpty()) { return row; } - List masks = - getMaskValues(row.getColumn(this.tenantColumnIndex).getString()); Row.Builder maskedRowBuilder = Row.newBuilder(row); @@ -111,9 +111,9 @@ private class TenantMasks { String startTimeAttributeName; List maskValues; - private TenantMasks(Config config, Optional startTimeAttributeName) { + private TenantMasks(Config config, String startTimeAttributeName) { this.tenantId = config.getString(TENANT_ID_CONFIG_KEY); - this.startTimeAttributeName = startTimeAttributeName.orElse(null); + this.startTimeAttributeName = startTimeAttributeName; this.maskValues = config.getConfigList(TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY).stream() .map(MaskValuesForTimeRange::new) 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 5d8f241a..72a1dbab 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 @@ -146,7 +146,7 @@ private void processConfig(Config config) { this.handlerScopedFiltersConfig = new HandlerScopedFiltersConfig(config, this.startTimeAttributeName); this.handlerScopedMaskingConfig = - new HandlerScopedMaskingConfig(config, this.startTimeAttributeName, tenantColumnName); + new HandlerScopedMaskingConfig(config, viewDefinition.getPhysicalColumnNames(this.startTimeAttributeName.orElse(null)).get(0), tenantColumnName); LOG.info( "Using {}ms as the threshold for logging slow queries of handler: {}", slowQueryThreshold, @@ -428,7 +428,8 @@ 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.getSelectedColumns(), executionContext.getTenantId()) .doOnComplete( () -> { long requestTimeMs = stopwatch.stop().elapsed(TimeUnit.MILLISECONDS); @@ -497,11 +498,12 @@ private Filter rewriteLeafFilter( return queryFilter; } - Observable convert(ResultSetGroup resultSetGroup, LinkedHashSet selectedAttributes) { + Observable convert( + ResultSetGroup resultSetGroup, LinkedHashSet selectedAttributes, String tenantId) { List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { ResultSet resultSet = resultSetGroup.getResultSet(0); - handlerScopedMaskingConfig.parseColumns(resultSet); + handlerScopedMaskingConfig.parseColumns(resultSet, tenantId); // 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 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..3b9302bf 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 @@ -1353,7 +1353,8 @@ public void testConvertSimpleSelectionsQueryResultSet() throws IOException { ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), resultTable); + pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + resultTable); } @Test @@ -1371,7 +1372,8 @@ public void testConvertAggregationColumnsQueryResultSet() throws IOException { ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), resultTable); + pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + resultTable); } @Test @@ -1432,7 +1434,8 @@ public void testConvertSelectionsWithMapKeysAndValuesQueryResultSet() throws IOE }; verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), expectedRows); + pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + expectedRows); } @Test @@ -1467,7 +1470,8 @@ public void testConvertMultipleResultSetsInFResultSetGroup() throws IOException }; verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>()), expectedRows); + pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + expectedRows); } @Test @@ -1756,6 +1760,81 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { } } + @Test + public void testMaskColumnValue() 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")) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + + // The query filter is based on both isEntrySpan and startTime. Since the viewFilter + // checks for both the true and false values of isEntrySpan and query filter only needs + // "true", isEntrySpan predicate is still passed to the store in the query. + String expectedQuery = + "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ?"; + Params params = + Params.newBuilder() + .addStringParam("maskTenant") + .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 testViewColumnFilterRemovalComplexCase() throws IOException { for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { @@ -2002,6 +2081,7 @@ private ResultSetGroup mockResultSetGroup(List resultSets) { private void verifyResponseRows(Observable rowObservable, String[][] expectedResultTable) throws IOException { + System.out.println(rowObservable); List rows = rowObservable.toList().blockingGet(); assertEquals(expectedResultTable.length, rows.size()); for (int rowIdx = 0; rowIdx < rows.size(); rowIdx++) { diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index f05d40c5..dca49e26 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -93,18 +93,14 @@ service.config = { ] tenantScopedMaskingCriteria = [ { - "tenantId": "testTenant", + "tenantId": "maskTenant", "timeRangeAndMaskValues": [ { - "startTimeMillis": 0, - "endTimeMillis": -1, + startTimeMillis = 1 + endTimeMillis = 100 "maskValues": [ { - "attributeId": "column_name1", - "maskedValue": "*" - }, - { - "attributeId": "column_name2", + "attributeId": "span_id", "maskedValue": "*" } ] From 4f178ae5ae96f3e18179ab95fdcc04bf26c478fd Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Mon, 30 Sep 2024 21:03:07 +0530 Subject: [PATCH 03/26] working --- .../service/HandlerScopedMaskingConfig.java | 116 ++++++++---------- .../pinot/PinotBasedRequestHandler.java | 21 ++-- .../pinot/PinotBasedRequestHandlerTest.java | 89 +++++++------- .../src/test/resources/application.conf | 4 +- .../resources/configs/common/application.conf | 34 ++--- 5 files changed, 128 insertions(+), 136 deletions(-) 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 index 76cd5076..73561a9e 100644 --- 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 @@ -1,41 +1,30 @@ 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.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Optional; import java.util.stream.Collectors; import lombok.Value; import lombok.experimental.NonFinal; import lombok.extern.slf4j.Slf4j; -import org.apache.pinot.client.ResultSet; -import org.hypertrace.core.query.service.api.Row; @Slf4j public class HandlerScopedMaskingConfig { private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; private final Map> tenantToMaskValuesMap; - private String tenantId; - private final String timeFilterColumnName; - private int timeFilterColumnIndex; - private final HashMap columnNameToIndexMap = new HashMap<>(); + private HashMap shouldMaskAttribute = new HashMap<>(); + private HashMap maskedValue = new HashMap<>(); - List getMaskValues(String tenantId) { - if (tenantToMaskValuesMap.containsKey(tenantId)) return tenantToMaskValuesMap.get(tenantId); - - return new ArrayList<>(); - } - - public HandlerScopedMaskingConfig( - Config config, String timeFilterColumnName, String tenantColumnName) { + public HandlerScopedMaskingConfig(Config config) { if (config.hasPath(TENANT_SCOPED_MASKS_CONFIG_KEY)) { this.tenantToMaskValuesMap = config.getConfigList(TENANT_SCOPED_MASKS_CONFIG_KEY).stream() - .map(maskConfig -> new TenantMasks(maskConfig, timeFilterColumnName)) + .map(maskConfig -> new TenantMasks(maskConfig)) .collect( Collectors.toMap( tenantFilters -> tenantFilters.tenantId, @@ -43,63 +32,64 @@ public HandlerScopedMaskingConfig( } else { this.tenantToMaskValuesMap = Collections.emptyMap(); } - - this.timeFilterColumnName = timeFilterColumnName; } - public void parseColumns(ResultSet resultSet, String tenantId) { - timeFilterColumnIndex = -1; - columnNameToIndexMap.clear(); - this.tenantId = tenantId; - - for (int colIdx = 0; colIdx < resultSet.getColumnCount(); colIdx++) { - String temp = resultSet.getColumnName(colIdx); - if (Objects.equals(this.timeFilterColumnName, resultSet.getColumnName(colIdx))) { - this.timeFilterColumnIndex = colIdx; - } + public void parseColumns(ExecutionContext executionContext) { + shouldMaskAttribute.clear(); + String tenantId = executionContext.getTenantId(); - columnNameToIndexMap.put(resultSet.getColumnName(colIdx), colIdx); + if (!tenantToMaskValuesMap.containsKey(tenantId)) { + return; } - } - public Row mask(Row row) { - if (this.timeFilterColumnIndex == -1) { - return row; + Optional queryTimeRange = executionContext.getQueryTimeRange(); + Instant queryStartTime, queryEndTime; + if (queryTimeRange.isPresent()) { + queryStartTime = queryTimeRange.get().getStartTime(); + queryEndTime = queryTimeRange.get().getEndTime(); + } else { + queryEndTime = Instant.MAX; + queryStartTime = Instant.MIN; } - - List masks = getMaskValues(this.tenantId); - if (masks.isEmpty()) { - return row; + for (MaskValuesForTimeRange timeRangeAndMasks : tenantToMaskValuesMap.get(tenantId)) { + boolean timeRangeOverlap = + isTimeRangeOverlap(timeRangeAndMasks, queryStartTime, queryEndTime); + + if (timeRangeOverlap) { + Map attributeToMaskedValue = + timeRangeAndMasks.maskValues.attributeToMaskedValue; + for (String attribute : attributeToMaskedValue.keySet()) { + shouldMaskAttribute.put(attribute, true); + maskedValue.put(attribute, attributeToMaskedValue.get(attribute)); + } + } } + } - Row.Builder maskedRowBuilder = Row.newBuilder(row); + private static boolean isTimeRangeOverlap( + MaskValuesForTimeRange timeRangeAndMasks, Instant queryStartTime, Instant queryEndTime) { + boolean timeRangeOverlap = true; - for (MaskValuesForTimeRange mask : masks) { - boolean toBeMasked = true; - if (mask.getEndTimeMillis().isPresent()) { - if (mask.getEndTimeMillis().get() < row.getColumn(this.timeFilterColumnIndex).getLong()) { - toBeMasked = false; - } - } - if (mask.getStartTimeMillis().isPresent()) { - if (mask.getStartTimeMillis().get() > row.getColumn(this.timeFilterColumnIndex).getLong()) { - toBeMasked = false; - } + if (timeRangeAndMasks.getStartTimeMillis().isPresent()) { + Instant startTimeInstant = Instant.ofEpochMilli(timeRangeAndMasks.getStartTimeMillis().get()); + if (startTimeInstant.isBefore(queryStartTime) || startTimeInstant.isAfter(queryEndTime)) { + timeRangeOverlap = false; } - if (toBeMasked) { - for (String columnName : mask.maskValues.getColumnToMaskedValue().keySet()) { - int colIdx = columnNameToIndexMap.get(columnName); - org.hypertrace.core.query.service.api.Value value = - org.hypertrace.core.query.service.api.Value.newBuilder() - .setString(mask.maskValues.getColumnToMaskedValue().get(columnName)) - .build(); - maskedRowBuilder.setColumn(colIdx, value); - } + Instant endTimeInstant = Instant.ofEpochMilli(timeRangeAndMasks.getStartTimeMillis().get()); + if (endTimeInstant.isBefore(queryStartTime) || endTimeInstant.isAfter(queryEndTime)) { + timeRangeOverlap = false; } } + return timeRangeOverlap; + } + + public boolean shouldMask(String attributeName) { + return this.maskedValue.containsKey(attributeName); + } - return maskedRowBuilder.build(); + public String getMaskedValue(String attributeName) { + return this.maskedValue.get(attributeName); } @Value @@ -108,12 +98,10 @@ private class TenantMasks { private static final String TENANT_ID_CONFIG_KEY = "tenantId"; private static final String TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY = "timeRangeAndMaskValues"; String tenantId; - String startTimeAttributeName; List maskValues; - private TenantMasks(Config config, String startTimeAttributeName) { + private TenantMasks(Config config) { this.tenantId = config.getString(TENANT_ID_CONFIG_KEY); - this.startTimeAttributeName = startTimeAttributeName; this.maskValues = config.getConfigList(TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY).stream() .map(MaskValuesForTimeRange::new) @@ -123,10 +111,10 @@ private TenantMasks(Config config, String startTimeAttributeName) { @Value private class MaskValues { - Map columnToMaskedValue; + Map attributeToMaskedValue; MaskValues(Map columnToMaskedValue) { - this.columnToMaskedValue = columnToMaskedValue; + this.attributeToMaskedValue = columnToMaskedValue; } } 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 72a1dbab..293ae327 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 @@ -145,8 +145,7 @@ private void processConfig(Config config) { this.handlerScopedFiltersConfig = new HandlerScopedFiltersConfig(config, this.startTimeAttributeName); - this.handlerScopedMaskingConfig = - new HandlerScopedMaskingConfig(config, viewDefinition.getPhysicalColumnNames(this.startTimeAttributeName.orElse(null)).get(0), tenantColumnName); + this.handlerScopedMaskingConfig = new HandlerScopedMaskingConfig(config); LOG.info( "Using {}ms as the threshold for logging slow queries of handler: {}", slowQueryThreshold, @@ -428,8 +427,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(), executionContext.getTenantId()) + return this.convert(resultSetGroup, executionContext) .doOnComplete( () -> { long requestTimeMs = stopwatch.stop().elapsed(TimeUnit.MILLISECONDS); @@ -498,12 +496,13 @@ private Filter rewriteLeafFilter( return queryFilter; } - Observable convert( - ResultSetGroup resultSetGroup, LinkedHashSet selectedAttributes, String tenantId) { + Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executionContext) { + String tenantId = executionContext.getTenantId(); + LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { ResultSet resultSet = resultSetGroup.getResultSet(0); - handlerScopedMaskingConfig.parseColumns(resultSet, tenantId); + handlerScopedMaskingConfig.parseColumns(executionContext); // 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 @@ -518,7 +517,7 @@ Observable convert( return Observable.fromIterable(rowBuilderList) .map(Builder::build) - .map(row -> handlerScopedMaskingConfig.mask(row)) + // .map(row -> handlerScopedMaskingConfig.mask(row)) .doOnNext(row -> LOG.debug("collect a row: {}", row)); } @@ -545,7 +544,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 = + !handlerScopedMaskingConfig.shouldMask(logicalName) + ? resultAnalyzer.getDataFromRow(rowId, logicalName) + : handlerScopedMaskingConfig.getMaskedValue(logicalName); + builder.addColumn(Value.newBuilder().setString(colVal).build()); } } 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 3b9302bf..9b4b5ded 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,6 @@ import com.typesafe.config.ConfigFactory; import io.reactivex.rxjava3.core.Observable; import java.io.IOException; -import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -1353,7 +1352,8 @@ public void testConvertSimpleSelectionsQueryResultSet() throws IOException { ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), resultTable); } @@ -1372,7 +1372,8 @@ public void testConvertAggregationColumnsQueryResultSet() throws IOException { ResultSetGroup resultSetGroup = mockResultSetGroup(List.of(resultSet)); verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), resultTable); } @@ -1434,7 +1435,8 @@ public void testConvertSelectionsWithMapKeysAndValuesQueryResultSet() throws IOE }; verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), expectedRows); } @@ -1470,7 +1472,8 @@ public void testConvertMultipleResultSetsInFResultSetGroup() throws IOException }; verifyResponseRows( - pinotBasedRequestHandler.convert(resultSetGroup, new LinkedHashSet<>(), "__default"), + pinotBasedRequestHandler.convert( + resultSetGroup, new ExecutionContext("__default", QueryRequest.newBuilder().build())), expectedRows); } @@ -1777,64 +1780,63 @@ public void testMaskColumnValue() throws IOException { 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"} - }; + 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); + 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")) - .build(); + QueryRequest.newBuilder() + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.id")) + .addSelection(QueryRequestBuilderUtils.createColumnExpression("EVENT.traceId")) + .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); // The query filter is based on both isEntrySpan and startTime. Since the viewFilter // checks for both the true and false values of isEntrySpan and query filter only needs // "true", isEntrySpan predicate is still passed to the store in the query. - String expectedQuery = - "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ?"; - Params params = - Params.newBuilder() - .addStringParam("maskTenant") - .build(); + String expectedQuery = "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ?"; + Params params = Params.newBuilder().addStringParam("maskTenant").build(); when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); String[][] expectedTable = - new String[][] { - {"*", "trace-id-1", }, - {"*", "trace-id-1"}, - {"*", "trace-id-1"}, - {"*", "trace-id-2"} - }; + new String[][] { + { + "*", "trace-id-1", + }, + {"*", "trace-id-1"}, + {"*", "trace-id-1"}, + {"*", "trace-id-2"} + }; verifyResponseRows(handler.handleRequest(request, context), expectedTable); } } - @Test public void testViewColumnFilterRemovalComplexCase() throws IOException { for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { @@ -2081,7 +2083,6 @@ private ResultSetGroup mockResultSetGroup(List resultSets) { private void verifyResponseRows(Observable rowObservable, String[][] expectedResultTable) throws IOException { - System.out.println(rowObservable); List rows = rowObservable.toList().blockingGet(); assertEquals(expectedResultTable.length, rows.size()); for (int rowIdx = 0; rowIdx < rows.size(); rowIdx++) { diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index dca49e26..25bc94d7 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -96,11 +96,9 @@ service.config = { "tenantId": "maskTenant", "timeRangeAndMaskValues": [ { - startTimeMillis = 1 - endTimeMillis = 100 "maskValues": [ { - "attributeId": "span_id", + "attributeId": "EVENT.id", "maskedValue": "*" } ] diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index 5233b3f5..25038550 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -99,22 +99,24 @@ service.config = { "tenantId": "testTenant", "timeRangeAndMaskValues": [ { - "startTimeMillis": 0, - "endTimeMillis": -1, - "maskValues": [ - { - "attributeId": "column_name1", - "maskedValue": "*" - }, - { - "attributeId": "column_name2", - "maskedValue": "*" - } - ] - }, - ] - } - ] +# "startTimeMillis": 0, +# # No startTimeMillis implies no filter on startTime +# "endTimeMillis": 1000, +# # No endTimeMillis implies no filter on endTime +# "maskValues": [ +# { +# "attributeId": "attribute_1", +# "maskedValue": "*" +# }, +# { +# "attributeId": "attribute_2", +# "maskedValue": "*" +# } +# ] +# }, +# ] +# } +# ] viewDefinition = { viewName = spanEventView mapFields = ["tags"] From a2a381c0cb058995a1f9a962099b3baf3613b538 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 1 Oct 2024 14:27:49 +0530 Subject: [PATCH 04/26] fix app.conf --- .../src/main/resources/configs/common/application.conf | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index 25038550..8222ce8b 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -94,11 +94,11 @@ service.config = { # ] # } # ] - tenantScopedMaskingCriteria = [ - { - "tenantId": "testTenant", - "timeRangeAndMaskValues": [ - { +# tenantScopedMaskingCriteria = [ +# { +# "tenantId": "testTenant", +# "timeRangeAndMaskValues": [ +# { # "startTimeMillis": 0, # # No startTimeMillis implies no filter on startTime # "endTimeMillis": 1000, From e1674117e5efcf6bea498fc0a39658c708ca41ea Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 1 Oct 2024 16:59:43 +0530 Subject: [PATCH 05/26] comments --- .../core/query/service/HandlerScopedMaskingConfig.java | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) 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 index 73561a9e..bee1250e 100644 --- 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 @@ -17,7 +17,6 @@ public class HandlerScopedMaskingConfig { private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; private final Map> tenantToMaskValuesMap; - private HashMap shouldMaskAttribute = new HashMap<>(); private HashMap maskedValue = new HashMap<>(); public HandlerScopedMaskingConfig(Config config) { @@ -35,9 +34,8 @@ public HandlerScopedMaskingConfig(Config config) { } public void parseColumns(ExecutionContext executionContext) { - shouldMaskAttribute.clear(); String tenantId = executionContext.getTenantId(); - + maskedValue.clear(); if (!tenantToMaskValuesMap.containsKey(tenantId)) { return; } @@ -59,7 +57,6 @@ public void parseColumns(ExecutionContext executionContext) { Map attributeToMaskedValue = timeRangeAndMasks.maskValues.attributeToMaskedValue; for (String attribute : attributeToMaskedValue.keySet()) { - shouldMaskAttribute.put(attribute, true); maskedValue.put(attribute, attributeToMaskedValue.get(attribute)); } } @@ -75,8 +72,10 @@ private static boolean isTimeRangeOverlap( if (startTimeInstant.isBefore(queryStartTime) || startTimeInstant.isAfter(queryEndTime)) { timeRangeOverlap = false; } + } - Instant endTimeInstant = Instant.ofEpochMilli(timeRangeAndMasks.getStartTimeMillis().get()); + if (timeRangeAndMasks.getEndTimeMillis().isPresent()) { + Instant endTimeInstant = Instant.ofEpochMilli(timeRangeAndMasks.getEndTimeMillis().get()); if (endTimeInstant.isBefore(queryStartTime) || endTimeInstant.isAfter(queryEndTime)) { timeRangeOverlap = false; } From 1167c62c788cafbdff2c100b1a9805db5568fb8a Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 1 Oct 2024 17:01:54 +0530 Subject: [PATCH 06/26] remove comment --- .../core/query/service/pinot/PinotBasedRequestHandler.java | 1 - 1 file changed, 1 deletion(-) 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 293ae327..910ef275 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 @@ -517,7 +517,6 @@ Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executio return Observable.fromIterable(rowBuilderList) .map(Builder::build) - // .map(row -> handlerScopedMaskingConfig.mask(row)) .doOnNext(row -> LOG.debug("collect a row: {}", row)); } From 3201839f0188691ac081c2b033252c0ce4e9d0e2 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Mon, 7 Oct 2024 13:47:04 +0530 Subject: [PATCH 07/26] remove state --- .../service/HandlerScopedMaskingConfig.java | 19 ++++++-------- .../pinot/PinotBasedRequestHandler.java | 25 +++++++++++-------- 2 files changed, 22 insertions(+), 22 deletions(-) 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 index bee1250e..bdb4215b 100644 --- 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 @@ -33,11 +33,12 @@ public HandlerScopedMaskingConfig(Config config) { } } - public void parseColumns(ExecutionContext executionContext) { + public HashMap getColumnToMaskedValuesMap(ExecutionContext executionContext) { String tenantId = executionContext.getTenantId(); - maskedValue.clear(); + HashMap columnMaskedValue = new HashMap<>(); + // maskedValue.clear(); if (!tenantToMaskValuesMap.containsKey(tenantId)) { - return; + return columnMaskedValue; } Optional queryTimeRange = executionContext.getQueryTimeRange(); @@ -57,10 +58,12 @@ public void parseColumns(ExecutionContext executionContext) { Map attributeToMaskedValue = timeRangeAndMasks.maskValues.attributeToMaskedValue; for (String attribute : attributeToMaskedValue.keySet()) { - maskedValue.put(attribute, attributeToMaskedValue.get(attribute)); + columnMaskedValue.put(attribute, attributeToMaskedValue.get(attribute)); } } } + + return columnMaskedValue; } private static boolean isTimeRangeOverlap( @@ -83,14 +86,6 @@ private static boolean isTimeRangeOverlap( return timeRangeOverlap; } - public boolean shouldMask(String attributeName) { - return this.maskedValue.containsKey(attributeName); - } - - public String getMaskedValue(String attributeName) { - return this.maskedValue.get(attributeName); - } - @Value @NonFinal private class TenantMasks { 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 910ef275..0fa26cfc 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 @@ -497,21 +497,21 @@ private Filter rewriteLeafFilter( } Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executionContext) { - String tenantId = executionContext.getTenantId(); LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { ResultSet resultSet = resultSetGroup.getResultSet(0); - handlerScopedMaskingConfig.parseColumns(executionContext); + HashMap columnToMaskedValue = + handlerScopedMaskingConfig.getColumnToMaskedValuesMap(executionContext); // 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, columnToMaskedValue); } else if (resultSetTypePredicateProvider.isResultTableResultSetType(resultSet)) { - handleTableFormatResultSet(resultSetGroup, rowBuilderList); + handleTableFormatResultSet(resultSetGroup, rowBuilderList, columnToMaskedValue); } else { - handleAggregationAndGroupBy(resultSetGroup, rowBuilderList); + handleAggregationAndGroupBy(resultSetGroup, rowBuilderList, columnToMaskedValue); } } @@ -523,7 +523,8 @@ Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executio private void handleSelection( ResultSetGroup resultSetGroup, List rowBuilderList, - LinkedHashSet selectedAttributes) { + LinkedHashSet selectedAttributes, + HashMap columnToMaskedValue) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); @@ -544,9 +545,9 @@ private void handleSelection( // colVal will never be null. But getDataRow can throw a runtime exception if it failed // to retrieve data String colVal = - !handlerScopedMaskingConfig.shouldMask(logicalName) + !columnToMaskedValue.containsKey(logicalName) ? resultAnalyzer.getDataFromRow(rowId, logicalName) - : handlerScopedMaskingConfig.getMaskedValue(logicalName); + : columnToMaskedValue.get(logicalName); builder.addColumn(Value.newBuilder().setString(colVal).build()); } @@ -555,7 +556,9 @@ private void handleSelection( } private void handleAggregationAndGroupBy( - ResultSetGroup resultSetGroup, List rowBuilderList) { + ResultSetGroup resultSetGroup, + List rowBuilderList, + HashMap columnToMaskedValue) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); Map groupKey2RowIdMap = new HashMap<>(); for (int i = 0; i < resultSetGroupCount; i++) { @@ -599,7 +602,9 @@ private void handleAggregationAndGroupBy( } private void handleTableFormatResultSet( - ResultSetGroup resultSetGroup, List rowBuilderList) { + ResultSetGroup resultSetGroup, + List rowBuilderList, + HashMap columnToMaskedValue) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); From 22a7c05bb875943e77c058cd11d6bcb6e1ab5fc0 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Mon, 7 Oct 2024 16:58:33 +0530 Subject: [PATCH 08/26] stateless changes --- .../service/HandlerScopedMaskingConfig.java | 47 ++++--------------- .../pinot/PinotBasedRequestHandler.java | 45 +++++++++++++----- .../service/pinot/PinotResultAnalyzer.java | 20 +++++++- .../pinot/PinotBasedRequestHandlerTest.java | 19 +++++--- .../src/test/resources/application.conf | 17 ++++--- .../resources/configs/common/application.conf | 12 ++--- 6 files changed, 85 insertions(+), 75 deletions(-) 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 index bdb4215b..c9466914 100644 --- 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 @@ -4,7 +4,6 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -17,7 +16,6 @@ public class HandlerScopedMaskingConfig { private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; private final Map> tenantToMaskValuesMap; - private HashMap maskedValue = new HashMap<>(); public HandlerScopedMaskingConfig(Config config) { if (config.hasPath(TENANT_SCOPED_MASKS_CONFIG_KEY)) { @@ -33,12 +31,12 @@ public HandlerScopedMaskingConfig(Config config) { } } - public HashMap getColumnToMaskedValuesMap(ExecutionContext executionContext) { + public List getMaskedAttributes(ExecutionContext executionContext) { String tenantId = executionContext.getTenantId(); - HashMap columnMaskedValue = new HashMap<>(); + List maskedAttributes = new ArrayList<>(); // maskedValue.clear(); if (!tenantToMaskValuesMap.containsKey(tenantId)) { - return columnMaskedValue; + return maskedAttributes; } Optional queryTimeRange = executionContext.getQueryTimeRange(); @@ -55,15 +53,11 @@ public HashMap getColumnToMaskedValuesMap(ExecutionContext execu isTimeRangeOverlap(timeRangeAndMasks, queryStartTime, queryEndTime); if (timeRangeOverlap) { - Map attributeToMaskedValue = - timeRangeAndMasks.maskValues.attributeToMaskedValue; - for (String attribute : attributeToMaskedValue.keySet()) { - columnMaskedValue.put(attribute, attributeToMaskedValue.get(attribute)); - } + maskedAttributes.addAll(timeRangeAndMasks.maskedAttributes); } } - return columnMaskedValue; + return maskedAttributes; } private static boolean isTimeRangeOverlap( @@ -103,26 +97,15 @@ private TenantMasks(Config config) { } } - @Value - private class MaskValues { - Map attributeToMaskedValue; - - MaskValues(Map columnToMaskedValue) { - this.attributeToMaskedValue = columnToMaskedValue; - } - } - @Value @NonFinal class MaskValuesForTimeRange { private static final String START_TIME_CONFIG_PATH = "startTimeMillis"; private static final String END_TIME_CONFIG_PATH = "endTimeMillis"; - private static final String MASK_VALUE_CONFIG_PATH = "maskValues"; - private static final String ATTRIBUTE_ID_CONFIG_PATH = "attributeId"; - private static final String MASKED_VALUE_CONFIG_PATH = "maskedValue"; + private static final String MASK_ATTRIBUTES_CONFIG_PATH = "maskedAttributes"; Optional startTimeMillis; Optional endTimeMillis; - MaskValues maskValues; + ArrayList maskedAttributes; private MaskValuesForTimeRange(Config config) { if (config.hasPath(START_TIME_CONFIG_PATH) && config.hasPath(END_TIME_CONFIG_PATH)) { @@ -132,20 +115,10 @@ private MaskValuesForTimeRange(Config config) { startTimeMillis = Optional.empty(); endTimeMillis = Optional.empty(); } - if (config.hasPath(MASK_VALUE_CONFIG_PATH)) { - List maskedValuesList = - new ArrayList<>(config.getConfigList(MASK_VALUE_CONFIG_PATH)); - HashMap maskedValuesMap = new HashMap<>(); - maskedValuesList.forEach( - maskedValue -> { - maskedValuesMap.put( - maskedValue.getString(ATTRIBUTE_ID_CONFIG_PATH), - maskedValue.getString(MASKED_VALUE_CONFIG_PATH)); - }); - - maskValues = new MaskValues(maskedValuesMap); + if (config.hasPath(MASK_ATTRIBUTES_CONFIG_PATH)) { + maskedAttributes = new ArrayList<>(config.getStringList(MASK_ATTRIBUTES_CONFIG_PATH)); } else { - maskValues = new MaskValues(new HashMap<>()); + maskedAttributes = new ArrayList<>(); } } } 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 0fa26cfc..d29586a9 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 @@ -59,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 MASKED_VALUE = "*"; + // This is how empty list is represented in Pinot + private static final String PINOT_EMPTY_LIST = "[\"\"]"; + private static final int DEFAULT_SLOW_QUERY_THRESHOLD_MS = 3000; private static final Set GTE_OPERATORS = Set.of(Operator.GE, Operator.GT, Operator.EQ); @@ -501,17 +505,18 @@ Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executio List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { ResultSet resultSet = resultSetGroup.getResultSet(0); - HashMap columnToMaskedValue = - handlerScopedMaskingConfig.getColumnToMaskedValuesMap(executionContext); + List maskedAttributes = + handlerScopedMaskingConfig.getMaskedAttributes(executionContext); // 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, columnToMaskedValue); + handleSelection(resultSetGroup, rowBuilderList, selectedAttributes, maskedAttributes); } else if (resultSetTypePredicateProvider.isResultTableResultSetType(resultSet)) { - handleTableFormatResultSet(resultSetGroup, rowBuilderList, columnToMaskedValue); + handleTableFormatResultSet( + resultSetGroup, rowBuilderList, selectedAttributes, maskedAttributes); } else { - handleAggregationAndGroupBy(resultSetGroup, rowBuilderList, columnToMaskedValue); + handleAggregationAndGroupBy(resultSetGroup, rowBuilderList); } } @@ -524,7 +529,7 @@ private void handleSelection( ResultSetGroup resultSetGroup, List rowBuilderList, LinkedHashSet selectedAttributes, - HashMap columnToMaskedValue) { + List maskedAttributes) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); @@ -545,9 +550,9 @@ private void handleSelection( // colVal will never be null. But getDataRow can throw a runtime exception if it failed // to retrieve data String colVal = - !columnToMaskedValue.containsKey(logicalName) + !maskedAttributes.contains(logicalName) ? resultAnalyzer.getDataFromRow(rowId, logicalName) - : columnToMaskedValue.get(logicalName); + : MASKED_VALUE; builder.addColumn(Value.newBuilder().setString(colVal).build()); } @@ -556,9 +561,7 @@ private void handleSelection( } private void handleAggregationAndGroupBy( - ResultSetGroup resultSetGroup, - List rowBuilderList, - HashMap columnToMaskedValue) { + ResultSetGroup resultSetGroup, List rowBuilderList) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); Map groupKey2RowIdMap = new HashMap<>(); for (int i = 0; i < resultSetGroupCount; i++) { @@ -604,10 +607,13 @@ private void handleAggregationAndGroupBy( private void handleTableFormatResultSet( ResultSetGroup resultSetGroup, List rowBuilderList, - HashMap columnToMaskedValue) { + LinkedHashSet selectedAttributes, + List 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(); @@ -620,6 +626,15 @@ private void handleTableFormatResultSet( // is structured String mapKeys = resultSet.getString(rowIdx, colIdx); String mapVals = resultSet.getString(rowIdx, colIdx + 1); + + String logicalNameKey = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + String logicalNameValue = resultAnalyzer.getLogicalNameFromColIdx(colIdx + 1); + + if (maskedAttributes.contains(logicalNameKey) + || maskedAttributes.contains(logicalNameValue)) { + mapVals = PINOT_EMPTY_LIST; + } + try { builder.addColumn( Value.newBuilder().setString(pinotMapConverter.merge(mapKeys, mapVals)).build()); @@ -632,6 +647,12 @@ private void handleTableFormatResultSet( colIdx++; } else { String val = resultSet.getString(rowIdx, colIdx); + String columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + + if (maskedAttributes.contains(columnLogicalName)) { + val = MASKED_VALUE; + } + 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..268a7bb7 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,6 +79,7 @@ static PinotResultAnalyzer create( String physName = resultSet.getColumnName(colIndex); if (physName.equalsIgnoreCase(names.get(0))) { logicalNameToPhysicalNameIndex.put(logicalName, colIndex); + indexToLogicalName.put(colIndex, logicalName); break; } } @@ -87,7 +94,8 @@ static PinotResultAnalyzer create( viewDefinition, mapLogicalNameToKeyIndex, mapLogicalNameToValueIndex, - logicalNameToPhysicalNameIndex); + logicalNameToPhysicalNameIndex, + indexToLogicalName); } @VisibleForTesting @@ -149,4 +157,12 @@ String getDataFromRow(int rowIndex, String logicalName) { } return result; } + + String getLogicalNameFromColIdx(Integer colIdx) { + if (indexToLogicalName.containsKey(colIdx)) { + return indexToLogicalName.get(colIdx); + } + + return null; + } } 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 9b4b5ded..e73dfb2e 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 @@ -1781,9 +1781,7 @@ public void testMaskColumnValue() throws IOException { String[][] resultTable = new String[][] { - { - "test-span-id-1", "trace-id-1", - }, + {"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"} @@ -1813,14 +1811,23 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { 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)))) .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); - + context.setTimeFilterColumn("EVENT.startTime"); // The query filter is based on both isEntrySpan and startTime. Since the viewFilter // checks for both the true and false values of isEntrySpan and query filter only needs // "true", isEntrySpan predicate is still passed to the store in the query. - String expectedQuery = "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ?"; - Params params = Params.newBuilder().addStringParam("maskTenant").build(); + String expectedQuery = "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; + Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); String[][] expectedTable = diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index 25bc94d7..693611f1 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -92,20 +92,19 @@ service.config = { } ] tenantScopedMaskingCriteria = [ - { - "tenantId": "maskTenant", - "timeRangeAndMaskValues": [ { - "maskValues": [ + tenantId = "maskTenant" + timeRangeAndMaskValues = [ { - "attributeId": "EVENT.id", - "maskedValue": "*" + startTimeMillis = 0 + endTimeMillis = 2728297299868 + maskedAttributes = [ + "EVENT.id" + ] } ] - }, + } ] - } - ] viewDefinition = { viewName = spanEventView mapFields = ["tags"] diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index 8222ce8b..d1b34727 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -103,15 +103,9 @@ service.config = { # # No startTimeMillis implies no filter on startTime # "endTimeMillis": 1000, # # No endTimeMillis implies no filter on endTime -# "maskValues": [ -# { -# "attributeId": "attribute_1", -# "maskedValue": "*" -# }, -# { -# "attributeId": "attribute_2", -# "maskedValue": "*" -# } +# "maskAttributes": [ +# "attribute_1", +# "attribute_2", # ] # }, # ] From e033b5b75bbc9cee8acb02949a60909ea081d9b6 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Mon, 7 Oct 2024 16:59:55 +0530 Subject: [PATCH 09/26] spotless --- .../query/service/pinot/PinotBasedRequestHandlerTest.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 e73dfb2e..da2e2a49 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 @@ -1818,15 +1818,15 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { QueryRequestBuilderUtils.createFilter( "EVENT.startTime", Operator.GT, - QueryRequestBuilderUtils.createLongLiteralValueExpression( - 99)))) + QueryRequestBuilderUtils.createLongLiteralValueExpression(99)))) .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); context.setTimeFilterColumn("EVENT.startTime"); // The query filter is based on both isEntrySpan and startTime. Since the viewFilter // checks for both the true and false values of isEntrySpan and query filter only needs // "true", isEntrySpan predicate is still passed to the store in the query. - String expectedQuery = "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; + String expectedQuery = + "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); when(pinotClient.executeQuery(expectedQuery, params)).thenReturn(resultSetGroup); From 0421c5df3f01d5169c0ad6c02f334a254a1b2f88 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 8 Oct 2024 11:25:45 +0530 Subject: [PATCH 10/26] comments --- .../service/HandlerScopedMaskingConfig.java | 7 ++++++ .../pinot/PinotBasedRequestHandler.java | 24 ++++++++++--------- .../service/pinot/PinotResultAnalyzer.java | 9 +++---- 3 files changed, 23 insertions(+), 17 deletions(-) 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 index c9466914..47fd7835 100644 --- 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 @@ -93,6 +93,7 @@ private TenantMasks(Config config) { this.maskValues = config.getConfigList(TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY).stream() .map(MaskValuesForTimeRange::new) + .filter(MaskValuesForTimeRange::isValid) .collect(Collectors.toList()); } } @@ -114,6 +115,8 @@ private MaskValuesForTimeRange(Config config) { } else { startTimeMillis = Optional.empty(); endTimeMillis = Optional.empty(); + log.warn( + "A masking filter is provided without startTimeMillis or endTimeMillis in tenantScopedMaskingCriteria. This filter will be ignored."); } if (config.hasPath(MASK_ATTRIBUTES_CONFIG_PATH)) { maskedAttributes = new ArrayList<>(config.getStringList(MASK_ATTRIBUTES_CONFIG_PATH)); @@ -121,5 +124,9 @@ private MaskValuesForTimeRange(Config config) { maskedAttributes = new ArrayList<>(); } } + + boolean isValid() { + return startTimeMillis.isPresent() && endTimeMillis.isPresent(); + } } } 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 d29586a9..cdd1aa1f 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 @@ -59,9 +59,9 @@ 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 MASKED_VALUE = "*"; + private static final String DEFAULT_MASKED_VALUE = "*"; // This is how empty list is represented in Pinot - private static final String PINOT_EMPTY_LIST = "[\"\"]"; + 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); @@ -552,7 +552,7 @@ private void handleSelection( String colVal = !maskedAttributes.contains(logicalName) ? resultAnalyzer.getDataFromRow(rowId, logicalName) - : MASKED_VALUE; + : DEFAULT_MASKED_VALUE; builder.addColumn(Value.newBuilder().setString(colVal).build()); } @@ -627,12 +627,13 @@ private void handleTableFormatResultSet( String mapKeys = resultSet.getString(rowIdx, colIdx); String mapVals = resultSet.getString(rowIdx, colIdx + 1); - String logicalNameKey = resultAnalyzer.getLogicalNameFromColIdx(colIdx); - String logicalNameValue = resultAnalyzer.getLogicalNameFromColIdx(colIdx + 1); + Optional logicalNameKey = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + Optional logicalNameValue = resultAnalyzer.getLogicalNameFromColIdx(colIdx + 1); - if (maskedAttributes.contains(logicalNameKey) - || maskedAttributes.contains(logicalNameValue)) { - mapVals = PINOT_EMPTY_LIST; + if ((logicalNameKey.isPresent() && maskedAttributes.contains(logicalNameKey.get())) + || (logicalNameValue.isPresent() + && maskedAttributes.contains(logicalNameValue.get()))) { + mapVals = ARRAY_TYPE_MASKED_VALUE; } try { @@ -647,10 +648,11 @@ private void handleTableFormatResultSet( colIdx++; } else { String val = resultSet.getString(rowIdx, colIdx); - String columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + Optional columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); - if (maskedAttributes.contains(columnLogicalName)) { - val = MASKED_VALUE; + if (columnLogicalName.isPresent() + && maskedAttributes.contains(columnLogicalName.get())) { + val = DEFAULT_MASKED_VALUE; } 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 268a7bb7..93506f1d 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 @@ -8,6 +8,7 @@ import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import javax.annotation.Nonnull; import org.apache.pinot.client.ResultSet; import org.slf4j.Logger; @@ -158,11 +159,7 @@ String getDataFromRow(int rowIndex, String logicalName) { return result; } - String getLogicalNameFromColIdx(Integer colIdx) { - if (indexToLogicalName.containsKey(colIdx)) { - return indexToLogicalName.get(colIdx); - } - - return null; + Optional getLogicalNameFromColIdx(Integer colIdx) { + return Optional.ofNullable(indexToLogicalName.get(colIdx)); } } From 345b5c5da5c5c65c594645f64c7ba02956035787 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 8 Oct 2024 12:24:58 +0530 Subject: [PATCH 11/26] comments --- .../core/query/service/HandlerScopedMaskingConfig.java | 5 +++-- .../core/query/service/pinot/PinotBasedRequestHandler.java | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) 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 index 47fd7835..7af9dd35 100644 --- 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 @@ -34,7 +34,6 @@ public HandlerScopedMaskingConfig(Config config) { public List getMaskedAttributes(ExecutionContext executionContext) { String tenantId = executionContext.getTenantId(); List maskedAttributes = new ArrayList<>(); - // maskedValue.clear(); if (!tenantToMaskValuesMap.containsKey(tenantId)) { return maskedAttributes; } @@ -126,7 +125,9 @@ private MaskValuesForTimeRange(Config config) { } boolean isValid() { - return startTimeMillis.isPresent() && endTimeMillis.isPresent(); + return startTimeMillis.isPresent() + && endTimeMillis.isPresent() + && !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 cdd1aa1f..450216a3 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 @@ -550,9 +550,9 @@ private void handleSelection( // colVal will never be null. But getDataRow can throw a runtime exception if it failed // to retrieve data String colVal = - !maskedAttributes.contains(logicalName) - ? resultAnalyzer.getDataFromRow(rowId, logicalName) - : DEFAULT_MASKED_VALUE; + maskedAttributes.contains(logicalName) + ? DEFAULT_MASKED_VALUE + : resultAnalyzer.getDataFromRow(rowId, logicalName); builder.addColumn(Value.newBuilder().setString(colVal).build()); } From 455945756e714e6b04a4f1c931fb58c860f6e20d Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 8 Oct 2024 14:44:19 +0530 Subject: [PATCH 12/26] tests --- .../pinot/PinotBasedRequestHandler.java | 7 +- .../pinot/PinotBasedRequestHandlerTest.java | 332 +++++++++++++++++- .../src/test/resources/application.conf | 5 +- 3 files changed, 336 insertions(+), 8 deletions(-) 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 450216a3..f279747d 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 @@ -627,12 +627,9 @@ private void handleTableFormatResultSet( String mapKeys = resultSet.getString(rowIdx, colIdx); String mapVals = resultSet.getString(rowIdx, colIdx + 1); - Optional logicalNameKey = resultAnalyzer.getLogicalNameFromColIdx(colIdx); - Optional logicalNameValue = resultAnalyzer.getLogicalNameFromColIdx(colIdx + 1); + Optional logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); - if ((logicalNameKey.isPresent() && maskedAttributes.contains(logicalNameKey.get())) - || (logicalNameValue.isPresent() - && maskedAttributes.contains(logicalNameValue.get()))) { + if (logicalName.isPresent() && maskedAttributes.contains(logicalName.get())) { mapVals = ARRAY_TYPE_MASKED_VALUE; } 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 da2e2a49..3b3a8aee 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,6 +18,8 @@ import com.typesafe.config.ConfigFactory; import io.reactivex.rxjava3.core.Observable; import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Objects; @@ -1764,7 +1766,7 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { } @Test - public void testMaskColumnValue() throws IOException { + public void testMaskColumnValueSelection() throws IOException { for (Config config : serviceConfig.getConfigList("queryRequestHandlersConfig")) { if (!isPinotConfig(config)) { continue; @@ -1844,6 +1846,334 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { } } + @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)))) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + // The query filter is based on both isEntrySpan and startTime. Since the viewFilter + // checks for both the true and false values of isEntrySpan and query filter only needs + // "true", isEntrySpan predicate is still passed to the store in the query. + String expectedQuery = + "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; + Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).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)))) + .build(); + ExecutionContext context = new ExecutionContext("maskTenant", request); + context.setTimeFilterColumn("EVENT.startTime"); + // The query filter is based on both isEntrySpan and startTime. Since the viewFilter + // checks for both the true and false values of isEntrySpan and query filter only needs + // "true", isEntrySpan predicate is still passed to the store in the query. + String expectedQuery = + "Select tags__KEYS, tags__VALUES FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; + Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).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"); + // The query filter is based on both isEntrySpan and startTime. Since the viewFilter + // checks for both the true and false values of isEntrySpan and query filter only needs + // "true", isEntrySpan predicate is still passed to the store in the query. + 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"); + // The query filter is based on both isEntrySpan and startTime. Since the viewFilter + // checks for both the true and false values of isEntrySpan and query filter only needs + // "true", isEntrySpan predicate is still passed to the store in the query. + 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 693611f1..a2e2b863 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -97,9 +97,10 @@ service.config = { timeRangeAndMaskValues = [ { startTimeMillis = 0 - endTimeMillis = 2728297299868 + endTimeMillis = 500 maskedAttributes = [ - "EVENT.id" + "EVENT.id", + "EVENT.spanTags" ] } ] From e89416fa6b4e0dda53db5dade3fc3f309b9cdfcb Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 8 Oct 2024 14:47:05 +0530 Subject: [PATCH 13/26] logical col idx --- .../query/service/pinot/PinotBasedRequestHandler.java | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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 f279747d..b4e3cd23 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 @@ -619,7 +619,9 @@ private void handleTableFormatResultSet( builder = Row.newBuilder(); rowBuilderList.add(builder); - for (int colIdx = 0; colIdx < resultSet.getColumnCount(); colIdx++) { + for (int colIdx = 0, logicalColIdx = 0; + colIdx < resultSet.getColumnCount(); + colIdx++, logicalColIdx++) { if (resultSet.getColumnName(colIdx).endsWith(ViewDefinition.MAP_KEYS_SUFFIX)) { // Read the key and value column values. The columns should be side by side. That's how // the Pinot query @@ -627,7 +629,7 @@ private void handleTableFormatResultSet( String mapKeys = resultSet.getString(rowIdx, colIdx); String mapVals = resultSet.getString(rowIdx, colIdx + 1); - Optional logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + Optional logicalName = resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); if (logicalName.isPresent() && maskedAttributes.contains(logicalName.get())) { mapVals = ARRAY_TYPE_MASKED_VALUE; @@ -645,7 +647,8 @@ private void handleTableFormatResultSet( colIdx++; } else { String val = resultSet.getString(rowIdx, colIdx); - Optional columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + Optional columnLogicalName = + resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); if (columnLogicalName.isPresent() && maskedAttributes.contains(columnLogicalName.get())) { From 19e48347cdd53e7fb1ae3bf686504636b4719f1c Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 8 Oct 2024 14:50:31 +0530 Subject: [PATCH 14/26] application.conf comments --- .../src/main/resources/configs/common/application.conf | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index d1b34727..f92ba29b 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -99,11 +99,11 @@ service.config = { # "tenantId": "testTenant", # "timeRangeAndMaskValues": [ # { +# #Start and end time filter must always be provided # "startTimeMillis": 0, -# # No startTimeMillis implies no filter on startTime # "endTimeMillis": 1000, -# # No endTimeMillis implies no filter on endTime # "maskAttributes": [ +# # Masking is only supported for attributes which have string data # "attribute_1", # "attribute_2", # ] From a055b3174de970adaa93521aa7ce657c18f33079 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Tue, 8 Oct 2024 15:01:50 +0530 Subject: [PATCH 15/26] remove stale comments --- .../pinot/PinotBasedRequestHandlerTest.java | 15 --------------- 1 file changed, 15 deletions(-) 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 3b3a8aee..0adafafe 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 @@ -1824,9 +1824,6 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); context.setTimeFilterColumn("EVENT.startTime"); - // The query filter is based on both isEntrySpan and startTime. Since the viewFilter - // checks for both the true and false values of isEntrySpan and query filter only needs - // "true", isEntrySpan predicate is still passed to the store in the query. String expectedQuery = "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); @@ -1905,9 +1902,6 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); context.setTimeFilterColumn("EVENT.startTime"); - // The query filter is based on both isEntrySpan and startTime. Since the viewFilter - // checks for both the true and false values of isEntrySpan and query filter only needs - // "true", isEntrySpan predicate is still passed to the store in the query. String expectedQuery = "Select span_id, trace_id FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); @@ -1987,9 +1981,6 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); context.setTimeFilterColumn("EVENT.startTime"); - // The query filter is based on both isEntrySpan and startTime. Since the viewFilter - // checks for both the true and false values of isEntrySpan and query filter only needs - // "true", isEntrySpan predicate is still passed to the store in the query. String expectedQuery = "Select tags__KEYS, tags__VALUES FROM spanEventView WHERE tenant_id = ? AND start_time_millis > ?"; Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); @@ -2066,9 +2057,6 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { .build(); ExecutionContext context = new ExecutionContext("noMaskTenant", request); context.setTimeFilterColumn("EVENT.startTime"); - // The query filter is based on both isEntrySpan and startTime. Since the viewFilter - // checks for both the true and false values of isEntrySpan and query filter only needs - // "true", isEntrySpan predicate is still passed to the store in the query. 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(); @@ -2150,9 +2138,6 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { .build(); ExecutionContext context = new ExecutionContext("maskTenant", request); context.setTimeFilterColumn("EVENT.startTime"); - // The query filter is based on both isEntrySpan and startTime. Since the viewFilter - // checks for both the true and false values of isEntrySpan and query filter only needs - // "true", isEntrySpan predicate is still passed to the store in the query. String expectedQuery = "Select tags__KEYS, tags__VALUES FROM spanEventView WHERE tenant_id = ? AND ( start_time_millis > ? AND start_time_millis < ? )"; Params params = From fdfcfcb21d2cf6f75e670b4e2097041e75dde403 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 12:57:46 +0530 Subject: [PATCH 16/26] comments --- .../service/pinot/PinotBasedRequestHandler.java | 12 +++++------- .../query/service/pinot/PinotResultAnalyzer.java | 5 ++--- .../service/pinot/PinotBasedRequestHandlerTest.java | 4 +--- 3 files changed, 8 insertions(+), 13 deletions(-) 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 b4e3cd23..092b62e1 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 @@ -501,9 +501,9 @@ private Filter rewriteLeafFilter( } Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executionContext) { - LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { + LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); ResultSet resultSet = resultSetGroup.getResultSet(0); List maskedAttributes = handlerScopedMaskingConfig.getMaskedAttributes(executionContext); @@ -629,9 +629,9 @@ private void handleTableFormatResultSet( String mapKeys = resultSet.getString(rowIdx, colIdx); String mapVals = resultSet.getString(rowIdx, colIdx + 1); - Optional logicalName = resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); + String logicalName = resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); - if (logicalName.isPresent() && maskedAttributes.contains(logicalName.get())) { + if (maskedAttributes.contains(logicalName)) { mapVals = ARRAY_TYPE_MASKED_VALUE; } @@ -647,11 +647,9 @@ private void handleTableFormatResultSet( colIdx++; } else { String val = resultSet.getString(rowIdx, colIdx); - Optional columnLogicalName = - resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); + String columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); - if (columnLogicalName.isPresent() - && maskedAttributes.contains(columnLogicalName.get())) { + if (maskedAttributes.contains(columnLogicalName)) { val = DEFAULT_MASKED_VALUE; } 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 93506f1d..35efb9dd 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 @@ -8,7 +8,6 @@ import java.util.LinkedHashSet; import java.util.List; import java.util.Map; -import java.util.Optional; import javax.annotation.Nonnull; import org.apache.pinot.client.ResultSet; import org.slf4j.Logger; @@ -159,7 +158,7 @@ String getDataFromRow(int rowIndex, String logicalName) { return result; } - Optional getLogicalNameFromColIdx(Integer colIdx) { - return Optional.ofNullable(indexToLogicalName.get(colIdx)); + 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 0adafafe..9cdc1057 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 @@ -1831,9 +1831,7 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { String[][] expectedTable = new String[][] { - { - "*", "trace-id-1", - }, + {"*", "trace-id-1"}, {"*", "trace-id-1"}, {"*", "trace-id-1"}, {"*", "trace-id-2"} From 8fd154fc154c2539f619190472bee1b8932e2d33 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 14:37:20 +0530 Subject: [PATCH 17/26] changes --- .../service/HandlerScopedMaskingConfig.java | 107 ++++++++---------- .../pinot/PinotBasedRequestHandler.java | 8 +- 2 files changed, 48 insertions(+), 67 deletions(-) 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 index 7af9dd35..6a9dea31 100644 --- 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 @@ -15,119 +15,102 @@ @Slf4j public class HandlerScopedMaskingConfig { private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; - private final Map> tenantToMaskValuesMap; + private Map> tenantToTimeRangeMaskedAttributes = + Collections.emptyMap(); public HandlerScopedMaskingConfig(Config config) { if (config.hasPath(TENANT_SCOPED_MASKS_CONFIG_KEY)) { - this.tenantToMaskValuesMap = + this.tenantToTimeRangeMaskedAttributes = config.getConfigList(TENANT_SCOPED_MASKS_CONFIG_KEY).stream() - .map(maskConfig -> new TenantMasks(maskConfig)) + .map(TenantMaskingConfig::new) .collect( Collectors.toMap( - tenantFilters -> tenantFilters.tenantId, - tenantFilters -> tenantFilters.maskValues)); - } else { - this.tenantToMaskValuesMap = Collections.emptyMap(); + TenantMaskingConfig::getTenantId, + TenantMaskingConfig::getTimeRangeToMaskedAttributes)); } } public List getMaskedAttributes(ExecutionContext executionContext) { String tenantId = executionContext.getTenantId(); List maskedAttributes = new ArrayList<>(); - if (!tenantToMaskValuesMap.containsKey(tenantId)) { + if (!tenantToTimeRangeMaskedAttributes.containsKey(tenantId)) { return maskedAttributes; } Optional queryTimeRange = executionContext.getQueryTimeRange(); - Instant queryStartTime, queryEndTime; + Instant queryStartTime = Instant.MIN, queryEndTime = Instant.MAX; if (queryTimeRange.isPresent()) { queryStartTime = queryTimeRange.get().getStartTime(); queryEndTime = queryTimeRange.get().getEndTime(); - } else { - queryEndTime = Instant.MAX; - queryStartTime = Instant.MIN; } - for (MaskValuesForTimeRange timeRangeAndMasks : tenantToMaskValuesMap.get(tenantId)) { - boolean timeRangeOverlap = - isTimeRangeOverlap(timeRangeAndMasks, queryStartTime, queryEndTime); - - if (timeRangeOverlap) { + for (TimeRangeToMaskedAttributes timeRangeAndMasks : + tenantToTimeRangeMaskedAttributes.get(tenantId)) { + if (isTimeRangeOverlap(timeRangeAndMasks, queryStartTime, queryEndTime)) { maskedAttributes.addAll(timeRangeAndMasks.maskedAttributes); } } - return maskedAttributes; } private static boolean isTimeRangeOverlap( - MaskValuesForTimeRange timeRangeAndMasks, Instant queryStartTime, Instant queryEndTime) { - boolean timeRangeOverlap = true; - - if (timeRangeAndMasks.getStartTimeMillis().isPresent()) { - Instant startTimeInstant = Instant.ofEpochMilli(timeRangeAndMasks.getStartTimeMillis().get()); - if (startTimeInstant.isBefore(queryStartTime) || startTimeInstant.isAfter(queryEndTime)) { - timeRangeOverlap = false; - } - } - - if (timeRangeAndMasks.getEndTimeMillis().isPresent()) { - Instant endTimeInstant = Instant.ofEpochMilli(timeRangeAndMasks.getEndTimeMillis().get()); - if (endTimeInstant.isBefore(queryStartTime) || endTimeInstant.isAfter(queryEndTime)) { - timeRangeOverlap = false; - } - } - return timeRangeOverlap; + TimeRangeToMaskedAttributes timeRangeAndMasks, Instant queryStartTime, Instant queryEndTime) { + return !(timeRangeAndMasks.startTimeMillis.isAfter(queryEndTime) + || timeRangeAndMasks.endTimeMillis.isBefore(queryStartTime)); } @Value @NonFinal - private class TenantMasks { + static class TenantMaskingConfig { private static final String TENANT_ID_CONFIG_KEY = "tenantId"; private static final String TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY = "timeRangeAndMaskValues"; String tenantId; - List maskValues; + List timeRangeToMaskedAttributes; - private TenantMasks(Config config) { + private TenantMaskingConfig(Config config) { this.tenantId = config.getString(TENANT_ID_CONFIG_KEY); - this.maskValues = + this.timeRangeToMaskedAttributes = config.getConfigList(TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY).stream() - .map(MaskValuesForTimeRange::new) - .filter(MaskValuesForTimeRange::isValid) + .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()); } } - @Value @NonFinal - class MaskValuesForTimeRange { + 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"; - Optional startTimeMillis; - Optional endTimeMillis; - ArrayList maskedAttributes; + Instant startTimeMillis = null; + Instant endTimeMillis = null; + ArrayList maskedAttributes = new ArrayList<>(); - private MaskValuesForTimeRange(Config config) { + private TimeRangeToMaskedAttributes(Config config) { if (config.hasPath(START_TIME_CONFIG_PATH) && config.hasPath(END_TIME_CONFIG_PATH)) { - this.startTimeMillis = Optional.of(config.getLong(START_TIME_CONFIG_PATH)); - this.endTimeMillis = Optional.of(config.getLong(END_TIME_CONFIG_PATH)); - } else { - startTimeMillis = Optional.empty(); - endTimeMillis = Optional.empty(); - log.warn( - "A masking filter is provided without startTimeMillis or endTimeMillis in tenantScopedMaskingCriteria. This filter will be ignored."); - } - if (config.hasPath(MASK_ATTRIBUTES_CONFIG_PATH)) { - maskedAttributes = new ArrayList<>(config.getStringList(MASK_ATTRIBUTES_CONFIG_PATH)); - } else { - maskedAttributes = new ArrayList<>(); + 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.isPresent() - && endTimeMillis.isPresent() - && !maskedAttributes.isEmpty(); + 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 092b62e1..2e45ae76 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 @@ -619,9 +619,7 @@ private void handleTableFormatResultSet( builder = Row.newBuilder(); rowBuilderList.add(builder); - for (int colIdx = 0, logicalColIdx = 0; - colIdx < resultSet.getColumnCount(); - colIdx++, logicalColIdx++) { + for (int colIdx = 0; colIdx < resultSet.getColumnCount(); colIdx++) { if (resultSet.getColumnName(colIdx).endsWith(ViewDefinition.MAP_KEYS_SUFFIX)) { // Read the key and value column values. The columns should be side by side. That's how // the Pinot query @@ -629,7 +627,7 @@ private void handleTableFormatResultSet( String mapKeys = resultSet.getString(rowIdx, colIdx); String mapVals = resultSet.getString(rowIdx, colIdx + 1); - String logicalName = resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); + String logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); if (maskedAttributes.contains(logicalName)) { mapVals = ARRAY_TYPE_MASKED_VALUE; @@ -647,7 +645,7 @@ private void handleTableFormatResultSet( colIdx++; } else { String val = resultSet.getString(rowIdx, colIdx); - String columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(logicalColIdx); + String columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); if (maskedAttributes.contains(columnLogicalName)) { val = DEFAULT_MASKED_VALUE; From 041e520d7d2e07cf3408aac7b01a8873f3775898 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 14:47:07 +0530 Subject: [PATCH 18/26] change config --- .../core/query/service/HandlerScopedMaskingConfig.java | 5 +++-- query-service-impl/src/test/resources/application.conf | 4 ++-- .../src/main/resources/configs/common/application.conf | 4 ++-- 3 files changed, 7 insertions(+), 6 deletions(-) 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 index 6a9dea31..1713f316 100644 --- 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 @@ -14,7 +14,7 @@ @Slf4j public class HandlerScopedMaskingConfig { - private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingCriteria"; + private static final String TENANT_SCOPED_MASKS_CONFIG_KEY = "tenantScopedMaskingConfig"; private Map> tenantToTimeRangeMaskedAttributes = Collections.emptyMap(); @@ -62,7 +62,8 @@ private static boolean isTimeRangeOverlap( @NonFinal static class TenantMaskingConfig { private static final String TENANT_ID_CONFIG_KEY = "tenantId"; - private static final String TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY = "timeRangeAndMaskValues"; + private static final String TIME_RANGE_AND_MASK_VALUES_CONFIG_KEY = + "timeRangeToMaskedAttributes"; String tenantId; List timeRangeToMaskedAttributes; diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index a2e2b863..680e299d 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -91,10 +91,10 @@ service.config = { ] } ] - tenantScopedMaskingCriteria = [ + tenantScopedMaskingConfig = [ { tenantId = "maskTenant" - timeRangeAndMaskValues = [ + timeRangeToMaskedAttributes = [ { startTimeMillis = 0 endTimeMillis = 500 diff --git a/query-service/src/main/resources/configs/common/application.conf b/query-service/src/main/resources/configs/common/application.conf index f92ba29b..e797d039 100644 --- a/query-service/src/main/resources/configs/common/application.conf +++ b/query-service/src/main/resources/configs/common/application.conf @@ -94,10 +94,10 @@ service.config = { # ] # } # ] -# tenantScopedMaskingCriteria = [ +# tenantScopedMaskingConfig = [ # { # "tenantId": "testTenant", -# "timeRangeAndMaskValues": [ +# "timeRangeToMaskedAttributes": [ # { # #Start and end time filter must always be provided # "startTimeMillis": 0, From 916076efa295a6ba53da3d507153754554af93e9 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 15:01:51 +0530 Subject: [PATCH 19/26] cleanup tests --- .../service/pinot/PinotResultAnalyzer.java | 6 +-- .../pinot/PinotBasedRequestHandlerTest.java | 52 ++++++++++++++----- .../src/test/resources/application.conf | 5 ++ 3 files changed, 48 insertions(+), 15 deletions(-) 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 35efb9dd..3aa73f10 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 @@ -85,9 +85,9 @@ static PinotResultAnalyzer create( } } } - 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); return new PinotResultAnalyzer( resultSet, selectedAttributes, 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 9cdc1057..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 @@ -1820,13 +1820,23 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { QueryRequestBuilderUtils.createFilter( "EVENT.startTime", Operator.GT, - QueryRequestBuilderUtils.createLongLiteralValueExpression(99)))) + 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 > ?"; - Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); + "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 = @@ -1896,20 +1906,28 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { QueryRequestBuilderUtils.createFilter( "EVENT.startTime", Operator.GT, - QueryRequestBuilderUtils.createLongLiteralValueExpression(99)))) + 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 > ?"; - Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); + "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-1"}, {"*", "trace-id-2"} @@ -1975,13 +1993,23 @@ public boolean isResultTableResultSetType(ResultSet resultSet) { QueryRequestBuilderUtils.createFilter( "EVENT.startTime", Operator.GT, - QueryRequestBuilderUtils.createLongLiteralValueExpression(99)))) + 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 > ?"; - Params params = Params.newBuilder().addStringParam("maskTenant").addLongParam(99).build(); + "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 = diff --git a/query-service-impl/src/test/resources/application.conf b/query-service-impl/src/test/resources/application.conf index 680e299d..832108c4 100644 --- a/query-service-impl/src/test/resources/application.conf +++ b/query-service-impl/src/test/resources/application.conf @@ -102,6 +102,11 @@ service.config = { "EVENT.id", "EVENT.spanTags" ] + }, + # Invalid mask added to test warning + { + maskedAttributes = [ + ] } ] } From 7b5e7cab3c45951b165f5027d0c7a55f6e16fc93 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 16:38:40 +0530 Subject: [PATCH 20/26] change --- .../core/query/service/pinot/PinotBasedRequestHandler.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 2e45ae76..2bd06902 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 @@ -503,10 +503,10 @@ private Filter rewriteLeafFilter( Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executionContext) { List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { - LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); - ResultSet resultSet = resultSetGroup.getResultSet(0); List maskedAttributes = handlerScopedMaskingConfig.getMaskedAttributes(executionContext); + ResultSet resultSet = resultSetGroup.getResultSet(0); + LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); // 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 From dc54e027dabe3ceef6689804ea418659ce1b4458 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 16:41:47 +0530 Subject: [PATCH 21/26] refactor --- .../core/query/service/pinot/PinotBasedRequestHandler.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2bd06902..8075216e 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 @@ -503,10 +503,10 @@ private Filter rewriteLeafFilter( Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executionContext) { List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { + LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); List maskedAttributes = handlerScopedMaskingConfig.getMaskedAttributes(executionContext); ResultSet resultSet = resultSetGroup.getResultSet(0); - LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); // 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 From 612f3bdc4c9fd06e831dbe255d592f26d062ac35 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 16:42:25 +0530 Subject: [PATCH 22/26] remove extra line --- .../core/query/service/pinot/PinotBasedRequestHandler.java | 1 - 1 file changed, 1 deletion(-) 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 8075216e..eee2f9c7 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 @@ -519,7 +519,6 @@ Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executio handleAggregationAndGroupBy(resultSetGroup, rowBuilderList); } } - return Observable.fromIterable(rowBuilderList) .map(Builder::build) .doOnNext(row -> LOG.debug("collect a row: {}", row)); From 02d2d728c7b3be225c8315007270cf05132bf7bd Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 16:47:04 +0530 Subject: [PATCH 23/26] refactor --- .../pinot/PinotBasedRequestHandler.java | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) 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 eee2f9c7..e361b505 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 @@ -623,14 +623,12 @@ 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 mapKeys = resultSet.getString(rowIdx, colIdx); - String mapVals = resultSet.getString(rowIdx, colIdx + 1); - String logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); - - if (maskedAttributes.contains(logicalName)) { - mapVals = ARRAY_TYPE_MASKED_VALUE; - } + String mapKeys = resultSet.getString(rowIdx, colIdx); + String mapVals = + maskedAttributes.contains(logicalName) + ? ARRAY_TYPE_MASKED_VALUE + : resultSet.getString(rowIdx, colIdx + 1); try { builder.addColumn( @@ -643,13 +641,11 @@ private void handleTableFormatResultSet( // advance colIdx by 1 since we have read 2 columns colIdx++; } else { - String val = resultSet.getString(rowIdx, colIdx); - String columnLogicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); - - if (maskedAttributes.contains(columnLogicalName)) { - val = DEFAULT_MASKED_VALUE; - } - + String logicalName = resultAnalyzer.getLogicalNameFromColIdx(colIdx); + String val = + maskedAttributes.contains(logicalName) + ? DEFAULT_MASKED_VALUE + : resultSet.getString(rowIdx, colIdx); builder.addColumn(Value.newBuilder().setString(val).build()); } } @@ -712,8 +708,4 @@ private boolean isInvalidExpression(Expression expression) { && viewDefinition.getColumnType(expression.getAttributeExpression().getAttributeId()) != ValueType.STRING_MAP; } - - HandlerScopedMaskingConfig getHandlerScopedMaskingConfig() { - return handlerScopedMaskingConfig; - } } From 4ebe89cd7831f0118d1b721eb603b5b7a1c6fc57 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 16:49:49 +0530 Subject: [PATCH 24/26] refactor --- .../hypertrace/core/query/service/pinot/PinotResultAnalyzer.java | 1 + 1 file changed, 1 insertion(+) 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 3aa73f10..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 @@ -88,6 +88,7 @@ static PinotResultAnalyzer create( 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, From 5375d854b634d6837ab929b8f027ad4148a60f47 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 16:57:32 +0530 Subject: [PATCH 25/26] list to set --- .../core/query/service/HandlerScopedMaskingConfig.java | 6 ++++-- .../core/query/service/pinot/PinotBasedRequestHandler.java | 6 +++--- 2 files changed, 7 insertions(+), 5 deletions(-) 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 index 1713f316..bb4e3bfc 100644 --- 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 @@ -4,9 +4,11 @@ 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; @@ -30,9 +32,9 @@ public HandlerScopedMaskingConfig(Config config) { } } - public List getMaskedAttributes(ExecutionContext executionContext) { + public Set getMaskedAttributes(ExecutionContext executionContext) { String tenantId = executionContext.getTenantId(); - List maskedAttributes = new ArrayList<>(); + HashSet maskedAttributes = new HashSet<>(); if (!tenantToTimeRangeMaskedAttributes.containsKey(tenantId)) { return maskedAttributes; } 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 e361b505..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 @@ -504,7 +504,7 @@ Observable convert(ResultSetGroup resultSetGroup, ExecutionContext executio List rowBuilderList = new ArrayList<>(); if (resultSetGroup.getResultSetCount() > 0) { LinkedHashSet selectedAttributes = executionContext.getSelectedColumns(); - List maskedAttributes = + Set maskedAttributes = handlerScopedMaskingConfig.getMaskedAttributes(executionContext); ResultSet resultSet = resultSetGroup.getResultSet(0); // Pinot has different Response format for selection and aggregation/group by query. @@ -528,7 +528,7 @@ private void handleSelection( ResultSetGroup resultSetGroup, List rowBuilderList, LinkedHashSet selectedAttributes, - List maskedAttributes) { + Set maskedAttributes) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); @@ -607,7 +607,7 @@ private void handleTableFormatResultSet( ResultSetGroup resultSetGroup, List rowBuilderList, LinkedHashSet selectedAttributes, - List maskedAttributes) { + Set maskedAttributes) { int resultSetGroupCount = resultSetGroup.getResultSetCount(); for (int i = 0; i < resultSetGroupCount; i++) { ResultSet resultSet = resultSetGroup.getResultSet(i); From d38e08c6a56b4e2d9fb0efd03f66fbb699ff4836 Mon Sep 17 00:00:00 2001 From: siddhant2001 Date: Thu, 10 Oct 2024 17:02:20 +0530 Subject: [PATCH 26/26] vuln fix --- query-service-impl/build.gradle.kts | 4 ++-- query-service/build.gradle.kts | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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/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")