Skip to content

Commit

Permalink
Fix behaviour of LazyDictWithCache wheni import fails (#32248)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
potiuk authored Jun 29, 2023
1 parent b6477d0 commit 462df0d
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 3 deletions.
3 changes: 1 addition & 2 deletions airflow/providers_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,7 @@ def __getitem__(self, key):
# 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):
Expand Down
31 changes: 30 additions & 1 deletion tests/always/test_providers_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
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:
Expand Down Expand Up @@ -373,3 +373,32 @@ def test_optional_feature_debug(self, mock_importlib_import_string):
assert [
"Optional provider feature disabled when importing 'HookClass' from 'test_package' package"
] == self._caplog.messages


@pytest.mark.parametrize(
"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"]

0 comments on commit 462df0d

Please sign in to comment.