Access Token support Added and Macro added for refresh token fields.

data-integrations · Mar 21, 2024 · 18fab3a · 18fab3a
1 parent c58f79f
commit 18fab3a
Show file tree

Hide file tree

Showing 15 changed files with 381 additions and 16 deletions.
diff --git a/docs/GoogleDrive-batchsink.md b/docs/GoogleDrive-batchsink.md
@@ -43,6 +43,11 @@ Make sure that:
 OAuth2 client credentials can be generated on Google Cloud
 [Credentials Page](https://console.cloud.google.com/apis/credentials)
 
+**OAuth Method:** The method used to get OAuth access tokens. Users have the option to either directly provide
+the OAuth access token or supply a client ID, client secret, and refresh token.
+
+**Access Token:** Short lived access token used for connecting.
+
 **Client ID:** OAuth2 client id used to identify the application.
 
 **Client Secret:** OAuth2 client secret used to access the authorization server.

diff --git a/docs/GoogleDrive-batchsource.md b/docs/GoogleDrive-batchsource.md
@@ -58,6 +58,11 @@ Make sure that:
 OAuth2 client credentials can be generated on Google Cloud
 [Credentials Page](https://console.cloud.google.com/apis/credentials)
 
+**OAuth Method:** The method used to get OAuth access tokens. Users have the option to either directly provide
+the OAuth access token or supply a client ID, client secret, and refresh token.
+
+**Access Token:** Short lived access token used for connecting.
+
 **Client ID:** OAuth2 client id used to identify the application.
 
 **Client Secret:** OAuth2 client secret used to access the authorization server.

diff --git a/docs/GoogleSheets-batchsink.md b/docs/GoogleSheets-batchsink.md
@@ -46,6 +46,11 @@ Make sure that:
 OAuth2 client credentials can be generated on Google Cloud
 [Credentials Page](https://console.cloud.google.com/apis/credentials)
 
+**OAuth Method:** The method used to get OAuth access tokens. Users have the option to either directly provide
+the OAuth access token or supply a client ID, client secret, and refresh token.
+
+**Access Token:** Short lived access token used for connecting.
+
 **Client ID:** OAuth2 client id used to identify the application.
 
 **Client Secret:** OAuth2 client secret used to access the authorization server.

diff --git a/docs/GoogleSheets-batchsource.md b/docs/GoogleSheets-batchsource.md
@@ -55,6 +55,11 @@ Make sure that:
 OAuth2 client credentials can be generated on Google Cloud 
 [Credentials Page](https://console.cloud.google.com/apis/credentials)
 
+**OAuth Method:** The method used to get OAuth access tokens. Users have the option to either directly provide
+the OAuth access token or supply a client ID, client secret, and refresh token.
+
+**Access Token:** Short lived access token used for connecting.
+
 **Client ID:** OAuth2 client id used to identify the application.
 
 **Client Secret:** OAuth2 client secret used to access the authorization server.

diff --git a/src/main/java/io/cdap/plugin/google/common/GoogleAuthBaseConfig.java b/src/main/java/io/cdap/plugin/google/common/GoogleAuthBaseConfig.java
@@ -40,20 +40,23 @@ public abstract class GoogleAuthBaseConfig extends PluginConfig {
   public static final String AUTO_DETECT_VALUE = "auto-detect";
   public static final String REFERENCE_NAME = "referenceName";
   public static final String AUTH_TYPE = "authType";
-  public static final String AUTH_TYPE_LABEL = "Auth type";
+  public static final String AUTH_TYPE_LABEL = "Authentication Type";
   public static final String CLIENT_ID = "clientId";
   public static final String CLIENT_ID_LABEL = "Client ID";
   public static final String CLIENT_SECRET = "clientSecret";
-  public static final String CLIENT_SECRET_LABEL = "Client secret";
+  public static final String CLIENT_SECRET_LABEL = "Client Secret";
   public static final String REFRESH_TOKEN = "refreshToken";
-  public static final String REFRESH_TOKEN_LABEL = "Refresh token";
+  public static final String REFRESH_TOKEN_LABEL = "Refresh Token";
   public static final String ACCOUNT_FILE_PATH = "accountFilePath";
   public static final String DIRECTORY_IDENTIFIER = "directoryIdentifier";
   public static final String NAME_SERVICE_ACCOUNT_TYPE = "serviceAccountType";
   public static final String NAME_SERVICE_ACCOUNT_JSON = "serviceAccountJSON";
   public static final String SERVICE_ACCOUNT_FILE_PATH = "filePath";
   public static final String SERVICE_ACCOUNT_JSON = "JSON";
   public static final String SCHEMA = "schema";
+  public static final String ACCESS_TOKEN = "accessToken";
+  public static final String ACCESS_TOKEN_LABEL = "Access Token";
+  public static final String OAUTH_METHOD = "oAuthMethod";
 
   private static final String IS_SET_FAILURE_MESSAGE_PATTERN = "'%s' property is empty or macro is not available.";
 
@@ -76,21 +79,37 @@ public abstract class GoogleAuthBaseConfig extends PluginConfig {
   @Macro
   protected String serviceAccountType;
 
+  @Macro
+  @Nullable
+  @Name(OAUTH_METHOD)
+  @Description("The method used to get OAuth access tokens. Users have the option to either directly provide " +
+    "the OAuth access token or supply a client ID, client secret, and refresh token.")
+  private String oAuthMethod;
+
   @Nullable
+  @Macro
   @Name(CLIENT_ID)
   @Description("OAuth2 client id.")
   private String clientId;
 
   @Nullable
+  @Macro
   @Name(CLIENT_SECRET)
   @Description("OAuth2 client secret.")
   private String clientSecret;
 
   @Nullable
+  @Macro
   @Name(REFRESH_TOKEN)
   @Description("OAuth2 refresh token.")
   private String refreshToken;
 
+  @Nullable
+  @Macro
+  @Name(ACCESS_TOKEN)
+  @Description("Short lived access token used for connecting.")
+  private String accessToken;
+
   @Nullable
   @Macro
   @Name(ACCOUNT_FILE_PATH)
@@ -177,9 +196,13 @@ private boolean validateAuthType(FailureCollector collector) {
   }
 
   private boolean validateOAuth2Properties(FailureCollector collector) {
-    return checkPropertyIsSet(collector, clientId, CLIENT_ID, CLIENT_ID_LABEL)
-      & checkPropertyIsSet(collector, clientSecret, CLIENT_SECRET, CLIENT_SECRET_LABEL)
-      & checkPropertyIsSet(collector, refreshToken, REFRESH_TOKEN, REFRESH_TOKEN_LABEL);
+    if (OAuthMethod.REFRESH_TOKEN.equals(getOAuthMethod())) {
+      return checkPropertyIsSet(collector, clientId, CLIENT_ID, CLIENT_ID_LABEL)
+        & checkPropertyIsSet(collector, clientSecret, CLIENT_SECRET, CLIENT_SECRET_LABEL)
+        & checkPropertyIsSet(collector, refreshToken, REFRESH_TOKEN, REFRESH_TOKEN_LABEL);
+    } else {
+      return checkPropertyIsSet(collector, accessToken, ACCESS_TOKEN, ACCESS_TOKEN_LABEL);
+    }
   }
 
   private boolean validateServiceAccount(FailureCollector collector) {
@@ -290,6 +313,14 @@ public void setRefreshToken(String refreshToken) {
     this.refreshToken = refreshToken;
   }
 
+  public void setoAuthMethod(String oAuthMethod) {
+    this.oAuthMethod = oAuthMethod;
+  }
+
+  public void setAccessToken(String accessToken) {
+    this.accessToken = accessToken;
+  }
+
   @Nullable
   public String getClientId() {
     return clientId;
@@ -305,6 +336,11 @@ public String getRefreshToken() {
     return refreshToken;
   }
 
+  @Nullable
+  public String getAccessToken() {
+    return accessToken;
+  }
+
   @Nullable
   public String getServiceAccountFilePath() {
     if (containsMacro(ACCOUNT_FILE_PATH) || Strings.isNullOrEmpty(accountFilePath)
@@ -344,4 +380,15 @@ public Boolean isServiceAccountFilePath() {
     String serviceAccountType = getServiceAccountType();
     return Strings.isNullOrEmpty(serviceAccountType) ? null : serviceAccountType.equals(SERVICE_ACCOUNT_FILE_PATH);
   }
+
+  public OAuthMethod getOAuthMethod() {
+    if (oAuthMethod == null) {
+      return OAuthMethod.REFRESH_TOKEN;
+    }
+    try {
+      return OAuthMethod.valueOf(oAuthMethod.toUpperCase());
+    } catch (IllegalArgumentException e) {
+      throw new IllegalArgumentException("Invalid oAuth method " + oAuthMethod);
+    }
+  }
 }
diff --git a/src/main/java/io/cdap/plugin/google/common/GoogleDriveClient.java b/src/main/java/io/cdap/plugin/google/common/GoogleDriveClient.java
@@ -88,13 +88,18 @@ protected Drive getDriveClient() throws IOException {
   }
 
   protected Credential getOAuth2Credential() {
-    GoogleCredential credential = new GoogleCredential.Builder()
-        .setTransport(httpTransport)
-        .setJsonFactory(JSON_FACTORY)
-        .setClientSecrets(config.getClientId(),
-            config.getClientSecret())
-        .build();
-    credential.createScoped(getRequiredScopes()).setRefreshToken(config.getRefreshToken());
+    GoogleCredential credential;
+    GoogleCredential.Builder builder = new GoogleCredential.Builder()
+      .setTransport(httpTransport)
+      .setJsonFactory(JSON_FACTORY);
+    if (OAuthMethod.ACCESS_TOKEN.equals(config.getOAuthMethod())) {
+      credential = builder.build();
+      credential.createScoped(getRequiredScopes()).setAccessToken(config.getAccessToken());
+    } else {
+      credential = builder.setClientSecrets(config.getClientId(),
+                                            config.getClientSecret()).build();
+      credential.createScoped(getRequiredScopes()).setRefreshToken(config.getRefreshToken());
+    }
     return credential;
   }
 

diff --git a/src/main/java/io/cdap/plugin/google/common/OAuthMethod.java b/src/main/java/io/cdap/plugin/google/common/OAuthMethod.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package io.cdap.plugin.google.common;
+
+/**
+ * Methods of getting an OAuth access token.
+ */
+public enum OAuthMethod {
+  ACCESS_TOKEN,
+  REFRESH_TOKEN;
+}
diff --git a/src/main/java/io/cdap/plugin/google/drive/source/GoogleDriveSourceConfig.java b/src/main/java/io/cdap/plugin/google/drive/source/GoogleDriveSourceConfig.java
@@ -364,6 +364,12 @@ public static GoogleDriveSourceConfig of(JsonObject properties) throws IOExcepti
     if (properties.has(GoogleDriveSourceConfig.REFRESH_TOKEN)) {
       googleDriveSourceConfig.setRefreshToken(properties.get(GoogleDriveSourceConfig.REFRESH_TOKEN).getAsString());
     }
+    if (properties.has(GoogleDriveSourceConfig.ACCESS_TOKEN)) {
+      googleDriveSourceConfig.setAccessToken(properties.get(GoogleDriveSourceConfig.ACCESS_TOKEN).getAsString());
+    }
+    if (properties.has(GoogleDriveSourceConfig.OAUTH_METHOD)) {
+      googleDriveSourceConfig.setoAuthMethod(properties.get(GoogleDriveSourceConfig.OAUTH_METHOD).getAsString());
+    }
     return googleDriveSourceConfig;
   }
 }
diff --git a/src/main/java/io/cdap/plugin/google/sheets/source/GoogleSheetsSourceConfig.java b/src/main/java/io/cdap/plugin/google/sheets/source/GoogleSheetsSourceConfig.java
@@ -282,7 +282,10 @@ private boolean shouldGetSchema() {
     return !containsMacro(SHEETS_TO_PULL) && !containsMacro(SHEETS_IDENTIFIERS) &&
       !containsMacro(COLUMN_NAMES_SELECTION) && !containsMacro(CUSTOM_COLUMN_NAMES_ROW) &&
       !containsMacro(LAST_DATA_COLUMN) && !containsMacro(NAME_SERVICE_ACCOUNT_TYPE) &&
-      !containsMacro(ACCOUNT_FILE_PATH) && !containsMacro(NAME_SERVICE_ACCOUNT_JSON);
+      !containsMacro(ACCOUNT_FILE_PATH) && !containsMacro(NAME_SERVICE_ACCOUNT_JSON) &&
+      !containsMacro(CLIENT_ID) && !containsMacro(CLIENT_SECRET) &&
+      !containsMacro(REFRESH_TOKEN) && !containsMacro(ACCESS_TOKEN) &&
+      !containsMacro(OAUTH_METHOD);
   }
 
   /**
@@ -1258,6 +1261,15 @@ public static GoogleSheetsSourceConfig of(JsonObject properties) throws IOExcept
         properties.get(GoogleSheetsSourceConfig.DIRECTORY_IDENTIFIER).getAsString());
     }
 
+    if (properties.has(GoogleSheetsSourceConfig.OAUTH_METHOD)) {
+      googleSheetsSourceConfig.setoAuthMethod(
+        properties.get(GoogleSheetsSourceConfig.OAUTH_METHOD).getAsString());
+    }
+
+    if (properties.has(GoogleSheetsSourceConfig.ACCESS_TOKEN)) {
+      googleSheetsSourceConfig.setAccessToken(
+        properties.get(GoogleSheetsSourceConfig.ACCESS_TOKEN).getAsString());
+    }
     return googleSheetsSourceConfig;
   }
 }
diff --git a/src/test/java/io/cdap/plugin/google/common/GoogleAuthBaseConfigTest.java b/src/test/java/io/cdap/plugin/google/common/GoogleAuthBaseConfigTest.java
@@ -84,4 +84,48 @@ public void testValidationErrorJSON() {
     Assert.assertEquals("serviceAccountJSON",
                         collector.getValidationFailures().get(0).getCauses().get(0).getAttribute("stageConfig"));
   }
+
+  @Test
+  public void testValidationOauthWithoutAccessToken() {
+    GoogleDriveSourceConfig config = new GoogleDriveSourceConfig("ref");
+    config.setReferenceName("validationErrorFilePath");
+    config.setAuthType("oAuth2");
+    config.setoAuthMethod(OAuthMethod.ACCESS_TOKEN.name());
+    config.setModificationDateRange("today");
+    config.getBodyFormat("string");
+    config.setStartDate("today");
+    FailureCollector collector = new DefaultFailureCollector("stageConfig", Collections.EMPTY_MAP);
+    config.validate(collector);
+    Assert.assertEquals(1, collector.getValidationFailures().size());
+    Assert.assertEquals("'Access Token' property is empty or macro is not available.",
+                        collector.getValidationFailures().get(0).getMessage());
+    Assert.assertEquals("accessToken",
+                        collector.getValidationFailures().get(0).getCauses().get(0).getAttribute("stageConfig"));
+  }
+
+  @Test
+  public void testValidationOauthWithoutRefreshToken() {
+    GoogleDriveSourceConfig config = new GoogleDriveSourceConfig(null);
+    config.setReferenceName("validationErrorFilePath");
+    config.setAuthType("oAuth2");
+    config.setoAuthMethod(OAuthMethod.REFRESH_TOKEN.name());
+    config.setModificationDateRange("today");
+    config.getBodyFormat("string");
+    config.setStartDate("today");
+    FailureCollector collector = new DefaultFailureCollector("stageConfig", Collections.EMPTY_MAP);
+    config.validate(collector);
+    Assert.assertEquals(3, collector.getValidationFailures().size());
+    Assert.assertEquals("'Client ID' property is empty or macro is not available.",
+                        collector.getValidationFailures().get(0).getMessage());
+    Assert.assertEquals("'Client Secret' property is empty or macro is not available.",
+                        collector.getValidationFailures().get(1).getMessage());
+    Assert.assertEquals("'Refresh Token' property is empty or macro is not available.",
+                        collector.getValidationFailures().get(2).getMessage());
+    Assert.assertEquals("clientId",
+                        collector.getValidationFailures().get(0).getCauses().get(0).getAttribute("stageConfig"));
+    Assert.assertEquals("clientSecret",
+                        collector.getValidationFailures().get(1).getCauses().get(0).getAttribute("stageConfig"));
+    Assert.assertEquals("refreshToken",
+                        collector.getValidationFailures().get(2).getCauses().get(0).getAttribute("stageConfig"));
+  }
 }
diff --git a/src/test/java/io/cdap/plugin/google/sheets/sink/GoogleSheetsSinkClientTest.java b/src/test/java/io/cdap/plugin/google/sheets/sink/GoogleSheetsSinkClientTest.java
@@ -21,6 +21,7 @@
 import com.google.api.services.sheets.v4.model.MergeCellsRequest;
 import com.google.api.services.sheets.v4.model.RowData;
 import io.cdap.plugin.google.common.AuthType;
+import io.cdap.plugin.google.common.OAuthMethod;
 import io.cdap.plugin.google.sheets.sink.utils.ComplexHeader;
 import io.cdap.plugin.google.sheets.sink.utils.FlatteredRowsRecord;
 import io.cdap.plugin.google.sheets.sink.utils.FlatteredRowsRequest;
@@ -49,7 +50,7 @@ public static void setupClient() throws IOException {
     EasyMock.expect(sinkConfig.getRefreshToken()).andReturn("dsfdsfdsfdsf").anyTimes();
     EasyMock.expect(sinkConfig.getClientSecret()).andReturn("dsfdsfdsfdsf").anyTimes();
     EasyMock.expect(sinkConfig.getClientId()).andReturn("dsdsrfegvrb").anyTimes();
-
+    EasyMock.expect(sinkConfig.getOAuthMethod()).andReturn(OAuthMethod.REFRESH_TOKEN).anyTimes();
     EasyMock.replay(sinkConfig);
     sinkClient = new GoogleSheetsSinkClient(sinkConfig);
 

diff --git a/widgets/GoogleDrive-batchsink.json b/widgets/GoogleDrive-batchsink.json
@@ -56,6 +56,33 @@
             ]
           }
         },
+        {
+          "name": "oAuthMethod",
+          "label": "OAuth Method",
+          "widget-type": "radio-group",
+          "widget-attributes": {
+            "layout": "inline",
+            "default": "REFRESH_TOKEN",
+            "options": [
+              {
+                "id": "REFRESH_TOKEN",
+                "label": "Refresh Token"
+              },
+              {
+                "id": "ACCESS_TOKEN",
+                "label": "Access Token"
+              }
+            ]
+          }
+        },
+        {
+          "name": "accessToken",
+          "label": "Access Token",
+          "widget-type": "textbox",
+          "widget-attributes": {
+            "placeholder": "${oauthAccessToken(provider,credential)}"
+          }
+        },
         {
           "widget-type": "textbox",
           "label": "Client ID",
@@ -115,6 +142,30 @@
         "operator": "equal to",
         "value": "oAuth2"
       },
+      "show": [
+        {
+          "name": "oAuthMethod",
+          "type": "property"
+        }
+      ]
+    },
+    {
+      "name": "Authenticate with OAuth2 Access Token",
+      "condition": {
+        "expression": "oAuthMethod == 'ACCESS_TOKEN' && authType == 'oAuth2'"
+      },
+      "show": [
+        {
+          "name": "accessToken",
+          "type": "property"
+        }
+      ]
+    },
+    {
+      "name": "Authenticate with OAuth2 Refresh Token",
+      "condition": {
+        "expression": "oAuthMethod == null || (oAuthMethod != 'ACCESS_TOKEN' && authType == 'oAuth2')"
+      },
       "show": [
         {
           "name": "clientId",