Add auth flow to declarative manifest schema (#24441)

* Add auth flow to declarative manifest schema

* Rename

* fix rename

* set advanced_auth

* Automated Commit - Formatting Changes

* update unit test

* format

* Add examples

* example -> examples

* add missing examples

* Automated Commit - Formatting Changes

---------

Co-authored-by: girarda <girarda@users.noreply.github.com>

airbyte-cdk/python/airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1043,6 +1043,149 @@ definitions:
         additionalProperties: true
       documentation_url:
         type: string
+      advanced_auth:
+        "$ref": "#/definitions/AuthFlow"
+  AuthFlow:
+    title: "Auth flow"
+    type: object
+    description: |-
+      Additional and optional specification object to describe what an 'advanced' Auth flow would need to function.
+        - A connector should be able to fully function with the configuration as described by the ConnectorSpecification in a 'basic' mode.
+        - The 'advanced' mode provides easier UX for the user with UI improvements and automations. However, this requires further setup on the
+        server side by instance or workspace admins beforehand. The trade-off is that the user does not have to provide as many technical
+        inputs anymore and the auth process is faster and easier to complete.
+    properties:
+      auth_flow_type:
+        title: "Auth flow type"
+        description: "The type of auth to use"
+        type: string
+        enum: ["oauth2.0", "oauth1.0"] # Future auth types should be added here
+      predicate_key:
+        title: "Predicate key"
+        description: Json Path to a field in the connectorSpecification that should exist for the advanced auth to be applicable.
+        type: array
+        items:
+          type: string
+        examples:
+          - ["credentials", "auth_type"]
+      predicate_value:
+        title: "Predicate value"
+        description: Value of the predicate_key fields for the advanced auth to be applicable.
+        type: string
+        examples:
+          - "Oauth"
+      oauth_config_specification:
+        "$ref": "#/definitions/OAuthConfigSpecification"
+  OAuthConfigSpecification:
+    type: object
+    additionalProperties: true
+    properties:
+      oauth_user_input_from_connector_config_specification:
+        title: "OAuth user input"
+        description: |-
+          OAuth specific blob. This is a Json Schema used to validate Json configurations used as input to OAuth.
+          Must be a valid non-nested JSON that refers to properties from ConnectorSpecification.connectionSpecification
+          using special annotation 'path_in_connector_config'.
+          These are input values the user is entering through the UI to authenticate to the connector, that might also shared
+          as inputs for syncing data via the connector.
+          Examples:
+          if no connector values is shared during oauth flow, oauth_user_input_from_connector_config_specification=[]
+          if connector values such as 'app_id' inside the top level are used to generate the API url for the oauth flow,
+            oauth_user_input_from_connector_config_specification={
+              app_id: {
+                type: string
+                path_in_connector_config: ['app_id']
+              }
+            }
+          if connector values such as 'info.app_id' nested inside another object are used to generate the API url for the oauth flow,
+            oauth_user_input_from_connector_config_specification={
+              app_id: {
+                type: string
+                path_in_connector_config: ['info', 'app_id']
+              }
+            }
+        type: object
+        examples:
+          - app_id:
+              type: string
+              path_in_connector_config: ["app_id"]
+          - app_id:
+              type: string
+              path_in_connector_config: ["info", "app_id"]
+      complete_oauth_output_specification:
+        title: "OAuth output specification"
+        description: |-
+          OAuth specific blob. This is a Json Schema used to validate Json configurations produced by the OAuth flows as they are
+          returned by the distant OAuth APIs.
+          Must be a valid JSON describing the fields to merge back to `ConnectorSpecification.connectionSpecification`.
+          For each field, a special annotation `path_in_connector_config` can be specified to determine where to merge it,
+          Examples:
+              complete_oauth_output_specification={
+                refresh_token: {
+                  type: string,
+                  path_in_connector_config: ['credentials', 'refresh_token']
+                }
+              }
+        type: object
+        additionalProperties: true
+        examples:
+          - refresh_token:
+              type: string,
+              path_in_connector_config: ["credentials", "refresh_token"]
+      complete_oauth_server_input_specification:
+        title: "OAuth input specification"
+        description: |-
+          OAuth specific blob. This is a Json Schema used to validate Json configurations persisted as Airbyte Server configurations.
+          Must be a valid non-nested JSON describing additional fields configured by the Airbyte Instance or Workspace Admins to be used by the
+          server when completing an OAuth flow (typically exchanging an auth code for refresh token).
+          Examples:
+              complete_oauth_server_input_specification={
+                client_id: {
+                  type: string
+                },
+                client_secret: {
+                  type: string
+                }
+              }
+        type: object
+        additionalProperties: true
+        examples:
+          - client_id:
+              type: string
+            client_secret:
+              type: string
+      complete_oauth_server_output_specification:
+        title: "OAuth server output specification"
+        description: |-
+          OAuth specific blob. This is a Json Schema used to validate Json configurations persisted as Airbyte Server configurations that
+          also need to be merged back into the connector configuration at runtime.
+          This is a subset configuration of `complete_oauth_server_input_specification` that filters fields out to retain only the ones that
+          are necessary for the connector to function with OAuth. (some fields could be used during oauth flows but not needed afterwards, therefore
+          they would be listed in the `complete_oauth_server_input_specification` but not `complete_oauth_server_output_specification`)
+          Must be a valid non-nested JSON describing additional fields configured by the Airbyte Instance or Workspace Admins to be used by the
+          connector when using OAuth flow APIs.
+          These fields are to be merged back to `ConnectorSpecification.connectionSpecification`.
+          For each field, a special annotation `path_in_connector_config` can be specified to determine where to merge it,
+          Examples:
+                complete_oauth_server_output_specification={
+                  client_id: {
+                    type: string,
+                    path_in_connector_config: ['credentials', 'client_id']
+                  },
+                  client_secret: {
+                    type: string,
+                    path_in_connector_config: ['credentials', 'client_secret']
+                  }
+                }
+        type: object
+        additionalProperties: true
+        examples:
+          - client_id:
+              type: string,
+              path_in_connector_config: ["credentials", "client_id"]
+            client_secret:
+              type: string,
+              path_in_connector_config: ["credentials", "client_secret"]
   SubstreamPartitionRouter:
     description: Partition router that is used to retrieve records that have been partitioned according to records from the specified parent streams
     type: object

airbyte-cdk/python/airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -309,10 +309,65 @@ class SessionTokenAuthenticator(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class Spec(BaseModel):
-    type: Literal["Spec"]
-    connection_specification: Dict[str, Any]
-    documentation_url: Optional[str] = None
+class AuthFlowType(Enum):
+    oauth2_0 = "oauth2.0"
+    oauth1_0 = "oauth1.0"
+
+
+class OAuthConfigSpecification(BaseModel):
+    class Config:
+        extra = Extra.allow
+
+    oauth_user_input_from_connector_config_specification: Optional[Dict[str, Any]] = Field(
+        None,
+        description="OAuth specific blob. This is a Json Schema used to validate Json configurations used as input to OAuth.\nMust be a valid non-nested JSON that refers to properties from ConnectorSpecification.connectionSpecification\nusing special annotation 'path_in_connector_config'.\nThese are input values the user is entering through the UI to authenticate to the connector, that might also shared\nas inputs for syncing data via the connector.\nExamples:\nif no connector values is shared during oauth flow, oauth_user_input_from_connector_config_specification=[]\nif connector values such as 'app_id' inside the top level are used to generate the API url for the oauth flow,\n  oauth_user_input_from_connector_config_specification={\n    app_id: {\n      type: string\n      path_in_connector_config: ['app_id']\n    }\n  }\nif connector values such as 'info.app_id' nested inside another object are used to generate the API url for the oauth flow,\n  oauth_user_input_from_connector_config_specification={\n    app_id: {\n      type: string\n      path_in_connector_config: ['info', 'app_id']\n    }\n  }",
+        examples=[
+            {"app_id": {"type": "string", "path_in_connector_config": ["app_id"]}},
+            {
+                "app_id": {
+                    "type": "string",
+                    "path_in_connector_config": ["info", "app_id"],
+                }
+            },
+        ],
+        title="OAuth user input",
+    )
+    complete_oauth_output_specification: Optional[Dict[str, Any]] = Field(
+        None,
+        description="OAuth specific blob. This is a Json Schema used to validate Json configurations produced by the OAuth flows as they are\nreturned by the distant OAuth APIs.\nMust be a valid JSON describing the fields to merge back to `ConnectorSpecification.connectionSpecification`.\nFor each field, a special annotation `path_in_connector_config` can be specified to determine where to merge it,\nExamples:\n    complete_oauth_output_specification={\n      refresh_token: {\n        type: string,\n        path_in_connector_config: ['credentials', 'refresh_token']\n      }\n    }",
+        examples=[
+            {
+                "refresh_token": {
+                    "type": "string,",
+                    "path_in_connector_config": ["credentials", "refresh_token"],
+                }
+            }
+        ],
+        title="OAuth output specification",
+    )
+    complete_oauth_server_input_specification: Optional[Dict[str, Any]] = Field(
+        None,
+        description="OAuth specific blob. This is a Json Schema used to validate Json configurations persisted as Airbyte Server configurations.\nMust be a valid non-nested JSON describing additional fields configured by the Airbyte Instance or Workspace Admins to be used by the\nserver when completing an OAuth flow (typically exchanging an auth code for refresh token).\nExamples:\n    complete_oauth_server_input_specification={\n      client_id: {\n        type: string\n      },\n      client_secret: {\n        type: string\n      }\n    }",
+        examples=[{"client_id": {"type": "string"}, "client_secret": {"type": "string"}}],
+        title="OAuth input specification",
+    )
+    complete_oauth_server_output_specification: Optional[Dict[str, Any]] = Field(
+        None,
+        description="OAuth specific blob. This is a Json Schema used to validate Json configurations persisted as Airbyte Server configurations that\nalso need to be merged back into the connector configuration at runtime.\nThis is a subset configuration of `complete_oauth_server_input_specification` that filters fields out to retain only the ones that\nare necessary for the connector to function with OAuth. (some fields could be used during oauth flows but not needed afterwards, therefore\nthey would be listed in the `complete_oauth_server_input_specification` but not `complete_oauth_server_output_specification`)\nMust be a valid non-nested JSON describing additional fields configured by the Airbyte Instance or Workspace Admins to be used by the\nconnector when using OAuth flow APIs.\nThese fields are to be merged back to `ConnectorSpecification.connectionSpecification`.\nFor each field, a special annotation `path_in_connector_config` can be specified to determine where to merge it,\nExamples:\n      complete_oauth_server_output_specification={\n        client_id: {\n          type: string,\n          path_in_connector_config: ['credentials', 'client_id']\n        },\n        client_secret: {\n          type: string,\n          path_in_connector_config: ['credentials', 'client_secret']\n        }\n      }",
+        examples=[
+            {
+                "client_id": {
+                    "type": "string,",
+                    "path_in_connector_config": ["credentials", "client_id"],
+                },
+                "client_secret": {
+                    "type": "string,",
+                    "path_in_connector_config": ["credentials", "client_secret"],
+                },
+            }
+        ],
+        title="OAuth server output specification",
+    )
 
 
 class WaitTimeFromHeader(BaseModel):
@@ -419,6 +474,23 @@ class RecordSelector(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class AuthFlow(BaseModel):
+    auth_flow_type: Optional[AuthFlowType] = Field(None, description="The type of auth to use", title="Auth flow type")
+    predicate_key: Optional[List[str]] = Field(
+        None,
+        description="Json Path to a field in the connectorSpecification that should exist for the advanced auth to be applicable.",
+        examples=[["credentials", "auth_type"]],
+        title="Predicate key",
+    )
+    predicate_value: Optional[str] = Field(
+        None,
+        description="Value of the predicate_key fields for the advanced auth to be applicable.",
+        examples=["Oauth"],
+        title="Predicate value",
+    )
+    oauth_config_specification: Optional[OAuthConfigSpecification] = None
+
+
 class CompositeErrorHandler(BaseModel):
     type: Literal["CompositeErrorHandler"]
     error_handlers: List[Union[CompositeErrorHandler, DefaultErrorHandler]]
@@ -469,6 +541,13 @@ class HttpRequester(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class Spec(BaseModel):
+    type: Literal["Spec"]
+    connection_specification: Dict[str, Any]
+    documentation_url: Optional[str] = None
+    advanced_auth: Optional[AuthFlow] = None
+
+
 class DeclarativeSource(BaseModel):
     class Config:
         extra = Extra.forbid

airbyte-cdk/python/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -759,7 +759,12 @@ def create_simple_retriever(
 
     @staticmethod
     def create_spec(model: SpecModel, config: Config, **kwargs) -> Spec:
-        return Spec(connection_specification=model.connection_specification, documentation_url=model.documentation_url, parameters={})
+        return Spec(
+            connection_specification=model.connection_specification,
+            documentation_url=model.documentation_url,
+            advanced_auth=model.advanced_auth,
+            parameters={},
+        )
 
     def create_substream_partition_router(self, model: SubstreamPartitionRouterModel, config: Config, **kwargs) -> SubstreamPartitionRouter:
         parent_stream_configs = []

airbyte-cdk/python/airbyte_cdk/sources/declarative/spec/spec.py

@@ -6,6 +6,7 @@
 from typing import Any, Mapping, Optional
 
 from airbyte_cdk.models.airbyte_protocol import ConnectorSpecification
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import AuthFlow
 
 
 @dataclass
@@ -21,6 +22,7 @@ class Spec:
     connection_specification: Mapping[str, Any]
     parameters: InitVar[Mapping[str, Any]]
     documentation_url: Optional[str] = None
+    advanced_auth: Optional[AuthFlow] = None
 
     def generate_spec(self) -> ConnectorSpecification:
         """
@@ -31,6 +33,9 @@ def generate_spec(self) -> ConnectorSpecification:
 
         if self.documentation_url:
             obj["documentationUrl"] = self.documentation_url
+        if self.advanced_auth:
+            obj["advanced_auth"] = self.advanced_auth
+            obj["advanced_auth"].auth_flow_type = obj["advanced_auth"].auth_flow_type.value  # Get enum value
 
         # We remap these keys to camel case because that's the existing format expected by the rest of the platform
         return ConnectorSpecification.parse_obj(obj)

airbyte-cdk/python/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -172,6 +172,8 @@ def test_full_config_stream():
         title: API Key
         description: Test API Key
         order: 0
+  advanced_auth:
+    auth_flow_type: "oauth2.0"
     """
     parsed_manifest = YamlDeclarativeSource._parse(content)
     resolved_manifest = resolver.preprocess_manifest(parsed_manifest)
@@ -257,6 +259,8 @@ def test_full_config_stream():
         "description": "Test API Key",
         "order": 0,
     }
+    advanced_auth = spec.advanced_auth
+    assert advanced_auth.auth_flow_type.value == "oauth2.0"
 
 
 def test_interpolate_config():

airbyte-cdk/python/unit_tests/sources/declarative/spec/test_spec.py

@@ -0,0 +1,20 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import pytest
+from airbyte_cdk.models.airbyte_protocol import AdvancedAuth, ConnectorSpecification
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import AuthFlow
+from airbyte_cdk.sources.declarative.spec.spec import Spec
+
+
+@pytest.mark.parametrize(
+    "test_name, spec, expected_connection_specification",
+    [
+        ("test_only_connection_specification", Spec(connection_specification={"client_id": "my_client_id"}, parameters={}), ConnectorSpecification(connectionSpecification={"client_id": "my_client_id"})),
+        ("test_with_doc_url", Spec(connection_specification={"client_id": "my_client_id"}, parameters={}, documentation_url="https://airbyte.io"), ConnectorSpecification(connectionSpecification={"client_id": "my_client_id"}, documentationUrl="https://airbyte.io")),
+        ("test_auth_flow", Spec(connection_specification={"client_id": "my_client_id"}, parameters={}, advanced_auth=AuthFlow(auth_flow_type="oauth2.0")), ConnectorSpecification(connectionSpecification={"client_id": "my_client_id"}, advanced_auth=AdvancedAuth(auth_flow_type="oauth2.0"))),
+    ],
+)
+def test_spec(test_name, spec, expected_connection_specification):
+    assert spec.generate_spec() == expected_connection_specification