airbyte_cdk.sources.declarative.parsers.custom_code_compiler

  1"""Contains functions to compile custom code from text."""
  2
  3import hashlib
  4import os
  5import sys
  6from collections.abc import Mapping
  7from types import ModuleType
  8from typing import Any, cast
  9
 10from typing_extensions import Literal
 11
 12ChecksumType = Literal["md5", "sha256"]
 13CHECKSUM_FUNCTIONS = {
 14    "md5": hashlib.md5,
 15    "sha256": hashlib.sha256,
 16}
 17COMPONENTS_MODULE_NAME = "components"
 18SDM_COMPONENTS_MODULE_NAME = "source_declarative_manifest.components"
 19INJECTED_MANIFEST = "__injected_declarative_manifest"
 20INJECTED_COMPONENTS_PY = "__injected_components_py"
 21INJECTED_COMPONENTS_PY_CHECKSUMS = "__injected_components_py_checksums"
 22ENV_VAR_ALLOW_CUSTOM_CODE = "AIRBYTE_ENABLE_UNSAFE_CODE"
 23
 24
 25class AirbyteCodeTamperedError(Exception):
 26    """Raised when the connector's components module does not match its checksum.
 27
 28    This is a fatal error, as it can be a sign of code tampering.
 29    """
 30
 31
 32class AirbyteCustomCodeNotPermittedError(Exception):
 33    """Raised when custom code is attempted to be run in an environment that does not support it."""
 34
 35    def __init__(self) -> None:
 36        super().__init__(
 37            "Custom connector code is not permitted in this environment. "
 38            "If you need to run custom code, please ask your administrator to set the `AIRBYTE_ENABLE_UNSAFE_CODE` "
 39            "environment variable to 'true' in your Airbyte environment. "
 40            "If you see this message in Airbyte Cloud, your workspace does not allow executing "
 41            "custom connector code."
 42        )
 43
 44
 45def _hash_text(input_text: str, hash_type: str = "md5") -> str:
 46    """Return the hash of the input text using the specified hash type."""
 47    if not input_text:
 48        raise ValueError("Hash input text cannot be empty.")
 49
 50    hash_object = CHECKSUM_FUNCTIONS[hash_type]()
 51    hash_object.update(input_text.encode())
 52    return hash_object.hexdigest()
 53
 54
 55def custom_code_execution_permitted() -> bool:
 56    """Return `True` if custom code execution is permitted, otherwise `False`.
 57
 58    Custom code execution is permitted if the `AIRBYTE_ENABLE_UNSAFE_CODE` environment variable is set to 'true'.
 59    """
 60    return os.environ.get(ENV_VAR_ALLOW_CUSTOM_CODE, "").lower() == "true"
 61
 62
 63def validate_python_code(
 64    code_text: str,
 65    checksums: dict[str, str] | None,
 66) -> None:
 67    """Validate the provided Python code text against the provided checksums.
 68
 69    Currently we fail if no checksums are provided, although this may change in the future.
 70    """
 71    if not code_text:
 72        # No code provided, nothing to validate.
 73        return
 74
 75    if not checksums:
 76        raise ValueError(f"A checksum is required to validate the code. Received: {checksums}")
 77
 78    for checksum_type, checksum in checksums.items():
 79        if checksum_type not in CHECKSUM_FUNCTIONS:
 80            raise ValueError(
 81                f"Unsupported checksum type: {checksum_type}. Supported checksum types are: {CHECKSUM_FUNCTIONS.keys()}"
 82            )
 83
 84        calculated_checksum = _hash_text(code_text, checksum_type)
 85        if calculated_checksum != checksum:
 86            raise AirbyteCodeTamperedError(
 87                f"{checksum_type} checksum does not match."
 88                + str(
 89                    {
 90                        "expected_checksum": checksum,
 91                        "actual_checksum": calculated_checksum,
 92                    }
 93                ),
 94            )
 95
 96
 97def get_registered_components_module(
 98    config: Mapping[str, Any] | None,
 99) -> ModuleType | None:
100    """Get a components module object based on the provided config.
101
102    If custom python components is provided, this will be loaded. Otherwise, we will
103    attempt to load from the `components` module already imported/registered in sys.modules.
104
105    If custom `components.py` text is provided in config, it will be registered with sys.modules
106    so that it can be later imported by manifest declarations which reference the provided classes.
107
108    Returns `None` if no components is provided and the `components` module is not found.
109    """
110    if config and config.get(INJECTED_COMPONENTS_PY, None):
111        if not custom_code_execution_permitted():
112            raise AirbyteCustomCodeNotPermittedError
113
114        # Create a new module object and execute the provided Python code text within it
115        python_text: str = config[INJECTED_COMPONENTS_PY]
116        return register_components_module_from_string(
117            components_py_text=python_text,
118            checksums=config.get(INJECTED_COMPONENTS_PY_CHECKSUMS, None),
119        )
120
121    # Check for `components` or `source_declarative_manifest.components`.
122    if SDM_COMPONENTS_MODULE_NAME in sys.modules:
123        return cast(ModuleType, sys.modules.get(SDM_COMPONENTS_MODULE_NAME))
124
125    if COMPONENTS_MODULE_NAME in sys.modules:
126        return cast(ModuleType, sys.modules.get(COMPONENTS_MODULE_NAME))
127
128    # Could not find module 'components' in `sys.modules`
129    # and INJECTED_COMPONENTS_PY was not provided in config.
130    return None
131
132
133def register_components_module_from_string(
134    components_py_text: str,
135    checksums: dict[str, Any] | None,
136) -> ModuleType:
137    """Load and return the components module from a provided string containing the python code."""
138    # First validate the code
139    validate_python_code(
140        code_text=components_py_text,
141        checksums=checksums,
142    )
143
144    # Create a new module object
145    components_module = ModuleType(name=COMPONENTS_MODULE_NAME)
146
147    # Execute the module text in the module's namespace
148    exec(components_py_text, components_module.__dict__)
149
150    # Register the module in `sys.modules`` so it can be imported as
151    # `source_declarative_manifest.components` and/or `components`.
152    sys.modules[SDM_COMPONENTS_MODULE_NAME] = components_module
153    sys.modules[COMPONENTS_MODULE_NAME] = components_module
154
155    # Now you can import and use the module
156    return components_module