def remove_product_from_product_set( self, product_set_id, product_id, location=None, project_id=None, retry=None, timeout=None, metadata=None, ): """ For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionRemoveProductFromProductSetOperator` """ client = self.get_conn() product_name = ProductSearchClient.product_path(project_id, location, product_id) product_set_name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Remove Product[name=%s] from Product Set[name=%s]', product_name, product_set_name) client.remove_product_from_product_set( name=product_set_name, product=product_name, retry=retry, timeout=timeout, metadata=metadata ) self.log.info('Product removed from Product Set')
def test_update_product_explicit_name_different_from_constructed(self, get_conn): # Given update_product_method = get_conn.return_value.update_product update_product_method.return_value = None explicit_p_name = ProductSearchClient.product_path( PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCT_ID_TEST_2 ) product = Product(name=explicit_p_name) template_p_name = ProductSearchClient.product_path(PROJECT_ID_TEST, LOC_ID_TEST, PRODUCT_ID_TEST) # When # Location and product_id are passed in addition to a Product with an explicit name, # but both names differ (constructed != explicit). # Should throw AirflowException in this case. with self.assertRaises(AirflowException) as cm: self.hook.update_product( location=LOC_ID_TEST, product_id=PRODUCT_ID_TEST, product=product, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) err = cm.exception self.assertTrue(err) self.assertIn( "The Product name provided in the object ({}) is different than the name created from the input " "parameters ({}). Please either: 1) Remove the Product name, 2) Remove the location and product_" "id parameters, 3) Unify the Product name and input parameters.".format( explicit_p_name, template_p_name ), str(err), ) update_product_method.assert_not_called()
def test_create_productset_autogenerated_id(self, get_conn): # Given autogenerated_id = 'autogen-id' response_product_set = ProductSet( name=ProductSearchClient.product_set_path(PROJECT_ID_TEST, LOC_ID_TEST, autogenerated_id) ) create_product_set_method = get_conn.return_value.create_product_set create_product_set_method.return_value = response_product_set parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) product_set = ProductSet() # When result = self.hook.create_product_set( location=LOC_ID_TEST, product_set_id=None, product_set=product_set, project_id=PROJECT_ID_TEST ) # Then # ProductSet ID was not provided in the method call above. Should be extracted from the API response # and returned. self.assertEqual(result, autogenerated_id) create_product_set_method.assert_called_once_with( parent=parent, product_set=product_set, product_set_id=None, retry=None, timeout=None, metadata=None, )
def create_product( self, location, product, project_id=None, product_id=None, retry=None, timeout=None, metadata=None ): """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductCreateOperator` """ client = self.get_conn() parent = ProductSearchClient.location_path(project_id, location) self.log.info('Creating a new Product under the parent: %s', parent) response = client.create_product( parent=parent, product=product, product_id=product_id, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('Product created: %s', response.name if response else '') self.log.debug('Product created:\n%s', response) if not product_id: # Product id was generated by the API product_id = self._get_autogenerated_id(response) self.log.info('Extracted autogenerated Product ID from the response: %s', product_id) return product_id
def get_product(self, location, product_id, project_id=None, retry=None, timeout=None, metadata=None): """ For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductGetOperator` """ client = self.get_conn() name = ProductSearchClient.product_path(project_id, location, product_id) self.log.info('Retrieving Product: %s', name) response = self._handle_request( lambda **kwargs: client.get_product(**kwargs), name=name, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('Product retrieved.') self.log.debug('Product retrieved:\n%s', response) return MessageToDict(response)
def test_create_productset_explicit_id(self, get_conn): # Given create_product_set_method = get_conn.return_value.create_product_set create_product_set_method.return_value = None parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) product_set = ProductSet() # When result = self.hook.create_product_set( location=LOC_ID_TEST, product_set_id=PRODUCTSET_ID_TEST, product_set=product_set, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) # Then # ProductSet ID was provided explicitly in the method call above, should be returned from the method self.assertEqual(result, PRODUCTSET_ID_TEST) create_product_set_method.assert_called_once_with( parent=parent, product_set=product_set, product_set_id=PRODUCTSET_ID_TEST, retry=None, timeout=None, metadata=None, )
def test_create_product_autogenerated_id_wrong_api_response( self, get_conn): # Given response_product = None create_product_method = get_conn.return_value.create_product create_product_method.return_value = response_product parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) product = Product() # When with self.assertRaises(AirflowException) as cm: self.hook.create_product(location=LOC_ID_TEST, product_id=None, product=product, project_id=PROJECT_ID_TEST) # Then # API response was wrong (None) and thus ProductSet ID extraction should fail. err = cm.exception self.assertIn('Unable to get name from response...', str(err)) create_product_method.assert_called_once_with(parent=parent, product=product, product_id=None, retry=None, timeout=None, metadata=None)
def test_update_productset_no_explicit_name(self, get_conn): # Given product_set = ProductSet() update_product_set_method = get_conn.return_value.update_product_set update_product_set_method.return_value = product_set productset_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST, LOC_ID_TEST, PRODUCTSET_ID_TEST) # When result = self.hook.update_product_set( location=LOC_ID_TEST, product_set_id=PRODUCTSET_ID_TEST, product_set=product_set, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) # Then self.assertEqual(result, MessageToDict(product_set)) update_product_set_method.assert_called_once_with( product_set=ProductSet(name=productset_name), metadata=None, retry=None, timeout=None, update_mask=None, )
def delete_reference_image( self, location: str, product_id: str, reference_image_id: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> Dict: """ 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) response = client.delete_reference_image( name=name, # pylint: disable=assignment-from-no-return retry=retry, timeout=timeout, metadata=metadata) self.log.info('ReferenceImage with the name [%s] deleted.', name) return MessageToDict(response)
def create_product(self, location: str, product: Union[dict, Product], project_id: str, product_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None): """ For the documentation see: :class:`~airflow.providers.google.cloud.operators.vision.CloudVisionCreateProductOperator` """ client = self.get_conn() parent = ProductSearchClient.location_path(project_id, location) self.log.info('Creating a new Product under the parent: %s', parent) response = client.create_product( parent=parent, product=product, product_id=product_id, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('Product created: %s', response.name if response else '') self.log.debug('Product created:\n%s', response) if not product_id: # Product id was generated by the API product_id = self._get_autogenerated_id(response) self.log.info( 'Extracted autogenerated Product ID from the response: %s', product_id) return product_id
def delete_reference_image( self, location: str, product_id: str, reference_image_id: str, project_id: str, retry: Union[Retry, _MethodDefault] = DEFAULT, timeout: Optional[float] = None, metadata: Sequence[Tuple[str, str]] = (), ) -> dict: """ For the documentation see: :py:class:`~airflow.providers.google.cloud.operators.vision.CloudVisionDeleteReferenceImageOperator` """ 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) response = client.delete_reference_image( name=name, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('ReferenceImage with the name [%s] deleted.', name) return MessageToDict(response)
def test_create_product_autogenerated_id_wrong_name_in_response( self, get_conn): # Given wrong_name = 'wrong_name_not_a_correct_path' response_product = Product(name=wrong_name) create_product_method = get_conn.return_value.create_product create_product_method.return_value = response_product parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) hook = self.vision_hook_default_project_id product = Product() # When with self.assertRaises(AirflowException) as cm: hook.create_product(location=LOC_ID_TEST, product_id=None, product=product, project_id=PROJECT_ID_TEST) # Then # API response was wrong (wrong name format) and thus ProductSet ID extraction should fail. err = cm.exception self.assertIn('Unable to get id from name', str(err)) create_product_method.assert_called_once_with(parent=parent, product=product, product_id=None, retry=None, timeout=None, metadata=None)
def test_update_productset_explicit_name_missing_params_for_constructed_name( self, location, product_set_id, get_conn): # Given explicit_ps_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCTSET_ID_TEST_2) product_set = ProductSet(name=explicit_ps_name) update_product_set_method = get_conn.return_value.update_product_set update_product_set_method.return_value = product_set # When result = self.hook.update_product_set( location=location, product_set_id=product_set_id, product_set=product_set, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) # Then self.assertEqual(result, MessageToDict(product_set)) update_product_set_method.assert_called_once_with( product_set=ProductSet(name=explicit_ps_name), metadata=None, retry=None, timeout=None, update_mask=None, )
def test_update_product_explicit_name_missing_params_for_constructed_name( self, location, product_id, get_conn ): # Given explicit_p_name = ProductSearchClient.product_path( PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCT_ID_TEST_2 ) product = Product(name=explicit_p_name) update_product_method = get_conn.return_value.update_product update_product_method.return_value = product # When result = self.hook.update_product( location=location, product_id=product_id, product=product, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) # Then self.assertEqual(result, MessageToDict(product)) update_product_method.assert_called_once_with( product=Product(name=explicit_p_name), metadata=None, retry=None, timeout=None, update_mask=None )
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) response = client.delete_reference_image(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ReferenceImage with the name [%s] deleted.', name) return MessageToDict(response)
def create_product( self, location, product, project_id=None, product_id=None, retry=None, timeout=None, metadata=None ): """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductCreateOperator` """ client = self.get_conn() parent = ProductSearchClient.location_path(project_id, location) self.log.info('Creating a new Product under the parent: %s', parent) response = self._handle_request( lambda **kwargs: client.create_product(**kwargs), parent=parent, product=product, product_id=product_id, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('Product created: %s', response.name if response else '') self.log.debug('Product created:\n%s', response) if not product_id: # Product id was generated by the API product_id = self._get_autogenerated_id(response) self.log.info('Extracted autogenerated Product ID from the response: %s', product_id) return product_id
def test_create_productset_autogenerated_id_wrong_api_response(self, get_conn): # Given response_product_set = None create_product_set_method = get_conn.return_value.create_product_set create_product_set_method.return_value = response_product_set parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) product_set = ProductSet() # When with self.assertRaises(AirflowException) as cm: self.hook.create_product_set( location=LOC_ID_TEST, product_set_id=None, product_set=product_set, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) # Then # API response was wrong (None) and thus ProductSet ID extraction should fail. err = cm.exception self.assertIn('Unable to get name from response...', str(err)) create_product_set_method.assert_called_once_with( parent=parent, product_set=product_set, product_set_id=None, retry=None, timeout=None, metadata=None, )
def create_product_set( self, location: str, product_set: Union[dict, ProductSet], project_id: str = None, product_set_id: str = None, retry: Retry = None, timeout: float = None, metadata: Sequence[Tuple[str, str]] = None, ) -> str: """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetCreateOperator` """ assert project_id is not None client = self.get_conn() parent = ProductSearchClient.location_path(project_id, location) self.log.info('Creating a new ProductSet under the parent: %s', parent) response = client.create_product_set( parent=parent, product_set=product_set, product_set_id=product_set_id, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('ProductSet created: %s', response.name if response else '') self.log.debug('ProductSet created:\n%s', response) if not product_set_id: # Product set id was generated by the API product_set_id = self._get_autogenerated_id(response) self.log.info('Extracted autogenerated ProductSet ID from the response: %s', product_set_id) return product_set_id
def test_update_productset_no_explicit_name(self, get_conn): # Given product_set = ProductSet() update_product_set_method = get_conn.return_value.update_product_set update_product_set_method.return_value = product_set productset_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST, LOC_ID_TEST, PRODUCTSET_ID_TEST ) # When result = self.hook.update_product_set( location=LOC_ID_TEST, product_set_id=PRODUCTSET_ID_TEST, product_set=product_set, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) # Then self.assertEqual(result, MessageToDict(product_set)) update_product_set_method.assert_called_once_with( product_set=ProductSet(name=productset_name), metadata=None, retry=None, timeout=None, update_mask=None, )
def get_conn(self) -> ProductSearchClient: """ Retrieves connection to Cloud Vision. :return: Google Cloud Vision client object. :rtype: google.cloud.vision_v1.ProductSearchClient """ if not self._client: self._client = ProductSearchClient(credentials=self._get_credentials(), client_info=CLIENT_INFO) return self._client
def delete_product(self, location, product_id, project_id=None, retry=None, timeout=None, metadata=None): """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductDeleteOperator` """ client = self.get_conn() name = ProductSearchClient.product_path(project_id, location, product_id) self.log.info('Deleting ProductSet: %s', name) client.delete_product(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('Product with the name [%s] deleted:', name)
def get_product(self, location, product_id, project_id=None, retry=None, timeout=None, metadata=None): """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductGetOperator` """ client = self.get_conn() name = ProductSearchClient.product_path(project_id, location, product_id) self.log.info('Retrieving Product: %s', name) response = client.get_product(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('Product retrieved.') self.log.debug('Product retrieved:\n%s', response) return MessageToDict(response)
def test_create_product_autogenerated_id(self, get_conn): # Given autogenerated_id = 'autogen-p-id' response_product = Product( name=ProductSearchClient.product_path(PROJECT_ID_TEST, LOC_ID_TEST, autogenerated_id) ) create_product_method = get_conn.return_value.create_product create_product_method.return_value = response_product parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) product = Product() # When result = self.hook.create_product( location=LOC_ID_TEST, product_id=None, product=product, project_id=PROJECT_ID_TEST ) # Then # Product ID was not provided in the method call above. Should be extracted from the API response # and returned. self.assertEqual(result, autogenerated_id) create_product_method.assert_called_once_with( parent=parent, product=product, product_id=None, retry=None, timeout=None, metadata=None )
def test_delete_product(self, get_conn): # Given delete_product_method = get_conn.return_value.delete_product delete_product_method.return_value = None name = ProductSearchClient.product_path(PROJECT_ID_TEST, LOC_ID_TEST, PRODUCT_ID_TEST) # When response = self.hook.delete_product( location=LOC_ID_TEST, product_id=PRODUCT_ID_TEST, project_id=PROJECT_ID_TEST ) # Then self.assertIsNone(response) delete_product_method.assert_called_once_with(name=name, retry=None, timeout=None, metadata=None)
def test_update_productset_explicit_name_different_from_constructed(self, get_conn): # Given update_product_set_method = get_conn.return_value.update_product_set update_product_set_method.return_value = None explicit_ps_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCTSET_ID_TEST_2 ) product_set = ProductSet(name=explicit_ps_name) template_ps_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST, LOC_ID_TEST, PRODUCTSET_ID_TEST ) # When # Location and product_set_id are passed in addition to a ProductSet with an explicit name, # but both names differ (constructed != explicit). # Should throw AirflowException in this case. with pytest.raises(AirflowException) as ctx: self.hook.update_product_set( location=LOC_ID_TEST, product_set_id=PRODUCTSET_ID_TEST, product_set=product_set, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) err = ctx.value # self.assertIn("The required parameter 'project_id' is missing", str(err)) assert err assert ( ERR_DIFF_NAMES.format( explicit_name=explicit_ps_name, constructed_name=template_ps_name, label="ProductSet", id_label="productset_id", ) in str(err) ) update_product_set_method.assert_not_called()
def test_update_productset_explicit_name_different_from_constructed( self, get_conn): # Given update_product_set_method = get_conn.return_value.update_product_set update_product_set_method.return_value = None hook = self.vision_hook_default_project_id explicit_ps_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCTSET_ID_TEST_2) product_set = ProductSet(name=explicit_ps_name) template_ps_name = ProductSearchClient.product_set_path( PROJECT_ID_TEST, LOC_ID_TEST, PRODUCTSET_ID_TEST) # When # Location and product_set_id are passed in addition to a ProductSet with an explicit name, # but both names differ (constructed != explicit). # Should throw AirflowException in this case. with self.assertRaises(AirflowException) as cm: hook.update_product_set( location=LOC_ID_TEST, product_set_id=PRODUCTSET_ID_TEST, product_set=product_set, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) err = cm.exception # self.assertIn("The required parameter 'project_id' is missing", str(err)) self.assertTrue(err) self.assertIn( "The ProductSet name provided in the object ({}) is different than the name " "created from the input parameters ({}). Please either: 1) Remove the ProductSet " "name, 2) Remove the location and productset_id parameters, 3) Unify the " "ProductSet name and input parameters.".format( explicit_ps_name, template_ps_name), str(err), ) update_product_set_method.assert_not_called()
def test_get_productset(self, get_conn): # Given name = ProductSearchClient.product_set_path(PROJECT_ID_TEST, LOC_ID_TEST, PRODUCTSET_ID_TEST) response_product_set = ProductSet(name=name) get_product_set_method = get_conn.return_value.get_product_set get_product_set_method.return_value = response_product_set # When response = self.hook.get_product_set( location=LOC_ID_TEST, product_set_id=PRODUCTSET_ID_TEST, project_id=PROJECT_ID_TEST ) # Then self.assertTrue(response) self.assertEqual(response, MessageToDict(response_product_set)) get_product_set_method.assert_called_once_with(name=name, retry=None, timeout=None, metadata=None)
def test_update_product_explicit_name_different_from_constructed( self, get_conn): # Given update_product_method = get_conn.return_value.update_product update_product_method.return_value = None explicit_p_name = ProductSearchClient.product_path( PROJECT_ID_TEST_2, LOC_ID_TEST_2, PRODUCT_ID_TEST_2) product = Product(name=explicit_p_name) template_p_name = ProductSearchClient.product_path( PROJECT_ID_TEST, LOC_ID_TEST, PRODUCT_ID_TEST) # When # Location and product_id are passed in addition to a Product with an explicit name, # but both names differ (constructed != explicit). # Should throw AirflowException in this case. with self.assertRaises(AirflowException) as cm: self.hook.update_product( location=LOC_ID_TEST, product_id=PRODUCT_ID_TEST, product=product, update_mask=None, project_id=PROJECT_ID_TEST, retry=None, timeout=None, metadata=None, ) err = cm.exception self.assertTrue(err) self.assertIn( ERR_DIFF_NAMES.format( explicit_name=explicit_p_name, constructed_name=template_p_name, label="Product", id_label="product_id", ), str(err), ) update_product_method.assert_not_called()
def delete_product( self, location: str, product_id: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> None: """ For the documentation see: :class:`~airflow.providers.google.cloud.operators.vision.CloudVisionDeleteProductOperator` """ client = self.get_conn() name = ProductSearchClient.product_path(project_id, location, product_id) self.log.info('Deleting ProductSet: %s', name) client.delete_product(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('Product with the name [%s] deleted:', name)
def delete_product_set( self, location: str, product_set_id: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None ): """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetDeleteOperator` """ client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Deleting ProductSet: %s', name) client.delete_product_set(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ProductSet with the name [%s] deleted.', name)
def delete_product_set( self, location: str, product_set_id: str, project_id: str = PROVIDE_PROJECT_ID, retry: Union[Retry, _MethodDefault] = DEFAULT, timeout: Optional[float] = None, metadata: Sequence[Tuple[str, str]] = (), ) -> None: """ For the documentation see: :class:`~airflow.providers.google.cloud.operators.vision.CloudVisionDeleteProductSetOperator` """ client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Deleting ProductSet: %s', name) client.delete_product_set(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ProductSet with the name [%s] deleted.', name)
def create_reference_image( self, location: str, product_id: str, reference_image: Union[Dict, ReferenceImage], reference_image_id: Optional[str] = None, project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, ) -> str: """ For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionReferenceImageCreateOperator` """ if not project_id: raise ValueError("Project ID should be set.") 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: # Reference 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', reference_image_id) return reference_image_id
def get_product_set( self, location: str, product_set_id: str, project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None ) -> Dict: """ For the documentation see: :class:`~airflow.providers.google.cloud.operators.vision.CloudVisionGetProductSetOperator` """ client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Retrieving ProductSet: %s', name) response = client.get_product_set(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ProductSet retrieved.') self.log.debug('ProductSet retrieved:\n%s', response) return MessageToDict(response)
def delete_product_set( self, location, product_set_id, project_id=None, retry=None, timeout=None, metadata=None ): """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetDeleteOperator` """ client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Deleting ProductSet: %s', name) response = self._handle_request( lambda **kwargs: client.delete_product_set(**kwargs), name=name, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('ProductSet with the name [%s] deleted.', name) return response
def create_reference_image( self, location: str, product_id: str, reference_image: Union[dict, ReferenceImage], project_id: str, reference_image_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Sequence[Tuple[str, str]] = (), ) -> str: """ For the documentation see: :py:class:`~airflow.providers.google.cloud.operators.vision.CloudVisionCreateReferenceImageOperator` """ 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: # Reference 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', reference_image_id) return reference_image_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', reference_image_id) return reference_image_id
def test_create_product_autogenerated_id_wrong_name_in_response(self, get_conn): # Given wrong_name = 'wrong_name_not_a_correct_path' response_product = Product(name=wrong_name) create_product_method = get_conn.return_value.create_product create_product_method.return_value = response_product parent = ProductSearchClient.location_path(PROJECT_ID_TEST, LOC_ID_TEST) product = Product() # When with self.assertRaises(AirflowException) as cm: self.hook.create_product( location=LOC_ID_TEST, product_id=None, product=product, project_id=PROJECT_ID_TEST ) # Then # API response was wrong (wrong name format) and thus ProductSet ID extraction should fail. err = cm.exception self.assertIn('Unable to get id from name', str(err)) create_product_method.assert_called_once_with( parent=parent, product=product, product_id=None, retry=None, timeout=None, metadata=None )
def get_product_set( self, location: str, product_set_id: str, project_id: str = None, retry: Retry = None, timeout: float = None, metadata: Sequence[Tuple[str, str]] = None ) -> Dict: """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetGetOperator` """ assert project_id is not None client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Retrieving ProductSet: %s', name) response = client.get_product_set(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ProductSet retrieved.') self.log.debug('ProductSet retrieved:\n%s', response) return MessageToDict(response)
def get_product_set( self, location: str, product_set_id: str, project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None ) -> Dict: """ For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetGetOperator` """ if not project_id: raise ValueError("Project ID should be set.") client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Retrieving ProductSet: %s', name) response = client.get_product_set(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ProductSet retrieved.') self.log.debug('ProductSet retrieved:\n%s', response) return MessageToDict(response)
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', reference_image_id ) return reference_image_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 ) response = client.delete_reference_image(name=name, retry=retry, timeout=timeout, metadata=metadata) self.log.info('ReferenceImage with the name [%s] deleted.', name) return MessageToDict(response)
def delete_product_set(self, location, product_set_id, project_id=None, retry=None, timeout=None, metadata=None): """ For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetDeleteOperator` """ client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Deleting ProductSet: %s', name) response = self._handle_request( lambda **kwargs: client.delete_product_set(**kwargs), name=name, retry=retry, timeout=timeout, metadata=metadata, ) self.log.info('ProductSet with the name [%s] deleted.', name) return response
def _get_entity_name(is_product, project_id, location, entity_id): if is_product: return ProductSearchClient.product_path(project_id, location, entity_id) else: return ProductSearchClient.product_set_path(project_id, location, entity_id)