This is an automated email from the ASF dual-hosted git repository. apucher pushed a commit to branch pinot-flexible-auth-provider in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 8ecba358ca24b4536cc4a42e5372cbda3d89df9f Author: Alexander Pucher <[email protected]> AuthorDate: Tue May 10 00:05:38 2022 -0700 auth provider round 2 --- .../{NullAuthProvider.java => AuthConfig.java} | 17 ++-- .../pinot/common/auth/AuthProviderUtils.java | 110 +++++++++++++++------ .../apache/pinot/common/auth/NullAuthProvider.java | 15 ++- .../pinot/common/auth/StaticTokenAuthProvider.java | 35 +++++-- .../apache/pinot/common/auth/UrlAuthProvider.java | 39 +++++--- .../common/utils/FileUploadDownloadClient.java | 4 +- .../common/utils/fetcher/BaseSegmentFetcher.java | 4 +- .../utils/fetcher/SegmentFetcherFactory.java | 19 ++-- .../apache/pinot/common/utils/http/HttpClient.java | 10 +- .../resources/PinotIngestionRestletResource.java | 13 +-- .../core/data/manager/BaseTableDataManager.java | 6 +- .../manager/offline/TableDataManagerProvider.java | 6 +- .../realtime/Server2ControllerSegmentUploader.java | 2 +- .../ServerSegmentCompletionProtocolHandler.java | 7 +- .../BaseTableDataManagerAcquireSegmentTest.java | 2 +- .../data/manager/BaseTableDataManagerTest.java | 2 +- .../offline/DimensionTableDataManagerTest.java | 2 +- .../realtime/LLRealtimeSegmentDataManagerTest.java | 2 +- .../executor/QueryExecutorExceptionsTest.java | 2 +- .../core/query/executor/QueryExecutorTest.java | 2 +- .../queries/SegmentWithNullValueVectorTest.java | 2 +- .../tests/BasicAuthBatchIntegrationTest.java | 3 +- .../BaseMultipleSegmentsConversionExecutor.java | 7 +- .../tasks/BaseSingleSegmentConversionExecutor.java | 4 +- .../local/data/manager/TableDataManager.java | 3 +- .../local/data/manager/TableDataManagerConfig.java | 20 ++-- .../pinot/segment/local/utils/IngestionUtils.java | 5 +- .../segment/local/utils/SegmentPushUtils.java | 14 +-- .../starter/helix/HelixInstanceDataManager.java | 10 +- .../helix/HelixInstanceDataManagerConfig.java | 29 +++--- .../apache/pinot/server/api/BaseResourceTest.java | 2 +- .../org/apache/pinot/spi/auth/AuthProvider.java | 4 +- .../config/instance/InstanceDataManagerConfig.java | 6 +- .../apache/pinot/spi/utils/CommonConstants.java | 8 +- .../org/apache/pinot/tools/AuthQuickstart.java | 6 +- .../org/apache/pinot/tools/BootstrapTableTool.java | 13 ++- .../org/apache/pinot/tools/EmptyQuickstart.java | 5 +- .../java/org/apache/pinot/tools/Quickstart.java | 5 +- .../admin/command/AbstractBaseAdminCommand.java | 32 ++++-- .../tools/admin/command/AddSchemaCommand.java | 12 ++- .../pinot/tools/admin/command/AddTableCommand.java | 15 ++- .../tools/admin/command/AddTenantCommand.java | 11 ++- .../tools/admin/command/BootstrapTableCommand.java | 14 ++- .../tools/admin/command/ChangeTableState.java | 12 ++- .../tools/admin/command/ImportDataCommand.java | 11 ++- .../command/LaunchDataIngestionJobCommand.java | 10 +- .../admin/command/OperateClusterConfigCommand.java | 11 ++- .../tools/admin/command/PostQueryCommand.java | 12 ++- .../tools/admin/command/QuickstartRunner.java | 28 ++---- .../tools/admin/command/UploadSegmentCommand.java | 16 ++- 50 files changed, 396 insertions(+), 233 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java b/pinot-common/src/main/java/org/apache/pinot/common/auth/AuthConfig.java similarity index 74% copy from pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java copy to pinot-common/src/main/java/org/apache/pinot/common/auth/AuthConfig.java index 80756f3aae..1da479b0dc 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/auth/AuthConfig.java @@ -18,14 +18,19 @@ */ package org.apache.pinot.common.auth; -import java.util.Collections; import java.util.Map; -import org.apache.pinot.spi.auth.AuthProvider; -public class NullAuthProvider implements AuthProvider { - @Override - public Map<String, Object> getHttpHeaders() { - return Collections.emptyMap(); +public class AuthConfig { + public static final String PROVIDER_CLASS = "provider.class"; + + protected Map<String, Object> _properties; + + public AuthConfig(Map<String, Object> properties) { + _properties = properties; + } + + public Map<String, Object> getProperties() { + return _properties; } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/auth/AuthProviderUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/auth/AuthProviderUtils.java index c8bac0593f..488c68841f 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/auth/AuthProviderUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/auth/AuthProviderUtils.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.common.auth; +import java.lang.reflect.Constructor; import java.util.Collections; import java.util.List; import java.util.Map; @@ -28,6 +29,7 @@ import org.apache.commons.lang3.StringUtils; import org.apache.http.Header; import org.apache.http.message.BasicHeader; import org.apache.pinot.spi.auth.AuthProvider; +import org.apache.pinot.spi.env.PinotConfiguration; /** @@ -39,55 +41,69 @@ public final class AuthProviderUtils { } /** - * Infer optimal auth provider based on the availability of static token, if any. + * Extract an AuthConfig from a pinot configuration subset namespace. * - * @param authToken static auth token - * @return auth provider + * @param pinotConfig pinot configuration + * @param namespace subset namespace + * @return auth config */ - public static AuthProvider inferProvider(String authToken) { - return inferProvider(authToken, null); + public static AuthConfig extractAuthConfig(PinotConfiguration pinotConfig, String namespace) { + return new AuthConfig(pinotConfig.subset(namespace).toMap()); } /** - * Infer optimal auth provider based on the availability of token and token url, if any. + * Create an AuthProvider after extracting a config from a pinot configuration subset namespace + * @see AuthProviderUtils#extractAuthConfig(PinotConfiguration, String) * - * @param authToken static auth token - * @param authTokenUrl dynamic token URL + * @param pinotConfig pinot configuration + * @param namespace subset namespace * @return auth provider */ - public static AuthProvider inferProvider(String authToken, String authTokenUrl) { - if (StringUtils.isNotBlank(authTokenUrl)) { - return new UrlAuthProvider(authTokenUrl); - } - if (StringUtils.isNotBlank(authToken)) { - return new StaticTokenAuthProvider(authToken); - } - return new NullAuthProvider(); + public static AuthProvider extractAuthProvider(PinotConfiguration pinotConfig, String namespace) { + return makeProvider(extractAuthConfig(pinotConfig, namespace)); } /** - * Resolve auth token right now, e.g. for job specs. + * Create auth provider based on the availability of a static token only, if any. * * @param authToken static auth token - * @param authTokenUrl dynamic token URL - * @return resolved static token + * @return auth provider */ - public static String resolveToToken(String authToken, String authTokenUrl) { - return resolveToToken(inferProvider(authToken, authTokenUrl)); + public static AuthProvider makeProvider(String authToken) { + if (StringUtils.isBlank(authToken)) { + return new NullAuthProvider(); + } + return new StaticTokenAuthProvider(authToken); } /** - * Resolve auth provider to token right now. + * Create auth provider based on an auth config. Mimics legacy behavior for static tokens if provided, or dynamic auth + * providers if additional configs are given. * - * @param authProvider - * @return + * @param authConfig auth config + * @return auth provider */ - public static String resolveToToken(AuthProvider authProvider) { - if (authProvider == null) { - return null; + public static AuthProvider makeProvider(AuthConfig authConfig) { + if (authConfig == null) { + return new NullAuthProvider(); + } + + Object providerClassValue = authConfig.getProperties().get(AuthConfig.PROVIDER_CLASS); + if (providerClassValue != null) { + try { + Class<?> providerClass = Class.forName(providerClassValue.toString()); + Constructor<?> constructor = providerClass.getConstructor(AuthConfig.class); + return (AuthProvider) constructor.newInstance(authConfig); + } catch (Exception e) { + throw new IllegalStateException("Could not create AuthProvider " + providerClassValue, e); + } + } + + if (authConfig.getProperties().containsKey(StaticTokenAuthProvider.TOKEN)) { + return new StaticTokenAuthProvider(authConfig); } - return authProvider.getHttpHeaders().entrySet().stream().findFirst().map(Map.Entry::getValue) - .filter(Objects::nonNull).map(Object::toString).orElse(null); + + return new NullAuthProvider(); } /** @@ -95,7 +111,7 @@ public final class AuthProviderUtils { * @param headers header map * @return list of http headers */ - public static List<Header> toHeaders(@Nullable Map<String, Object> headers) { + public static List<Header> toRequestHeaders(@Nullable Map<String, Object> headers) { if (headers == null) { return Collections.emptyList(); } @@ -108,10 +124,40 @@ public final class AuthProviderUtils { * @param authProvider auth provider * @return list of http headers */ - public static List<Header> toHeaders(@Nullable AuthProvider authProvider) { + public static List<Header> toRequestHeaders(@Nullable AuthProvider authProvider) { if (authProvider == null) { return Collections.emptyList(); } - return toHeaders(authProvider.getHttpHeaders()); + return toRequestHeaders(authProvider.getRequestHeaders()); + } + + /** + * Convenience helper to convert an optional authProvider to a static job spec token + * @param authProvider auth provider + * @return static token + */ + public static String toTaskToken(@Nullable AuthProvider authProvider) { + if (authProvider == null) { + return null; + } + return authProvider.getTaskToken(); + } + + /** + * Helper to extract string values from complex AuthConfig instance. + * + * @param config auth config + * @param key config key + * @param defaultValue default value + * @return config value + */ + static String getOrDefault(AuthConfig config, String key, String defaultValue) { + if (config == null || !config.getProperties().containsKey(key)) { + return defaultValue; + } + if (config.getProperties().get(key) instanceof String) { + return (String) config.getProperties().get(key); + } + throw new IllegalArgumentException("Expected String but got " + config.getProperties().get(key).getClass()); } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java b/pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java index 80756f3aae..e3de707d93 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/auth/NullAuthProvider.java @@ -24,8 +24,21 @@ import org.apache.pinot.spi.auth.AuthProvider; public class NullAuthProvider implements AuthProvider { + public NullAuthProvider() { + // left blank + } + + public NullAuthProvider(AuthConfig ignore) { + // left blank + } + @Override - public Map<String, Object> getHttpHeaders() { + public Map<String, Object> getRequestHeaders() { return Collections.emptyMap(); } + + @Override + public String getTaskToken() { + return null; + } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/auth/StaticTokenAuthProvider.java b/pinot-common/src/main/java/org/apache/pinot/common/auth/StaticTokenAuthProvider.java index fe0406598d..f686b732d0 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/auth/StaticTokenAuthProvider.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/auth/StaticTokenAuthProvider.java @@ -25,9 +25,13 @@ import org.apache.pinot.spi.auth.AuthProvider; public class StaticTokenAuthProvider implements AuthProvider { - final String _header; - final String _prefix; - final String _token; + public static final String HEADER = "header"; + public static final String PREFIX = "prefix"; + public static final String TOKEN = "token"; + + protected final String _header; + protected final String _prefix; + protected final String _token; public StaticTokenAuthProvider(String token) { _header = HttpHeaders.AUTHORIZATION; @@ -35,14 +39,27 @@ public class StaticTokenAuthProvider implements AuthProvider { _token = token; } - public StaticTokenAuthProvider(String header, String prefix, String token) { - _header = header; - _prefix = prefix; - _token = token; + public StaticTokenAuthProvider(AuthConfig authConfig) { + _header = AuthProviderUtils.getOrDefault(authConfig, HEADER, HttpHeaders.AUTHORIZATION); + _prefix = AuthProviderUtils.getOrDefault(authConfig, PREFIX, "Basic"); + _token = authConfig.getProperties().get(TOKEN).toString(); } @Override - public Map<String, Object> getHttpHeaders() { - return Collections.singletonMap(_header, _prefix + _token); + public Map<String, Object> getRequestHeaders() { + return Collections.singletonMap(_header, makeToken()); + } + + @Override + public String getTaskToken() { + return makeToken(); + } + + private String makeToken() { + String token = _token; + if (token.startsWith(_prefix)) { + return token; + } + return _prefix + " " + _token; } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/auth/UrlAuthProvider.java b/pinot-common/src/main/java/org/apache/pinot/common/auth/UrlAuthProvider.java index 0c8653453c..961e0ab69b 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/auth/UrlAuthProvider.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/auth/UrlAuthProvider.java @@ -30,36 +30,53 @@ import org.apache.pinot.spi.auth.AuthProvider; public class UrlAuthProvider implements AuthProvider { - final String _header; - final String _prefix; - final URL _url; + public static final String HEADER = "header"; + public static final String PREFIX = "prefix"; + public static final String URL = "url"; + + protected final String _header; + protected final String _prefix; + protected final URL _url; public UrlAuthProvider(String url) { try { _header = HttpHeaders.AUTHORIZATION; - _prefix = "Bearer "; + _prefix = "Bearer"; _url = new URL(url); } catch (MalformedURLException e) { throw new IllegalArgumentException(e); } } - public UrlAuthProvider(String header, String prefix, String url) { + public UrlAuthProvider(AuthConfig authConfig) { try { - _header = header; - _prefix = prefix; - _url = new URL(url); + _header = AuthProviderUtils.getOrDefault(authConfig, HEADER, HttpHeaders.AUTHORIZATION); + _prefix = AuthProviderUtils.getOrDefault(authConfig, PREFIX, "Bearer"); + _url = new URL(authConfig.getProperties().get(URL).toString()); } catch (MalformedURLException e) { throw new IllegalArgumentException(e); } } @Override - public Map<String, Object> getHttpHeaders() { + public Map<String, Object> getRequestHeaders() { + return Collections.singletonMap(_header, makeToken()); + } + + @Override + public String getTaskToken() { + return makeToken(); + } + + private String makeToken() { try { - return Collections.singletonMap(_header, _prefix + IOUtils.toString(_url, StandardCharsets.UTF_8)); + String token = IOUtils.toString(_url, StandardCharsets.UTF_8); + if (token.startsWith(_prefix)) { + return token; + } + return _prefix + " " + token; } catch (IOException e) { - throw new IllegalArgumentException("Could not access auth url", e); + throw new IllegalArgumentException("Could not resolve auth url " + _url, e); } } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java index 10429387a5..01d461782e 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java @@ -408,7 +408,7 @@ public class FileUploadDownloadClient implements AutoCloseable { RequestBuilder requestBuilder = RequestBuilder.post(uri).setVersion(HttpVersion.HTTP_1_1) .setHeader(HttpHeaders.CONTENT_TYPE, HttpClient.JSON_CONTENT_TYPE) .setEntity(new StringEntity(jsonRequestBody, ContentType.APPLICATION_JSON)); - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); HttpClient.setTimeout(requestBuilder, socketTimeoutMs); return requestBuilder.build(); } @@ -417,7 +417,7 @@ public class FileUploadDownloadClient implements AutoCloseable { @Nullable AuthProvider authProvider) { RequestBuilder requestBuilder = RequestBuilder.post(uri).setVersion(HttpVersion.HTTP_1_1) .setHeader(HttpHeaders.CONTENT_TYPE, HttpClient.JSON_CONTENT_TYPE); - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); HttpClient.setTimeout(requestBuilder, socketTimeoutMs); return requestBuilder.build(); } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/BaseSegmentFetcher.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/BaseSegmentFetcher.java index 21479feece..d99c256233 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/BaseSegmentFetcher.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/BaseSegmentFetcher.java @@ -38,8 +38,6 @@ public abstract class BaseSegmentFetcher implements SegmentFetcher { public static final String RETRY_COUNT_CONFIG_KEY = "retry.count"; public static final String RETRY_WAIT_MS_CONFIG_KEY = "retry.wait.ms"; public static final String RETRY_DELAY_SCALE_FACTOR_CONFIG_KEY = "retry.delay.scale.factor"; - public static final String AUTH_TOKEN = CommonConstants.KEY_OF_AUTH_TOKEN; - public static final String AUTH_TOKEN_URL = CommonConstants.KEY_OF_AUTH_TOKEN_URL; public static final int DEFAULT_RETRY_COUNT = 3; public static final int DEFAULT_RETRY_WAIT_MS = 100; public static final int DEFAULT_RETRY_DELAY_SCALE_FACTOR = 5; @@ -56,7 +54,7 @@ public abstract class BaseSegmentFetcher implements SegmentFetcher { _retryCount = config.getProperty(RETRY_COUNT_CONFIG_KEY, DEFAULT_RETRY_COUNT); _retryWaitMs = config.getProperty(RETRY_WAIT_MS_CONFIG_KEY, DEFAULT_RETRY_WAIT_MS); _retryDelayScaleFactor = config.getProperty(RETRY_DELAY_SCALE_FACTOR_CONFIG_KEY, DEFAULT_RETRY_DELAY_SCALE_FACTOR); - _authProvider = AuthProviderUtils.inferProvider(config.getProperty(AUTH_TOKEN), config.getProperty(AUTH_TOKEN_URL)); + _authProvider = AuthProviderUtils.extractAuthProvider(config, CommonConstants.KEY_OF_AUTH); doInit(config); _logger .info("Initialized with retryCount: {}, retryWaitMs: {}, retryDelayScaleFactor: {}", _retryCount, _retryWaitMs, diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/SegmentFetcherFactory.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/SegmentFetcherFactory.java index 577991193f..f9835206d7 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/SegmentFetcherFactory.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/fetcher/SegmentFetcherFactory.java @@ -24,7 +24,8 @@ import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; -import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.common.auth.AuthConfig; +import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.spi.crypt.PinotCrypter; import org.apache.pinot.spi.crypt.PinotCrypterFactory; import org.apache.pinot.spi.env.PinotConfiguration; @@ -38,9 +39,8 @@ public class SegmentFetcherFactory { static final String SEGMENT_FETCHER_CLASS_KEY_SUFFIX = ".class"; private static final String PROTOCOLS_KEY = "protocols"; - private static final String AUTH_TOKEN_KEY = CommonConstants.KEY_OF_AUTH_TOKEN; - private static final String AUTH_TOKEN_URL_KEY = CommonConstants.KEY_OF_AUTH_TOKEN_URL; private static final String ENCODED_SUFFIX = ".enc"; + private static final String AUTH_KEY = CommonConstants.KEY_OF_AUTH; private static final Logger LOGGER = LoggerFactory.getLogger(SegmentFetcherFactory.class); @@ -91,15 +91,14 @@ public class SegmentFetcherFactory { segmentFetcher = (SegmentFetcher) Class.forName(segmentFetcherClassName).newInstance(); } - String authToken = config.getProperty(AUTH_TOKEN_KEY); - String authTokenUrl = config.getProperty(AUTH_TOKEN_URL_KEY); + AuthConfig authConfig = AuthProviderUtils.extractAuthConfig(config, AUTH_KEY); + + PinotConfiguration subConfig = config.subset(protocol); + AuthConfig subAuthConfig = AuthProviderUtils.extractAuthConfig(subConfig, AUTH_KEY); Map<String, Object> subConfigMap = config.subset(protocol).toMap(); - if (!subConfigMap.containsKey(AUTH_TOKEN_KEY) && StringUtils.isNotBlank(authToken)) { - subConfigMap.put(AUTH_TOKEN_KEY, authToken); - } - if (!subConfigMap.containsKey(AUTH_TOKEN_URL_KEY) && StringUtils.isNotBlank(authTokenUrl)) { - subConfigMap.put(AUTH_TOKEN_URL_KEY, authTokenUrl); + if (subAuthConfig.getProperties().isEmpty() && !authConfig.getProperties().isEmpty()) { + authConfig.getProperties().forEach((key, value) -> subConfigMap.put(AUTH_KEY + "." + key, value)); } segmentFetcher.init(new PinotConfiguration(subConfigMap)); diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/http/HttpClient.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/http/HttpClient.java index 2327066656..79b050bf68 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/http/HttpClient.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/http/HttpClient.java @@ -126,7 +126,7 @@ public class HttpClient implements AutoCloseable { @Nullable AuthProvider authProvider) throws IOException { RequestBuilder requestBuilder = RequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1); - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); if (MapUtils.isNotEmpty(headers)) { for (Map.Entry<String, String> header : headers.entrySet()) { requestBuilder.addHeader(header.getKey(), header.getValue()); @@ -155,7 +155,7 @@ public class HttpClient implements AutoCloseable { @Nullable AuthProvider authProvider) throws IOException { RequestBuilder requestBuilder = RequestBuilder.delete(uri).setVersion(HttpVersion.HTTP_1_1); - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); if (MapUtils.isNotEmpty(headers)) { for (Map.Entry<String, String> header : headers.entrySet()) { requestBuilder.addHeader(header.getKey(), header.getValue()); @@ -183,7 +183,7 @@ public class HttpClient implements AutoCloseable { if (payload != null) { requestBuilder.setEntity(payload); } - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); if (MapUtils.isNotEmpty(headers)) { for (Map.Entry<String, String> header : headers.entrySet()) { requestBuilder.addHeader(header.getKey(), header.getValue()); @@ -210,7 +210,7 @@ public class HttpClient implements AutoCloseable { if (payload != null) { requestBuilder.setEntity(payload); } - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); if (MapUtils.isNotEmpty(headers)) { for (Map.Entry<String, String> header : headers.entrySet()) { requestBuilder.addHeader(header.getKey(), header.getValue()); @@ -453,7 +453,7 @@ public class HttpClient implements AutoCloseable { private static HttpUriRequest getDownloadFileRequest(URI uri, int socketTimeoutMs, AuthProvider authProvider, List<Header> httpHeaders) { RequestBuilder requestBuilder = RequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1); - AuthProviderUtils.toHeaders(authProvider).forEach(requestBuilder::addHeader); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); HttpClient.setTimeout(requestBuilder, socketTimeoutMs); String userInfo = uri.getUserInfo(); if (userInfo != null) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotIngestionRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotIngestionRestletResource.java index bc332e0678..564d0ed566 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotIngestionRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotIngestionRestletResource.java @@ -196,20 +196,15 @@ public class PinotIngestionRestletResource { }); Schema schema = _pinotHelixResourceManager.getTableSchema(tableNameWithType); + AuthProvider authProvider = AuthProviderUtils.extractAuthProvider(_controllerConf, + CommonConstants.Controller.PREFIX_OF_CONFIG_OF_SEGMENT_FETCHER_FACTORY + ".auth"); + FileIngestionHelper fileIngestionHelper = new FileIngestionHelper(tableConfig, schema, batchConfigMap, getControllerUri(), - new File(_controllerConf.getDataDir(), UPLOAD_DIR), getAuthProvider()); + new File(_controllerConf.getDataDir(), UPLOAD_DIR), authProvider); return fileIngestionHelper.buildSegmentAndPush(payload); } - private AuthProvider getAuthProvider() { - String authToken = _controllerConf - .getProperty(CommonConstants.Controller.PREFIX_OF_CONFIG_OF_SEGMENT_FETCHER_FACTORY + ".auth.token"); - String authTokenUrl = _controllerConf - .getProperty(CommonConstants.Controller.PREFIX_OF_CONFIG_OF_SEGMENT_FETCHER_FACTORY + ".auth.token.url"); - return AuthProviderUtils.inferProvider(authToken, authTokenUrl); - } - private URI getControllerUri() { try { return new URI(_controllerConf.generateVipUrl()); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java index 9c8b301275..830332edc3 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java @@ -36,7 +36,6 @@ import org.apache.commons.io.FileUtils; import org.apache.helix.HelixManager; import org.apache.helix.ZNRecord; import org.apache.helix.store.zk.ZkHelixPropertyStore; -import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.common.metadata.ZKMetadataProvider; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.ServerGauge; @@ -94,7 +93,7 @@ public abstract class BaseTableDataManager implements TableDataManager { @Override public void init(TableDataManagerConfig tableDataManagerConfig, String instanceId, ZkHelixPropertyStore<ZNRecord> propertyStore, ServerMetrics serverMetrics, HelixManager helixManager, - @Nullable LoadingCache<Pair<String, String>, SegmentErrorInfo> errorCache) { + @Nullable LoadingCache<Pair<String, String>, SegmentErrorInfo> errorCache, AuthProvider authProvider) { LOGGER.info("Initializing table data manager for table: {}", tableDataManagerConfig.getTableName()); _tableDataManagerConfig = tableDataManagerConfig; @@ -102,8 +101,7 @@ public abstract class BaseTableDataManager implements TableDataManager { _propertyStore = propertyStore; _serverMetrics = serverMetrics; _helixManager = helixManager; - _authProvider = AuthProviderUtils.inferProvider(tableDataManagerConfig.getAuthToken(), - tableDataManagerConfig.getAuthTokenUrl()); + _authProvider = authProvider; _tableNameWithType = tableDataManagerConfig.getTableName(); _tableDataDir = tableDataManagerConfig.getDataDir(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/TableDataManagerProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/TableDataManagerProvider.java index 42994edcfd..0603b2f05a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/TableDataManagerProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/TableDataManagerProvider.java @@ -28,6 +28,7 @@ import org.apache.pinot.common.restlet.resources.SegmentErrorInfo; import org.apache.pinot.core.data.manager.realtime.RealtimeTableDataManager; import org.apache.pinot.segment.local.data.manager.TableDataManager; import org.apache.pinot.segment.local.data.manager.TableDataManagerConfig; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.utils.Pair; @@ -51,7 +52,7 @@ public class TableDataManagerProvider { public static TableDataManager getTableDataManager(TableDataManagerConfig tableDataManagerConfig, String instanceId, ZkHelixPropertyStore<ZNRecord> propertyStore, ServerMetrics serverMetrics, HelixManager helixManager, - LoadingCache<Pair<String, String>, SegmentErrorInfo> errorCache) { + LoadingCache<Pair<String, String>, SegmentErrorInfo> errorCache, AuthProvider authProvider) { TableDataManager tableDataManager; switch (TableType.valueOf(tableDataManagerConfig.getTableDataManagerType())) { case OFFLINE: @@ -67,7 +68,8 @@ public class TableDataManagerProvider { default: throw new IllegalStateException(); } - tableDataManager.init(tableDataManagerConfig, instanceId, propertyStore, serverMetrics, helixManager, errorCache); + tableDataManager.init(tableDataManagerConfig, instanceId, propertyStore, serverMetrics, helixManager, errorCache, + authProvider); return tableDataManager; } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/Server2ControllerSegmentUploader.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/Server2ControllerSegmentUploader.java index 023e60399e..2559616b49 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/Server2ControllerSegmentUploader.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/Server2ControllerSegmentUploader.java @@ -74,7 +74,7 @@ public class Server2ControllerSegmentUploader implements SegmentUploader { try { String responseStr = _fileUploadDownloadClient .uploadSegment(_controllerSegmentUploadCommitUrl, _segmentName, segmentFile, - AuthProviderUtils.toHeaders(_authProvider), null, _segmentUploadRequestTimeoutMs).getResponse(); + AuthProviderUtils.toRequestHeaders(_authProvider), null, _segmentUploadRequestTimeoutMs).getResponse(); response = SegmentCompletionProtocol.Response.fromJsonString(responseStr); _segmentLogger.info("Controller response {} for {}", response.toJsonString(), _controllerSegmentUploadCommitUrl); if (response.getStatus().equals(SegmentCompletionProtocol.ControllerResponseStatus.NOT_LEADER)) { diff --git a/pinot-core/src/main/java/org/apache/pinot/server/realtime/ServerSegmentCompletionProtocolHandler.java b/pinot-core/src/main/java/org/apache/pinot/server/realtime/ServerSegmentCompletionProtocolHandler.java index cd5275d367..644ec3fbdb 100644 --- a/pinot-core/src/main/java/org/apache/pinot/server/realtime/ServerSegmentCompletionProtocolHandler.java +++ b/pinot-core/src/main/java/org/apache/pinot/server/realtime/ServerSegmentCompletionProtocolHandler.java @@ -75,8 +75,7 @@ public class ServerSegmentCompletionProtocolHandler { _segmentUploadRequestTimeoutMs = uploaderConfig .getProperty(CONFIG_OF_SEGMENT_UPLOAD_REQUEST_TIMEOUT_MS, DEFAULT_SEGMENT_UPLOAD_REQUEST_TIMEOUT_MS); - _authProvider = AuthProviderUtils.inferProvider(uploaderConfig.getProperty(CONFIG_OF_SEGMENT_UPLOADER_AUTH_TOKEN), - uploaderConfig.getProperty(CONFIG_OF_SEGMENT_UPLOADER_AUTH_TOKEN_URL)); + _authProvider = AuthProviderUtils.extractAuthProvider(uploaderConfig, CONFIG_OF_SEGMENT_UPLOADER_AUTH); } public ServerSegmentCompletionProtocolHandler(ServerMetrics serverMetrics, String tableNameWithType) { @@ -215,7 +214,7 @@ public class ServerSegmentCompletionProtocolHandler { SegmentCompletionProtocol.Response response; try { String responseStr = _fileUploadDownloadClient - .sendSegmentCompletionProtocolRequest(new URI(url), AuthProviderUtils.toHeaders(_authProvider), null, + .sendSegmentCompletionProtocolRequest(new URI(url), AuthProviderUtils.toRequestHeaders(_authProvider), null, DEFAULT_OTHER_REQUESTS_TIMEOUT).getResponse(); response = SegmentCompletionProtocol.Response.fromJsonString(responseStr); LOGGER.info("Controller response {} for {}", response.toJsonString(), url); @@ -240,7 +239,7 @@ public class ServerSegmentCompletionProtocolHandler { SegmentCompletionProtocol.Response response; try { String responseStr = _fileUploadDownloadClient - .uploadSegmentMetadataFiles(new URI(url), metadataFiles, AuthProviderUtils.toHeaders(_authProvider), + .uploadSegmentMetadataFiles(new URI(url), metadataFiles, AuthProviderUtils.toRequestHeaders(_authProvider), null, _segmentUploadRequestTimeoutMs).getResponse(); response = SegmentCompletionProtocol.Response.fromJsonString(responseStr); LOGGER.info("Controller response {} for {}", response.toJsonString(), url); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerAcquireSegmentTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerAcquireSegmentTest.java index 86299bfb92..882e935708 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerAcquireSegmentTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerAcquireSegmentTest.java @@ -119,7 +119,7 @@ public class BaseTableDataManagerAcquireSegmentTest { when(config.getDataDir()).thenReturn(_tmpDir.getAbsolutePath()); } tableDataManager.init(config, "dummyInstance", mock(ZkHelixPropertyStore.class), - new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()), mock(HelixManager.class), null); + new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()), mock(HelixManager.class), null, null); tableDataManager.start(); Field segsMapField = BaseTableDataManager.class.getDeclaredField("_segmentDataManagerMap"); segsMapField.setAccessible(true); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java index fac8977c5f..d8104c4957 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/BaseTableDataManagerTest.java @@ -529,7 +529,7 @@ public class BaseTableDataManagerTest { OfflineTableDataManager tableDataManager = new OfflineTableDataManager(); tableDataManager.init(config, "dummyInstance", mock(ZkHelixPropertyStore.class), - new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()), mock(HelixManager.class), null); + new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()), mock(HelixManager.class), null, null); tableDataManager.start(); return tableDataManager; } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManagerTest.java index 7eb74e5fc7..3d73e20b33 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManagerTest.java @@ -124,7 +124,7 @@ public class DimensionTableDataManagerTest { when(config.getDataDir()).thenReturn(INDEX_DIR.getAbsolutePath()); } tableDataManager.init(config, "dummyInstance", mockPropertyStore(), - new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()), helixManager, null); + new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()), helixManager, null, null); tableDataManager.start(); return tableDataManager; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/LLRealtimeSegmentDataManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/LLRealtimeSegmentDataManagerTest.java index 6a76f4cedb..c77aea3e8a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/LLRealtimeSegmentDataManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/LLRealtimeSegmentDataManagerTest.java @@ -726,7 +726,7 @@ public class LLRealtimeSegmentDataManagerTest { TableDataManager tableDataManager = TableDataManagerProvider.getTableDataManager(tableDataManagerConfig, "testInstance", propertyStore, - mock(ServerMetrics.class), mock(HelixManager.class), null); + mock(ServerMetrics.class), mock(HelixManager.class), null, null); tableDataManager.start(); tableDataManager.shutDown(); Assert.assertFalse(SegmentBuildTimeLeaseExtender.isExecutorShutdown()); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorExceptionsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorExceptionsTest.java index 9ab49c0497..0c5e7a68c1 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorExceptionsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorExceptionsTest.java @@ -136,7 +136,7 @@ public class QueryExecutorExceptionsTest { @SuppressWarnings("unchecked") TableDataManager tableDataManager = TableDataManagerProvider.getTableDataManager(tableDataManagerConfig, "testInstance", - mock(ZkHelixPropertyStore.class), mock(ServerMetrics.class), mock(HelixManager.class), null); + mock(ZkHelixPropertyStore.class), mock(ServerMetrics.class), mock(HelixManager.class), null, null); tableDataManager.start(); //we don't add index segments to the data manager to simulate numSegmentsAcquired < numSegmentsQueried InstanceDataManager instanceDataManager = mock(InstanceDataManager.class); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorTest.java index b87502b220..242def7df1 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/executor/QueryExecutorTest.java @@ -130,7 +130,7 @@ public class QueryExecutorTest { @SuppressWarnings("unchecked") TableDataManager tableDataManager = TableDataManagerProvider.getTableDataManager(tableDataManagerConfig, "testInstance", - mock(ZkHelixPropertyStore.class), mock(ServerMetrics.class), mock(HelixManager.class), null); + mock(ZkHelixPropertyStore.class), mock(ServerMetrics.class), mock(HelixManager.class), null, null); tableDataManager.start(); for (ImmutableSegment indexSegment : _indexSegments) { tableDataManager.addSegment(indexSegment); diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/SegmentWithNullValueVectorTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/SegmentWithNullValueVectorTest.java index d7c115f448..4cf75f97e5 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/SegmentWithNullValueVectorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/SegmentWithNullValueVectorTest.java @@ -145,7 +145,7 @@ public class SegmentWithNullValueVectorTest { TableDataManager tableDataManager = TableDataManagerProvider.getTableDataManager(tableDataManagerConfig, "testInstance", Mockito.mock(ZkHelixPropertyStore.class), Mockito.mock(ServerMetrics.class), - Mockito.mock(HelixManager.class), null); + Mockito.mock(HelixManager.class), null, null); tableDataManager.start(); tableDataManager.addSegment(_segment); _instanceDataManager = Mockito.mock(InstanceDataManager.class); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BasicAuthBatchIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BasicAuthBatchIntegrationTest.java index 3332f9d257..c9056b03c6 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BasicAuthBatchIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BasicAuthBatchIntegrationTest.java @@ -27,6 +27,7 @@ import java.io.IOException; import java.util.Map; import org.apache.commons.io.FileUtils; import org.apache.commons.io.IOUtils; +import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.controller.helix.core.minion.generator.PinotTaskGenerator; import org.apache.pinot.minion.executor.PinotTaskExecutor; import org.apache.pinot.spi.env.PinotConfiguration; @@ -159,7 +160,7 @@ public class BasicAuthBatchIntegrationTest extends ClusterTest { .write(jobFileContents.replaceAll("9000", String.valueOf(getControllerPort())), new FileOutputStream(jobFile)); new BootstrapTableTool("http", "localhost", getControllerPort(), baseDir.getAbsolutePath(), - AUTH_TOKEN, null).execute(); + AuthProviderUtils.makeProvider(AUTH_TOKEN)).execute(); Thread.sleep(5000); diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java index 0ce4b7177c..f3eca88572 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java @@ -34,7 +34,6 @@ import org.apache.http.NameValuePair; import org.apache.http.message.BasicHeader; import org.apache.http.message.BasicNameValuePair; import org.apache.pinot.common.auth.AuthProviderUtils; -import org.apache.pinot.common.auth.StaticTokenAuthProvider; import org.apache.pinot.common.metadata.segment.SegmentZKMetadataCustomMapModifier; import org.apache.pinot.common.restlet.resources.StartReplaceSegmentsRequest; import org.apache.pinot.common.utils.FileUploadDownloadClient; @@ -134,7 +133,7 @@ public abstract class BaseMultipleSegmentsConversionExecutor extends BaseTaskExe String downloadURLString = configs.get(MinionConstants.DOWNLOAD_URL_KEY); String[] downloadURLs = downloadURLString.split(MinionConstants.URL_SEPARATOR); String uploadURL = configs.get(MinionConstants.UPLOAD_URL_KEY); - AuthProvider authProvider = AuthProviderUtils.inferProvider(configs.get(MinionConstants.AUTH_TOKEN)); + AuthProvider authProvider = AuthProviderUtils.makeProvider(configs.get(MinionConstants.AUTH_TOKEN)); LOGGER.info("Start executing {} on table: {}, input segments: {} with downloadURLs: {}, uploadURL: {}", taskType, tableNameWithType, inputSegmentNames, downloadURLString, uploadURL); @@ -217,7 +216,7 @@ public abstract class BaseMultipleSegmentsConversionExecutor extends BaseTaskExe List<Header> httpHeaders = new ArrayList<>(); httpHeaders.add(segmentZKMetadataCustomMapModifierHeader); - httpHeaders.addAll(AuthProviderUtils.toHeaders(authProvider.getHttpHeaders())); + httpHeaders.addAll(AuthProviderUtils.toRequestHeaders(authProvider)); // Set parameters for upload request NameValuePair enableParallelPushProtectionParameter = @@ -270,7 +269,7 @@ public abstract class BaseMultipleSegmentsConversionExecutor extends BaseTaskExe Map<String, String> configs = pinotTaskConfig.getConfigs(); _tableNameWithType = configs.get(MinionConstants.TABLE_NAME_KEY); _uploadURL = configs.get(MinionConstants.UPLOAD_URL_KEY); - _authProvider = new StaticTokenAuthProvider(configs.get(MinionConstants.AUTH_TOKEN)); + _authProvider = AuthProviderUtils.makeProvider(configs.get(MinionConstants.AUTH_TOKEN)); _inputSegmentNames = configs.get(MinionConstants.SEGMENT_NAME_KEY); String replaceSegmentsString = configs.get(MinionConstants.ENABLE_REPLACE_SEGMENTS_KEY); _replaceSegmentsEnabled = Boolean.parseBoolean(replaceSegmentsString); diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java index e8eb4e1ba3..6c323a372c 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java @@ -70,7 +70,7 @@ public abstract class BaseSingleSegmentConversionExecutor extends BaseTaskExecut String downloadURL = configs.get(MinionConstants.DOWNLOAD_URL_KEY); String uploadURL = configs.get(MinionConstants.UPLOAD_URL_KEY); String originalSegmentCrc = configs.get(MinionConstants.ORIGINAL_SEGMENT_CRC_KEY); - AuthProvider authProvider = AuthProviderUtils.inferProvider(configs.get(MinionConstants.AUTH_TOKEN)); + AuthProvider authProvider = AuthProviderUtils.makeProvider(configs.get(MinionConstants.AUTH_TOKEN)); long currentSegmentCrc = getSegmentCrc(tableNameWithType, segmentName); if (Long.parseLong(originalSegmentCrc) != currentSegmentCrc) { @@ -151,7 +151,7 @@ public abstract class BaseSingleSegmentConversionExecutor extends BaseTaskExecut httpHeaders.add(ifMatchHeader); httpHeaders.add(refreshOnlyHeader); httpHeaders.add(segmentZKMetadataCustomMapModifierHeader); - httpHeaders.addAll(AuthProviderUtils.toHeaders(authProvider)); + httpHeaders.addAll(AuthProviderUtils.toRequestHeaders(authProvider)); // Set parameters for upload request. NameValuePair enableParallelPushProtectionParameter = diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManager.java index 3ae95a6452..205fdd796c 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManager.java @@ -33,6 +33,7 @@ import org.apache.pinot.common.restlet.resources.SegmentErrorInfo; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; import org.apache.pinot.segment.spi.ImmutableSegment; import org.apache.pinot.segment.spi.SegmentMetadata; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.Pair; @@ -49,7 +50,7 @@ public interface TableDataManager { */ void init(TableDataManagerConfig tableDataManagerConfig, String instanceId, ZkHelixPropertyStore<ZNRecord> propertyStore, ServerMetrics serverMetrics, HelixManager helixManager, - LoadingCache<Pair<String, String>, SegmentErrorInfo> errorCache); + LoadingCache<Pair<String, String>, SegmentErrorInfo> errorCache, AuthProvider authProvider); /** * Starts the table data manager. Should be called only once after table data manager gets initialized but before diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManagerConfig.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManagerConfig.java index e085a72abf..fdd6f6f72f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManagerConfig.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/data/manager/TableDataManagerConfig.java @@ -36,8 +36,7 @@ public class TableDataManagerConfig { private static final String TABLE_DATA_MANAGER_CONSUMER_DIRECTORY = "consumerDirectory"; private static final String TABLE_DATA_MANAGER_NAME = "name"; private static final String TABLE_IS_DIMENSION = "isDimTable"; - private static final String TABLE_DATA_MANGER_AUTH_TOKEN = "authToken"; - private static final String TABLE_DATA_MANGER_AUTH_TOKEN_URL = "authTokenUrl"; + private static final String TABLE_DATA_MANGER_AUTH = "auth"; private final Configuration _tableDataManagerConfig; @@ -69,13 +68,9 @@ public class TableDataManagerConfig { return _tableDataManagerConfig.getBoolean(TABLE_IS_DIMENSION); } - public String getAuthToken() { - return _tableDataManagerConfig.getString(TABLE_DATA_MANGER_AUTH_TOKEN); - } - - public String getAuthTokenUrl() { - return _tableDataManagerConfig.getString(TABLE_DATA_MANGER_AUTH_TOKEN_URL); - } +// public AuthConfig getAuthConfig() { +// return AuthProviderUtils.extractAuthConfig(_tableDataManagerConfig, TABLE_DATA_MANGER_AUTH); +// } public static TableDataManagerConfig getDefaultHelixTableDataManagerConfig( InstanceDataManagerConfig instanceDataManagerConfig, String tableNameWithType) { @@ -87,17 +82,16 @@ public class TableDataManagerConfig { TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableNameWithType); Preconditions.checkNotNull(tableType); defaultConfig.addProperty(TABLE_DATA_MANAGER_TYPE, tableType.name()); - defaultConfig.addProperty(TABLE_DATA_MANGER_AUTH_TOKEN, instanceDataManagerConfig.getAuthToken()); - defaultConfig.addProperty(TABLE_DATA_MANGER_AUTH_TOKEN_URL, instanceDataManagerConfig.getAuthTokenUrl()); +// instanceDataManagerConfig.getAuthConfig().forEach((key, value) -> defaultConfig +// .addProperty(TABLE_DATA_MANGER_AUTH + "." + key, value)); return new TableDataManagerConfig(defaultConfig); } - public void overrideConfigs(TableConfig tableConfig, String authToken) { + public void overrideConfigs(TableConfig tableConfig) { // Override table level configs _tableDataManagerConfig.addProperty(TABLE_IS_DIMENSION, tableConfig.isDimTable()); - _tableDataManagerConfig.addProperty(TABLE_DATA_MANGER_AUTH_TOKEN, authToken); // If we wish to override some table level configs using table config, override them here // Note: the configs in TableDataManagerConfig is immutable once the table is created, which mean it will not pick diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/IngestionUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/IngestionUtils.java index a67c0e6151..0f1949f4cd 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/IngestionUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/IngestionUtils.java @@ -271,9 +271,8 @@ public final class IngestionUtils { spec.setPushJobSpec(pushJobSpec); spec.setTableSpec(tableSpec); spec.setPinotClusterSpecs(pinotClusterSpecs); - if (authProvider != null) { - spec.setAuthToken(AuthProviderUtils.resolveToToken(authProvider)); - } + spec.setAuthToken(AuthProviderUtils.toTaskToken(authProvider)); + return spec; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java index 68067f34f1..07d1558780 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java @@ -107,7 +107,7 @@ public class SegmentPushUtils implements Serializable { String fileName = tarFile.getName(); Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT)); String segmentName = fileName.substring(0, fileName.length() - Constants.TAR_GZ_FILE_EXT.length()); - AuthProvider authProvider = AuthProviderUtils.inferProvider(spec.getAuthToken()); + AuthProvider authProvider = AuthProviderUtils.makeProvider(spec.getAuthToken()); for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) { URI controllerURI; try { @@ -128,7 +128,7 @@ public class SegmentPushUtils implements Serializable { try (InputStream inputStream = fileSystem.open(tarFileURI)) { SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.uploadSegment(FileUploadDownloadClient.getUploadSegmentURI(controllerURI), - segmentName, inputStream, AuthProviderUtils.toHeaders(authProvider), + segmentName, inputStream, AuthProviderUtils.toRequestHeaders(authProvider), FileUploadDownloadClient.makeTableParam(tableName), tableName, tableType); LOGGER.info("Response for pushing table {} segment {} to location {} - {}: {}", tableName, segmentName, controllerURI, response.getStatusCode(), response.getResponse()); @@ -165,7 +165,7 @@ public class SegmentPushUtils implements Serializable { for (String segmentUri : segmentUris) { URI segmentURI = URI.create(segmentUri); PinotFS outputDirFS = PinotFSFactory.create(segmentURI.getScheme()); - AuthProvider authProvider = AuthProviderUtils.inferProvider(spec.getAuthToken()); + AuthProvider authProvider = AuthProviderUtils.makeProvider(spec.getAuthToken()); for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) { URI controllerURI; try { @@ -186,8 +186,8 @@ public class SegmentPushUtils implements Serializable { try { SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT .sendSegmentUri(FileUploadDownloadClient.getUploadSegmentURI(controllerURI), segmentUri, - AuthProviderUtils.toHeaders(authProvider), FileUploadDownloadClient.makeTableParam(tableName), - HttpClient.DEFAULT_SOCKET_TIMEOUT_MS); + AuthProviderUtils.toRequestHeaders(authProvider), + FileUploadDownloadClient.makeTableParam(tableName), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS); LOGGER.info("Response for pushing table {} segment uri {} to location {} - {}: {}", tableName, segmentUri, controllerURI, response.getStatusCode(), response.getResponse()); return true; @@ -240,7 +240,7 @@ public class SegmentPushUtils implements Serializable { Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT)); String segmentName = fileName.substring(0, fileName.length() - Constants.TAR_GZ_FILE_EXT.length()); File segmentMetadataFile = generateSegmentMetadataFile(fileSystem, URI.create(tarFilePath)); - AuthProvider authProvider = AuthProviderUtils.inferProvider(spec.getAuthToken()); + AuthProvider authProvider = AuthProviderUtils.makeProvider(spec.getAuthToken()); try { for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) { URI controllerURI; @@ -264,7 +264,7 @@ public class SegmentPushUtils implements Serializable { headers.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI, segmentUriPath)); headers.add(new BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE, FileUploadDownloadClient.FileUploadType.METADATA.toString())); - headers.addAll(AuthProviderUtils.toHeaders(authProvider)); + headers.addAll(AuthProviderUtils.toRequestHeaders(authProvider)); SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT .uploadSegmentMetadata(FileUploadDownloadClient.getUploadSegmentURI(controllerURI), segmentName, diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java index 3cc12961ca..998fcf6d62 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManager.java @@ -42,6 +42,7 @@ import org.apache.helix.HelixManager; import org.apache.helix.ZNRecord; import org.apache.helix.model.ExternalView; import org.apache.helix.store.zk.ZkHelixPropertyStore; +import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.common.metadata.ZKMetadataProvider; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.ServerMetrics; @@ -60,6 +61,7 @@ import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; import org.apache.pinot.segment.spi.ImmutableSegment; import org.apache.pinot.segment.spi.IndexSegment; import org.apache.pinot.segment.spi.SegmentMetadata; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.env.PinotConfiguration; @@ -86,7 +88,7 @@ public class HelixInstanceDataManager implements InstanceDataManager { private HelixManager _helixManager; private ServerMetrics _serverMetrics; private ZkHelixPropertyStore<ZNRecord> _propertyStore; - private String _authToken; + private AuthProvider _authProvider; private SegmentUploader _segmentUploader; // Fixed size LRU cache for storing last N errors on the instance. @@ -103,7 +105,7 @@ public class HelixInstanceDataManager implements InstanceDataManager { _instanceId = _instanceDataManagerConfig.getInstanceId(); _helixManager = helixManager; _serverMetrics = serverMetrics; - _authToken = config.getProperty(CommonConstants.Server.CONFIG_OF_AUTH_TOKEN); + _authProvider = AuthProviderUtils.extractAuthProvider(config, CommonConstants.Server.CONFIG_OF_AUTH); _segmentUploader = new PinotFSSegmentUploader(_instanceDataManagerConfig.getSegmentStoreUri(), PinotFSSegmentUploader.DEFAULT_SEGMENT_UPLOAD_TIMEOUT_MILLIS); @@ -180,10 +182,10 @@ public class HelixInstanceDataManager implements InstanceDataManager { LOGGER.info("Creating table data manager for table: {}", tableNameWithType); TableDataManagerConfig tableDataManagerConfig = TableDataManagerConfig.getDefaultHelixTableDataManagerConfig(_instanceDataManagerConfig, tableNameWithType); - tableDataManagerConfig.overrideConfigs(tableConfig, _authToken); + tableDataManagerConfig.overrideConfigs(tableConfig); TableDataManager tableDataManager = TableDataManagerProvider.getTableDataManager(tableDataManagerConfig, _instanceId, _propertyStore, - _serverMetrics, _helixManager, _errorCache); + _serverMetrics, _helixManager, _errorCache, _authProvider); tableDataManager.start(); LOGGER.info("Created table data manager for table: {}", tableNameWithType); return tableDataManager; diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java index 14b193e4dd..650b2ddc1a 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixInstanceDataManagerConfig.java @@ -62,10 +62,10 @@ public class HelixInstanceDataManagerConfig implements InstanceDataManagerConfig public static final String SEGMENT_FORMAT_VERSION = "segment.format.version"; // Key of whether to enable reloading consuming segments public static final String INSTANCE_RELOAD_CONSUMING_SEGMENT = "reload.consumingSegment"; - // Key of the auth token - public static final String AUTH_TOKEN = "auth.token"; - // Key of the auth token url - public static final String AUTH_TOKEN_URL = "auth.token.url"; + // // (legacy) key of the auth token + //public static final String LEGACY_AUTH_TOKEN = "auth.token"; + // // Key of the auth configs + //public static final String AUTH = "auth"; // Key of segment directory loader public static final String SEGMENT_DIRECTORY_LOADER = "segment.directory.loader"; @@ -213,15 +213,18 @@ public class HelixInstanceDataManagerConfig implements InstanceDataManagerConfig .getProperty(MAX_PARALLEL_SEGMENT_BUILDS, DEFAULT_MAX_PARALLEL_SEGMENT_BUILDS); } - @Override - public String getAuthToken() { - return _instanceDataManagerConfiguration.getProperty(AUTH_TOKEN); - } - - @Override - public String getAuthTokenUrl() { - return _instanceDataManagerConfiguration.getProperty(AUTH_TOKEN_URL); - } +// @Override +// public Map<String, Object> getAuthConfig() { +// Map<String, Object> prop = new HashMap<>(_instanceDataManagerConfiguration.getProperty(AUTH, Map.class)); +// +// // legacy compatibility +// String authToken = _instanceDataManagerConfiguration.getProperty(LEGACY_AUTH_TOKEN); +// if (StringUtils.isNotBlank(authToken)) { +// prop.put(LEGACY_AUTH_TOKEN, authToken); +// } +// +// return prop; +// } @Override public String getSegmentDirectoryLoader() { diff --git a/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java b/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java index 2708373a08..56026734de 100644 --- a/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java +++ b/pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java @@ -193,7 +193,7 @@ public abstract class BaseResourceTest { TableDataManager tableDataManager = new OfflineTableDataManager(); tableDataManager .init(tableDataManagerConfig, "testInstance", mock(ZkHelixPropertyStore.class), mock(ServerMetrics.class), - mock(HelixManager.class), null); + mock(HelixManager.class), null, null); tableDataManager.start(); _tableDataManagerMap.put(tableNameWithType, tableDataManager); } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/auth/AuthProvider.java b/pinot-spi/src/main/java/org/apache/pinot/spi/auth/AuthProvider.java index 1ecbcc1188..a646a33a71 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/auth/AuthProvider.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/auth/AuthProvider.java @@ -29,5 +29,7 @@ import java.util.Map; * expiring JWTs and other token rotation mechanisms. */ public interface AuthProvider { - Map<String, Object> getHttpHeaders(); + Map<String, Object> getRequestHeaders(); + + String getTaskToken(); } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java index 422add0b37..402ced8a53 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/config/instance/InstanceDataManagerConfig.java @@ -53,10 +53,8 @@ public interface InstanceDataManagerConfig { int getMaxParallelSegmentBuilds(); - String getAuthToken(); - - String getAuthTokenUrl(); - +// Map<String, Object> getAuthConfig(); +// String getSegmentDirectoryLoader(); long getErrorCacheSize(); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 37fc66fab4..d799476b27 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -33,8 +33,7 @@ public class CommonConstants { public static final String HTTP_PROTOCOL = "http"; public static final String HTTPS_PROTOCOL = "https"; - public static final String KEY_OF_AUTH_TOKEN = "auth.token"; - public static final String KEY_OF_AUTH_TOKEN_URL = "auth.token.url"; + public static final String KEY_OF_AUTH = "auth"; public static final String TABLE_NAME = "tableName"; @@ -332,7 +331,7 @@ public class CommonConstants { * Service token for accessing protected controller APIs. * E.g. null (auth disabled), "Basic abcdef..." (basic auth), "Bearer 123def..." (oauth2) */ - public static final String CONFIG_OF_AUTH_TOKEN = KEY_OF_AUTH_TOKEN; + public static final String CONFIG_OF_AUTH = KEY_OF_AUTH; // Configuration to consider the server ServiceStatus as being STARTED if the percent of resources (tables) that // are ONLINE for this this server has crossed the threshold percentage of the total number of tables @@ -433,8 +432,7 @@ public class CommonConstants { * Service token for accessing protected controller APIs. * E.g. null (auth disabled), "Basic abcdef..." (basic auth), "Bearer 123def..." (oauth2) */ - public static final String CONFIG_OF_SEGMENT_UPLOADER_AUTH_TOKEN = KEY_OF_AUTH_TOKEN; - public static final String CONFIG_OF_SEGMENT_UPLOADER_AUTH_TOKEN_URL = KEY_OF_AUTH_TOKEN_URL; + public static final String CONFIG_OF_SEGMENT_UPLOADER_AUTH = KEY_OF_AUTH; public static final int DEFAULT_SEGMENT_UPLOAD_REQUEST_TIMEOUT_MS = 300_000; public static final int DEFAULT_OTHER_REQUESTS_TIMEOUT = 10_000; diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/AuthQuickstart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/AuthQuickstart.java index 00ec73d5af..f186f95338 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/AuthQuickstart.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/AuthQuickstart.java @@ -22,7 +22,9 @@ import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.core.auth.BasicAuthUtils; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.plugin.PluginManager; @@ -33,8 +35,8 @@ public class AuthQuickstart extends Quickstart { } @Override - public String getAuthToken() { - return BasicAuthUtils.toBasicAuthToken("admin", "verysecret"); + public AuthProvider getAuthProvider() { + return AuthProviderUtils.makeProvider(BasicAuthUtils.toBasicAuthToken("admin", "verysecret")); } @Override diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/BootstrapTableTool.java b/pinot-tools/src/main/java/org/apache/pinot/tools/BootstrapTableTool.java index 5e0a4ac218..9aedbdf8f5 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/BootstrapTableTool.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/BootstrapTableTool.java @@ -34,6 +34,7 @@ import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.common.minion.MinionClient; import org.apache.pinot.common.utils.TlsUtils; import org.apache.pinot.core.common.MinionConstants; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties; import org.apache.pinot.spi.ingestion.batch.IngestionJobLauncher; @@ -54,13 +55,12 @@ public class BootstrapTableTool { private final String _controllerProtocol; private final String _controllerHost; private final int _controllerPort; - private final String _authToken; - private final String _authTokenUrl; + private final AuthProvider _authProvider; private final String _tableDir; private final MinionClient _minionClient; public BootstrapTableTool(String controllerProtocol, String controllerHost, int controllerPort, String tableDir, - String authToken, String authTokenUrl) { + AuthProvider authProvider) { Preconditions.checkNotNull(controllerProtocol); Preconditions.checkNotNull(controllerHost); Preconditions.checkNotNull(tableDir); @@ -69,8 +69,7 @@ public class BootstrapTableTool { _controllerPort = controllerPort; _tableDir = tableDir; _minionClient = new MinionClient(controllerHost, String.valueOf(controllerPort)); - _authToken = authToken; - _authTokenUrl = authTokenUrl; + _authProvider = authProvider; } public boolean execute() @@ -120,7 +119,7 @@ public class BootstrapTableTool { return new AddTableCommand().setSchemaFile(schemaFile.getAbsolutePath()) .setTableConfigFile(tableConfigFile.getAbsolutePath()).setControllerProtocol(_controllerProtocol) .setControllerHost(_controllerHost).setControllerPort(String.valueOf(_controllerPort)).setExecute(true) - .setAuthToken(_authToken).setAuthTokenUrl(_authTokenUrl).execute(); + .setAuthProvider(_authProvider).execute(); } private boolean bootstrapOfflineTable(File setupTableTmpDir, String tableName, File schemaFile, @@ -183,7 +182,7 @@ public class BootstrapTableTool { } // url-based token needs to be resolved before job run - spec.setAuthToken(AuthProviderUtils.resolveToToken(_authToken, _authTokenUrl)); + spec.setAuthToken(AuthProviderUtils.toTaskToken(_authProvider)); IngestionJobLauncher.runIngestionJob(spec); } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/EmptyQuickstart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/EmptyQuickstart.java index 08586ed5d5..5fc3e8ce82 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/EmptyQuickstart.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/EmptyQuickstart.java @@ -22,6 +22,7 @@ import java.io.File; import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.tools.admin.PinotAdministrator; import org.apache.pinot.tools.admin.command.QuickstartRunner; @@ -32,7 +33,7 @@ public class EmptyQuickstart extends QuickStartBase { return Arrays.asList("EMPTY", "DEFAULT"); } - public String getAuthToken() { + public AuthProvider getAuthProvider() { return null; } @@ -48,7 +49,7 @@ public class EmptyQuickstart extends QuickStartBase { QuickstartRunner runner = new QuickstartRunner(new ArrayList<>(), 1, 1, 1, 0, - dataDir, true, getAuthToken(), getConfigOverrides(), _zkExternalAddress, false); + dataDir, true, getAuthProvider(), getConfigOverrides(), _zkExternalAddress, false); if (_zkExternalAddress != null) { printStatus(Quickstart.Color.CYAN, "***** Starting controller, broker and server *****"); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/Quickstart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/Quickstart.java index 01fb8fb716..6baf5c16e9 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/Quickstart.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/Quickstart.java @@ -28,6 +28,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; import org.apache.commons.io.FileUtils; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.tools.admin.PinotAdministrator; import org.apache.pinot.tools.admin.command.QuickstartRunner; @@ -60,7 +61,7 @@ public class Quickstart extends QuickStartBase { return 0; } - public String getAuthToken() { + public AuthProvider getAuthProvider() { return null; } @@ -104,7 +105,7 @@ public class Quickstart extends QuickStartBase { QuickstartTableRequest request = new QuickstartTableRequest(baseDir.getAbsolutePath()); QuickstartRunner runner = new QuickstartRunner(Lists.newArrayList(request), 1, 1, 1, - getNumMinions(), dataDir, true, getAuthToken(), + getNumMinions(), dataDir, true, getAuthProvider(), getConfigOverrides(), null, true); printStatus(Color.CYAN, "***** Starting Zookeeper, controller, broker and server *****"); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AbstractBaseAdminCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AbstractBaseAdminCommand.java index 9b463ca652..300a49b6bd 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AbstractBaseAdminCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AbstractBaseAdminCommand.java @@ -36,7 +36,11 @@ import org.apache.commons.configuration.ConfigurationException; import org.apache.commons.lang3.StringUtils; import org.apache.http.Header; import org.apache.pinot.common.auth.AuthProviderUtils; +import org.apache.pinot.common.auth.NullAuthProvider; +import org.apache.pinot.common.auth.StaticTokenAuthProvider; +import org.apache.pinot.common.auth.UrlAuthProvider; import org.apache.pinot.core.auth.BasicAuthUtils; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.tools.AbstractBaseCommand; import org.apache.pinot.tools.utils.PinotConfigUtils; @@ -125,13 +129,13 @@ public class AbstractBaseAdminCommand extends AbstractBaseCommand { } /** - * Generate an (optional) HTTP Authorization header given an auth token + * Generate an (optional) HTTP Authorization header given an auth config * - * @param authToken auth token - * @return list of 0 or 1 "Authorization" headers + * @param authProvider auth provider + * @return list of headers */ - static List<Header> makeAuthHeaders(String authToken, String authTokenUrl) { - return AuthProviderUtils.toHeaders(AuthProviderUtils.inferProvider(authToken, authTokenUrl)); + static List<Header> makeAuthHeaders(AuthProvider authProvider) { + return AuthProviderUtils.toRequestHeaders(authProvider); } /** @@ -140,18 +144,28 @@ public class AbstractBaseAdminCommand extends AbstractBaseCommand { * @param authToken optional pass-thru token * @param user optional username * @param password optional password + * @param tokenUrl optional token url * @return auth token, or null if neither pass-thru token nor user info available */ @Nullable - static String makeAuthToken(String authToken, String user, String password) { + static AuthProvider makeAuthProvider(AuthProvider provider, String tokenUrl, String authToken, String user, + String password) { + if (provider != null) { + return provider; + } + + if (StringUtils.isNotBlank(tokenUrl)) { + return new UrlAuthProvider(tokenUrl); + } + if (StringUtils.isNotBlank(authToken)) { - return authToken; + return new StaticTokenAuthProvider(authToken); } if (StringUtils.isNotBlank(user)) { - return BasicAuthUtils.toBasicAuthToken(user, password); + return new StaticTokenAuthProvider(BasicAuthUtils.toBasicAuthToken(user, password)); } - return null; + return new NullAuthProvider(); } } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java index cac2c295e1..4cde2cde1f 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddSchemaCommand.java @@ -22,6 +22,7 @@ import java.io.File; import java.io.FileNotFoundException; import java.util.Collections; import org.apache.pinot.common.utils.FileUploadDownloadClient; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.NetUtils; @@ -66,6 +67,8 @@ public class AddSchemaCommand extends AbstractBaseAdminCommand implements Comman description = "Print this message.") private boolean _help = false; + private AuthProvider _authProvider; + @Override public boolean getHelp() { return _help; @@ -126,6 +129,10 @@ public class AddSchemaCommand extends AbstractBaseAdminCommand implements Comman _authToken = authToken; } + public void setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + } + public AddSchemaCommand setExecute(boolean exec) { _exec = exec; return this; @@ -154,9 +161,8 @@ public class AddSchemaCommand extends AbstractBaseAdminCommand implements Comman try (FileUploadDownloadClient fileUploadDownloadClient = new FileUploadDownloadClient()) { fileUploadDownloadClient.addSchema(FileUploadDownloadClient .getUploadSchemaURI(_controllerProtocol, _controllerHost, Integer.parseInt(_controllerPort)), - schema.getSchemaName(), schemaFile, makeAuthHeaders(makeAuthToken(_authToken, _user, _password), - _authTokenUrl), - Collections.emptyList()); + schema.getSchemaName(), schemaFile, makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, + _user, _password)), Collections.emptyList()); } catch (Exception e) { LOGGER.error("Got Exception to upload Pinot Schema: " + schema.getSchemaName(), e); return false; diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTableCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTableCommand.java index 871a4de15b..5eb61c2704 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTableCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTableCommand.java @@ -23,6 +23,7 @@ import java.io.File; import java.io.IOException; import java.util.Collections; import org.apache.pinot.common.utils.FileUploadDownloadClient; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.JsonUtils; @@ -81,6 +82,8 @@ public class AddTableCommand extends AbstractBaseAdminCommand implements Command private String _controllerAddress; + private AuthProvider _authProvider; + @Override public boolean getHelp() { return _help; @@ -159,6 +162,11 @@ public class AddTableCommand extends AbstractBaseAdminCommand implements Command return this; } + public AddTableCommand setAuthProvider(AuthProvider authProvider) { + _authProvider = _authProvider; + return this; + } + public void uploadSchema() throws Exception { File schemaFile; @@ -173,9 +181,8 @@ public class AddTableCommand extends AbstractBaseAdminCommand implements Command try (FileUploadDownloadClient fileUploadDownloadClient = new FileUploadDownloadClient()) { fileUploadDownloadClient.addSchema(FileUploadDownloadClient .getUploadSchemaURI(_controllerProtocol, _controllerHost, Integer.parseInt(_controllerPort)), - schema.getSchemaName(), schemaFile, makeAuthHeaders(makeAuthToken(_authToken, _user, _password), - _authTokenUrl), - Collections.emptyList()); + schema.getSchemaName(), schemaFile, makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, + _user, _password)), Collections.emptyList()); } catch (Exception e) { LOGGER.error("Got Exception to upload Pinot Schema: " + schema.getSchemaName(), e); throw e; @@ -186,7 +193,7 @@ public class AddTableCommand extends AbstractBaseAdminCommand implements Command throws IOException { String res = AbstractBaseAdminCommand .sendRequest("POST", ControllerRequestURLBuilder.baseUrl(_controllerAddress).forTableCreate(), node.toString(), - makeAuthHeaders(makeAuthToken(_authToken, _user, _password), _authTokenUrl)); + makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, _password))); LOGGER.info(res); return res.contains("succesfully added"); } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTenantCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTenantCommand.java index 3cb5bd7cbb..6c80817215 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTenantCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/AddTenantCommand.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.tools.admin.command; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.tenant.Tenant; import org.apache.pinot.spi.config.tenant.TenantRole; import org.apache.pinot.spi.utils.CommonConstants; @@ -80,6 +81,8 @@ public class AddTenantCommand extends AbstractBaseAdminCommand implements Comman private String _controllerAddress; + private AuthProvider _authProvider; + public AddTenantCommand setControllerUrl(String url) { _controllerAddress = url; return this; @@ -130,6 +133,11 @@ public class AddTenantCommand extends AbstractBaseAdminCommand implements Comman return this; } + public AddTenantCommand setAuthPRovider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + @Override public boolean execute() throws Exception { @@ -150,7 +158,8 @@ public class AddTenantCommand extends AbstractBaseAdminCommand implements Comman Tenant tenant = new Tenant(_role, _name, _instanceCount, _offlineInstanceCount, _realtimeInstanceCount); String res = AbstractBaseAdminCommand .sendRequest("POST", ControllerRequestURLBuilder.baseUrl(_controllerAddress).forTenantCreate(), - tenant.toJsonString(), makeAuthHeaders(makeAuthToken(_authToken, _user, _password), _authTokenUrl)); + tenant.toJsonString(), makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, + _password))); LOGGER.info(res); System.out.print(res); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/BootstrapTableCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/BootstrapTableCommand.java index 72057df286..2e6b95b9c7 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/BootstrapTableCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/BootstrapTableCommand.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.tools.admin.command; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.plugin.PluginManager; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.NetUtils; @@ -94,6 +95,8 @@ public class BootstrapTableCommand extends AbstractBaseAdminCommand implements C description = "Print this message.") private boolean _help = false; + private AuthProvider _authProvider; + @Override public boolean getHelp() { return _help; @@ -109,6 +112,11 @@ public class BootstrapTableCommand extends AbstractBaseAdminCommand implements C return this; } + public BootstrapTableCommand setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + @Override public String toString() { return ("BootstrapTable -dir " + _dir); @@ -130,9 +138,7 @@ public class BootstrapTableCommand extends AbstractBaseAdminCommand implements C if (_controllerHost == null) { _controllerHost = NetUtils.getHostAddress(); } - String token = makeAuthToken(_authToken, _user, _password); - return new BootstrapTableTool(_controllerProtocol, _controllerHost, Integer.parseInt(_controllerPort), _dir, token, - _authTokenUrl) - .execute(); + return new BootstrapTableTool(_controllerProtocol, _controllerHost, Integer.parseInt(_controllerPort), _dir, + makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, _password)).execute(); } } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ChangeTableState.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ChangeTableState.java index 381b3b2e1e..855cda9704 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ChangeTableState.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ChangeTableState.java @@ -21,6 +21,7 @@ package org.apache.pinot.tools.admin.command; import java.net.URI; import org.apache.commons.httpclient.HttpClient; import org.apache.commons.httpclient.methods.GetMethod; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.NetUtils; import org.apache.pinot.tools.Command; @@ -64,6 +65,13 @@ public class ChangeTableState extends AbstractBaseAdminCommand implements Comman description = "Print this message.") private boolean _help = false; + private AuthProvider _authProvider; + + public ChangeTableState setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + @Override public boolean execute() throws Exception { @@ -80,10 +88,8 @@ public class ChangeTableState extends AbstractBaseAdminCommand implements Comman URI uri = new URI(_controllerProtocol, null, _controllerHost, Integer.parseInt(_controllerPort), URI_TABLES_PATH + _tableName, "state=" + stateValue, null); - String token = makeAuthToken(_authToken, _user, _password); - GetMethod httpGet = new GetMethod(uri.toString()); - makeAuthHeaders(makeAuthToken(_authToken, _user, _password), _authTokenUrl) + makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, _password)) .forEach(header -> httpGet.addRequestHeader(header.getName(), header.getValue())); int status = httpClient.executeMethod(httpGet); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ImportDataCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ImportDataCommand.java index 0e135d8e57..d40069a299 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ImportDataCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/ImportDataCommand.java @@ -30,7 +30,7 @@ import java.util.List; import java.util.Map; import org.apache.commons.codec.digest.DigestUtils; import org.apache.commons.io.FileUtils; -import org.apache.pinot.common.auth.AuthProviderUtils; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.data.readers.FileFormat; import org.apache.pinot.spi.filesystem.PinotFSFactory; import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties; @@ -100,6 +100,8 @@ public class ImportDataCommand extends AbstractBaseAdminCommand implements Comma @CommandLine.Option(names = {"-help", "-h", "--h", "--help"}, help = true, description = "Print this message.") private boolean _help = false; + private AuthProvider _authProvider; + public ImportDataCommand setDataFilePath(String dataFilePath) { _dataFilePath = dataFilePath; return this; @@ -157,6 +159,11 @@ public class ImportDataCommand extends AbstractBaseAdminCommand implements Comma return this; } + public ImportDataCommand setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + public List<String> getAdditionalConfigs() { return _additionalConfigs; } @@ -262,7 +269,7 @@ public class ImportDataCommand extends AbstractBaseAdminCommand implements Comma spec.setCleanUpOutputDir(true); spec.setOverwriteOutput(true); spec.setJobType("SegmentCreationAndTarPush"); - spec.setAuthToken(AuthProviderUtils.resolveToToken(makeAuthToken(_authToken, _user, _password), _authTokenUrl)); + spec.setAuthToken(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, _password).getTaskToken()); // set ExecutionFrameworkSpec ExecutionFrameworkSpec executionFrameworkSpec = new ExecutionFrameworkSpec(); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchDataIngestionJobCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchDataIngestionJobCommand.java index c1a8b0f8df..c21f312104 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchDataIngestionJobCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchDataIngestionJobCommand.java @@ -21,8 +21,8 @@ package org.apache.pinot.tools.admin.command; import java.util.Arrays; import java.util.List; import org.apache.commons.lang3.StringUtils; -import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.common.utils.TlsUtils; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.ingestion.batch.IngestionJobLauncher; import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec; import org.apache.pinot.spi.ingestion.batch.spec.TlsSpec; @@ -62,6 +62,8 @@ public class LaunchDataIngestionJobCommand extends AbstractBaseAdminCommand impl @CommandLine.Option(names = {"-authTokenUrl"}, required = false, description = "Http auth token url.") private String _authTokenUrl; + private AuthProvider _authProvider; + public String getJobSpecFile() { return _jobSpecFile; } @@ -86,6 +88,10 @@ public class LaunchDataIngestionJobCommand extends AbstractBaseAdminCommand impl _propertyFile = propertyFile; } + public void setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + } + @Override public boolean getHelp() { return _help; @@ -117,7 +123,7 @@ public class LaunchDataIngestionJobCommand extends AbstractBaseAdminCommand impl } if (StringUtils.isBlank(spec.getAuthToken())) { - spec.setAuthToken(AuthProviderUtils.resolveToToken(makeAuthToken(_authToken, _user, _password), _authTokenUrl)); + spec.setAuthToken(makeAuthProvider(_authProvider, _authToken, _user, _password, _authTokenUrl).getTaskToken()); } try { diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/OperateClusterConfigCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/OperateClusterConfigCommand.java index b8647ddf64..fba7b4fd20 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/OperateClusterConfigCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/OperateClusterConfigCommand.java @@ -24,6 +24,7 @@ import java.util.Iterator; import java.util.List; import org.apache.commons.lang.StringUtils; import org.apache.http.Header; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.NetUtils; @@ -69,6 +70,8 @@ public class OperateClusterConfigCommand extends AbstractBaseAdminCommand implem description = "Print this message.") private boolean _help = false; + private AuthProvider _authProvider; + @Override public boolean getHelp() { return _help; @@ -140,6 +143,11 @@ public class OperateClusterConfigCommand extends AbstractBaseAdminCommand implem return this; } + public OperateClusterConfigCommand setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + public String run() throws Exception { if (_controllerHost == null) { @@ -151,7 +159,8 @@ public class OperateClusterConfigCommand extends AbstractBaseAdminCommand implem } String clusterConfigUrl = _controllerProtocol + "://" + _controllerHost + ":" + _controllerPort + "/cluster/configs"; - List<Header> headers = makeAuthHeaders(makeAuthToken(_authToken, _user, _password), _authTokenUrl); + List<Header> headers = makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, + _password)); switch (_operation.toUpperCase()) { case "ADD": case "UPDATE": diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/PostQueryCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/PostQueryCommand.java index e8f6194360..a9629c0d6c 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/PostQueryCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/PostQueryCommand.java @@ -19,6 +19,7 @@ package org.apache.pinot.tools.admin.command; import java.util.Collections; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.CommonConstants.Broker.Request; import org.apache.pinot.spi.utils.JsonUtils; @@ -61,6 +62,8 @@ public class PostQueryCommand extends AbstractBaseAdminCommand implements Comman + "this message.") private boolean _help = false; + private AuthProvider _authProvider; + @Override public boolean getHelp() { return _help; @@ -126,6 +129,11 @@ public class PostQueryCommand extends AbstractBaseAdminCommand implements Comman return this; } + public PostQueryCommand setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + public String run() throws Exception { if (_brokerHost == null) { @@ -134,8 +142,8 @@ public class PostQueryCommand extends AbstractBaseAdminCommand implements Comman LOGGER.info("Executing command: " + this); String url = _brokerProtocol + "://" + _brokerHost + ":" + _brokerPort + "/query/sql"; String request = JsonUtils.objectToString(Collections.singletonMap(Request.SQL, _query)); - return sendRequest("POST", url, request, makeAuthHeaders(makeAuthToken(_authToken, _user, _password), - _authTokenUrl)); + return sendRequest("POST", url, request, makeAuthHeaders(makeAuthProvider(_authProvider, + _authTokenUrl, _authToken, _user, _password))); } @Override diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/QuickstartRunner.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/QuickstartRunner.java index 3f241666c6..a59291dd99 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/QuickstartRunner.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/QuickstartRunner.java @@ -31,6 +31,7 @@ import java.util.List; import java.util.Map; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.config.tenant.TenantRole; import org.apache.pinot.spi.env.PinotConfiguration; @@ -72,8 +73,7 @@ public class QuickstartRunner { private final int _numMinions; private final File _tempDir; private final boolean _enableTenantIsolation; - private final String _authToken; - private final String _authTokenUrl; + private final AuthProvider _authProvider; private final Map<String, Object> _configOverrides; private final boolean _deleteExistingData; @@ -87,21 +87,12 @@ public class QuickstartRunner { public QuickstartRunner(List<QuickstartTableRequest> tableRequests, int numControllers, int numBrokers, int numServers, int numMinions, File tempDir, Map<String, Object> configOverrides) throws Exception { - this(tableRequests, numControllers, numBrokers, numServers, numMinions, tempDir, true, null, - null, configOverrides, null, + this(tableRequests, numControllers, numBrokers, numServers, numMinions, tempDir, true, null, configOverrides, null, true); } public QuickstartRunner(List<QuickstartTableRequest> tableRequests, int numControllers, int numBrokers, - int numServers, int numMinions, File tempDir, boolean enableIsolation, String authToken, - Map<String, Object> configOverrides, String zkExternalAddress, boolean deleteExistingData) - throws Exception { - this(tableRequests, numControllers, numBrokers, numServers, numMinions, tempDir, enableIsolation, authToken, null, - configOverrides, zkExternalAddress, deleteExistingData); - } - - public QuickstartRunner(List<QuickstartTableRequest> tableRequests, int numControllers, int numBrokers, - int numServers, int numMinions, File tempDir, boolean enableIsolation, String authToken, String authTokenUrl, + int numServers, int numMinions, File tempDir, boolean enableIsolation, AuthProvider authProvider, Map<String, Object> configOverrides, String zkExternalAddress, boolean deleteExistingData) throws Exception { _tableRequests = tableRequests; @@ -111,8 +102,7 @@ public class QuickstartRunner { _numMinions = numMinions; _tempDir = tempDir; _enableTenantIsolation = enableIsolation; - _authToken = authToken; - _authTokenUrl = authTokenUrl; + _authProvider = authProvider; _configOverrides = configOverrides; _zkExternalAddress = zkExternalAddress; _deleteExistingData = deleteExistingData; @@ -238,8 +228,8 @@ public class QuickstartRunner { public void bootstrapTable() throws Exception { for (QuickstartTableRequest request : _tableRequests) { - if (!new BootstrapTableTool("http", "localhost", _controllerPorts.get(0), request.getBootstrapTableDir(), - _authToken, _authTokenUrl).execute()) { + if (!new BootstrapTableTool("http", "localhost", _controllerPorts.get(0), + request.getBootstrapTableDir(), _authProvider).execute()) { throw new RuntimeException("Failed to bootstrap table with request - " + request); } } @@ -284,8 +274,8 @@ public class QuickstartRunner { throws Exception { int brokerPort = _brokerPorts.get(RANDOM.nextInt(_brokerPorts.size())); return JsonUtils.stringToJsonNode( - new PostQueryCommand().setBrokerPort(String.valueOf(brokerPort)).setAuthToken(_authToken) - .setAuthTokenUrl(_authTokenUrl).setQuery(query).run()); + new PostQueryCommand().setBrokerPort(String.valueOf(brokerPort)).setAuthProvider(_authProvider) + .setQuery(query).run()); } public static void registerDefaultPinotFS() { diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java index 64610150db..c6953956c9 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java @@ -24,10 +24,10 @@ import java.net.URI; import java.util.Collections; import org.apache.commons.io.FileUtils; import org.apache.http.message.BasicNameValuePair; -import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.common.utils.FileUploadDownloadClient; import org.apache.pinot.common.utils.TarGzCompressionUtils; import org.apache.pinot.common.utils.http.HttpClient; +import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.NetUtils; import org.apache.pinot.tools.Command; @@ -78,6 +78,8 @@ public class UploadSegmentCommand extends AbstractBaseAdminCommand implements Co description = "Print this message.") private boolean _help = false; + private AuthProvider _authProvider; + @Override public boolean getHelp() { return _help; @@ -143,6 +145,11 @@ public class UploadSegmentCommand extends AbstractBaseAdminCommand implements Co return this; } + public UploadSegmentCommand setAuthProvider(AuthProvider authProvider) { + _authProvider = authProvider; + return this; + } + @Override public boolean execute() throws Exception { @@ -177,10 +184,9 @@ public class UploadSegmentCommand extends AbstractBaseAdminCommand implements Co LOGGER.info("Uploading segment tar file: {}", segmentTarFile); fileUploadDownloadClient.uploadSegment(uploadSegmentHttpURI, segmentTarFile.getName(), segmentTarFile, - AuthProviderUtils.toHeaders(AuthProviderUtils.inferProvider(makeAuthToken(_authToken, _user, _password), - _authTokenUrl)), Collections.singletonList(new BasicNameValuePair( - FileUploadDownloadClient.QueryParameters.TABLE_NAME, _tableName)), - HttpClient.DEFAULT_SOCKET_TIMEOUT_MS); + makeAuthHeaders(makeAuthProvider(_authProvider, _authTokenUrl, _authToken, _user, _password)), + Collections.singletonList(new BasicNameValuePair(FileUploadDownloadClient.QueryParameters.TABLE_NAME, + _tableName)), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS); } } finally { // Delete the temporary working directory. --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
