OptionalMangledDLTContext.py

import os
from typing import Optional


class BothENVAndConfigPresent(Exception):
    def __init__(self, msg: str) -> None:
        """

        :type msg: str
        """
        super().__init__(msg)


class OptionalMangledDLTContext:
    def _is_mangleddlt_envs_present(self) -> bool:
        return all(os.environ.get(var) for var in ["DATABRICKS_HOST", "DATABRICKS_TOKEN", "DATABRICKS_WAREHOUSE_ID"])

    def _is_databricks_runtime(self) -> bool:
        """Check if running on Databricks runtime."""
        return 'DATABRICKS_RUNTIME_VERSION' in os.environ

    def __init__(self, config: Optional[dict] = None) -> None:
        self.config = config

    def __enter__(self):
        if self._is_databricks_runtime():
            return self
        from mangledlt import MangledDLT

        if self.config is not None and self._is_mangleddlt_envs_present():
            raise BothENVAndConfigPresent("Provide either config or ENV vars, not both")

        if self.config is not None:
            self.mdlt = MangledDLT(config=self.config)
            self.mdlt.enable()
        elif self._is_mangleddlt_envs_present():
            self.mdlt = MangledDLT()
            self.mdlt.enable()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.mdlt is not None:
            self.mdlt.disable()
        return False # Propagate any exception