coufon commented on a change in pull request #4791:  [AIRFLOW-3908] Add more 
Google Cloud Vision operators
URL: https://github.com/apache/airflow/pull/4791#discussion_r261411536
 
 

 ##########
 File path: airflow/contrib/hooks/gcp_vision_hook.py
 ##########
 @@ -354,3 +307,140 @@ def _get_autogenerated_id(response):
         if '/' not in name:
             raise AirflowException('Unable to get id from name... 
[{}]'.format(name))
         return name.rsplit('/', 1)[1]
+
+    @GoogleCloudBaseHook.catch_http_exception
+    @GoogleCloudBaseHook.fallback_to_default_project_id
+    def create_reference_image(
+        self,
+        location,
+        product_id,
+        reference_image,
+        reference_image_id=None,
+        project_id=None,
+        retry=None,
+        timeout=None,
+        metadata=None,
+    ):
+        """
+        For the documentation see:
+        
:py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionReferenceImageCreateOperator`
+        """
+        client = self.get_conn()
+        self.log.info('Creating ReferenceImage')
+        parent = ProductSearchClient.product_path(project=project_id, 
location=location, product=product_id)
+
+        response = client.create_reference_image(
+            parent=parent,
+            reference_image=reference_image,
+            reference_image_id=reference_image_id,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+        self.log.info('ReferenceImage created: %s', response.name if response 
else '')
+        self.log.debug('ReferenceImage created:\n%s', response)
+
+        if not reference_image_id:
+            # Refernece image  id was generated by the API
+            reference_image_id = self._get_autogenerated_id(response)
+            self.log.info('Extracted autogenerated ReferenceImage ID from the 
response: %s', product_set_id)
+
+        return reference_image_id
+
+    @GoogleCloudBaseHook.catch_http_exception
+    @GoogleCloudBaseHook.fallback_to_default_project_id
+    def delete_reference_image(
+        self,
+        location,
+        product_id,
+        reference_image_id,
+        project_id=None,
+        retry=None,
+        timeout=None,
+        metadata=None,
+    ):
+        """
+        For the documentation see:
+        
:py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionReferenceImageCreateOperator`
+        """
+        client = self.get_conn()
+        self.log.info('Deleting ReferenceImage')
+        name = ProductSearchClient.reference_image_path(
+            project=project_id, location=location, product=product_id, 
reference_image=reference_image_id
+        )
+        client.delete_reference_image(name=name, retry=retry, timeout=timeout, 
metadata=metadata)
 
 Review comment:
   Maybe add a log.info here?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


With regards,
Apache Git Services

Reply via email to