🎉 square: added oauth support

airbyte-integrations/connectors/source-square/acceptance-test-config.yml

@@ -5,17 +5,23 @@ tests:
   connection:
     - config_path: "secrets/config.json"
       status: "succeed"
+    - config_path: "secrets/config_oauth.json"
+      status: "succeed"
     - config_path: "integration_tests/invalid_config.json"
       status: "failed"
   discovery:
     - config_path: "secrets/config.json"
+    - config_path: "secrets/config_oauth.json"
   basic_read:
     - config_path: "secrets/config.json"
       configured_catalog_path: "integration_tests/configured_catalog.json"
+    - config_path: "secrets/config_oauth.json"
+      configured_catalog_path: "integration_tests/configured_catalog_oauth.json"
   incremental:
     - config_path: "secrets/config.json"
       configured_catalog_path: "integration_tests/configured_catalog.json"
       future_state_path: "integration_tests/abnormal_state.json"
   full_refresh:
     - config_path: "secrets/config.json"
       configured_catalog_path: "integration_tests/configured_catalog.json"
+

airbyte-integrations/connectors/source-square/integration_tests/configured_catalog_oauth.json

@@ -0,0 +1,40 @@
+{
+  "streams": [
+    {
+      "stream": {
+        "name": "locations",
+        "json_schema": {},
+        "supported_sync_modes": ["full_refresh"],
+        "source_defined_cursor": true,
+        "default_cursor_field": ["id"]
+      },
+      "sync_mode": "full_refresh",
+      "cursor_field": ["id"],
+      "destination_sync_mode": "overwrite"
+    },
+    {
+      "stream": {
+        "name": "team_members",
+        "json_schema": {},
+        "supported_sync_modes": ["full_refresh"],
+        "source_defined_cursor": true,
+        "default_cursor_field": ["id"]
+      },
+      "sync_mode": "full_refresh",
+      "cursor_field": ["id"],
+      "destination_sync_mode": "overwrite"
+    },
+    {
+      "stream": {
+        "name": "team_member_wages",
+        "json_schema": {},
+        "supported_sync_modes": ["full_refresh"],
+        "source_defined_cursor": true,
+        "default_cursor_field": ["id"]
+      },
+      "sync_mode": "full_refresh",
+      "cursor_field": ["id"],
+      "destination_sync_mode": "overwrite"
+    }
+  ]
+}

airbyte-integrations/connectors/source-square/source_square/source.py

@@ -4,15 +4,18 @@
 
 import json
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
 import pendulum
 import requests
+from airbyte_cdk.logger import AirbyteLogger
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.http import HttpStream
-from airbyte_cdk.sources.streams.http.auth import TokenAuthenticator
+from airbyte_cdk.sources.streams.http.auth.core import HttpAuthenticator
+from airbyte_cdk.sources.streams.http.requests_native_auth import Oauth2Authenticator, TokenAuthenticator
+from requests.auth import AuthBase
 from source_square.utils import separate_items_by_count
 
 
@@ -35,8 +38,16 @@ def parse_square_error_response(error: requests.exceptions.HTTPError) -> SquareE
 
 
 class SquareStream(HttpStream, ABC):
-    def __init__(self, is_sandbox: bool, api_version: str, start_date: str, include_deleted_objects: bool, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        is_sandbox: bool,
+        api_version: str,
+        start_date: str,
+        include_deleted_objects: bool,
+        authenticator: Union[AuthBase, HttpAuthenticator],
+    ):
+        super().__init__(authenticator)
+        self._authenticator = authenticator
         self.is_sandbox = is_sandbox
         self.api_version = api_version
         # Converting users ISO 8601 format (YYYY-MM-DD) to RFC 3339 (2021-06-14T13:47:56.799Z)
@@ -358,16 +369,75 @@ def stream_slices(self, **kwargs) -> Iterable[Optional[Mapping[str, Any]]]:
             yield {"location_ids": location}
 
 
+class Oauth2AuthenticatorSquare(Oauth2Authenticator):
+    def refresh_access_token(self) -> Tuple[str, int]:
+        """Handle differences in expiration attr:
+        from API: "expires_at": "2021-11-05T14:26:57Z"
+        expected: "expires_in": number of seconds
+        """
+        token, expires_at = super().refresh_access_token()
+        expires_in = pendulum.parse(expires_at) - pendulum.now()
+        return token, expires_in.seconds
+
+
 class SourceSquare(AbstractSource):
-    api_version = "2021-06-16"  # Latest Stable Release
+    api_version = "2021-09-15"  # Latest Stable Release
+
+    @staticmethod
+    def get_auth(config: Mapping[str, Any]) -> AuthBase:
+
+        credential = config.get("credentials", {})
+        auth_type = credential.get("auth_type")
+        if auth_type == "Oauth":
+            # scopes needed for all currently supported streams:
+            scopes = [
+                "CUSTOMERS_READ",
+                "EMPLOYEES_READ",
+                "ITEMS_READ",
+                "MERCHANT_PROFILE_READ",
+                "ORDERS_READ",
+                "PAYMENTS_READ",
+                "TIMECARDS_READ",
+                # OAuth Permissions:
+                # https://developer.squareup.com/docs/oauth-api/square-permissions
+                # https://developer.squareup.com/reference/square/enums/OAuthPermission
+                # "DISPUTES_READ",
+                # "GIFTCARDS_READ",
+                # "INVENTORY_READ",
+                # "INVOICES_READ",
+                # "TIMECARDS_SETTINGS_READ",
+                # "LOYALTY_READ",
+                # "ONLINE_STORE_SITE_READ",
+                # "ONLINE_STORE_SNIPPETS_READ",
+                # "SUBSCRIPTIONS_READ",
+            ]
+
+            auth = Oauth2AuthenticatorSquare(
+                token_refresh_endpoint="https://connect.squareup.com/oauth2/token",
+                client_secret=credential.get("client_secret"),
+                client_id=credential.get("client_id"),
+                refresh_token=credential.get("refresh_token"),
+                scopes=scopes,
+                expires_in_name="expires_at",
+            )
+        elif auth_type == "Apikey":
+            auth = TokenAuthenticator(token=credential.get("api_key"))
+        elif not auth_type and config.get("api_key"):
+            auth = TokenAuthenticator(token=config.get("api_key"))
+        else:
+            raise Exception(f"Invalid auth type: {auth_type}")
+
+        return auth
 
-    def check_connection(self, logger, config) -> Tuple[bool, any]:
+    def check_connection(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> Tuple[bool, any]:
 
         headers = {
             "Square-Version": self.api_version,
-            "Authorization": "Bearer {}".format(config["api_key"]),
             "Content-Type": "application/json",
         }
+        auth = self.get_auth(config)
+        headers.update(auth.get_auth_header())
+
         url = "https://connect.squareup{}.com/v2/catalog/info".format("sandbox" if config["is_sandbox"] else "")
 
         try:
@@ -383,9 +453,8 @@ def check_connection(self, logger, config) -> Tuple[bool, any]:
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
 
-        auth = TokenAuthenticator(token=config["api_key"])
         args = {
-            "authenticator": auth,
+            "authenticator": self.get_auth(config),
             "is_sandbox": config["is_sandbox"],
             "api_version": self.api_version,
             "start_date": config["start_date"],

airbyte-integrations/connectors/source-square/source_square/spec.json

@@ -4,37 +4,146 @@
     "$schema": "http://json-schema.org/draft-07/schema#",
     "title": "Square Source CDK Specifications",
     "type": "object",
-    "required": ["api_key", "is_sandbox"],
-    "additionalProperties": false,
+    "required": ["is_sandbox"],
+    "additionalProperties": true,
     "properties": {
-      "api_key": {
-        "type": "string",
-        "description": "The API key for a Square application.",
-        "title": "API Key",
-        "airbyte_secret": true
-      },
       "is_sandbox": {
         "type": "boolean",
         "description": "Determines whether to use the sandbox or production environment.",
         "title": "Sandbox",
         "examples": [true, false],
-        "default": true
+        "default": false
       },
       "start_date": {
         "type": "string",
         "description": "UTC date in the format YYYY-MM-DD. Any data before this date will not be replicated. If not set, all data will be replicated.",
         "title": "Start Date",
         "examples": ["2021-01-01"],
-        "default": "1970-01-01",
+        "default": "2021-01-01",
         "pattern": "^[0-9]{4}-[0-9]{2}-[0-9]{2}$"
       },
       "include_deleted_objects": {
         "type": "boolean",
         "description": "In some streams there is an option to include deleted objects (Items, Categories, Discounts, Taxes)",
-        "title": "Include Deleded Objects",
+        "title": "Include Deleted Objects",
         "examples": [true, false],
         "default": false
+      },
+      "credentials": {
+        "type": "object",
+        "title": "Credential Type",
+        "oneOf": [
+          {
+            "title": "Oauth authentication",
+            "type": "object",
+            "required": [
+              "auth_type",
+              "client_id",
+              "client_secret",
+              "refresh_token"
+            ],
+            "properties": {
+              "auth_type": {
+                "type": "string",
+                "const": "Oauth",
+                "enum": ["Oauth"],
+                "default": "Oauth",
+                "order": 0
+              },
+              "client_id": {
+                "title": "Client ID",
+                "type": "string",
+                "description": "The Square-issued ID of your application",
+                "airbyte_secret": true
+              },
+              "client_secret": {
+                "title": "Client Secret",
+                "type": "string",
+                "description": "The Square-issued application secret for your application",
+                "airbyte_secret": true
+              },
+              "refresh_token": {
+                "title": "Refresh Token",
+                "type": "string",
+                "description": "A refresh token generated using the above client ID and secret",
+                "airbyte_secret": true
+              }
+            }
+          },
+          {
+            "type": "object",
+            "title": "API Key",
+            "required": ["auth_type", "api_key"],
+            "properties": {
+              "auth_type": {
+                "type": "string",
+                "const": "Apikey",
+                "enum": ["Apikey"],
+                "default": "Apikey",
+                "order": 1
+              },
+              "api_key": {
+                "title": "API key token",
+                "type": "string",
+                "description": "The API key for a Square application",
+                "airbyte_secret": true
+              }
+            }
+          }
+        ]
+      }
+    }
+  },
+  "authSpecification": {
+    "auth_type": "oauth2.0",
+    "oauth2Specification": {
+      "rootObject": ["credentials", 0],
+      "oauthFlowInitParameters": [["client_id"], ["client_secret"]],
+      "oauthFlowOutputParameters": [["refresh_token"]]
+    }
+  },
+  "advanced_auth": {
+    "auth_flow_type": "oauth2.0",
+    "predicate_key": ["credentials", "auth_type"],
+    "predicate_value": "Oauth",
+    "oauth_config_specification": {
+      "complete_oauth_output_specification": {
+        "type": "object",
+        "additionalProperties": false,
+        "properties": {
+          "refresh_token": {
+            "type": "string",
+            "path_in_connector_config": ["credentials", "refresh_token"]
+          }
+        }
+      },
+      "complete_oauth_server_input_specification": {
+        "type": "object",
+        "additionalProperties": false,
+        "properties": {
+          "client_id": {
+            "type": "string"
+          },
+          "client_secret": {
+            "type": "string"
+          }
+        }
+      },
+      "complete_oauth_server_output_specification": {
+        "type": "object",
+        "additionalProperties": false,
+        "properties": {
+          "client_id": {
+            "type": "string",
+            "path_in_connector_config": ["credentials", "client_id"]
+          },
+          "client_secret": {
+            "type": "string",
+            "path_in_connector_config": ["credentials", "client_secret"]
+          }
+        }
       }
     }
   }
+
 }

airbyte-integrations/connectors/source-square/unit_tests/connection_test.py

@@ -8,5 +8,11 @@
 
 def test_source_wrong_credentials():
     source = SourceSquare()
-    status, error = source.check_connection(logger=AirbyteLogger(), config={"api_key": "wrong.api.key", "is_sandbox": True})
+    config = {
+        "credentials": {"auth_type": "Apikey", "api_key": "bla"},
+        "is_sandbox": True,
+        "start_date": "2021-06-01",
+        "include_deleted_objects": False,
+    }
+    status, error = source.check_connection(logger=AirbyteLogger(), config=config)
     assert not status

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

@@ -47,6 +47,7 @@ public OAuthImplementationFactory(final ConfigRepository configRepository, final
         .put("airbyte/source-salesforce", new SalesforceOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-slack", new SlackOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-snapchat-marketing", new SnapchatMarketingOAuthFlow(configRepository, httpClient))
+        .put("airbyte/source-square", new SquareOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-strava", new StravaOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-surveymonkey", new SurveymonkeyOAuthFlow(configRepository, httpClient))
         .put("airbyte/source-trello", new TrelloOAuthFlow(configRepository, httpClient))