[airbyte-ci] Implement pre/post build hooks

airbyte-ci/connectors/pipelines/README.md

@@ -395,6 +395,7 @@ This command runs the Python tests for a airbyte-ci poetry package.
 ## Changelog
 | Version | PR                                                         | Description                                                                                               |
 | ------- | ---------------------------------------------------------- | --------------------------------------------------------------------------------------------------------- |
+| 1.7.0   | [#30526](https://github.com/airbytehq/airbyte/pull/30526)  | Implement pre/post install hooks support.                                                                  |
 | 1.6.0   | [#30474](https://github.com/airbytehq/airbyte/pull/30474)  | Test connector inside their containers.                                                                  |
 | 1.5.1   | [#31227](https://github.com/airbytehq/airbyte/pull/31227)  | Use python 3.11 in amazoncorretto-bazed gradle containers, run 'test' gradle task instead of 'check'.     |
 | 1.5.0   | [#30456](https://github.com/airbytehq/airbyte/pull/30456)  | Start building Python connectors using our base images.                                                   |

airbyte-ci/connectors/pipelines/pipelines/builds/build_customization.py

@@ -0,0 +1,88 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import importlib
+from logging import Logger
+from types import ModuleType
+from typing import List, Optional
+
+from connector_ops.utils import Connector
+from dagger import Container
+
+BUILD_CUSTOMIZATION_MODULE_NAME = "build_customization"
+BUILD_CUSTOMIZATION_SPEC_NAME = f"{BUILD_CUSTOMIZATION_MODULE_NAME}.py"
+DEFAULT_MAIN_FILE_NAME = "main.py"
+
+
+def get_build_customization_module(connector: Connector) -> Optional[ModuleType]:
+    """Import the build_customization.py file from the connector directory if it exists.
+    Returns:
+        Optional[ModuleType]: The build_customization.py module if it exists, None otherwise.
+    """
+    build_customization_spec_path = connector.code_directory / BUILD_CUSTOMIZATION_SPEC_NAME
+    if not build_customization_spec_path.exists():
+        return None
+    build_customization_spec = importlib.util.spec_from_file_location(
+        f"{connector.code_directory.name}_{BUILD_CUSTOMIZATION_MODULE_NAME}", build_customization_spec_path
+    )
+    build_customization_module = importlib.util.module_from_spec(build_customization_spec)
+    build_customization_spec.loader.exec_module(build_customization_module)
+    return build_customization_module
+
+
+def get_main_file_name(connector: Connector) -> str:
+    """Get the main file name from the build_customization.py module if it exists, DEFAULT_MAIN_FILE_NAME otherwise.
+
+    Args:
+        connector (Connector): The connector to build.
+
+    Returns:
+        str: The main file name.
+    """
+    build_customization_module = get_build_customization_module(connector)
+    if hasattr(build_customization_module, "MAIN_FILE_NAME"):
+        return build_customization_module.MAIN_FILE_NAME
+    return DEFAULT_MAIN_FILE_NAME
+
+
+def get_entrypoint(connector: Connector) -> List[str]:
+    main_file_name = get_main_file_name(connector)
+    return ["python", f"/airbyte/integration_code/{main_file_name}"]
+
+
+async def pre_install_hooks(connector: Connector, base_container: Container, logger: Logger) -> Container:
+    """Run the pre_connector_install hook if it exists in the build_customization.py module.
+    It will mutate the base_container and return it.
+
+    Args:
+        connector (Connector): The connector to build.
+        base_container (Container): The base container to mutate.
+        logger (Logger): The logger to use.
+
+    Returns:
+        Container: The mutated base_container.
+    """
+    build_customization_module = get_build_customization_module(connector)
+    if hasattr(build_customization_module, "pre_connector_install"):
+        base_container = await build_customization_module.pre_connector_install(base_container)
+        logger.info(f"Connector {connector.technical_name} pre install hook executed.")
+    return base_container
+
+
+async def post_install_hooks(connector: Connector, connector_container: Container, logger: Logger) -> Container:
+    """Run the post_connector_install hook if it exists in the build_customization.py module.
+    It will mutate the connector_container and return it.
+
+    Args:
+        connector (Connector): The connector to build.
+        connector_container (Container): The connector container to mutate.
+        logger (Logger): The logger to use.
+
+    Returns:
+        Container: The mutated connector_container.
+    """
+    build_customization_module = get_build_customization_module(connector)
+    if hasattr(build_customization_module, "post_connector_install"):
+        connector_container = await build_customization_module.post_connector_install(connector_container)
+        logger.info(f"Connector {connector.technical_name} post install hook executed.")
+    return connector_container

airbyte-ci/connectors/pipelines/pipelines/builds/python_connectors.py

@@ -6,6 +6,7 @@
 from dagger import Container, Platform
 from pipelines.actions.environments import apply_python_development_overrides, with_python_connector_installed
 from pipelines.bases import StepResult
+from pipelines.builds import build_customization
 from pipelines.builds.common import BuildConnectorImagesBase
 from pipelines.contexts import ConnectorContext
 
@@ -16,7 +17,6 @@ class BuildConnectorImages(BuildConnectorImagesBase):
     A spec command is run on the container to validate it was built successfully.
     """
 
-    DEFAULT_ENTRYPOINT = ["python", "/airbyte/integration_code/main.py"]
     PATH_TO_INTEGRATION_CODE = "/airbyte/integration_code"
 
     async def _build_connector(self, platform: Platform):
@@ -35,8 +35,6 @@ def _get_base_container(self, platform: Platform) -> Container:
 
     async def _create_builder_container(self, base_container: Container) -> Container:
         """Pre install the connector dependencies in a builder container.
-        If a python connectors depends on another local python connector, we need to mount its source in the container
-        This occurs for the source-file-secure connector for example, which depends on source-file
 
         Args:
             base_container (Container): The base container to use to build the connector.
@@ -62,27 +60,32 @@ async def _build_from_base_image(self, platform: Platform) -> Container:
         """
         self.logger.info("Building connector from base image in metadata")
         base = self._get_base_container(platform)
-        builder = await self._create_builder_container(base)
+        customized_base = await build_customization.pre_install_hooks(self.context.connector, base, self.logger)
+        entrypoint = build_customization.get_entrypoint(self.context.connector)
+        main_file_name = build_customization.get_main_file_name(self.context.connector)
+
+        builder = await self._create_builder_container(customized_base)
 
         # The snake case name of the connector corresponds to the python package name of the connector
         # We want to mount it to the container under PATH_TO_INTEGRATION_CODE/connector_snake_case_name
         connector_snake_case_name = self.context.connector.technical_name.replace("-", "_")
 
         connector_container = (
             # copy python dependencies from builder to connector container
-            base.with_directory("/usr/local", builder.directory("/usr/local"))
+            customized_base.with_directory("/usr/local", builder.directory("/usr/local"))
             .with_workdir(self.PATH_TO_INTEGRATION_CODE)
-            .with_file("main.py", (await self.context.get_connector_dir(include="main.py")).file("main.py"))
+            .with_file(main_file_name, (await self.context.get_connector_dir(include=main_file_name)).file(main_file_name))
             .with_directory(
                 connector_snake_case_name,
                 (await self.context.get_connector_dir(include=connector_snake_case_name)).directory(connector_snake_case_name),
             )
-            .with_env_variable("AIRBYTE_ENTRYPOINT", " ".join(self.DEFAULT_ENTRYPOINT))
-            .with_entrypoint(self.DEFAULT_ENTRYPOINT)
+            .with_env_variable("AIRBYTE_ENTRYPOINT", " ".join(entrypoint))
+            .with_entrypoint(entrypoint)
             .with_label("io.airbyte.version", self.context.connector.metadata["dockerImageTag"])
             .with_label("io.airbyte.name", self.context.connector.metadata["dockerRepository"])
         )
-        return connector_container
+        customized_connector = await build_customization.post_install_hooks(self.context.connector, connector_container, self.logger)
+        return customized_connector
 
     async def _build_from_dockerfile(self, platform: Platform) -> Container:
         """Build the connector container using its Dockerfile.

airbyte-ci/connectors/pipelines/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "pipelines"
-version = "1.6.0"
+version = "1.7.0"
 description = "Packaged maintained by the connector operations team to perform CI for connectors' pipelines"
 authors = ["Airbyte <contact@airbyte.io>"]
 

airbyte-ci/connectors/pipelines/tests/test_builds/dummy_build_customization.py

@@ -0,0 +1,35 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from dagger import Container
+
+
+async def pre_connector_install(base_image_container: Container) -> Container:
+    """This function will run before the connector installation.
+    It can mutate the base image container.
+
+    Args:
+        base_image_container (Container): The base image container to mutate.
+
+    Returns:
+        Container: The mutated base image container.
+    """
+    return await base_image_container.with_env_variable("MY_PRE_BUILD_ENV_VAR", "my_pre_build_env_var_value")
+
+
+async def post_connector_install(connector_container: Container) -> Container:
+    """This function will run after the connector installation during the build process.
+    It can mutate the connector container.
+
+    Args:
+        connector_container (Container): The connector container to mutate.
+
+    Returns:
+        Container: The mutated connector container.
+    """
+    return await connector_container.with_env_variable("MY_POST_BUILD_ENV_VAR", "my_post_build_env_var_value")

airbyte-ci/connectors/pipelines/tests/test_builds/test_python_connectors.py

@@ -2,9 +2,11 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from pathlib import Path
+
 import pytest
 from pipelines.bases import StepStatus
-from pipelines.builds import python_connectors
+from pipelines.builds import build_customization, python_connectors
 from pipelines.contexts import ConnectorContext
 
 pytestmark = [
@@ -31,9 +33,17 @@ def test_context_with_connector_without_base_image(self, test_context):
     def connector_with_base_image(self, all_connectors):
         for connector in all_connectors:
             if connector.metadata and connector.metadata.get("connectorBuildOptions", {}).get("baseImage"):
-                return connector
+                if not (connector.code_directory / "build_customization.py").exists():
+                    return connector
         pytest.skip("No connector with a connectorBuildOptions.baseImage metadata found")
 
+    @pytest.fixture
+    def connector_with_base_image_with_build_customization(self, connector_with_base_image):
+        dummy_build_customization = (Path(__file__).parent / "dummy_build_customization.py").read_text()
+        (connector_with_base_image.code_directory / "build_customization.py").write_text(dummy_build_customization)
+        yield connector_with_base_image
+        (connector_with_base_image.code_directory / "build_customization.py").unlink()
+
     @pytest.fixture
     def test_context_with_real_connector_using_base_image(self, connector_with_base_image, dagger_client):
         context = ConnectorContext(
@@ -48,6 +58,22 @@ def test_context_with_real_connector_using_base_image(self, connector_with_base_
         context.dagger_client = dagger_client
         return context
 
+    @pytest.fixture
+    def test_context_with_real_connector_using_base_image_with_build_customization(
+        self, connector_with_base_image_with_build_customization, dagger_client
+    ):
+        context = ConnectorContext(
+            pipeline_name="test build",
+            connector=connector_with_base_image_with_build_customization,
+            git_branch="test",
+            git_revision="test",
+            report_output_prefix="test",
+            is_local=True,
+            use_remote_secrets=True,
+        )
+        context.dagger_client = dagger_client
+        return context
+
     @pytest.fixture
     def connector_without_base_image(self, all_connectors):
         for connector in all_connectors:
@@ -87,9 +113,11 @@ async def test_building_from_base_image_for_real(self, test_context_with_real_co
         step_result = await step._run()
         step_result.status is StepStatus.SUCCESS
         built_container = step_result.output_artifact[current_platform]
-        assert await built_container.env_variable("AIRBYTE_ENTRYPOINT") == " ".join(step.DEFAULT_ENTRYPOINT)
+        assert await built_container.env_variable("AIRBYTE_ENTRYPOINT") == " ".join(
+            build_customization.get_entrypoint(step.context.connector)
+        )
         assert await built_container.workdir() == step.PATH_TO_INTEGRATION_CODE
-        assert await built_container.entrypoint() == step.DEFAULT_ENTRYPOINT
+        assert await built_container.entrypoint() == build_customization.get_entrypoint(step.context.connector)
         assert (
             await built_container.label("io.airbyte.version")
             == test_context_with_real_connector_using_base_image.connector.metadata["dockerImageTag"]
@@ -99,6 +127,18 @@ async def test_building_from_base_image_for_real(self, test_context_with_real_co
             == test_context_with_real_connector_using_base_image.connector.metadata["dockerRepository"]
         )
 
+    async def test_building_from_base_image_with_customization_for_real(
+        self, test_context_with_real_connector_using_base_image_with_build_customization, current_platform
+    ):
+        step = python_connectors.BuildConnectorImages(
+            test_context_with_real_connector_using_base_image_with_build_customization, current_platform
+        )
+        step_result = await step._run()
+        step_result.status is StepStatus.SUCCESS
+        built_container = step_result.output_artifact[current_platform]
+        assert await built_container.env_variable("MY_PRE_BUILD_ENV_VAR") == "my_pre_build_env_var_value"
+        assert await built_container.env_variable("MY_POST_BUILD_ENV_VAR") == "my_post_build_env_var_value"
+
     async def test__run_using_base_dockerfile_with_mocks(self, mocker, test_context_with_connector_without_base_image, current_platform):
         container_built_from_dockerfile = mocker.AsyncMock()
         mocker.patch.object(