🎉 Source Drift: add OAuth 2.0 authentication option

airbyte-config/init/src/main/resources/config/STANDARD_SOURCE_DEFINITION/445831eb-78db-4b1f-8f1f-0d96ad8739e2.json

@@ -2,7 +2,7 @@
   "sourceDefinitionId": "445831eb-78db-4b1f-8f1f-0d96ad8739e2",
   "name": "Drift",
   "dockerRepository": "airbyte/source-drift",
-  "dockerImageTag": "0.2.3",
+  "dockerImageTag": "0.2.4",
   "documentationUrl": "https://docs.airbyte.io/integrations/sources/drift",
   "icon": "drift.svg"
 }

airbyte-config/init/src/main/resources/seed/source_definitions.yaml

@@ -123,7 +123,7 @@
 - name: Drift
   sourceDefinitionId: 445831eb-78db-4b1f-8f1f-0d96ad8739e2
   dockerRepository: airbyte/source-drift
-  dockerImageTag: 0.2.3
+  dockerImageTag: 0.2.4
   documentationUrl: https://docs.airbyte.io/integrations/sources/drift
   icon: drift.svg
   sourceType: api

airbyte-integrations/connectors/source-drift/Dockerfile

@@ -34,5 +34,5 @@ COPY source_drift ./source_drift
 ENV AIRBYTE_ENTRYPOINT "python /airbyte/integration_code/main.py"
 ENTRYPOINT ["python", "/airbyte/integration_code/main.py"]
 
-LABEL io.airbyte.version=0.2.3
+LABEL io.airbyte.version=0.2.4
 LABEL io.airbyte.name=airbyte/source-drift

airbyte-integrations/connectors/source-drift/README.md

@@ -47,10 +47,10 @@ and place them into `secrets/config.json`.
 
 ### Locally running the connector
 ```
-python main_dev.py spec
-python main_dev.py check --config secrets/config.json
-python main_dev.py discover --config secrets/config.json
-python main_dev.py read --config secrets/config.json --catalog sample_files/configured_catalog.json
+python main.py spec
+python main.py check --config secrets/config.json
+python main.py discover --config secrets/config.json
+python main.py read --config secrets/config.json --catalog sample_files/configured_catalog.json
 ```
 
 ### Unit Tests
@@ -88,6 +88,16 @@ docker run --rm -v $(pwd)/secrets:/secrets -v $(pwd)/sample_files:/sample_files
 1. To run additional integration tests, place your integration tests in a new directory `integration_tests` and run them with `python -m pytest -s integration_tests`.
    Make sure to familiarize yourself with [pytest test discovery](https://docs.pytest.org/en/latest/goodpractices.html#test-discovery) to know how your test files and methods should be named.
 
+#### Acceptance Tests
+Customize `acceptance-test-config.yml` file to configure tests. See [Source Acceptance Tests](https://docs.airbyte.io/connector-development/testing-connectors/source-acceptance-tests-reference) for more information.
+If your connector requires to create or destroy resources for use during acceptance tests create fixtures for it and place them inside integration_tests/acceptance.py.
+To run your integration tests with acceptance tests, from the connector root, run
+```
+docker build . --no-cache -t airbyte/source-drift:dev \
+&& python -m pytest -p source_acceptance_test.plugin
+```
+To run your integration tests with docker
+
 ## Dependency Management
 All of your dependencies should go in `setup.py`, NOT `requirements.txt`. The requirements file is only used to connect internal Airbyte dependencies in the monorepo for local development.
 

airbyte-integrations/connectors/source-drift/integration_tests/invalid_config.json

@@ -1,3 +1,7 @@
 {
-  "access_token": "invalid_access_token"
-}
+  "access_token": "invalid_access_token",
+  "credentials": {
+      "credentials": "access_token",
+      "access_token": "migrated_from_old_config"
+  }
+}

airbyte-integrations/connectors/source-drift/sample_files/sample_config.json

@@ -1,3 +1,5 @@
 {
-  "access_token": "1234567890abcdefghijk"
-}
+  "credentials": {
+      "access_token": "123412341234sfsdfs"
+  }
+}

airbyte-integrations/connectors/source-drift/source_drift/client/client.py

@@ -3,18 +3,30 @@
 #
 
 
-from typing import Iterator, Tuple
+from typing import Dict, Iterator, Tuple
 
 from airbyte_cdk.sources.deprecated.client import BaseClient
 
 from .api import APIClient
 from .common import AuthError, ValidationError
 
 
+class DriftAuthenticator:
+    def __init__(self, config: Dict):
+        self.config = config
+
+    def get_token(self) -> str:
+        access_token = self.config.get("access_token")
+        if access_token:
+            return access_token
+        else:
+            return self.config.get("credentials").get("access_token")
+
+
 class Client(BaseClient):
-    def __init__(self, access_token: str):
+    def __init__(self, **config: Dict):
         super().__init__()
-        self._client = APIClient(access_token)
+        self._client = APIClient(access_token=DriftAuthenticator(config).get_token())
 
     def stream__accounts(self, **kwargs) -> Iterator[dict]:
         yield from self._client.accounts.list()

airbyte-integrations/connectors/source-drift/source_drift/spec.json

@@ -4,14 +4,83 @@
     "$schema": "http://json-schema.org/draft-07/schema#",
     "title": "Drift Spec",
     "type": "object",
-    "required": ["access_token"],
-    "additionalProperties": false,
+    "required": [],
+    "additionalProperties": true,
     "properties": {
-      "access_token": {
-        "type": "string",
-        "description": "Drift Access Token. See the <a href=\"https://docs.airbyte.io/integrations/sources/drift\">docs</a> for more information on how to generate this key.",
-        "airbyte_secret": true
+      "credentials": {
+        "title": "Authorization Method",
+        "type": "object",
+        "oneOf": [
+          {
+            "type": "object",
+            "title": "OAuth2.0",
+            "required": ["client_id", "client_secret", "access_token", "refresh_token"],
+            "properties": {
+              "credentials": {
+                "type": "string",
+                "const": "oauth2.0",
+                "enum": ["oauth2.0"],
+                "default": "oauth2.0",
+                "order": 0
+              },
+              "client_id": {
+                "type": "string",
+                "title": "Client ID",
+                "description": "The Client ID of your Drift developer application.",
+                "airbyte_secret": true
+              },
+              "client_secret": {
+                "type": "string",
+                "title": "Client Secret",
+                "description": "The Client Secret of your Drift developer application.",
+                "airbyte_secret": true
+              },
+              "access_token": {
+                "type": "string",
+                "title": "Access Token",
+                "description": "Access Token for making authenticated requests.",
+                "airbyte_secret": true
+              },
+              "refresh_token": {
+                "type": "string",
+                "title": "Refresh Token",
+                "description": "Refresh Token to renew the expired access_token.",
+                "default": "",
+                "airbyte_secret": true
+              }
+            }
+          },
+          {
+            "title": "Access Token",
+            "type": "object",
+            "required": ["access_token"],
+            "properties": {
+              "credentials": {
+                "type": "string",
+                "const": "access_token",
+                "enum": ["access_token"],
+                "default": "access_token",
+                "order": 0
+              },
+              "access_token": {
+                "type": "string",
+                "title": "Access Token",
+                "description": "Drift Access Token. See the <a href=\"https://docs.airbyte.io/integrations/sources/drift\">docs</a> for more information on how to generate this key.",
+                "airbyte_secret": true
+              }
+            }
+          }
+        ]
       }
     }
+  },
+  "authSpecification": {
+    "auth_type": "oauth2.0",
+    "oauth2Specification": {
+      "rootObject": ["credentials", "0"],
+      "oauthFlowInitParameters": [["client_id"], ["client_secret"]],
+      "oauthFlowOutputParameters": [["access_token"], ["refresh_token"]]
+    }
   }
 }
+

airbyte-integrations/connectors/source-drift/unit_tests/test_client.py

@@ -6,16 +6,18 @@
 import pytest
 from source_drift.client import AuthError, Client
 
+config = {"credentials": {"access_token": "wrong_key"}}
+
 
 def test__heal_check_with_wrong_token():
-    client = Client(access_token="wrong_key")
+    client = Client(**config)
     alive, error = client.health_check()
 
     assert not alive
     assert error == "(401, 'The access token is invalid or has expired')"
 
 
 def test__users_with_wrong_token():
-    client = Client(access_token="wrong_key")
+    client = Client(**config)
     with pytest.raises(AuthError, match="(401, 'The access token is invalid or has expired')"):
         next(client.stream__users())

airbyte-oauth/src/main/java/io/airbyte/oauth/OAuthImplementationFactory.java

@@ -9,6 +9,7 @@
 import io.airbyte.config.StandardSourceDefinition;
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.oauth.flows.AsanaOAuthFlow;
+import io.airbyte.oauth.flows.DriftOAuthFlow;
 import io.airbyte.oauth.flows.GithubOAuthFlow;
 import io.airbyte.oauth.flows.HubspotOAuthFlow;
 import io.airbyte.oauth.flows.IntercomOAuthFlow;
@@ -54,6 +55,7 @@ public OAuthImplementationFactory(final ConfigRepository configRepository, final
         .put("airbyte/source-snapchat-marketing", new SnapchatMarketingOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-surveymonkey", new SurveymonkeyOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-trello", new TrelloOAuthFlow(configRepository, httpClient))
+        .put("airbyte/source-drift", new DriftOAuthFlow(configRepository))
         .build();
   }
 

airbyte-oauth/src/main/java/io/airbyte/oauth/flows/DriftOAuthFlow.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.oauth.flows;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import io.airbyte.config.persistence.ConfigRepository;
+import io.airbyte.oauth.BaseOAuthFlow;
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.http.HttpClient;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+import java.util.function.Supplier;
+import org.apache.http.client.utils.URIBuilder;
+
+/**
+ * Following docs from
+ * https://devdocs.drift.com/docs/authentication-and-scopes#1-direct-the-user-to-the-drift-oauth-url-
+ */
+public class DriftOAuthFlow extends BaseOAuthFlow {
+
+  private static final String ACCESS_TOKEN_URL = "https://driftapi.com/oauth2/token";
+
+  public DriftOAuthFlow(ConfigRepository configRepository) {
+    super(configRepository);
+  }
+
+  @VisibleForTesting
+  DriftOAuthFlow(ConfigRepository configRepository, HttpClient httpClient, Supplier<String> stateSupplier) {
+    super(configRepository, httpClient, stateSupplier);
+  }
+
+  @Override
+  protected String formatConsentUrl(UUID definitionId, String clientId, String redirectUrl) throws IOException {
+    final URIBuilder builder = new URIBuilder()
+        .setScheme("https")
+        .setHost("dev.drift.com")
+        .setPath("authorize")
+        .addParameter("response_type", "code")
+        .addParameter("client_id", clientId)
+        .addParameter("redirect_uri", redirectUrl)
+        .addParameter("state", getState());
+    try {
+      return builder.build().toString();
+    } catch (URISyntaxException e) {
+      throw new IOException("Failed to format Consent URL for OAuth flow", e);
+    }
+  }
+
+  @Override
+  protected String extractCodeParameter(Map<String, Object> queryParams) throws IOException {
+    if (queryParams.containsKey("code")) {
+      return (String) queryParams.get("code");
+    } else {
+      throw new IOException("Undefined 'code' from consent redirected url.");
+    }
+  }
+
+  @Override
+  protected String getClientIdUnsafe(JsonNode config) {
+    // the config object containing client ID is nested inside the "credentials" object
+    Preconditions.checkArgument(config.hasNonNull("credentials"));
+    return super.getClientIdUnsafe(config.get("credentials"));
+  }
+
+  @Override
+  protected String getClientSecretUnsafe(JsonNode config) {
+    // the config object containing client SECRET is nested inside the "credentials" object
+    Preconditions.checkArgument(config.hasNonNull("credentials"));
+    return super.getClientSecretUnsafe(config.get("credentials"));
+  }
+
+  @Override
+  protected String getAccessTokenUrl() {
+    return ACCESS_TOKEN_URL;
+  }
+
+  @Override
+  protected Map<String, String> getAccessTokenQueryParameters(String clientId, String clientSecret, String authCode, String redirectUrl) {
+    return ImmutableMap.<String, String>builder()
+        .put("client_id", clientId)
+        .put("client_secret", clientSecret)
+        .put("code", authCode)
+        .put("grant_type", "authorization_code")
+        .build();
+  }
+
+  protected Map<String, Object> extractRefreshToken(final JsonNode data, String accessTokenUrl) throws IOException {
+    final Map<String, Object> result = new HashMap<>();
+    if (data.has("access_token")) {
+      result.put("access_token", data.get("access_token").asText());
+    } else {
+      throw new IOException(String.format("Missing 'access_token' in query params from %s", accessTokenUrl));
+    }
+    return Map.of("credentials", result);
+  }
+
+}