Skip to content

Commit f489a70

Browse files
committed
Fix behaviour of LazyDictWithCache wheni import fails
When apache#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.
1 parent b6477d0 commit f489a70

File tree

2 files changed

+31
-3
lines changed

2 files changed

+31
-3
lines changed

airflow/providers_manager.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,7 @@ def __getitem__(self, key):
109109
# callable itself
110110
value = value()
111111
self._resolved.add(key)
112-
if value:
113-
self._raw_dict.__setitem__(key, value)
112+
self._raw_dict.__setitem__(key, value)
114113
return value
115114

116115
def __delitem__(self, key):

tests/always/test_providers_manager.py

+30-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
from wtforms import BooleanField, Field, StringField
2929

3030
from airflow.exceptions import AirflowOptionalProviderFeatureException
31-
from airflow.providers_manager import HookClassProvider, ProviderInfo, ProvidersManager
31+
from airflow.providers_manager import HookClassProvider, LazyDictWithCache, ProviderInfo, ProvidersManager
3232

3333

3434
class TestProviderManager:
@@ -373,3 +373,32 @@ def test_optional_feature_debug(self, mock_importlib_import_string):
373373
assert [
374374
"Optional provider feature disabled when importing 'HookClass' from 'test_package' package"
375375
] == self._caplog.messages
376+
377+
378+
@pytest.mark.parametrize(
379+
"value, expected_outputs,",
380+
[
381+
("a", "a"),
382+
(1, 1),
383+
(None, None),
384+
(lambda: 0, 0),
385+
(lambda: None, None),
386+
(lambda: "z", "z"),
387+
],
388+
)
389+
def test_lazy_cache_dict_resolving(value, expected_outputs):
390+
lazy_cache_dict = LazyDictWithCache()
391+
lazy_cache_dict["key"] = value
392+
assert lazy_cache_dict["key"] == expected_outputs
393+
# Retrieve it again to see if it is correctly returned again
394+
assert lazy_cache_dict["key"] == expected_outputs
395+
396+
397+
def test_lazy_cache_dict_raises_error():
398+
def raise_method():
399+
raise Exception("test")
400+
401+
lazy_cache_dict = LazyDictWithCache()
402+
lazy_cache_dict["key"] = raise_method
403+
with pytest.raises(Exception, match="test"):
404+
_ = lazy_cache_dict["key"]

0 commit comments

Comments
 (0)