Skip to content

Support num_retries field in env var for GCP connection #8700

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 5, 2020
Merged

Support num_retries field in env var for GCP connection #8700

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3e2377b
Support num_retries field in env var for GCP connection
May 4, 2020
988fe4e
fixup! Support num_retries field in env var for GCP connection
May 4, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion airflow/providers/google/common/hooks/base_google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,19 @@ def num_retries(self) -> int:
:return: the number of times each API request should be retried
:rtype: int
"""
return self._get_field('num_retries') or 5
field_value = self._get_field('num_retries', default=5)
if field_value is None:
return 5
if isinstance(field_value, str) and field_value.strip() == '':
return 5
try:
return int(field_value)
except ValueError:
raise AirflowException(
f"The num_retries field should be a integer. "
f"Current value: \"{field_value}\" (type: {type(field_value)}). "
f"Please check the connection configuration."
)

@property
def client_info(self) -> ClientInfo:
Expand Down
46 changes: 46 additions & 0 deletions tests/providers/google/common/hooks/test_base_google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -724,3 +724,49 @@ def test_provide_authorized_gcloud_via_gcloud_application_default(
],
any_order=False
)


class TestNumRetry(unittest.TestCase):

def test_should_return_int_when_set_int_via_connection(self):
instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default")
instance.extras = {
'extra__google_cloud_platform__num_retries': 10,
}

self.assertIsInstance(instance.num_retries, int)
self.assertEqual(10, instance.num_retries)

@mock.patch.dict(
'os.environ',
AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=(
'google-cloud-platform://?extra__google_cloud_platform__num_retries=5'
)
)
def test_should_return_int_when_set_via_env_var(self):
instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default")
self.assertIsInstance(instance.num_retries, int)

@mock.patch.dict(
'os.environ',
AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=(
'google-cloud-platform://?extra__google_cloud_platform__num_retries=cat'
)
)
def test_should_raise_when_invalid_value_via_env_var(self):
instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default")
with self.assertRaisesRegex(
AirflowException, re.escape("The num_retries field should be a integer.")
):
self.assertIsInstance(instance.num_retries, int)

@mock.patch.dict(
'os.environ',
AIRFLOW_CONN_GOOGLE_CLOUD_DEFAULT=(
'google-cloud-platform://?extra__google_cloud_platform__num_retries='
)
)
def test_should_fallback_when_empty_string_in_env_var(self):
instance = hook.GoogleBaseHook(gcp_conn_id="google_cloud_default")
self.assertIsInstance(instance.num_retries, int)
self.assertEqual(5, instance.num_retries)