From b76b73bbfbf6f0351f142e075dd975018eab7b66 Mon Sep 17 00:00:00 2001 From: Augustin Date: Thu, 21 Jul 2022 16:34:20 +0200 Subject: [PATCH] cdk: do not call `init_uncaught_exception_handler` from modules' root (#14892) --- airbyte-cdk/python/CHANGELOG.md | 3 ++- .../python/airbyte_cdk/destinations/destination.py | 3 +++ airbyte-cdk/python/airbyte_cdk/entrypoint.py | 2 +- airbyte-cdk/python/setup.py | 2 +- .../python/unit_tests/destinations/test_destination.py | 8 ++++++++ airbyte-cdk/python/unit_tests/test_entrypoint.py | 7 +++++++ 6 files changed, 22 insertions(+), 3 deletions(-) diff --git a/airbyte-cdk/python/CHANGELOG.md b/airbyte-cdk/python/CHANGELOG.md index 3cda55145ac..daf785e5dd5 100644 --- a/airbyte-cdk/python/CHANGELOG.md +++ b/airbyte-cdk/python/CHANGELOG.md @@ -1,6 +1,7 @@ # Changelog -## Unreleased +## 0.1.66 +- Call init_uncaught_exception_handler from AirbyteEntrypoint.__init__ and Destination.run_cmd - Add the ability to remove & add records in YAML-based sources ## 0.1.65 diff --git a/airbyte-cdk/python/airbyte_cdk/destinations/destination.py b/airbyte-cdk/python/airbyte_cdk/destinations/destination.py index 9c39a573ca5..97bb381113b 100644 --- a/airbyte-cdk/python/airbyte_cdk/destinations/destination.py +++ b/airbyte-cdk/python/airbyte_cdk/destinations/destination.py @@ -10,6 +10,7 @@ from abc import ABC, abstractmethod from typing import Any, Iterable, List, Mapping from airbyte_cdk.connector import Connector +from airbyte_cdk.exception_handler import init_uncaught_exception_handler from airbyte_cdk.models import AirbyteMessage, ConfiguredAirbyteCatalog, Type from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit from pydantic import ValidationError @@ -83,6 +84,7 @@ class Destination(Connector, ABC): return parsed_args def run_cmd(self, parsed_args: argparse.Namespace) -> Iterable[AirbyteMessage]: + cmd = parsed_args.command if cmd not in self.VALID_CMDS: raise Exception(f"Unrecognized command: {cmd}") @@ -103,6 +105,7 @@ class Destination(Connector, ABC): yield from self._run_write(config=config, configured_catalog_path=parsed_args.catalog, input_stream=wrapped_stdin) def run(self, args: List[str]): + init_uncaught_exception_handler(logger) parsed_args = self.parse_args(args) output_messages = self.run_cmd(parsed_args) for message in output_messages: diff --git a/airbyte-cdk/python/airbyte_cdk/entrypoint.py b/airbyte-cdk/python/airbyte_cdk/entrypoint.py index ad7ade90a93..7d57dc7e9e2 100644 --- a/airbyte-cdk/python/airbyte_cdk/entrypoint.py +++ b/airbyte-cdk/python/airbyte_cdk/entrypoint.py @@ -20,11 +20,11 @@ from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_o from airbyte_cdk.utils.airbyte_secrets_utils import get_secrets, update_secrets logger = init_logger("airbyte") -init_uncaught_exception_handler(logger) class AirbyteEntrypoint(object): def __init__(self, source: Source): + init_uncaught_exception_handler(logger) self.source = source self.logger = logging.getLogger(f"airbyte.{getattr(source, 'name', '')}") diff --git a/airbyte-cdk/python/setup.py b/airbyte-cdk/python/setup.py index b15db905a95..903006f23c7 100644 --- a/airbyte-cdk/python/setup.py +++ b/airbyte-cdk/python/setup.py @@ -15,7 +15,7 @@ README = (HERE / "README.md").read_text() setup( name="airbyte-cdk", - version="0.1.65", + version="0.1.66", description="A framework for writing Airbyte Connectors.", long_description=README, long_description_content_type="text/markdown", diff --git a/airbyte-cdk/python/unit_tests/destinations/test_destination.py b/airbyte-cdk/python/unit_tests/destinations/test_destination.py index 7f6b900a071..66043b0365b 100644 --- a/airbyte-cdk/python/unit_tests/destinations/test_destination.py +++ b/airbyte-cdk/python/unit_tests/destinations/test_destination.py @@ -11,6 +11,7 @@ from unittest.mock import ANY import pytest from airbyte_cdk.destinations import Destination +from airbyte_cdk.destinations import destination as destination_module from airbyte_cdk.models import ( AirbyteCatalog, AirbyteConnectionStatus, @@ -136,6 +137,13 @@ class OrderedIterableMatcher(Iterable): class TestRun: + def test_run_initializes_exception_handler(self, mocker, destination: Destination): + mocker.patch.object(destination_module, "init_uncaught_exception_handler") + mocker.patch.object(destination, "parse_args") + mocker.patch.object(destination, "run_cmd") + destination.run(["dummy"]) + destination_module.init_uncaught_exception_handler.assert_called_once_with(destination_module.logger) + def test_run_spec(self, mocker, destination: Destination): args = {"command": "spec"} parsed_args = argparse.Namespace(**args) diff --git a/airbyte-cdk/python/unit_tests/test_entrypoint.py b/airbyte-cdk/python/unit_tests/test_entrypoint.py index 168539b5d2c..26bc5b85541 100644 --- a/airbyte-cdk/python/unit_tests/test_entrypoint.py +++ b/airbyte-cdk/python/unit_tests/test_entrypoint.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock import pytest from airbyte_cdk import AirbyteEntrypoint +from airbyte_cdk import entrypoint as entrypoint_module from airbyte_cdk.models import ( AirbyteCatalog, AirbyteConnectionStatus, @@ -56,6 +57,12 @@ def entrypoint() -> AirbyteEntrypoint: return AirbyteEntrypoint(MockSource()) +def test_airbyte_entrypoint_init(mocker): + mocker.patch.object(entrypoint_module, "init_uncaught_exception_handler") + AirbyteEntrypoint(MockSource()) + entrypoint_module.init_uncaught_exception_handler.assert_called_once_with(entrypoint_module.logger) + + @pytest.mark.parametrize( ["cmd", "args", "expected_args"], [