airbyte-ci: Add upgrade cdk command

airbyte-ci/connectors/pipelines/README.md

@@ -109,6 +109,7 @@ At this point you can run `airbyte-ci` commands.
 - [Examples](#examples)
 - [Options](#options-2)
 - [`connectors bump_version` command](#connectors-bump_version)
+- [`connectors upgrade_cdk` command](#connectors-upgrade_cdk)
 - [`connectors upgrade_base_image` command](#connectors-upgrade_base_image)
 - [`connectors migrate_to_base_image` command](#connectors-migrate_to_base_image)
 - [`format` command subgroup](#format-subgroup)
@@ -390,6 +391,20 @@ Bump source-openweather: `airbyte-ci connectors --name=source-openweather bump_v
 | `PULL_REQUEST_NUMBER` | The GitHub pull request number, used in the changelog entry            |
 | `CHANGELOG_ENTRY`     | The changelog entry that will get added to the connector documentation |
 
+### <a id="connectors-upgrade_cdk"></a>`connectors upgrade_cdk` command
+
+Upgrade the CDK version of the selected connectors by updating the dependency in the setup.py file.
+
+### Examples
+
+Upgrade for source-openweather: `airbyte-ci connectors --name=source-openweather upgrade_cdk <new-cdk-version>`
+
+#### Arguments
+
+| Argument              | Description                                                            |
+| --------------------- | ---------------------------------------------------------------------- |
+| `CDK_VERSION`         | CDK version to set (default to the most recent version)                |
+
 ### <a id="connectors-upgrade_base_image"></a>`connectors upgrade_base_image` command
 
 Modify the selected connector metadata to use the latest base image version.

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

@@ -153,6 +153,7 @@ def should_use_remote_secrets(use_remote_secrets: Optional[bool]) -> bool:
         "bump_version": "pipelines.airbyte_ci.connectors.bump_version.commands.bump_version",
         "migrate_to_base_image": "pipelines.airbyte_ci.connectors.migrate_to_base_image.commands.migrate_to_base_image",
         "upgrade_base_image": "pipelines.airbyte_ci.connectors.upgrade_base_image.commands.upgrade_base_image",
+        "upgrade_cdk": "pipelines.airbyte_ci.connectors.upgrade_cdk.commands.bump_version",
     },
 )
 @click.option(

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/upgrade_cdk/__init__.py

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#

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

@@ -0,0 +1,67 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import asyncclick as click
+import requests
+from pipelines.airbyte_ci.connectors.context import ConnectorContext
+from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
+from pipelines.airbyte_ci.connectors.upgrade_cdk.pipeline import run_connector_cdk_upgrade_pipeline
+from pipelines.cli.dagger_pipeline_command import DaggerPipelineCommand
+
+
+def latest_cdk_version():
+    """
+    Get the latest version of airbyte-cdk from pypi
+    """
+    cdk_pypi_url = "https://pypi.org/pypi/airbyte-cdk/json"
+    response = requests.get(cdk_pypi_url)
+    response.raise_for_status()
+    package_info = response.json()
+    return package_info["info"]["version"]
+
+
+@click.command(cls=DaggerPipelineCommand, short_help="Upgrade CDK version")
+@click.argument("target-cdk-version", type=str, default=latest_cdk_version)
+@click.pass_context
+async def bump_version(
+    ctx: click.Context,
+    target_cdk_version: str,
+) -> bool:
+    """Upgrade CDK version"""
+
+    connectors_contexts = [
+        ConnectorContext(
+            pipeline_name=f"Upgrade CDK version of connector {connector.technical_name}",
+            connector=connector,
+            is_local=ctx.obj["is_local"],
+            git_branch=ctx.obj["git_branch"],
+            git_revision=ctx.obj["git_revision"],
+            ci_report_bucket=ctx.obj["ci_report_bucket_name"],
+            report_output_prefix=ctx.obj["report_output_prefix"],
+            use_remote_secrets=ctx.obj["use_remote_secrets"],
+            gha_workflow_run_url=ctx.obj.get("gha_workflow_run_url"),
+            dagger_logs_url=ctx.obj.get("dagger_logs_url"),
+            pipeline_start_timestamp=ctx.obj.get("pipeline_start_timestamp"),
+            ci_context=ctx.obj.get("ci_context"),
+            ci_gcs_credentials=ctx.obj["ci_gcs_credentials"],
+            ci_git_user=ctx.obj["ci_git_user"],
+            ci_github_access_token=ctx.obj["ci_github_access_token"],
+            enable_report_auto_open=False,
+            s3_build_cache_access_key_id=ctx.obj.get("s3_build_cache_access_key_id"),
+            s3_build_cache_secret_key=ctx.obj.get("s3_build_cache_secret_key"),
+        )
+        for connector in ctx.obj["selected_connectors_with_modified_files"]
+    ]
+
+    await run_connectors_pipelines(
+        connectors_contexts,
+        run_connector_cdk_upgrade_pipeline,
+        "Upgrade CDK version pipeline",
+        ctx.obj["concurrency"],
+        ctx.obj["dagger_logs_path"],
+        ctx.obj["execute_timeout"],
+        target_cdk_version,
+    )
+
+    return True

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/upgrade_cdk/pipeline.py

@@ -0,0 +1,91 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import os
+import re
+
+from pipelines.airbyte_ci.connectors.context import ConnectorContext
+from pipelines.airbyte_ci.connectors.reports import ConnectorReport
+from pipelines.helpers import git
+from pipelines.models.steps import Step, StepResult, StepStatus
+
+
+class SetCDKVersion(Step):
+    title = "Set CDK Version"
+
+    def __init__(
+        self,
+        context: ConnectorContext,
+        new_version: str,
+    ):
+        super().__init__(context)
+        self.new_version = new_version
+
+    async def _run(self) -> StepResult:
+        context: ConnectorContext = self.context
+        og_connector_dir = await context.get_connector_dir()
+        if not "setup.py" in await og_connector_dir.entries():
+            return StepResult(
+                self,
+                StepStatus.FAILURE,
+                stdout="Connector does not have a setup.py file.",
+            )
+        setup_py = og_connector_dir.file("setup.py")
+        setup_py_content = await setup_py.contents()
+        try:
+            updated_setup_py = self.update_cdk_version(setup_py_content)
+            updated_connector_dir = og_connector_dir.with_new_file("setup.py", updated_setup_py)
+            diff = og_connector_dir.diff(updated_connector_dir)
+            exported_successfully = await diff.export(os.path.join(git.get_git_repo_path(), context.connector.code_directory))
+            if not exported_successfully:
+                return StepResult(
+                    self,
+                    StepStatus.FAILURE,
+                    stdout="Could not export diff to local git repo.",
+                )
+            return StepResult(self, StepStatus.SUCCESS, stdout=f"Updated CDK version to {self.new_version}", output_artifact=diff)
+        except ValueError as e:
+            return StepResult(
+                self,
+                StepStatus.FAILURE,
+                stderr=f"Could not set CDK version: {e}",
+                exc_info=e,
+            )
+
+    def update_cdk_version(self, og_setup_py_content: str) -> str:
+        airbyte_cdk_dependency = re.search(
+            r"airbyte-cdk(?P<extra>\[[a-zA-Z0-9-]*\])?(?P<version>[<>=!~]+[0-9]*\.[0-9]*\.[0-9]*)?", og_setup_py_content
+        )
+        # If there is no airbyte-cdk dependency, add the version
+        if airbyte_cdk_dependency is not None:
+            new_version = f"airbyte-cdk{airbyte_cdk_dependency.group('extra') or ''}>={self.new_version}"
+            return og_setup_py_content.replace(airbyte_cdk_dependency.group(), new_version)
+        else:
+            raise ValueError("Could not find airbyte-cdk dependency in setup.py")
+
+
+async def run_connector_cdk_upgrade_pipeline(
+    context: ConnectorContext,
+    semaphore,
+    target_version: str,
+) -> ConnectorReport:
+    """Run a pipeline to upgrade the CDK version for a single connector.
+
+    Args:
+        context (ConnectorContext): The initialized connector context.
+
+    Returns:
+        ConnectorReport: The reports holding the CDK version set results.
+    """
+    async with semaphore:
+        steps_results = []
+        async with context:
+            set_cdk_version = SetCDKVersion(
+                context,
+                target_version,
+            )
+            set_cdk_version_result = await set_cdk_version.run()
+            steps_results.append(set_cdk_version_result)
+            context.report = ConnectorReport(context, steps_results, name="CONNECTOR VERSION CDK UPGRADE RESULTS")
+    return context.report

airbyte-ci/connectors/pipelines/tests/test_upgrade_cdk.py

@@ -0,0 +1,121 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import json
+import random
+from pathlib import Path
+from typing import List
+from unittest.mock import AsyncMock, MagicMock
+
+import anyio
+import pytest
+from connector_ops.utils import Connector, ConnectorLanguage
+from dagger import Directory
+from pipelines.airbyte_ci.connectors.context import ConnectorContext
+from pipelines.airbyte_ci.connectors.publish import pipeline as publish_pipeline
+from pipelines.airbyte_ci.connectors.upgrade_cdk import pipeline as upgrade_cdk_pipeline
+from pipelines.models.steps import StepStatus
+
+pytestmark = [
+    pytest.mark.anyio,
+]
+
+
+@pytest.fixture
+def sample_connector():
+    return Connector("source-pokeapi")
+
+
+def get_sample_setup_py(airbyte_cdk_dependency: str):
+    return f"""from setuptools import find_packages, setup
+
+MAIN_REQUIREMENTS = [
+    "{airbyte_cdk_dependency}",
+]
+
+setup(
+    name="source_pokeapi",
+    description="Source implementation for Pokeapi.",
+    author="Airbyte",
+    author_email="contact@airbyte.io",
+    packages=find_packages(),
+    install_requires=MAIN_REQUIREMENTS,
+)
+"""
+
+
+@pytest.fixture
+def connector_context(sample_connector, dagger_client, current_platform):
+    context = ConnectorContext(
+        pipeline_name="test",
+        connector=sample_connector,
+        git_branch="test",
+        git_revision="test",
+        report_output_prefix="test",
+        is_local=True,
+        use_remote_secrets=True,
+        targeted_platforms=[current_platform],
+    )
+    context.dagger_client = dagger_client
+    return context
+
+
+@pytest.mark.parametrize(
+    "setup_py_content, expected_setup_py_content",
+    [
+        (get_sample_setup_py("airbyte-cdk"), get_sample_setup_py("airbyte-cdk>=6.6.6")),
+        (get_sample_setup_py("airbyte-cdk[file-based]"), get_sample_setup_py("airbyte-cdk[file-based]>=6.6.6")),
+        (get_sample_setup_py("airbyte-cdk==1.2.3"), get_sample_setup_py("airbyte-cdk>=6.6.6")),
+        (get_sample_setup_py("airbyte-cdk>=1.2.3"), get_sample_setup_py("airbyte-cdk>=6.6.6")),
+        (get_sample_setup_py("airbyte-cdk[file-based]>=1.2.3"), get_sample_setup_py("airbyte-cdk[file-based]>=6.6.6")),
+    ],
+)
+async def test_run_connector_cdk_upgrade_pipeline(
+    connector_context: ConnectorContext, setup_py_content: str, expected_setup_py_content: str
+):
+    full_og_connector_dir = await connector_context.get_connector_dir()
+    updated_connector_dir = full_og_connector_dir.with_new_file("setup.py", setup_py_content)
+
+    # For this test, replace the actual connector dir with an updated version that sets the setup.py contents
+    connector_context.get_connector_dir = AsyncMock(return_value=updated_connector_dir)
+
+    # Mock the diff method to record the resulting directory and return a mock to not actually export the diff to the repo
+    updated_connector_dir.diff = MagicMock(return_value=AsyncMock())
+    step = upgrade_cdk_pipeline.SetCDKVersion(connector_context, "6.6.6")
+    step_result = await step.run()
+    assert step_result.status == StepStatus.SUCCESS
+
+    # Check that the resulting directory that got passed to the mocked diff method looks as expected
+    resulting_directory: Directory = await full_og_connector_dir.diff(updated_connector_dir.diff.call_args[0][0])
+    files = await resulting_directory.entries()
+    # validate only setup.py is changed
+    assert files == ["setup.py"]
+    setup_py = resulting_directory.file("setup.py")
+    actual_setup_py_content = await setup_py.contents()
+    assert expected_setup_py_content == actual_setup_py_content
+
+    # Assert that the diff was exported to the repo
+    assert updated_connector_dir.diff.return_value.export.call_count == 1
+
+
+async def test_fail_connector_cdk_upgrade_pipeline_on_missing_setup_py(connector_context: ConnectorContext):
+    full_og_connector_dir = await connector_context.get_connector_dir()
+    updated_connector_dir = full_og_connector_dir.without_file("setup.py")
+
+    connector_context.get_connector_dir = AsyncMock(return_value=updated_connector_dir)
+
+    step = upgrade_cdk_pipeline.SetCDKVersion(connector_context, "6.6.6")
+    step_result = await step.run()
+    assert step_result.status == StepStatus.FAILURE
+
+
+async def test_fail_connector_cdk_upgrade_pipeline_on_missing_airbyte_cdk(connector_context: ConnectorContext):
+    full_og_connector_dir = await connector_context.get_connector_dir()
+    updated_connector_dir = full_og_connector_dir.with_new_file("setup.py", get_sample_setup_py("another-lib==1.2.3"))
+
+    connector_context.get_connector_dir = AsyncMock(return_value=updated_connector_dir)
+
+    step = upgrade_cdk_pipeline.SetCDKVersion(connector_context, "6.6.6")
+    step_result = await step.run()
+    assert step_result.status == StepStatus.FAILURE