This is an automated email from the ASF dual-hosted git repository. ephraimanierobi pushed a commit to branch v2-6-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 38df98bac8101ac61d1a27bcaa07b82a24b1487e Author: Jarek Potiuk <[email protected]> AuthorDate: Thu Jun 29 09:09:19 2023 +0200 Fix behaviour of LazyDictWithCache wheni import fails (#32248) When #15330 added docker.task, it also optimized replacement of the callable with it's result in LazyDictWithCache. LazyDictWithCache is used by Provider's Manager to optimize access to hooks - basically hook is only actually imported, when it is accessed. This helps with speeding up importing of connection information The optimization added result of running callable to _resolved set, but it missed the case when None was returned. Previously, when None was returned, the callable was not replaced and it was called again. After the change - the _resolved set was updated with the key and None was returned. But since the key has not been replaced, next time when the same key was retrieved, the original "callable" was returned, not the None value. So if callable returned None, and the same key was retrieved twice, the second time, instead of None, the dictionary returned Callable. This PR fixes it by setting the value to dictionary even if it was None. (cherry picked from commit 462df0d05abbf023ea367ad2591cc889a49c815a) --- airflow/providers_manager.py | 3 +-- tests/always/test_providers_manager.py | 31 ++++++++++++++++++++++++++++++- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index 84fd957a0e..dd214c755c 100644 --- a/airflow/providers_manager.py +++ b/airflow/providers_manager.py @@ -110,8 +110,7 @@ class LazyDictWithCache(MutableMapping): # callable itself value = value() self._resolved.add(key) - if value: - self._raw_dict.__setitem__(key, value) + self._raw_dict.__setitem__(key, value) return value def __delitem__(self, key): diff --git a/tests/always/test_providers_manager.py b/tests/always/test_providers_manager.py index 31a6f70823..ab834ec5a9 100644 --- a/tests/always/test_providers_manager.py +++ b/tests/always/test_providers_manager.py @@ -28,7 +28,7 @@ from flask_babel import lazy_gettext from wtforms import BooleanField, Field, StringField from airflow.exceptions import AirflowOptionalProviderFeatureException -from airflow.providers_manager import HookClassProvider, ProviderInfo, ProvidersManager +from airflow.providers_manager import HookClassProvider, LazyDictWithCache, ProviderInfo, ProvidersManager class TestProviderManager: @@ -372,3 +372,32 @@ class TestProviderManager: assert [ "Optional provider feature disabled when importing 'HookClass' from 'test_package' package" ] == self._caplog.messages + + [email protected]( + "value, expected_outputs,", + [ + ("a", "a"), + (1, 1), + (None, None), + (lambda: 0, 0), + (lambda: None, None), + (lambda: "z", "z"), + ], +) +def test_lazy_cache_dict_resolving(value, expected_outputs): + lazy_cache_dict = LazyDictWithCache() + lazy_cache_dict["key"] = value + assert lazy_cache_dict["key"] == expected_outputs + # Retrieve it again to see if it is correctly returned again + assert lazy_cache_dict["key"] == expected_outputs + + +def test_lazy_cache_dict_raises_error(): + def raise_method(): + raise Exception("test") + + lazy_cache_dict = LazyDictWithCache() + lazy_cache_dict["key"] = raise_method + with pytest.raises(Exception, match="test"): + _ = lazy_cache_dict["key"]
