Airbyte-ci: Use async click (#31748)

Co-authored-by: Ben Church <ben@airbyte.io>
Co-authored-by: erohmensing <erohmensing@users.noreply.github.com>

airbyte-ci/connectors/pipelines/README.md

@@ -400,6 +400,7 @@ This command runs the Python tests for a airbyte-ci poetry package.
 ## Changelog
 | Version | PR                                                         | Description                                                                                               |
 | ------- | ---------------------------------------------------------- | --------------------------------------------------------------------------------------------------------- |
+| 2.3.1   | [#31748](https://github.com/airbytehq/airbyte/pull/31748)  | Use AsyncClick library instead of base Click.                                                             |
 | 2.3.0   | [#31699](https://github.com/airbytehq/airbyte/pull/31699)  | Support optional concurrent CAT execution.                                                                |
 | 2.2.6   | [#31752](https://github.com/airbytehq/airbyte/pull/31752)  | Only authenticate when secrets are available.
 | 2.2.5   | [#31718](https://github.com/airbytehq/airbyte/pull/31718)  | Authenticate the sidecar docker daemon to DockerHub.                                                      |

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/build_image/commands.py

@@ -2,8 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines.airbyte_ci.connectors.build_image.steps import run_connector_build_pipeline
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
@@ -19,7 +18,7 @@
     type=bool,
 )
 @click.pass_context
-def build(ctx: click.Context, use_host_gradle_dist_tar: bool) -> bool:
+async def build(ctx: click.Context, use_host_gradle_dist_tar: bool) -> bool:
     """Runs a build pipeline for the selected connectors."""
 
     connectors_contexts = [
@@ -47,8 +46,7 @@ def build(ctx: click.Context, use_host_gradle_dist_tar: bool) -> bool:
     ]
     if use_host_gradle_dist_tar and not ctx.obj["is_local"]:
         raise Exception("flag --use-host-gradle-dist-tar requires --is-local")
-    anyio.run(
-        run_connectors_pipelines,
+    await run_connectors_pipelines(
         connectors_contexts,
         run_connector_build_pipeline,
         "Build Pipeline",

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/bump_version/commands.py

@@ -2,8 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines.airbyte_ci.connectors.bump_version.pipeline import run_connector_version_bump_pipeline
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
@@ -15,7 +14,7 @@
 @click.argument("pull-request-number", type=str)
 @click.argument("changelog-entry", type=str)
 @click.pass_context
-def bump_version(
+async def bump_version(
     ctx: click.Context,
     bump_type: str,
     pull_request_number: str,
@@ -47,8 +46,7 @@ def bump_version(
         for connector in ctx.obj["selected_connectors_with_modified_files"]
     ]
 
-    anyio.run(
-        run_connectors_pipelines,
+    await run_connectors_pipelines(
         connectors_contexts,
         run_connector_version_bump_pipeline,
         "Version bump pipeline pipeline",

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/commands.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 from typing import List, Set, Tuple
 
-import click
+import asyncclick as click
 from connector_ops.utils import ConnectorLanguage, SupportLevelEnum, get_all_connectors_in_repo
 from pipelines import main_logger
 from pipelines.cli.lazy_group import LazyGroup
@@ -192,7 +192,7 @@ def validate_environment(is_local: bool, use_remote_secrets: bool):
     envvar="DOCKER_HUB_PASSWORD",
 )
 @click.pass_context
-def connectors(
+async def connectors(
     ctx: click.Context,
     use_remote_secrets: bool,
     names: Tuple[str],

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/list/commands.py

@@ -2,7 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import click
+import asyncclick as click
 from connector_ops.utils import console
 from pipelines.cli.dagger_pipeline_command import DaggerPipelineCommand
 from rich.table import Table
@@ -11,7 +11,7 @@
 
 @click.command(cls=DaggerPipelineCommand, help="List all selected connectors.")
 @click.pass_context
-def list(
+async def list(
     ctx: click.Context,
 ):
     selected_connectors = sorted(ctx.obj["selected_connectors_with_modified_files"], key=lambda x: x.technical_name)

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/migrate_to_base_image/commands.py

@@ -2,8 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
 from pipelines.airbyte_ci.connectors.migrate_to_base_image.pipeline import run_connector_migration_to_base_image_pipeline
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
@@ -17,7 +16,7 @@
 )
 @click.argument("pull-request-number", type=str)
 @click.pass_context
-def migrate_to_base_image(
+async def migrate_to_base_image(
     ctx: click.Context,
     pull_request_number: str,
 ) -> bool:
@@ -51,8 +50,7 @@ def migrate_to_base_image(
         for connector in ctx.obj["selected_connectors_with_modified_files"]
     ]
 
-    anyio.run(
-        run_connectors_pipelines,
+    await run_connectors_pipelines(
         connectors_contexts,
         run_connector_migration_to_base_image_pipeline,
         "Migration to base image pipeline",

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/publish/commands.py

@@ -2,8 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines import main_logger
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
 from pipelines.airbyte_ci.connectors.publish.context import PublishConnectorContext
@@ -57,7 +56,7 @@
     default="#connector-publish-updates",
 )
 @click.pass_context
-def publish(
+async def publish(
     ctx: click.Context,
     pre_release: bool,
     spec_cache_gcs_credentials: str,
@@ -113,8 +112,7 @@ def publish(
     main_logger.warn("Concurrency is forced to 1. For stability reasons we disable parallel publish pipelines.")
     ctx.obj["concurrency"] = 1
 
-    publish_connector_contexts = anyio.run(
-        run_connectors_pipelines,
+    publish_connector_contexts = await run_connectors_pipelines(
         publish_connector_contexts,
         run_connector_publish_pipeline,
         "Publishing connectors",

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/test/commands.py

@@ -4,8 +4,7 @@
 
 import sys
 
-import anyio
-import click
+import asyncclick as click
 from pipelines import main_logger
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
@@ -46,7 +45,7 @@
     is_flag=True,
 )
 @click.pass_context
-def test(
+async def test(
     ctx: click.Context,
     code_tests_only: bool,
     fail_fast: bool,
@@ -100,8 +99,7 @@ def test(
         for connector in ctx.obj["selected_connectors_with_modified_files"]
     ]
     try:
-        anyio.run(
-            run_connectors_pipelines,
+        await run_connectors_pipelines(
             [connector_context for connector_context in connectors_tests_contexts],
             run_connector_test_pipeline,
             "Test Pipeline",

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/upgrade_base_image/commands.py

@@ -2,8 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
 from pipelines.airbyte_ci.connectors.migrate_to_base_image.pipeline import run_connector_base_image_upgrade_pipeline
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
@@ -14,7 +13,7 @@
 @click.command(cls=DaggerPipelineCommand, short_help="Upgrades the base image version used by the selected connectors.")
 @click.option("--set-if-not-exists", default=True)
 @click.pass_context
-def upgrade_base_image(ctx: click.Context, set_if_not_exists: bool, docker_hub_username: str, docker_hub_password: str) -> bool:
+async def upgrade_base_image(ctx: click.Context, set_if_not_exists: bool, docker_hub_username: str, docker_hub_password: str) -> bool:
     """Upgrades the base image version used by the selected connectors."""
 
     fail_if_missing_docker_hub_creds(ctx)
@@ -45,8 +44,7 @@ def upgrade_base_image(ctx: click.Context, set_if_not_exists: bool, docker_hub_u
         for connector in ctx.obj["selected_connectors_with_modified_files"]
     ]
 
-    anyio.run(
-        run_connectors_pipelines,
+    await run_connectors_pipelines(
         connectors_contexts,
         run_connector_base_image_upgrade_pipeline,
         "Upgrade base image pipeline",

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/metadata/commands.py

@@ -2,8 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines.airbyte_ci.metadata.pipeline import run_metadata_orchestrator_deploy_pipeline
 from pipelines.cli.dagger_pipeline_command import DaggerPipelineCommand
 
@@ -24,9 +23,8 @@ def deploy(ctx: click.Context):
 
 @deploy.command(cls=DaggerPipelineCommand, name="orchestrator", help="Deploy the metadata service orchestrator to production")
 @click.pass_context
-def deploy_orchestrator(ctx: click.Context) -> bool:
-    return anyio.run(
-        run_metadata_orchestrator_deploy_pipeline,
+async def deploy_orchestrator(ctx: click.Context) -> bool:
+    await run_metadata_orchestrator_deploy_pipeline(
         ctx.obj["is_local"],
         ctx.obj["git_branch"],
         ctx.obj["git_revision"],

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/test/commands.py

@@ -2,15 +2,14 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import anyio
-import click
+import asyncclick as click
 from pipelines.airbyte_ci.test.pipeline import run_test
 
 
 @click.command()
 @click.argument("poetry_package_path")
 @click.option("--test-directory", default="tests", help="The directory containing the tests to run.")
-def test(
+async def test(
     poetry_package_path: str,
     test_directory: str,
 ):
@@ -20,6 +19,6 @@ def test(
         poetry_package_path (str): Path to the poetry package to test, relative to airbyte-ci directory.
         test_directory (str): The directory containing the tests to run.
     """
-    success = anyio.run(run_test, poetry_package_path, test_directory)
+    success = await run_test(poetry_package_path, test_directory)
     if not success:
         click.Abort()

airbyte-ci/connectors/pipelines/pipelines/cli/airbyte_ci.py

@@ -10,7 +10,7 @@
 from pathlib import Path
 from typing import List
 
-import click
+import asyncclick as click
 import git
 from github import PullRequest
 from pipelines import main_logger
@@ -112,7 +112,7 @@ def set_working_directory_to_root() -> None:
     os.chdir(working_dir)
 
 
-def get_modified_files(
+async def get_modified_files(
     git_branch: str, git_revision: str, diffed_branch: str, is_local: bool, ci_context: CIContext, pull_request: PullRequest
 ) -> List[str]:
     """Get the list of modified files in the current git branch.
@@ -125,15 +125,15 @@ def get_modified_files(
     This latest case is the one we encounter when running the pipeline locally, on a local branch, or manually on GHA with a workflow dispatch event.
     """
     if ci_context is CIContext.MASTER or ci_context is CIContext.NIGHTLY_BUILDS:
-        return get_modified_files_in_commit(git_branch, git_revision, is_local)
+        return await get_modified_files_in_commit(git_branch, git_revision, is_local)
     if ci_context is CIContext.PULL_REQUEST and pull_request is not None:
         return get_modified_files_in_pull_request(pull_request)
     if ci_context is CIContext.MANUAL:
         if git_branch == "master":
-            return get_modified_files_in_commit(git_branch, git_revision, is_local)
+            return await get_modified_files_in_commit(git_branch, git_revision, is_local)
         else:
-            return get_modified_files_in_branch(git_branch, git_revision, diffed_branch, is_local)
-    return get_modified_files_in_branch(git_branch, git_revision, diffed_branch, is_local)
+            return await get_modified_files_in_branch(git_branch, git_revision, diffed_branch, is_local)
+    return await get_modified_files_in_branch(git_branch, git_revision, diffed_branch, is_local)
 
 
 # COMMANDS
@@ -178,7 +178,7 @@ def get_modified_files(
 @click.option("--show-dagger-logs/--hide-dagger-logs", default=False, type=bool)
 @click.pass_context
 @track_command
-def airbyte_ci(
+async def airbyte_ci(
     ctx: click.Context,
     is_local: bool,
     git_branch: str,
@@ -224,7 +224,7 @@ def airbyte_ci(
         ctx.obj["pull_request"] = None
 
     ctx.obj["modified_files"] = transform_strs_to_paths(
-        get_modified_files(git_branch, git_revision, diffed_branch, is_local, ci_context, ctx.obj["pull_request"])
+        await get_modified_files(git_branch, git_revision, diffed_branch, is_local, ci_context, ctx.obj["pull_request"])
     )
 
     if not is_local:

airbyte-ci/connectors/pipelines/pipelines/cli/dagger_pipeline_command.py

@@ -10,7 +10,7 @@
 from pathlib import Path
 from typing import Any
 
-import click
+import asyncclick as click
 from dagger import DaggerError
 from pipelines import consts, main_logger
 from pipelines.consts import GCS_PUBLIC_DOMAIN, STATIC_REPORT_PREFIX
@@ -21,7 +21,7 @@
 
 class DaggerPipelineCommand(click.Command):
     @sentry_utils.with_command_context
-    def invoke(self, ctx: click.Context) -> Any:
+    async def invoke(self, ctx: click.Context) -> Any:
         """Wrap parent invoke in a try catch suited to handle pipeline failures.
         Args:
             ctx (click.Context): The invocation context.
@@ -49,7 +49,7 @@ def invoke(self, ctx: click.Context) -> Any:
                     ctx.obj["dagger_logs_url"] = None
             else:
                 ctx.obj["dagger_logs_path"] = None
-            pipeline_success = super().invoke(ctx)
+            pipeline_success = await super().invoke(ctx)
             if not pipeline_success:
                 raise DaggerError(f"Dagger Command {command_name} failed.")
         except DaggerError as e:

airbyte-ci/connectors/pipelines/pipelines/cli/lazy_group.py

@@ -3,7 +3,7 @@
 import importlib
 from typing import Dict, List, Optional
 
-import click
+import asyncclick as click
 
 
 class LazyGroup(click.Group):
@@ -39,5 +39,6 @@ def _lazy_load(self, cmd_name: str) -> click.BaseCommand:
         cmd_object = getattr(mod, cmd_object_name)
         # check the result to make debugging easier
         if not isinstance(cmd_object, click.BaseCommand):
+            print(f"{cmd_object} is of instance {type(cmd_object)}")
             raise ValueError(f"Lazy loading of {import_path} failed by returning " "a non-command object")
         return cmd_object