Add CloudVisionDeleteReferenceImageOperator (#9698)

apache · Jul 15, 2020 · 2d8dbac · 2d8dbac
1 parent 52b6efe
commit 2d8dbac
Show file tree

Hide file tree

Showing 5 changed files with 204 additions and 12 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_vision.py b/airflow/providers/google/cloud/example_dags/example_vision.py
@@ -39,10 +39,11 @@
  CloudVisionAddProductToProductSetOperator, CloudVisionCreateProductOperator,
  CloudVisionCreateProductSetOperator, CloudVisionCreateReferenceImageOperator,
  CloudVisionDeleteProductOperator, CloudVisionDeleteProductSetOperator,
- CloudVisionDetectImageLabelsOperator, CloudVisionDetectImageSafeSearchOperator,
- CloudVisionDetectTextOperator, CloudVisionGetProductOperator, CloudVisionGetProductSetOperator,
- CloudVisionImageAnnotateOperator, CloudVisionRemoveProductFromProductSetOperator,
- CloudVisionTextDetectOperator, CloudVisionUpdateProductOperator, CloudVisionUpdateProductSetOperator,
+ CloudVisionDeleteReferenceImageOperator, CloudVisionDetectImageLabelsOperator,
+ CloudVisionDetectImageSafeSearchOperator, CloudVisionDetectTextOperator, CloudVisionGetProductOperator,
+ CloudVisionGetProductSetOperator, CloudVisionImageAnnotateOperator,
+ CloudVisionRemoveProductFromProductSetOperator, CloudVisionTextDetectOperator,
+ CloudVisionUpdateProductOperator, CloudVisionUpdateProductSetOperator,
 )
 from airflow.utils.dates import days_ago
 
@@ -185,6 +186,17 @@
  )
  # [END howto_operator_vision_reference_image_create]
 
+ # [START howto_operator_vision_reference_image_delete]
+ reference_image_delete = CloudVisionDeleteReferenceImageOperator(
+ location=GCP_VISION_LOCATION,
+ product_id="{{ task_instance.xcom_pull('product_create') }}",
+ reference_image_id=GCP_VISION_REFERENCE_IMAGE_ID,
+ retry=Retry(maximum=10.0),
+ timeout=5,
+ task_id='reference_image_delete',
+ )
+ # [END howto_operator_vision_reference_image_delete]
+
  # [START howto_operator_vision_add_product_to_product_set]
  add_product_to_product_set = CloudVisionAddProductToProductSetOperator(
  location=GCP_VISION_LOCATION,
@@ -214,7 +226,7 @@
  product_set_create >> product_set_get >> product_set_update >> product_set_delete
 
  # ReferenceImage path
- product_create >> reference_image_create >> product_delete
+ product_create >> reference_image_create >> reference_image_delete >> product_delete
 
  # Product/ProductSet path
  product_create >> add_product_to_product_set
@@ -326,6 +338,17 @@
  )
  # [END howto_operator_vision_reference_image_create_2]
 
+ # [START howto_operator_vision_reference_image_delete_2]
+ reference_image_delete_2 = CloudVisionDeleteReferenceImageOperator(
+ location=GCP_VISION_LOCATION,
+ reference_image_id=GCP_VISION_REFERENCE_IMAGE_ID,
+ product_id=GCP_VISION_PRODUCT_ID,
+ retry=Retry(maximum=10.0),
+ timeout=5,
+ task_id='reference_image_delete_2',
+ )
+ # [END howto_operator_vision_reference_image_delete_2]
+
  # Second 'create' task with the same product_id to demonstrate idempotence
  reference_image_create_2_idempotence = CloudVisionCreateReferenceImageOperator(
  location=GCP_VISION_LOCATION,
@@ -367,7 +390,8 @@
  product_set_create_2 >> product_set_create_2_idempotence >> product_set_delete_2
 
  # ReferenceImage path
- product_create_2 >> reference_image_create_2 >> reference_image_create_2_idempotence >> product_delete_2
+ product_create_2 >> reference_image_create_2 >> reference_image_create_2_idempotence
+ reference_image_create_2_idempotence >> reference_image_delete_2 >> product_delete_2
 
  # Product/ProductSet path
  add_product_to_product_set_2 >> remove_product_from_product_set_2

diff --git a/airflow/providers/google/cloud/hooks/vision.py b/airflow/providers/google/cloud/hooks/vision.py
@@ -425,7 +425,7 @@ def delete_reference_image(
  ) -> Dict:
  """
  For the documentation see:
- :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionReferenceImageCreateOperator`
+ :py:class:`~airflow.providers.google.cloud.operators.vision.CloudVisionDeleteReferenceImageOperator`
  """
  client = self.get_conn()
  self.log.info('Deleting ReferenceImage')

diff --git a/airflow/providers/google/cloud/operators/vision.py b/airflow/providers/google/cloud/operators/vision.py
@@ -828,6 +828,86 @@ def execute(self, context):
  return self.reference_image_id
 
 
+class CloudVisionDeleteReferenceImageOperator(BaseOperator):
+ """
+ Deletes a ReferenceImage ID resource.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:CloudVisionDeleteReferenceImageOperator`
+
+ :param location: (Required) The region where the Product is located. Valid regions (as of 2019-02-05) are:
+ us-east1, us-west1, europe-west1, asia-east1
+ :type location: str
+ :param reference_image_id: (Optional) A user-supplied resource id for the ReferenceImage to be added.
+ If set, the server will attempt to use this value as the resource id. If it is already in use, an
+ error is returned with code ALREADY_EXISTS. Must be at most 128 characters long. It cannot contain
+ the character `/`.
+ :type reference_image_id: str
+ :param product_id: (Optional) The resource id of this Product.
+ :type product_id: str
+ :param project_id: (Optional) The project in which the Product is located. If set to None or
+ missing, the default project_id from the GCP connection is used.
+ :type project_id: str
+ :param retry: (Optional) A retry object used 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 metadata: (Optional) Additional metadata that is provided to the method.
+ :type metadata: sequence[tuple[str, str]]
+ :param gcp_conn_id: (Optional) The connection ID used to connect to Google Cloud Platform.
+ :type gcp_conn_id: str
+ """
+ # [START vision_reference_image_create_template_fields]
+ template_fields = (
+ "location",
+ "product_id",
+ "reference_image_id",
+ "project_id",
+ "gcp_conn_id",
+ )
+ # [END vision_reference_image_create_template_fields]
+
+ @apply_defaults
+ def __init__(
+ self,
+ location: str,
+ product_id: str,
+ reference_image_id: str,
+ project_id: Optional[str] = None,
+ retry: Optional[Retry] = None,
+ timeout: Optional[float] = None,
+ metadata: Optional[MetaData] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ *args,
+ **kwargs
+ ) -> None:
+ super().__init__(*args, **kwargs)
+ self.location = location
+ self.product_id = product_id
+ self.reference_image_id = reference_image_id
+ self.project_id = project_id
+ self.retry = retry
+ self.timeout = timeout
+ self.metadata = metadata
+ self.gcp_conn_id = gcp_conn_id
+
+ def execute(self, context):
+ hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)
+ hook.delete_reference_image(
+ location=self.location,
+ product_id=self.product_id,
+ reference_image_id=self.reference_image_id,
+ project_id=self.project_id,
+ retry=self.retry,
+ timeout=self.timeout,
+ metadata=self.metadata,
+ )
+
+
 class CloudVisionAddProductToProductSetOperator(BaseOperator):
  """
  Adds a Product to the specified ProductSet. If the Product is already present, no change is made.

diff --git a/docs/howto/operator/google/cloud/vision.rst b/docs/howto/operator/google/cloud/vision.rst
@@ -680,9 +680,71 @@ Templating
 More information
 """"""""""""""""
 
-See `Google Cloud Vision ReferenceImage create documentation
+See `Google Cloud Vision ReferenceImage delete documentation
 <https://googleapis.github.io/google-cloud-python/latest/vision/gapic/v1/api.html#google.cloud.vision_v1.ProductSearchClient.create_reference_image>`_.
 
+.. _howto/operator:CloudVisionDeleteReferenceImageOperator:
+
+CloudVisionDeleteReferenceImageOperator
+---------------------------------------
+
+Deletes a :code:`ReferenceImage` resource.
+
+For parameter definition, take a look at
+:class:`~airflow.providers.google.cloud.operators.vision.CloudVisionDeleteReferenceImageOperator`
+
+Using the operator
+""""""""""""""""""
+
+We are using the :class:`~google.cloud.vision_v1.types.ReferenceImage` and :class:`~google.api_core.retry.Retry` objects from Google libraries:
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_vision.py
+ :language: python
+ :start-after: [START howto_operator_vision_reference_image_import]
+ :end-before: [END howto_operator_vision_reference_image_import]
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_vision.py
+ :language: python
+ :start-after: [START howto_operator_vision_retry_import]
+ :end-before: [END howto_operator_vision_retry_import]
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_vision.py
+ :language: python
+ :start-after: [START howto_operator_vision_reference_image]
+ :end-before: [END howto_operator_vision_reference_image]
+
+The ``product_set_id`` argument can be omitted (it will be generated by the API):
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_vision.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_vision_reference_image_delete]
+ :end-before: [END howto_operator_vision_reference_image_delete]
+
+Or it can be specified explicitly:
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_vision.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_vision_reference_image_delete_2]
+ :end-before: [END howto_operator_vision_reference_image_delete_2]
+
+
+Templating
+""""""""""
+
+.. literalinclude:: /../airflow/providers/google/cloud/operators/vision.py
+ :language: python
+ :dedent: 4
+ :start-after: [START vision_reference_image_create_template_fields]
+ :end-before: [END vision_reference_image_create_template_fields]
+
+More information
+""""""""""""""""
+
+See `Google Cloud Vision ReferenceImage create documentation
+<https://googleapis.github.io/google-cloud-python/latest/vision/gapic/v1/api.html#google.cloud.vision_v1.ProductSearchClient.delete_reference_image>`_.
+
 .. _howto/operator:CloudVisionRemoveProductFromProductSetOperator:
 
 CloudVisionRemoveProductFromProductSetOperator

diff --git a/tests/providers/google/cloud/operators/test_vision.py b/tests/providers/google/cloud/operators/test_vision.py
@@ -26,10 +26,11 @@
  CloudVisionAddProductToProductSetOperator, CloudVisionCreateProductOperator,
  CloudVisionCreateProductSetOperator, CloudVisionCreateReferenceImageOperator,
  CloudVisionDeleteProductOperator, CloudVisionDeleteProductSetOperator,
- CloudVisionDetectImageLabelsOperator, CloudVisionDetectImageSafeSearchOperator,
- CloudVisionDetectTextOperator, CloudVisionGetProductOperator, CloudVisionGetProductSetOperator,
- CloudVisionImageAnnotateOperator, CloudVisionRemoveProductFromProductSetOperator,
- CloudVisionTextDetectOperator, CloudVisionUpdateProductOperator, CloudVisionUpdateProductSetOperator,
+ CloudVisionDeleteReferenceImageOperator, CloudVisionDetectImageLabelsOperator,
+ CloudVisionDetectImageSafeSearchOperator, CloudVisionDetectTextOperator, CloudVisionGetProductOperator,
+ CloudVisionGetProductSetOperator, CloudVisionImageAnnotateOperator,
+ CloudVisionRemoveProductFromProductSetOperator, CloudVisionTextDetectOperator,
+ CloudVisionUpdateProductOperator, CloudVisionUpdateProductSetOperator,
 )
 
 PRODUCTSET_TEST = ProductSet(display_name='Test Product Set')
@@ -279,6 +280,31 @@ def test_already_exists(self, mock_hook):
  )
 
 
+class TestCloudVisionReferenceImageDelete(unittest.TestCase):
+ @mock.patch(
+ 'airflow.providers.google.cloud.operators.vision.CloudVisionHook',
+ )
+ def test_minimal_green_path(self, mock_hook):
+ mock_hook.return_value.delete_reference_image.return_value = {}
+ op = CloudVisionDeleteReferenceImageOperator(
+ location=LOCATION_TEST,
+ product_id=PRODUCT_ID_TEST,
+ reference_image_id=REFERENCE_IMAGE_ID_TEST,
+ task_id='id',
+ )
+ op.execute(context=None)
+ mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID)
+ mock_hook.return_value.delete_reference_image.assert_called_once_with(
+ location=LOCATION_TEST,
+ product_id=PRODUCT_ID_TEST,
+ reference_image_id=REFERENCE_IMAGE_ID_TEST,
+ project_id=None,
+ retry=None,
+ timeout=None,
+ metadata=None,
+ )
+
+
 class TestCloudVisionAddProductToProductSetOperator(unittest.TestCase):
  @mock.patch('airflow.providers.google.cloud.operators.vision.CloudVisionHook')
  def test_minimal_green_path(self, mock_hook):