🎉 square: added oauth support (airbytehq#6842)

* fixed test which check incorrect cred config

* Added oauth2 authentication

* Added oauth creds

* fixed formatting

* added oauth2 spec section, added missing type hints

* Added java part of Square OAuth

* fixed checkstyle

* removed commented code

* added support for old format of spec.json files, updated change logs docs

* renamed spec property 'authentication' to default 'credentials'. fixed changes in java part

* recovered empty files

* updated OAuthImplementationFactory.java

* fixed issue with autheticator for sub streams, added config catalog with all streams, updated docs

* use advanced_auth

* added advanced_auth

* moved scopes to private property

* updated source version

* Revert "updated source version"

This reverts commit ce3d061.

* updated source version

* added new version for airbyte index

Co-authored-by: ievgeniit <[email protected]>

airbyte-config/init/src/main/resources/config/STANDARD_SOURCE_DEFINITION/77225a51-cd15-4a13-af02-65816bd0ecf4.json

@@ -2,7 +2,7 @@
   "sourceDefinitionId": "77225a51-cd15-4a13-af02-65816bd0ecf4",
   "name": "Square",
   "dockerRepository": "airbyte/source-square",
-  "dockerImageTag": "0.1.3",
+  "dockerImageTag": "0.1.4",
   "documentationUrl": "https://docs.airbyte.io/integrations/sources/square",
   "icon": "square.svg"
 }

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

@@ -676,7 +676,7 @@
 - name: Square
   sourceDefinitionId: 77225a51-cd15-4a13-af02-65816bd0ecf4
   dockerRepository: airbyte/source-square
-  dockerImageTag: 0.1.3
+  dockerImageTag: 0.1.4
   documentationUrl: https://docs.airbyte.io/integrations/sources/square
   icon: square.svg
   sourceType: api

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

@@ -7049,52 +7049,148 @@
     supportsNormalization: false
     supportsDBT: false
     supported_destination_sync_modes: []
-- dockerImage: "airbyte/source-square:0.1.3"
+- dockerImage: "airbyte/source-square:0.1.4"
   spec:
     documentationUrl: "https://docs.airbyte.io/integrations/sources/square"
     connectionSpecification:
       $schema: "http://json-schema.org/draft-07/schema#"
       title: "Square Source CDK Specifications"
       type: "object"
       required:
-      - "api_key"
       - "is_sandbox"
-      additionalProperties: false
+      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
     supportsNormalization: false
     supportsDBT: false
     supported_destination_sync_modes: []
+    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"
 - dockerImage: "airbyte/source-strava:0.1.2"
   spec:
     documentationUrl: "https://docs.airbyte.io/integrations/sources/strava"

airbyte-integrations/connectors/source-square/Dockerfile

@@ -12,5 +12,5 @@ RUN pip install .
 ENV AIRBYTE_ENTRYPOINT "python /airbyte/integration_code/main.py"
 ENTRYPOINT ["python", "/airbyte/integration_code/main.py"]
 
-LABEL io.airbyte.version=0.1.3
+LABEL io.airbyte.version=0.1.4
 LABEL io.airbyte.name=airbyte/source-square

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"],