docs(samples): Added extra exception handling to operation samples

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
     #

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
     #

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

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]

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

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)