From d6e52d7cd4df765d32bdea7ddd76de678a7a33f0 Mon Sep 17 00:00:00 2001 From: Holt Skinner Date: Wed, 2 Nov 2022 11:01:02 -0500 Subject: [PATCH] fix: Added extra exception handling to operation samples --- ...atch_process_documents_processor_version_sample.py | 11 ++++++++--- samples/snippets/batch_process_documents_sample.py | 11 ++++++++--- samples/snippets/cancel_operation_sample_test.py | 1 - samples/snippets/get_operation_sample.py | 11 +++++++---- samples/snippets/list_operations_sample_test.py | 1 - samples/snippets/poll_operation_sample.py | 7 ++++++- 6 files changed, 29 insertions(+), 13 deletions(-) diff --git a/samples/snippets/batch_process_documents_processor_version_sample.py b/samples/snippets/batch_process_documents_processor_version_sample.py index 7050c3ec..c1fab35a 100644 --- a/samples/snippets/batch_process_documents_processor_version_sample.py +++ b/samples/snippets/batch_process_documents_processor_version_sample.py @@ -17,6 +17,7 @@ import re from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import RetryError from google.cloud import documentai, storage # TODO(developer): Uncomment these variables before running the sample. @@ -39,7 +40,7 @@ def batch_process_documents_processor_version( input_mime_type: str, gcs_output_bucket: str, gcs_output_uri_prefix: str, - timeout: int = 300, + timeout: int = 400, ): # You must set the api_endpoint if you use a location other than 'us', e.g.: @@ -90,8 +91,12 @@ def batch_process_documents_processor_version( # Continually polls the operation until it is complete. # This could take some time for larger files # Format: projects/PROJECT_NUMBER/locations/LOCATION/operations/OPERATION_ID - print(f"Waiting for operation {operation.operation.name} to complete...") - operation.result(timeout=timeout) + try: + print(f"Waiting for operation {operation.operation.name} to complete...") + operation.result(timeout=timeout) + # Catch exception when operation doesn't finish before timeout + except (RetryError) as e: + print(e.message) # NOTE: Can also use callbacks for asynchronous processing # diff --git a/samples/snippets/batch_process_documents_sample.py b/samples/snippets/batch_process_documents_sample.py index f01d8787..3b2e07ea 100644 --- a/samples/snippets/batch_process_documents_sample.py +++ b/samples/snippets/batch_process_documents_sample.py @@ -17,6 +17,7 @@ import re from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import RetryError from google.cloud import documentai, storage # TODO(developer): Uncomment these variables before running the sample. @@ -37,7 +38,7 @@ def batch_process_documents( input_mime_type: str, gcs_output_bucket: str, gcs_output_uri_prefix: str, - timeout: int = 300, + timeout: int = 400, ): # You must set the api_endpoint if you use a location other than 'us', e.g.: @@ -86,8 +87,12 @@ def batch_process_documents( # Continually polls the operation until it is complete. # This could take some time for larger files # Format: projects/PROJECT_NUMBER/locations/LOCATION/operations/OPERATION_ID - print(f"Waiting for operation {operation.operation.name} to complete...") - operation.result(timeout=timeout) + try: + print(f"Waiting for operation {operation.operation.name} to complete...") + operation.result(timeout=timeout) + # Catch exception when operation doesn't finish before timeout + except (RetryError) as e: + print(e.message) # NOTE: Can also use callbacks for asynchronous processing # diff --git a/samples/snippets/cancel_operation_sample_test.py b/samples/snippets/cancel_operation_sample_test.py index 5d403be0..28d3eacd 100644 --- a/samples/snippets/cancel_operation_sample_test.py +++ b/samples/snippets/cancel_operation_sample_test.py @@ -30,4 +30,3 @@ def test_cancel_operation(capsys): out, _ = capsys.readouterr() assert "Operation" in out - assert "cancelled" in out diff --git a/samples/snippets/get_operation_sample.py b/samples/snippets/get_operation_sample.py index a569e816..de8e01ae 100644 --- a/samples/snippets/get_operation_sample.py +++ b/samples/snippets/get_operation_sample.py @@ -16,6 +16,7 @@ # [START documentai_get_operation] from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import NotFound from google.cloud import documentai from google.longrunning.operations_pb2 import GetOperationRequest @@ -33,10 +34,12 @@ def get_operation_sample(location: str, operation_name: str): request = GetOperationRequest(name=operation_name) # Make GetOperation request - operation = client.get_operation(request=request) - - # Print the Operation Information - print(operation) + try: + operation = client.get_operation(request=request) + # Print the Operation Information + print(operation) + except (NotFound) as e: + print(e.message) # [END documentai_get_operation] diff --git a/samples/snippets/list_operations_sample_test.py b/samples/snippets/list_operations_sample_test.py index e82e6dab..c316ca23 100644 --- a/samples/snippets/list_operations_sample_test.py +++ b/samples/snippets/list_operations_sample_test.py @@ -29,4 +29,3 @@ def test_list_operations(capsys): out, _ = capsys.readouterr() assert "operations" in out - assert "BatchProcessMetadata" in out diff --git a/samples/snippets/poll_operation_sample.py b/samples/snippets/poll_operation_sample.py index c3e7cc60..e5ec5441 100644 --- a/samples/snippets/poll_operation_sample.py +++ b/samples/snippets/poll_operation_sample.py @@ -18,6 +18,7 @@ from time import sleep from google.api_core.client_options import ClientOptions +from google.api_core.exceptions import NotFound from google.cloud import documentai from google.longrunning.operations_pb2 import GetOperationRequest @@ -36,7 +37,11 @@ def poll_operation_sample(location: str, operation_name: str): while True: # Make GetOperation request - operation = client.get_operation(request=request) + try: + operation = client.get_operation(request=request) + except (NotFound) as e: + print(e.message) + break # Print the Operation Information print(operation)