diff --git a/shared/module_loading/src/airflow_shared/module_loading/__init__.py b/shared/module_loading/src/airflow_shared/module_loading/__init__.py index 238506ae52dae..20f6bf8d2be46 100644 --- a/shared/module_loading/src/airflow_shared/module_loading/__init__.py +++ b/shared/module_loading/src/airflow_shared/module_loading/__init__.py @@ -81,18 +81,39 @@ def import_string(dotted_path: str): Raise ImportError if the import failed. """ - # TODO: Add support for nested classes. Currently, it only works for top-level classes. try: module_path, class_name = dotted_path.rsplit(".", 1) except ValueError: raise ImportError(f"{dotted_path} doesn't look like a module path") - module = import_module(module_path) - try: - return getattr(module, class_name) + module = import_module(module_path) + except ImportError as e: + # Fallback for nested classes + parts = dotted_path.split(".") + # Iterate backwards to find the longest valid module path + for i in range(len(parts) - 1, 0, -1): + module_path = ".".join(parts[:i]) + class_path = parts[i:] + try: + module = import_module(module_path) + break + except ImportError: + if i == 1: + raise ImportError(f"{dotted_path} doesn't look like a module path") from e + continue + else: + raise ImportError(f"{dotted_path} doesn't look like a module path") from e + else: + class_path = [class_name] + + obj = module + try: + for attr in class_path: + obj = getattr(obj, attr) + return obj except AttributeError: - raise ImportError(f'Module "{module_path}" does not define a "{class_name}" attribute/class') + raise ImportError(f'Module "{module_path}" does not define a "{".".join(class_path)}" attribute/class') def qualname(o: object | Callable, use_qualname: bool = False, exclude_module: bool = False) -> str: diff --git a/shared/module_loading/tests/module_loading/test_module_loading.py b/shared/module_loading/tests/module_loading/test_module_loading.py index 7c26d8caece08..c41d6ee13d35d 100644 --- a/shared/module_loading/tests/module_loading/test_module_loading.py +++ b/shared/module_loading/tests/module_loading/test_module_loading.py @@ -32,19 +32,30 @@ def _sample_function(): pass +class TopLevelClass: + class NestedClass: + pass + + class TestModuleImport: def test_import_string(self): cls = import_string("module_loading.test_module_loading._import_string") assert cls == _import_string + def test_import_nested_class(self): + cls = import_string("module_loading.test_module_loading.TopLevelClass.NestedClass") + assert cls == TopLevelClass.NestedClass + + def test_import_string_exceptions(self): # Test exceptions raised with pytest.raises(ImportError): import_string("no_dots_in_path") - msg = 'Module "module_loading.test_module_loading" does not define a "nonexistent" attribute' + msg = 'Module "module_loading.test_module_loading" does not define a "nonexistent" attribute/class' with pytest.raises(ImportError, match=msg): import_string("module_loading.test_module_loading.nonexistent") + class TestModuleLoading: @pytest.mark.parametrize( ("path", "expected"), diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 87311f02da7a1..f2d435d72e1eb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -418,7 +418,12 @@ def _fork_main( - Catch un-handled exceptions and attempt to show _something_ in case of error - Finally, run the actual task runner code (``target`` argument, defaults to ``.task_runner:main`) """ - # TODO: Make this process a session leader + # Make this process a session leader + if hasattr(os, "setsid"): + try: + os.setsid() + except OSError: + pass # Store original stderr for last-chance exception handling last_chance_stderr = _get_last_chance_stderr()