merobi-hub/code-agent
0
1import logging2import os3import attr4 5from packaging.version import Version6 7from smolagents.tools import Tool8 9# Set this to True to bypass the latest version check for OpenLineage package.10# Version check will be skipped if unable to access PyPI URL11BYPASS_LATEST_VERSION_CHECK = False12# Update this to `CUSTOM` if using any other backend for OpenLineage events ingestion13# When using custom transport - implement custom checks in _verify_custom_backend function14LINEAGE_BACKEND = "MARQUEZ"15 16log = logging.getLogger(__name__) 17 18# Tool to check the configuration of OpenLineage in the local Apache Airflow instance.19class CheckOpenLineage(Tool):20 name = "check_openlineage_config_tool"21 description = """22 The preflight CheckOpenLineage class has been created to enable verifying of the setup 23 of OpenLineage within an Apache Airflow environment. It checks the Airflow version, the 24 version of the installed OpenLineage package, and the configuration settings read by the 25 OpenLineage listener. This validation is crucial because, after setting up OpenLineage 26 with Airflow and configuring necessary environment variables, users need confirmation 27 that the setup is correctly done to start receiving OpenLineage events.28 29 This is a slightly modified version of the Preflight Check DAG in the OpenLineage docs30 at https://openlineage.io/docs/integrations/airflow/preflight-check-dag."""31 inputs = {}32 output_type = "string"33 34 35 def _get_latest_package_version(self, library_name: str) -> Version | None:36 """37 Get the latest available version of the Apache Airflow OpenLineage Provider package38 from the PyPI.org API.39 """40 try:41 import requests42 43 response = requests.get(f"https://pypi.org/pypi/{library_name}/json")44 response.raise_for_status()45 version_string = response.json()["info"]["version"]46 return Version(version_string)47 except Exception as e:48 log.error(f"Failed to fetch latest version for `{library_name}` from PyPI: {e}")49 return None50 51 52 def _get_installed_package_version(self, library_name) -> Version | None:53 """54 Get the version of Apache Airflow OpenLineage Provider installed locally.55 """56 try:57 from importlib.metadata import version58 59 version = Version(version(library_name)) 60 log.info(f"Installed {library_name} version is {version}.")61 return version62 except Exception as e:63 raise ModuleNotFoundError(f"`{library_name}` is not installed") from e64 65 66 def _provider_can_be_used(self) -> [bool, str]:67 """68 Get the version of the locally installed Apache Airflow instance to determine if the69 Apache Airflow OpenLineage Provider can be used.70 """71 import subprocess72 73 app_name = "airflow"74 version_flag = "version"75 process = subprocess.run([app_name, version_flag], capture_output=True, text=True, check=True)76 version_output = process.stdout.strip()77 parsed_version = Version(version_output)78 if parsed_version < Version("2.1"):79 raise RuntimeError("OpenLineage is not supported in Airflow versions <2.1")80 elif parsed_version >= Version("2.7"):81 log.info("Provider can be used.")82 return True, version_output83 return False84 85 86 def validate_ol_installation(self) -> None:87 """88 Validate the OpenLineage installation by verifying the compatibility of the OpenLineage Provider and the89 locally installed copy of Apache Airflow.90 """91 library_name = "openlineage-airflow"92 provider_status = self._provider_can_be_used()93 if provider_status[0]:94 library_name = "apache-airflow-providers-openlineage"95 96 library_version = self._get_installed_package_version(library_name)97 if Version(provider_status[1]) >= Version("2.10.0") and library_version < Version("1.8.0"):98 raise ValueError(99 f"Airflow version `{provider_status[1]}` requires `{library_name}` version >=1.8.0. "100 f"Installed version: `{library_version}` "101 f"Please upgrade the package using `pip install --upgrade {library_name}`"102 )103 if BYPASS_LATEST_VERSION_CHECK:104 log.info(f"Bypassing the latest version check for `{library_name}`")105 return106 107 latest_version = self._get_latest_package_version(library_name)108 if latest_version is None:109 log.warning(f"Failed to fetch the latest version for `{library_name}`. Skipping version check.")110 return111 112 if library_version < latest_version:113 raise ValueError(114 f"`{library_name}` is out of date. "115 f"Installed version: `{library_version}`, "116 f"Required version: `{latest_version}`"117 f"Please upgrade the package using `pip install --upgrade {library_name}` or set BYPASS_LATEST_VERSION_CHECK to True"118 )119 120 121 def _is_transport_set(self) -> None:122 """Check if an OpenLineage transport has been set."""123 transport = conf.get("openlineage", "transport", fallback="")124 log.info(f"Transport: {transport}")125 if transport:126 raise ValueError(127 "Transport value found: `%s`\n"128 "Please check the format at "129 "https://openlineage.io/docs/client/python/#built-in-transport-types",130 transport,131 )132 log.info("Airflow OL transport is not set.")133 return134 135 136 def _is_config_set(self, provider: bool = True) -> None:137 """Check if an OpenLineage config exists."""138 if provider:139 config_path = conf.get("openlineage", "config_path", fallback="")140 else:141 config_path = os.getenv("OPENLINEAGE_CONFIG", "")142 143 log.info("OL config is not set.")144 return145 146 147 def _check_http_env_vars(self) -> None:148 """Check environment for OpenLineage URL and endpoint environment variables."""149 from urllib.parse import urljoin150 151 try:152 final_url = urljoin(os.getenv("OPENLINEAGE_URL", ""), os.getenv("OPENLINEAGE_ENDPOINT", ""))153 log.info("OPENLINEAGE_URL and OPENLINEAGE_ENDPOINT are set to: %s", final_url)154 except:155 raise ValueError(156 "OPENLINEAGE_URL and OPENLINEAGE_ENDPOINT are not set. "157 "Please set up OpenLineage using documentation at "158 "https://airflow.apache.org/docs/apache-airflow-providers-openlineage/stable/guides/user.html"159 )160 return161 162 163 def _debug_missing_transport(self):164 """Debug a missing transport."""165 if self._provider_can_be_used():166 self._is_config_set(provider=True)167 self._is_transport_set()168 self._is_config_set(provider=False)169 # self._check_openlineage_yml("openlineage.yml")170 # self._check_openlineage_yml("~/.openlineage/openlineage.yml")171 self._check_http_env_vars()172 raise ValueError("OpenLineage is missing configuration, please refer to the OL setup docs.")173 174 175 def _is_listener_accessible(self):176 """Check if an OpenLineage listener is accessible."""177 if self._provider_can_be_used():178 try:179 from airflow.providers.openlineage.plugins.openlineage import OpenLineageProviderPlugin as plugin180 except ImportError as e:181 raise ValueError("OpenLineage provider is not accessible") from e182 else:183 try:184 from openlineage.airflow.plugin import OpenLineagePlugin as plugin185 except ImportError as e:186 raise ValueError("OpenLineage is not accessible") from e187 188 if len(plugin.listeners) == 1:189 return True190 191 return False192 193 194 def _is_ol_disabled(self):195 """Confirm that OpenLineage is not disabled and inspect the configuration to suggest a fix."""196 if self._provider_can_be_used():197 try:198 # apache-airflow-providers-openlineage >= 1.7.0199 from airflow.providers.openlineage.conf import is_disabled200 except ImportError:201 # apache-airflow-providers-openlineage < 1.7.0202 from airflow.providers.openlineage.plugins.openlineage import _is_disabled as is_disabled203 else:204 from openlineage.airflow.plugin import _is_disabled as is_disabled205 206 if is_disabled():207 if self._provider_can_be_used() and conf.getboolean("openlineage", "disabled", fallback=False):208 raise ValueError("OpenLineage is disabled in airflow.cfg: openlineage.disabled")209 elif os.getenv("OPENLINEAGE_DISABLED", "false").lower() == "true":210 raise ValueError(211 "OpenLineage is disabled due to the environment variable OPENLINEAGE_DISABLED"212 )213 raise ValueError(214 "OpenLineage is disabled because required config/env variables are not set. "215 "Please refer to "216 "https://airflow.apache.org/docs/apache-airflow-providers-openlineage/stable/guides/user.html"217 )218 log.info("OpenLineage is not disabled.")219 return False220 221 222 def _get_transport(self):223 """Get the configured transport from the OpenLineage plugin."""224 if self._provider_can_be_used():225 from airflow.providers.openlineage.plugins.openlineage import OpenLineageProviderPlugin226 transport = OpenLineageProviderPlugin().listeners[0].adapter.get_or_create_openlineage_client().transport227 else:228 from openlineage.airflow.plugin import OpenLineagePlugin229 transport = (230 OpenLineagePlugin.listeners[0].adapter.get_or_create_openlineage_client().transport231 )232 return transport233 234 def is_ol_accessible_and_enabled(self):235 """Confirm that OpenLineage is accessible and enabled by attempting to build the transport."""236 if not self._is_listener_accessible():237 self._is_ol_disabled()238 239 try:240 transport = self._get_transport()241 except Exception as e:242 raise ValueError("There was an error when trying to build transport.") from e243 244 if transport is None or transport.kind in ("noop", "console"):245 _debug_missing_transport()246 247 248 def validate_connection(self):249 """Validate the connection to the lineage backend."""250 transport = self._get_transport()251 config = attr.asdict(transport.config)252 self._verify_backend(LINEAGE_BACKEND, config)253 254 255 def _verify_backend(self, backend_type: str, config: dict):256 """Verify the lineage backed."""257 backend_type = backend_type.lower()258 if backend_type == "marquez":259 log.info("Backend type: Marquez")260 elif backend_type == "atlan":261 log.info("Backend type: Atlan")262 return self._verify_atlan_http_backend(config)263 elif backend_type == "custom":264 log.info("Backend type: custom")265 return self._verify_custom_backend(config)266 else:267 raise ValueError(f"Unsupported backend type: {backend_type}")268 269 270 def _verify_atlan_http_backend(self, config):271 raise NotImplementedError("This feature is not implemented yet")272 273 274 def _verify_custom_backend(self, config):275 raise NotImplementedError("This feature is not implemented yet")276 277 278 def forward(self) -> str:279 try:280 self.validate_ol_installation()281 self.is_ol_accessible_and_enabled()282 self.validate_connection()283 return "OpenLineage is configured correctly."284 except:285 return "There is a problem with the OpenLineage configuration."286 287 def __init__(self, *args, **kwargs):288 self.is_initialized = False289