Skip to content
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

Jump to bottom

[AIRFLOW-4768] Add timeout parameter to Cloud Video Intelligence operators #5862

Merged
merged 1 commit into from
Aug 22, 2019
Merged

[AIRFLOW-4768] Add timeout parameter to Cloud Video Intelligence operators #5862

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
18 changes: 18 additions & 0 deletions airflow/gcp/operators/video_intelligence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ class CloudVideoIntelligenceDetectVideoLabelsOperator(BaseOperator):
:param retry: Retry object used to determine when/if to retry requests.
If None is specified, requests will not be retried.
:type retry: google.api_core.retry.Retry
:param timeout: Optional, The amount of time, in seconds, to wait for the request to complete.
Note that if retry is specified, the timeout applies to each individual attempt.
:type timeout: float
:param gcp_conn_id: Optional, The connection ID used to connect to Google Cloud
Platform. Defaults to ``google_cloud_default``.
:type gcp_conn_id: str
Expand All @@ -71,6 +74,7 @@ def __init__(
video_context=None,
location=None,
retry=None,
timeout=None,
gcp_conn_id="google_cloud_default",
*args,
**kwargs
Expand All @@ -83,6 +87,7 @@ def __init__(
self.location = location
self.retry = retry
self.gcp_conn_id = gcp_conn_id
self.timeout = timeout

def execute(self, context):
hook = CloudVideoIntelligenceHook(gcp_conn_id=self.gcp_conn_id)
Expand All @@ -93,6 +98,7 @@ def execute(self, context):
location=self.location,
retry=self.retry,
features=[enums.Feature.LABEL_DETECTION],
timeout=self.timeout
)
self.log.info("Processing video for label annotations")
result = MessageToDict(operation.result())
Expand Down Expand Up @@ -128,6 +134,9 @@ class CloudVideoIntelligenceDetectVideoExplicitContentOperator(BaseOperator):
:param retry: Retry object used to determine when/if to retry requests.
If None is specified, requests will not be retried.
:type retry: google.api_core.retry.Retry
:param timeout: Optional, The amount of time, in seconds, to wait for the request to complete.
Note that if retry is specified, the timeout applies to each individual attempt.
:type timeout: float
:param gcp_conn_id: Optional, The connection ID used to connect to Google Cloud
Platform. Defaults to ``google_cloud_default``.
:type gcp_conn_id: str
Expand All @@ -144,6 +153,7 @@ def __init__(
video_context=None,
location=None,
retry=None,
timeout=None,
gcp_conn_id="google_cloud_default",
*args,
**kwargs
Expand All @@ -156,6 +166,7 @@ def __init__(
self.location = location
self.retry = retry
self.gcp_conn_id = gcp_conn_id
self.timeout = timeout

def execute(self, context):
hook = CloudVideoIntelligenceHook(gcp_conn_id=self.gcp_conn_id)
Expand All @@ -166,6 +177,7 @@ def execute(self, context):
location=self.location,
retry=self.retry,
features=[enums.Feature.EXPLICIT_CONTENT_DETECTION],
timeout=self.timeout
)
self.log.info("Processing video for explicit content annotations")
result = MessageToDict(operation.result())
Expand Down Expand Up @@ -201,6 +213,9 @@ class CloudVideoIntelligenceDetectVideoShotsOperator(BaseOperator):
:param retry: Retry object used to determine when/if to retry requests.
If None is specified, requests will not be retried.
:type retry: google.api_core.retry.Retry
:param timeout: Optional, The amount of time, in seconds, to wait for the request to complete.
Note that if retry is specified, the timeout applies to each individual attempt.
:type timeout: float
:param gcp_conn_id: Optional, The connection ID used to connect to Google Cloud
Platform. Defaults to ``google_cloud_default``.
:type gcp_conn_id: str
Expand All @@ -217,6 +232,7 @@ def __init__(
video_context=None,
location=None,
retry=None,
timeout=None,
gcp_conn_id="google_cloud_default",
*args,
**kwargs
Expand All @@ -229,6 +245,7 @@ def __init__(
self.location = location
self.retry = retry
self.gcp_conn_id = gcp_conn_id
self.timeout = timeout

def execute(self, context):
hook = CloudVideoIntelligenceHook(gcp_conn_id=self.gcp_conn_id)
Expand All @@ -239,6 +256,7 @@ def execute(self, context):
location=self.location,
retry=self.retry,
features=[enums.Feature.SHOT_CHANGE_DETECTION],
timeout=self.timeout
)
self.log.info("Processing video for video shots annotations")
result = MessageToDict(operation.result())
Expand Down
3 changes: 3 additions & 0 deletions tests/gcp/operators/test_video_intelligence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def test_detect_video_labels_green_path(self, mock_hook):
video_context=None,
location=None,
retry=None,
timeout=None
)

@mock.patch("airflow.gcp.operators.video_intelligence.CloudVideoIntelligenceHook")
Expand All @@ -76,6 +77,7 @@ def test_detect_video_explicit_content_green_path(self, mock_hook):
video_context=None,
location=None,
retry=None,
timeout=None
)

@mock.patch("airflow.gcp.operators.video_intelligence.CloudVideoIntelligenceHook")
Expand All @@ -96,4 +98,5 @@ def test_detect_video_shots_green_path(self, mock_hook):
video_context=None,
location=None,
retry=None,
timeout=None
)