gortiz commented on code in PR #15240: URL: https://github.com/apache/pinot/pull/15240#discussion_r1988613220
########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java: ########## @@ -317,4 +323,66 @@ protected void onQueryFinish(long requestId) { protected boolean isQueryCancellationEnabled() { return _queriesById != null; } + + protected void updatePhaseTimingForTables(Set<String> tableNames, BrokerQueryPhase phase, long time) { + for (String tableName : tableNames) { + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + _brokerMetrics.addPhaseTiming(rawTableName, phase, time); + } + } + + /** + * Validates whether the requester has access to all the tables. + */ + protected TableAuthorizationResult hasTableAccess(RequesterIdentity requesterIdentity, Set<String> tableNames, + RequestContext requestContext, HttpHeaders httpHeaders) { + final long startTimeNs = System.nanoTime(); + AccessControl accessControl = _accessControlFactory.create(); + + TableAuthorizationResult tableAuthorizationResult = accessControl.authorize(requesterIdentity, tableNames); + + Set<String> failedTables = tableNames.stream() + .filter(table -> !accessControl.hasAccess(httpHeaders, TargetType.TABLE, table, Actions.Table.QUERY)) + .collect(Collectors.toSet()); + + failedTables.addAll(tableAuthorizationResult.getFailedTables()); + + if (!failedTables.isEmpty()) { + tableAuthorizationResult = new TableAuthorizationResult(failedTables); + } else { + tableAuthorizationResult = TableAuthorizationResult.success(); + } + + if (!tableAuthorizationResult.hasAccess()) { + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); + LOGGER.warn("Access denied for requestId {}", requestContext.getRequestId()); Review Comment: Are we sure we want to log this as a warning? It could be problematic, especially when it is more needed (e.g., when receiving an attack). The metric should be good enough. BTW, we could also make that metric table based instead of global for more information. ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java: ########## @@ -317,4 +323,66 @@ protected void onQueryFinish(long requestId) { protected boolean isQueryCancellationEnabled() { return _queriesById != null; } + + protected void updatePhaseTimingForTables(Set<String> tableNames, BrokerQueryPhase phase, long time) { + for (String tableName : tableNames) { + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + _brokerMetrics.addPhaseTiming(rawTableName, phase, time); + } + } + + /** + * Validates whether the requester has access to all the tables. + */ + protected TableAuthorizationResult hasTableAccess(RequesterIdentity requesterIdentity, Set<String> tableNames, + RequestContext requestContext, HttpHeaders httpHeaders) { + final long startTimeNs = System.nanoTime(); + AccessControl accessControl = _accessControlFactory.create(); + + TableAuthorizationResult tableAuthorizationResult = accessControl.authorize(requesterIdentity, tableNames); + + Set<String> failedTables = tableNames.stream() + .filter(table -> !accessControl.hasAccess(httpHeaders, TargetType.TABLE, table, Actions.Table.QUERY)) + .collect(Collectors.toSet()); + + failedTables.addAll(tableAuthorizationResult.getFailedTables()); + + if (!failedTables.isEmpty()) { + tableAuthorizationResult = new TableAuthorizationResult(failedTables); + } else { + tableAuthorizationResult = TableAuthorizationResult.success(); + } + + if (!tableAuthorizationResult.hasAccess()) { + _brokerMetrics.addMeteredGlobalValue(BrokerMeter.REQUEST_DROPPED_DUE_TO_ACCESS_ERROR, 1); + LOGGER.warn("Access denied for requestId {}", requestContext.getRequestId()); + requestContext.setErrorCode(QueryErrorCode.ACCESS_DENIED); + } + + updatePhaseTimingForTables(tableNames, BrokerQueryPhase.AUTHORIZATION, System.nanoTime() - startTimeNs); + + return tableAuthorizationResult; + } + + /** + * Returns true if the QPS quota of query tables, database or application has been exceeded. + */ + protected boolean hasExceededQPSQuota(@Nullable String database, Set<String> tableNames, + RequestContext requestContext) { + if (database != null && !_queryQuotaManager.acquireDatabase(database)) { + LOGGER.warn("Request {}: query exceeds quota for database: {}", requestContext.getRequestId(), database); + requestContext.setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS); + return true; + } + for (String tableName : tableNames) { + if (!_queryQuotaManager.acquire(tableName)) { + LOGGER.warn("Request {}: query exceeds quota for table: {}", requestContext.getRequestId(), tableName); + requestContext.setErrorCode(QueryErrorCode.TOO_MANY_REQUESTS); + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + _brokerMetrics.addMeteredTableValue(rawTableName, BrokerMeter.QUERY_QUOTA_EXCEEDED, 1); Review Comment: Same here. Logging as WARN seems to be overkill. We have the metric for that. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org