add the ability to upgrade CDK for java connectors

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

@@ -3,26 +3,14 @@
 #
 
 import asyncclick as click
-import requests  # type: ignore
 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() -> str:
-    """
-    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.argument("target-cdk-version", type=str, default="latest")
 @click.pass_context
 async def bump_version(
     ctx: click.Context,

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

@@ -7,6 +7,8 @@
 import re
 from typing import TYPE_CHECKING
 
+import requests  # type: ignore
+from connector_ops.utils import ConnectorLanguage
 from pipelines.airbyte_ci.connectors.context import ConnectorContext
 from pipelines.airbyte_ci.connectors.reports import ConnectorReport, Report
 from pipelines.helpers import git
@@ -30,14 +32,22 @@ def __init__(
 
     async def _run(self) -> StepResult:
         context = self.context
-        og_connector_dir = await context.get_connector_dir()
-        if "setup.py" not in await og_connector_dir.entries():
-            return self.skip("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)
+            og_connector_dir = await context.get_connector_dir()
+            if self.context.connector.language == ConnectorLanguage.PYTHON or self.context.connector.language == ConnectorLanguage.LOW_CODE:
+                updated_connector_dir = await self.upgrade_cdk_version_for_python_connector(og_connector_dir)
+            elif self.context.connector.language == ConnectorLanguage.JAVA:
+                updated_connector_dir = await self.upgrade_cdk_version_for_java_connector(og_connector_dir)
+            else:
+                return StepResult(
+                    self,
+                    StepStatus.FAILURE,
+                    f"Can't upgrade the CDK for connector {self.context.connector.technical_name} of written in {self.context.connector.language}",
+                )
+
+            if updated_connector_dir is None:
+                return self.skip(self.skip_reason)
             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:
@@ -55,17 +65,69 @@ async def _run(self) -> StepResult:
                 exc_info=e,
             )
 
-    def update_cdk_version(self, og_setup_py_content: str) -> str:
+    async def upgrade_cdk_version_for_java_connector(self, og_connector_dir: Directory) -> Directory:
+        context = self.context
+
+        if "build.gradle" not in await og_connector_dir.entries():
+            raise ValueError(f"Java connector {self.context.connector.technical_name} does not have a build.gradle file.")
+
+        build_gradle = og_connector_dir.file("build.gradle")
+        build_gradle_content = await build_gradle.contents()
+
+        old_cdk_version_required = re.search(r"cdkVersionRequired *= *'(?P<version>[0-9]*\.[0-9]*\.[0-9]*)?'", build_gradle_content)
+        # If there is no airbyte-cdk dependency, add the version
+        if old_cdk_version_required is None:
+            raise ValueError("Could not find airbyte-cdk dependency in build.gradle")
+
+        if self.new_version == "latest":
+            version_file_content = await self.context.get_repo_file(
+                "airbyte-cdk/java/airbyte-cdk/core/src/main/resources/version.properties"
+            ).contents()
+            new_version = re.search(r"version *= *(?P<version>[0-9]*\.[0-9]*\.[0-9]*)", version_file_content).group("version")
+        else:
+            new_version = self.new_version
+
+        updated_build_gradle = build_gradle_content.replace(old_cdk_version_required.group("version"), new_version)
+
+        use_local_cdk = re.search(r"useLocalCdk *=.*", updated_build_gradle)
+        if use_local_cdk is not None:
+            updated_build_gradle = updated_build_gradle.replace(use_local_cdk.group(), "useLocalCdk = false")
+
+        return og_connector_dir.with_new_file("build.gradle", updated_build_gradle)
+
+    async def upgrade_cdk_version_for_python_connector(self, og_connector_dir: Directory) -> Directory:
+        context = self.context
+        og_connector_dir = await context.get_connector_dir()
+        if "setup.py" not in await og_connector_dir.entries():
+            self.skip_reason = f"Python connector {self.context.connector.technical_name} does not have a setup.py file."
+            return None
+        setup_py = og_connector_dir.file("setup.py")
+        setup_py_content = await setup_py.contents()
+
         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
+            r"airbyte-cdk(?P<extra>\[[a-zA-Z0-9-]*\])?(?P<version>[<>=!~]+[0-9]*(?:\.[0-9]*)?(?:\.[0-9]*)?)?", 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:
+        if airbyte_cdk_dependency is None:
             raise ValueError("Could not find airbyte-cdk dependency in setup.py")
 
+        if self.new_version == "latest":
+            """
+            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()
+            new_version = package_info["info"]["version"]
+        else:
+            new_version = self.new_version
+
+        new_version_str = f"airbyte-cdk{airbyte_cdk_dependency.group('extra') or ''}>={new_version}"
+        updated_setup_py = setup_py_content.replace(airbyte_cdk_dependency.group(), new_version_str)
+
+        return og_connector_dir.with_new_file("setup.py", updated_setup_py)
+
 
 async def run_connector_cdk_upgrade_pipeline(
     context: ConnectorContext,

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

@@ -0,0 +1,133 @@
+#
+# 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-postgres")
+
+
+def get_sample_build_gradle(airbyte_cdk_version: str, useLocalCdk: str):
+    return f"""import org.jsonschema2pojo.SourceType
+
+plugins {{
+    id 'application'
+    id 'airbyte-java-connector'
+    id "org.jsonschema2pojo" version "1.2.1"
+}}
+
+java {{
+    compileJava {{
+        options.compilerArgs += "-Xlint:-try,-rawtypes,-unchecked"
+    }}
+}}
+
+airbyteJavaConnector {{
+    cdkVersionRequired = '{airbyte_cdk_version}'
+    features = ['db-sources']
+    useLocalCdk = {useLocalCdk}
+}}
+
+
+application {{
+    mainClass = 'io.airbyte.integrations.source.postgres.PostgresSource'
+    applicationDefaultJvmArgs = ['-XX:+ExitOnOutOfMemoryError', '-XX:MaxRAMPercentage=75.0']
+}}
+"""
+
+
+@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(
+    "build_gradle_content, expected_build_gradle_content",
+    [
+        (get_sample_build_gradle("1.2.3", "false"), get_sample_build_gradle("6.6.6", "false")),
+        (get_sample_build_gradle("1.2.3", "true"), get_sample_build_gradle("6.6.6", "false")),
+        (get_sample_build_gradle("6.6.6", "false"), get_sample_build_gradle("6.6.6", "false")),
+        (get_sample_build_gradle("6.6.6", "true"), get_sample_build_gradle("6.6.6", "false")),
+        (get_sample_build_gradle("7.0.0", "false"), get_sample_build_gradle("6.6.6", "false")),
+        (get_sample_build_gradle("7.0.0", "true"), get_sample_build_gradle("6.6.6", "false")),
+    ],
+)
+async def test_run_connector_cdk_upgrade_pipeline(
+    connector_context: ConnectorContext, build_gradle_content: str, expected_build_gradle_content: str
+):
+    full_og_connector_dir = await connector_context.get_connector_dir()
+    updated_connector_dir = full_og_connector_dir.with_new_file("build.gradle", build_gradle_content)
+
+    # For this test, replace the actual connector dir with an updated version that sets the build.gradle 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 build.gradle is changed
+    assert files == ["build.gradle"]
+    build_gradle = resulting_directory.file("build.gradle")
+    actual_build_gradle_content = await build_gradle.contents()
+    assert expected_build_gradle_content == actual_build_gradle_content
+
+    # Assert that the diff was exported to the repo
+    assert updated_connector_dir.diff.return_value.export.call_count == 1
+
+
+async def test_skip_connector_cdk_upgrade_pipeline_on_missing_build_gradle(connector_context: ConnectorContext):
+    full_og_connector_dir = await connector_context.get_connector_dir()
+    updated_connector_dir = full_og_connector_dir.without_file("build.gradle")
+
+    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("build.gradle", get_sample_build_gradle("abc", "false"))
+
+    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

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

@@ -69,6 +69,9 @@ def connector_context(sample_connector, dagger_client, current_platform):
         (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")),
+        (get_sample_setup_py("airbyte-cdk==1.2"), get_sample_setup_py("airbyte-cdk>=6.6.6")),
+        (get_sample_setup_py("airbyte-cdk>=1.2"), get_sample_setup_py("airbyte-cdk>=6.6.6")),
+        (get_sample_setup_py("airbyte-cdk[file-based]>=1.2"), get_sample_setup_py("airbyte-cdk[file-based]>=6.6.6")),
     ],
 )
 async def test_run_connector_cdk_upgrade_pipeline(