Skip to content
Permalink
Browse files
DRILL-8220: Add User Translation Support for OAuth Enabled Plugins (#…
  • Loading branch information
cgivre committed May 16, 2022
1 parent bd5b386 commit 58391fecd165665141402ccedf8d2789a1014751
Showing 16 changed files with 625 additions and 257 deletions.
@@ -59,7 +59,7 @@ credentialProvider.
To use OAuth2.0, you will have to create an `oAuthConfig` in the plugin configuration. Within the `oAuthConfig`, define the `callbackURL` and `authorizationURL` parameters:
* The `authorizationURL` is provided by the API and is the URL where the authorization code is obtained.
* The `callbackURL` parameter is the URL where the API will send the access token. You must provide this when you register and obtain your client ID and client secret. This
will be in the format: `http(s)://<your drill host>/storage/<storage plugin name>update_oauth2_authtoken`
will be in the format: `http(s)://<your drill host>/credentials/<storage plugin name>update_oauth2_authtoken`
* (Optional)`scope`: The scope parameter limits the scope of your access. This is something which can be found in the remote API documentation.

### The Credential Provider
@@ -104,7 +104,7 @@ The example configuration below demonstrates how to connect Drill to the API ava
},
"proxyType": "direct",
"oAuthConfig": {
"callbackURL": "http://localhost:8047/storage/clickup/update_oath2_authtoken",
"callbackURL": "http://localhost:8047/credentials/clickup/update_oath2_authtoken",
"authorizationURL": "https://app.clickup.com/api"
},
"credentialsProvider": {
@@ -125,3 +125,20 @@ There are a few optional parameters in the OAuth config which you may need to se

* `tokenType`: Some OAuth enabled APIs provide a `Bearer` token. If that is the case, this should be set to `Bearer`.
* `authorizationParams`: A key value parameters which are sent during the authentication process.

## Enabling Individual User Credentials with OAuth 2.0
Drill recently introduced the `USER_TRANSLATION` authorization mode, which is useful for plugins that do not have the concept of user impersonation. This is very much the
case for OAuth enabled APIs. When you configure an OAuth enabled API, the client secret keys belong to the application. Following this design pattern, each individual user
should authorize (or not) the application. Thus the `clientID` and `client_secret` tokens really belong to the application and the `access_token` and `refresh_token` belong
to the individual user.

Enabling user translation is quite simple. In the configuration for the storage plugin simply add the key below to your
plugin configuration. Note that for user translation to work, user impersonation and authentication must both be enabled globally.

```json
"authMode":"USER_TRANSLATION"
```

Once you've done this, when a user logs in, they will see a new menu option at the top bar called `Credentials`. This will contain a listing
of storage plugin instances that require credentials. For OAuth enabled plugins, there will be an `Authorize` button next to the plugin name.
Each user will have to authorize Drill to access the plugin.
@@ -21,6 +21,7 @@
import org.apache.calcite.schema.SchemaPlus;
import org.apache.drill.common.JSONOptions;
import org.apache.drill.common.expression.SchemaPath;
import org.apache.drill.common.logical.StoragePluginConfig.AuthMode;
import org.apache.drill.exec.metastore.MetadataProviderManager;
import org.apache.drill.exec.oauth.OAuthTokenProvider;
import org.apache.drill.exec.oauth.PersistentTokenTable;
@@ -34,6 +35,7 @@
import org.apache.drill.exec.store.SchemaConfig;
import org.apache.drill.exec.store.StoragePluginRegistry;
import org.apache.drill.exec.store.base.filter.FilterPushDownUtils;
import org.apache.drill.shaded.guava.com.google.common.annotations.VisibleForTesting;
import org.apache.drill.shaded.guava.com.google.common.collect.ImmutableSet;

import com.fasterxml.jackson.core.type.TypeReference;
@@ -46,25 +48,38 @@ public class HttpStoragePlugin extends AbstractStoragePlugin {
private final HttpStoragePluginConfig config;
private final HttpSchemaFactory schemaFactory;
private final StoragePluginRegistry registry;
private final TokenRegistry tokenRegistry;
private TokenRegistry tokenRegistry;

public HttpStoragePlugin(HttpStoragePluginConfig configuration, DrillbitContext context, String name) {
super(context, name);
this.config = configuration;
this.registry = context.getStorage();
this.schemaFactory = new HttpSchemaFactory(this);

// Get OAuth Token Provider if needed
OAuthTokenProvider tokenProvider = context.getoAuthTokenProvider();
tokenRegistry = tokenProvider.getOauthTokenRegistry();
tokenRegistry.createTokenTable(getName());
if (config.getAuthMode() != AuthMode.USER_TRANSLATION) {
initializeOauthTokenTable(null);
}
}

@Override
public void registerSchemas(SchemaConfig schemaConfig, SchemaPlus parent) {

// For user translation mode, this is moved here because we don't have the
// active username in the constructor. Removing it from the constructor makes
// it difficult to test, so we do the check and leave it in both places.
if (config.getAuthMode() == AuthMode.USER_TRANSLATION) {
initializeOauthTokenTable(schemaConfig.getUserName());
}
schemaFactory.registerSchemas(schemaConfig, parent);
}

@VisibleForTesting
public void initializeOauthTokenTable(String username) {
OAuthTokenProvider tokenProvider = context.getoAuthTokenProvider();
tokenRegistry = tokenProvider.getOauthTokenRegistry(username);
tokenRegistry.createTokenTable(getName());
}

@Override
public HttpStoragePluginConfig getConfig() {
return config;
@@ -78,6 +93,18 @@ public TokenRegistry getTokenRegistry() {
return tokenRegistry;
}

/**
* This method returns the {@link TokenRegistry} for a given user. It is only used for testing user translation
* with OAuth 2.0.
* @param username A {@link String} of the current active user.
* @return A {@link TokenRegistry} for the given user.
*/
@VisibleForTesting
public TokenRegistry getTokenRegistry(String username) {
initializeOauthTokenTable(username);
return tokenRegistry;
}

public PersistentTokenTable getTokenTable() { return tokenRegistry.getTokenTable(getName()); }

@Override
@@ -19,6 +19,7 @@

import org.apache.drill.common.PlanStringBuilder;
import org.apache.drill.common.exceptions.UserException;
import org.apache.drill.common.logical.OAuthConfig;
import org.apache.drill.common.map.CaseInsensitiveMap;
import org.apache.drill.common.logical.CredentialedStoragePluginConfig;

@@ -50,8 +51,6 @@ public class HttpStoragePluginConfig extends CredentialedStoragePluginConfig {
public final String proxyHost;
public final int proxyPort;
public final String proxyType;
public final HttpOAuthConfig oAuthConfig;

/**
* Timeout in {@link TimeUnit#SECONDS}.
*/
@@ -68,7 +67,7 @@ public HttpStoragePluginConfig(@JsonProperty("cacheResults") Boolean cacheResult
@JsonProperty("proxyType") String proxyType,
@JsonProperty("proxyUsername") String proxyUsername,
@JsonProperty("proxyPassword") String proxyPassword,
@JsonProperty("oAuthConfig") HttpOAuthConfig oAuthConfig,
@JsonProperty("oAuthConfig") OAuthConfig oAuthConfig,
@JsonProperty("credentialsProvider") CredentialsProvider credentialsProvider,
@JsonProperty("authMode") String authMode
) {
@@ -82,8 +81,8 @@ public HttpStoragePluginConfig(@JsonProperty("cacheResults") Boolean cacheResult
normalize(proxyPassword),
credentialsProvider),
credentialsProvider == null,
AuthMode.parseOrDefault(authMode, AuthMode.SHARED_USER)
);
AuthMode.parseOrDefault(authMode, AuthMode.SHARED_USER),
oAuthConfig);
this.cacheResults = cacheResults != null && cacheResults;

this.connections = CaseInsensitiveMap.newHashMap();
@@ -94,7 +93,6 @@ public HttpStoragePluginConfig(@JsonProperty("cacheResults") Boolean cacheResult
this.timeout = timeout == null ? 0 : timeout;
this.proxyHost = normalize(proxyHost);
this.proxyPort = proxyPort == null ? 0 : proxyPort;
this.oAuthConfig = oAuthConfig;

proxyType = normalize(proxyType);
this.proxyType = proxyType == null
@@ -130,7 +128,7 @@ private HttpStoragePluginConfig(HttpStoragePluginConfig that, CredentialsProvide
* @param that The current HTTP Plugin Config
* @param oAuthConfig The updated OAuth config
*/
public HttpStoragePluginConfig(HttpStoragePluginConfig that, HttpOAuthConfig oAuthConfig) {
public HttpStoragePluginConfig(HttpStoragePluginConfig that, OAuthConfig oAuthConfig) {
super(CredentialProviderUtils.getCredentialsProvider(that.proxyUsername(), that.proxyPassword(), that.credentialsProvider),
that.credentialsProvider == null);

@@ -140,7 +138,7 @@ public HttpStoragePluginConfig(HttpStoragePluginConfig that, HttpOAuthConfig oAu
this.proxyHost = that.proxyHost;
this.proxyPort = that.proxyPort;
this.proxyType = that.proxyType;
this.oAuthConfig = oAuthConfig;
this.oAuthConfig = that.oAuthConfig;
}

private static String normalize(String value) {
@@ -234,11 +232,6 @@ public int hashCode() {
@JsonProperty("proxyPort")
public int proxyPort() { return proxyPort; }

@JsonProperty("oAuthConfig")
public HttpOAuthConfig oAuthConfig() {
return oAuthConfig;
}

@JsonProperty("username")
public String username() {
if (!directCredentials) {
@@ -23,9 +23,9 @@
import okhttp3.Request;

import org.apache.drill.common.exceptions.UserException;
import org.apache.drill.common.logical.OAuthConfig;
import org.apache.drill.common.logical.security.CredentialsProvider;
import org.apache.drill.exec.oauth.PersistentTokenTable;
import org.apache.drill.exec.store.http.HttpOAuthConfig;
import org.apache.drill.exec.store.http.HttpStoragePluginConfig;
import org.apache.drill.exec.store.http.util.HttpProxyConfig;
import org.apache.drill.exec.store.http.util.SimpleHttp;
@@ -69,7 +69,7 @@ public AccessTokenRepository(HttpProxyConfig proxyConfig,
client = builder.build();
}

public HttpOAuthConfig getOAuthConfig() {
public OAuthConfig getOAuthConfig() {
return pluginConfig.oAuthConfig();
}

@@ -32,6 +32,7 @@

import org.apache.commons.lang3.StringUtils;
import org.apache.drill.common.exceptions.EmptyErrorContext;
import org.apache.drill.common.logical.OAuthConfig;
import org.apache.drill.common.logical.StoragePluginConfig.AuthMode;
import org.apache.drill.common.map.CaseInsensitiveMap;
import org.apache.drill.common.exceptions.CustomErrorContext;
@@ -47,7 +48,6 @@
import org.apache.drill.exec.store.http.HttpApiConfig;
import org.apache.drill.exec.store.http.HttpApiConfig.HttpMethod;
import org.apache.drill.exec.store.http.HttpApiConfig.PostLocation;
import org.apache.drill.exec.store.http.HttpOAuthConfig;
import org.apache.drill.exec.store.http.HttpStoragePlugin;
import org.apache.drill.exec.store.http.HttpStoragePluginConfig;
import org.apache.drill.exec.store.http.HttpSubScan;
@@ -114,7 +114,7 @@ public class SimpleHttp {
private final String connection;
private final HttpStoragePluginConfig pluginConfig;
private final HttpApiConfig apiConfig;
private final HttpOAuthConfig oAuthConfig;
private final OAuthConfig oAuthConfig;
private String responseMessage;
private int responseCode;
private String responseProtocol;
@@ -967,7 +967,7 @@ public static class SimpleHttpBuilder {
private PersistentTokenTable tokenTable;
private HttpStoragePluginConfig pluginConfig;
private HttpApiConfig endpointConfig;
private HttpOAuthConfig oAuthConfig;
private OAuthConfig oAuthConfig;
private Map<String,String> filters;
private String connection;
private String username;
@@ -23,6 +23,7 @@
import okhttp3.Response;
import okhttp3.mockwebserver.MockResponse;
import okhttp3.mockwebserver.MockWebServer;
import org.apache.drill.common.logical.OAuthConfig;
import org.apache.drill.common.logical.StoragePluginConfig.AuthMode;
import org.apache.drill.common.logical.security.CredentialsProvider;
import org.apache.drill.common.logical.security.PlainCredentialsProvider;
@@ -58,7 +59,7 @@
public class TestOAuthProcess extends ClusterTest {

private static final Logger logger = LoggerFactory.getLogger(TestOAuthProcess.class);
private static final int MOCK_SERVER_PORT = 47770;
private static final int MOCK_SERVER_PORT = 47775;

private static final int TIMEOUT = 30;
private static final String CONNECTION_NAME = "localOauth";
@@ -101,7 +102,7 @@ public static void setup() throws Exception {
.inputType("json")
.build();

HttpOAuthConfig oAuthConfig = HttpOAuthConfig.builder()
OAuthConfig oAuthConfig = OAuthConfig.builder()
.callbackURL(hostname + "/update_oauth2_authtoken")
.build();

@@ -24,6 +24,7 @@
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import org.apache.drill.common.logical.OAuthConfig;
import org.apache.drill.common.logical.StoragePluginConfig.AuthMode;
import org.apache.drill.common.logical.security.CredentialsProvider;
import org.apache.drill.common.logical.security.PlainCredentialsProvider;
@@ -81,7 +82,7 @@ public static void setup() throws Exception {
.inputType("json")
.build();

HttpOAuthConfig oAuthConfig = HttpOAuthConfig.builder()
OAuthConfig oAuthConfig = OAuthConfig.builder()
.callbackURL(hostname + "/update_ouath2_authtoken")
.build();

0 comments on commit 58391fe

Please sign in to comment.