241 lines
8.6 KiB
Python
241 lines
8.6 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
|
|
import logging
|
|
import sys
|
|
from argparse import Namespace
|
|
from typing import Any, Iterable, Mapping, MutableMapping
|
|
|
|
import pytest
|
|
from airbyte_cdk import AirbyteEntrypoint
|
|
from airbyte_cdk.logger import AirbyteLogFormatter
|
|
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: logging.Logger,
|
|
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,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_config():
|
|
yield {
|
|
SECRET_PROPERTY: I_AM_A_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, config, caplog):
|
|
caplog.set_level(logging.DEBUG, logger="airbyte.test")
|
|
caplog.handler.setFormatter(AirbyteLogFormatter())
|
|
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 = caplog.text
|
|
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])
|
|
|
|
|
|
def test_airbyte_secrets_are_masked_on_uncaught_exceptions(mocker, caplog, capsys):
|
|
caplog.set_level(logging.DEBUG, logger="airbyte.test")
|
|
caplog.handler.setFormatter(AirbyteLogFormatter())
|
|
|
|
class BrokenSource(MockSource):
|
|
def read(
|
|
self,
|
|
logger: logging.Logger,
|
|
config: Mapping[str, Any],
|
|
catalog: ConfiguredAirbyteCatalog,
|
|
state: MutableMapping[str, Any] = None,
|
|
):
|
|
raise Exception("Exception:" + I_AM_A_SECRET_VALUE)
|
|
|
|
entrypoint = AirbyteEntrypoint(BrokenSource())
|
|
parsed_args = Namespace(command="read", config="", state="", catalog="")
|
|
source_spec = {
|
|
"type": "object",
|
|
"required": ["api_token"],
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
SECRET_PROPERTY: {"type": "string", "airbyte_secret": True},
|
|
NOT_SECRET_PROPERTY: {"type": "string", "airbyte_secret": False},
|
|
},
|
|
}
|
|
simple_config = {
|
|
SECRET_PROPERTY: I_AM_A_SECRET_VALUE,
|
|
NOT_SECRET_PROPERTY: NOT_A_SECRET_VALUE,
|
|
}
|
|
mocker.patch.object(
|
|
MockSource,
|
|
"spec",
|
|
return_value=ConnectorSpecification(connectionSpecification=source_spec),
|
|
)
|
|
mocker.patch.object(MockSource, "configure", return_value=simple_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={})
|
|
|
|
try:
|
|
list(entrypoint.run(parsed_args))
|
|
except Exception:
|
|
sys.excepthook(*sys.exc_info())
|
|
assert I_AM_A_SECRET_VALUE not in capsys.readouterr().out, "Should have filtered non-secret value from exception trace message"
|
|
assert I_AM_A_SECRET_VALUE not in caplog.text, "Should have filtered secret value from exception log message"
|
|
|
|
|
|
def test_non_airbyte_secrets_are_not_masked_on_uncaught_exceptions(mocker, caplog, capsys):
|
|
caplog.set_level(logging.DEBUG, logger="airbyte.test")
|
|
caplog.handler.setFormatter(AirbyteLogFormatter())
|
|
|
|
class BrokenSource(MockSource):
|
|
def read(
|
|
self,
|
|
logger: logging.Logger,
|
|
config: Mapping[str, Any],
|
|
catalog: ConfiguredAirbyteCatalog,
|
|
state: MutableMapping[str, Any] = None,
|
|
):
|
|
raise Exception("Exception:" + NOT_A_SECRET_VALUE)
|
|
|
|
entrypoint = AirbyteEntrypoint(BrokenSource())
|
|
parsed_args = Namespace(command="read", config="", state="", catalog="")
|
|
source_spec = {
|
|
"type": "object",
|
|
"required": ["api_token"],
|
|
"additionalProperties": False,
|
|
"properties": {
|
|
SECRET_PROPERTY: {"type": "string", "airbyte_secret": True},
|
|
NOT_SECRET_PROPERTY: {"type": "string", "airbyte_secret": False},
|
|
},
|
|
}
|
|
simple_config = {
|
|
SECRET_PROPERTY: I_AM_A_SECRET_VALUE,
|
|
NOT_SECRET_PROPERTY: NOT_A_SECRET_VALUE,
|
|
}
|
|
mocker.patch.object(
|
|
MockSource,
|
|
"spec",
|
|
return_value=ConnectorSpecification(connectionSpecification=source_spec),
|
|
)
|
|
mocker.patch.object(MockSource, "configure", return_value=simple_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={})
|
|
mocker.patch.object(MockSource, "read", side_effect=Exception("Exception:" + NOT_A_SECRET_VALUE))
|
|
|
|
try:
|
|
list(entrypoint.run(parsed_args))
|
|
except Exception:
|
|
sys.excepthook(*sys.exc_info())
|
|
assert NOT_A_SECRET_VALUE in capsys.readouterr().out, "Should not have filtered non-secret value from exception trace message"
|
|
assert NOT_A_SECRET_VALUE in caplog.text, "Should not have filtered non-secret value from exception log message"
|