Secure logger implementation and unit tests

airbyte-cdk/python/airbyte_cdk/secure_logger.py

@@ -9,44 +9,41 @@
 
 from airbyte_cdk.sources import Source
 
-from .logger import LOGGING_CONFIG, TRACE_LEVEL_NUM, AirbyteLogFormatter, AirbyteNativeLogger
+from .logger import LOGGING_CONFIG, TRACE_LEVEL_NUM, AirbyteNativeLogger
 
 
 def init_secure_logger(name: str = None, source: Optional[Source] = None, config: Optional[Mapping] = None):
     """Initial set up of logger"""
     logging.setLoggerClass(AirbyteNativeLogger)
     logging.addLevelName(TRACE_LEVEL_NUM, "TRACE")
     logger = logging.getLogger(name)
+    logger.propagate = True
 
     logger.setLevel(TRACE_LEVEL_NUM)
     logging.config.dictConfig(LOGGING_CONFIG)
-    for handler in logger.handlers:
-        handler.setFormatter(AirbyteSecureLogFormatter(source, config))
-
+    logger.addFilter(AirbyteLogFilter(source, config))
     return logger
 
 
-class AirbyteSecureLogFormatter(AirbyteLogFormatter):
-    def __init__(
-        self,
-        fmt: Optional[str] = ...,
-        datefmt: Optional[str] = ...,
-        style: str = ...,
-        validate: bool = ...,
-        source: Source = None,
-        config: Mapping = None,
-    ) -> None:
-        super().__init__(fmt, datefmt, style, validate)
+class AirbyteLogFilter(logging.Filter):
+    def __init__(self, source: Optional[Source] = None, config: Optional[Mapping] = None):
+        super().__init__()
         if source and config:
             self.secrets_strings = self.get_secrets(source, config)
         else:
             self.secrets_strings = []
 
-    def format(self, record: logging.LogRecord) -> str:
-        reduce(lambda log_msg, secret: log_msg.replace(secret, "****"), self.secrets_strings, record.msg)
-        return super(AirbyteSecureLogFormatter, self).format(record)
+    def filter(self, record: logging.LogRecord) -> str:
+        record.msg = reduce(
+            lambda log_msg, secret: log_msg.replace(secret, "****"),
+            self.secrets_strings,
+            record.msg,
+        )
+        return True
 
     @staticmethod
     def get_secrets(source, config):
-        secret_key_names = [k for k, v in source.spec().connectionSpecification["properties"].items() if v.get("airbyte_secret", False)]
-        return [config.get(k) for k in secret_key_names if config.get(k)]
+        secret_key_names = [
+            k for k, v in source.spec().connectionSpecification.get("properties", {}).items() if v.get("airbyte_secret", False)
+        ]
+        return [str(config.get(k)) for k in secret_key_names if config.get(k)]

airbyte-cdk/python/setup.py

@@ -15,7 +15,7 @@
 
 setup(
     name="airbyte-cdk",
-    version="0.1.36",
+    version="0.1.37",
     description="A framework for writing Airbyte Connectors.",
     long_description=README,
     long_description_content_type="text/markdown",

airbyte-cdk/python/unit_tests/test_entrypoint.py

@@ -124,6 +124,7 @@ def config_mock(mocker, request):
         ({"username": "fake"}, {"type": "object", "properties": {"name": {"type": "string"}}, "additionalProperties": False}, False),
         ({"username": "fake"}, {"type": "object", "properties": {"username": {"type": "string"}}, "additionalProperties": False}, True),
         ({"username": "fake"}, {"type": "object", "properties": {"user": {"type": "string"}}}, True),
+        ({"username": "fake"}, {"type": "object", "properties": {"user": {"type": "string", "airbyte_secret": True}}}, True),
         (
             {"username": "fake", "_limit": 22},
             {"type": "object", "properties": {"username": {"type": "string"}}, "additionalProperties": False},

airbyte-cdk/python/unit_tests/test_secure_logger.py

@@ -0,0 +1,140 @@
+#
+# Copyright (c) 2021 Airbyte, Inc., all rights reserved.
+#
+
+from argparse import Namespace
+from typing import Any, Iterable, Mapping, MutableMapping
+
+import pytest
+from airbyte_cdk import AirbyteEntrypoint, AirbyteLogger
+from airbyte_cdk.models import AirbyteMessage, AirbyteRecordMessage, ConfiguredAirbyteCatalog, ConnectorSpecification, Type
+from airbyte_cdk.sources import Source
+
+SECRET_PROPERTY = "api_token"
+ANOTHER_SECRET_PROPERTY = "another_api_token"
+ANOTHER_NOT_SECRET_PROPERTY = "not_secret_property"
+
+NOT_SECRET_PROPERTY = "explicitly_not_secret_property"
+
+I_AM_A_SECRET_VALUE = "I am a secret"
+ANOTHER_SECRET_VALUE = "Another secret"
+SECRET_INTEGER_VALUE = 123456789
+NOT_A_SECRET_VALUE = "I am not a secret"
+ANOTHER_NOT_SECRET_VALUE = "I am not a secret"
+
+
+class MockSource(Source):
+    def read(
+        self,
+        logger: AirbyteLogger,
+        config: Mapping[str, Any],
+        catalog: ConfiguredAirbyteCatalog,
+        state: MutableMapping[str, Any] = None,
+    ) -> Iterable[AirbyteMessage]:
+        logger.info(I_AM_A_SECRET_VALUE)
+        logger.info(I_AM_A_SECRET_VALUE + " plus Some non secret Value in the same log record" + NOT_A_SECRET_VALUE)
+        logger.info(NOT_A_SECRET_VALUE)
+        yield AirbyteMessage(
+            record=AirbyteRecordMessage(stream="stream", data={"data": "stuff"}, emitted_at=1),
+            type=Type.RECORD,
+        )
+
+    def discover(self, **kwargs):
+        pass
+
+    def check(self, **kwargs):
+        pass
+
+
+spec_with_airbyte_secrets = {
+    "type": "object",
+    "required": ["api_token"],
+    "additionalProperties": False,
+    "properties": {
+        SECRET_PROPERTY: {"type": "string", "airbyte_secret": True},
+        NOT_SECRET_PROPERTY: {"type": "string", "airbyte_secret": False},
+    },
+}
+
+spec_with_airbyte_secrets_config = {
+    SECRET_PROPERTY: I_AM_A_SECRET_VALUE,
+    NOT_SECRET_PROPERTY: NOT_A_SECRET_VALUE,
+}
+
+spec_with_multiple_airbyte_secrets = {
+    "type": "object",
+    "required": ["api_token"],
+    "additionalProperties": True,
+    "properties": {
+        SECRET_PROPERTY: {"type": "string", "airbyte_secret": True},
+        ANOTHER_SECRET_PROPERTY: {"type": "string", "airbyte_secret": True},
+        NOT_SECRET_PROPERTY: {"type": "string", "airbyte_secret": False},
+        ANOTHER_NOT_SECRET_PROPERTY: {"type": "string"},
+    },
+}
+
+spec_with_multiple_airbyte_secrets_config = {
+    SECRET_PROPERTY: I_AM_A_SECRET_VALUE,
+    NOT_SECRET_PROPERTY: NOT_A_SECRET_VALUE,
+    ANOTHER_SECRET_PROPERTY: ANOTHER_SECRET_VALUE,
+    ANOTHER_NOT_SECRET_PROPERTY: ANOTHER_NOT_SECRET_VALUE,
+}
+
+spec_with_airbyte_secrets_not_string = {
+    "type": "object",
+    "required": ["api_token"],
+    "additionalProperties": True,
+    "properties": {
+        SECRET_PROPERTY: {"type": "string", "airbyte_secret": True},
+        ANOTHER_SECRET_PROPERTY: {"type": "integer", "airbyte_secret": True},
+    },
+}
+
+spec_with_airbyte_secrets_not_string_config = {
+    SECRET_PROPERTY: I_AM_A_SECRET_VALUE,
+    ANOTHER_SECRET_PROPERTY: SECRET_INTEGER_VALUE,
+}
+
+
+def simple_config():
+    yield {
+        SECRET_PROPERTY: I_AM_A_SECRET_VALUE,
+        ANOTHER_NOT_SECRET_PROPERTY: ANOTHER_NOT_SECRET_VALUE,
+        ANOTHER_SECRET_PROPERTY: ANOTHER_SECRET_VALUE,
+    }
+
+
+@pytest.mark.parametrize(
+    "source_spec, config",
+    [
+        [spec_with_airbyte_secrets, spec_with_airbyte_secrets_config],
+        [spec_with_multiple_airbyte_secrets, spec_with_multiple_airbyte_secrets_config],
+        [
+            spec_with_airbyte_secrets_not_string,
+            spec_with_airbyte_secrets_not_string_config,
+        ],
+    ],
+    ids=[
+        "spec_with_airbyte_secrets",
+        "spec_with_multiple_airbyte_secrets",
+        "spec_with_airbyte_secrets_not_string",
+    ],
+)
+def test_airbyte_secret_is_masked_on_logger_output(source_spec, mocker, capsys, config):
+    entrypoint = AirbyteEntrypoint(MockSource())
+    parsed_args = Namespace(command="read", config="", state="", catalog="")
+    mocker.patch.object(
+        MockSource,
+        "spec",
+        return_value=ConnectorSpecification(connectionSpecification=source_spec),
+    )
+    mocker.patch.object(MockSource, "configure", return_value=config)
+    mocker.patch.object(MockSource, "read_config", return_value=None)
+    mocker.patch.object(MockSource, "read_state", return_value={})
+    mocker.patch.object(MockSource, "read_catalog", return_value={})
+    list(entrypoint.run(parsed_args))
+    log_result = capsys.readouterr().out + capsys.readouterr().err
+    expected_secret_values = [config[k] for k, v in source_spec["properties"].items() if v.get("airbyte_secret")]
+    expected_plain_text_values = [config[k] for k, v in source_spec["properties"].items() if not v.get("airbyte_secret")]
+    assert all([str(v) not in log_result for v in expected_secret_values])
+    assert all([str(v) in log_result for v in expected_plain_text_values])