Python gcp_vision_hook.CloudVisionHook类代码示例

 def execute(self, context):
     hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)
     return hook.safe_search_detection(
         image=self.image,
         max_results=self.max_results,
         retry=self.retry,
         timeout=self.timeout,
         additional_properties=self.additional_properties,
     )

 def __init__(
     self,
     product,
     location=None,
     product_id=None,
     project_id=None,
     update_mask=None,
     retry=None,
     timeout=None,
     metadata=None,
     gcp_conn_id='google_cloud_default',
     *args,
     **kwargs
 ):
     super(CloudVisionProductUpdateOperator, self).__init__(*args, **kwargs)
     self.product = product
     self.location = location
     self.product_id = product_id
     self.project_id = project_id
     self.update_mask = update_mask
     self.retry = retry
     self.timeout = timeout
     self.metadata = metadata
     self.gcp_conn_id = gcp_conn_id
     self._hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)

 def __init__(
     self,
     location,
     product_id,
     project_id=None,
     retry=None,
     timeout=None,
     metadata=None,
     gcp_conn_id="google_cloud_default",
     *args,
     **kwargs
 ):
     super(CloudVisionProductGetOperator, self).__init__(*args, **kwargs)
     self.location = location
     self.product_id = product_id
     self.project_id = project_id
     self.retry = retry
     self.timeout = timeout
     self.metadata = metadata
     self.gcp_conn_id = gcp_conn_id
     self._hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)

 def setUp(self):
     with mock.patch(
         'airflow.contrib.hooks.gcp_vision_hook.CloudVisionHook.__init__',
         new=mock_base_gcp_hook_default_project_id,
     ):
         self.hook = CloudVisionHook(gcp_conn_id='test')

class TestGcpVisionHook(unittest.TestCase):
    def setUp(self):
        with mock.patch(
            'airflow.contrib.hooks.gcp_vision_hook.CloudVisionHook.__init__',
            new=mock_base_gcp_hook_default_project_id,
        ):
            self.hook = CloudVisionHook(gcp_conn_id='test')

    @mock.patch('airflow.contrib.hooks.gcp_vision_hook.CloudVisionHook.get_conn')
    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,
        )

    @mock.patch('airflow.contrib.hooks.gcp_vision_hook.CloudVisionHook.get_conn')
    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,
        )

    @mock.patch('airflow.contrib.hooks.gcp_vision_hook.CloudVisionHook.get_conn')
    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,
        )

    @mock.patch('airflow.contrib.hooks.gcp_vision_hook.CloudVisionHook.get_conn')
    def test_get_productset(self, get_conn):
#.........这里部分代码省略.........

class CloudVisionProductDeleteOperator(BaseOperator):
    """
    Permanently deletes a product and its reference images.

    Metadata of the product and all its images will be deleted right away, but search queries against
    ProductSets containing the product may still work until all related caches are refreshed.

    Possible errors:

    - Returns NOT_FOUND if the product does not exist.

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:CloudVisionProductDeleteOperator`

    :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 product_id: (Required) 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: The connection ID used to connect to Google Cloud Platform.
    :type gcp_conn_id: str
    """

    # [START vision_product_delete_template_fields]
    template_fields = ('location', 'project_id', 'product_id', 'gcp_conn_id')
    # [END vision_product_delete_template_fields]

    @apply_defaults
    def __init__(
        self,
        location,
        product_id,
        project_id=None,
        retry=None,
        timeout=None,
        metadata=None,
        gcp_conn_id='google_cloud_default',
        *args,
        **kwargs
    ):
        super(CloudVisionProductDeleteOperator, self).__init__(*args, **kwargs)
        self.location = location
        self.product_id = product_id
        self.project_id = project_id
        self.retry = retry
        self.timeout = timeout
        self.metadata = metadata
        self.gcp_conn_id = gcp_conn_id
        self._hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)

    def execute(self, context):
        return self._hook.delete_product(
            location=self.location,
            product_id=self.product_id,
            project_id=self.project_id,
            retry=self.retry,
            timeout=self.timeout,
            metadata=self.metadata,
        )

class CloudVisionProductUpdateOperator(BaseOperator):
    """
    Makes changes to a Product resource. Only the display_name, description, and labels fields can be
    updated right now.

    If labels are updated, the change will not be reflected in queries until the next index time.

    .. note:: To locate the `Product` resource, its `name` in the form
        `projects/PROJECT_ID/locations/LOC_ID/products/PRODUCT_ID` is necessary.

    You can provide the `name` directly as an attribute of the `product` object. However, you can leave it
    blank and provide `location` and `product_id` instead (and optionally `project_id` - if not present,
    the connection default will be used) and the `name` will be created by the operator itself.

    This mechanism exists for your convenience, to allow leaving the `project_id` empty and having Airflow
    use the connection default `project_id`.

    Possible errors related to the provided `Product`:

    - Returns NOT_FOUND if the Product does not exist.
    - Returns INVALID_ARGUMENT if display_name is present in update_mask but is missing from the request or
        longer than 4096 characters.
    - Returns INVALID_ARGUMENT if description is present in update_mask but is longer than 4096 characters.
    - Returns INVALID_ARGUMENT if product_category is present in update_mask.

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:CloudVisionProductUpdateOperator`

    :param product: (Required) The Product resource which replaces the one on the server. product.name is
        immutable. If a dict is provided, it must be of the same form as the protobuf message `Product`.
    :type product: dict or google.cloud.vision_v1.types.ProductSet
    :param location: (Optional) 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 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 update_mask: (Optional) The `FieldMask` that specifies which fields to update. If update_mask
        isn’t specified, all mutable fields are to be updated. Valid mask paths include product_labels,
        display_name, and description. If a dict is provided, it must be of the same form as the protobuf
        message `FieldMask`.
    :type update_mask: dict or google.cloud.vision_v1.types.FieldMask
    :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: The connection ID used to connect to Google Cloud Platform.
    :type gcp_conn_id: str
    """

    # [START vision_product_update_template_fields]
    template_fields = ('location', 'project_id', 'product_id', 'gcp_conn_id')
    # [END vision_product_update_template_fields]

    @apply_defaults
    def __init__(
        self,
        product,
        location=None,
        product_id=None,
        project_id=None,
        update_mask=None,
        retry=None,
        timeout=None,
        metadata=None,
        gcp_conn_id='google_cloud_default',
        *args,
        **kwargs
    ):
        super(CloudVisionProductUpdateOperator, self).__init__(*args, **kwargs)
        self.product = product
        self.location = location
        self.product_id = product_id
        self.project_id = project_id
        self.update_mask = update_mask
        self.retry = retry
        self.timeout = timeout
        self.metadata = metadata
        self.gcp_conn_id = gcp_conn_id
        self._hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)

    def execute(self, context):
        return self._hook.update_product(
            product=self.product,
            location=self.location,
            product_id=self.product_id,
            project_id=self.project_id,
            update_mask=self.update_mask,
            retry=self.retry,
            timeout=self.timeout,
            metadata=self.metadata,
        )

class CloudVisionProductCreateOperator(BaseOperator):
    """
    Creates and returns a new product resource.

    Possible errors regarding the `Product` object provided:

    - Returns INVALID_ARGUMENT if `display_name` is missing or longer than 4096 characters.
    - Returns INVALID_ARGUMENT if `description` is longer than 4096 characters.
    - Returns INVALID_ARGUMENT if `product_category` is missing or invalid.

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:CloudVisionProductCreateOperator`

    :param location: (Required) The region where the Product should be created. Valid regions
        (as of 2019-02-05) are: us-east1, us-west1, europe-west1, asia-east1
    :type location: str
    :param product: (Required) The product to create. If a dict is provided, it must be of the same form as
        the protobuf message `Product`.
    :type product: dict or google.cloud.vision_v1.types.Product
    :param project_id: (Optional) The project in which the Product should be created. If set to None or
        missing, the default project_id from the GCP connection is used.
    :type project_id: str
    :param product_id: (Optional) A user-supplied resource id for this Product.
        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 product_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: The connection ID used to connect to Google Cloud Platform.
    :type gcp_conn_id: str

    """

    # [START vision_product_create_template_fields]
    template_fields = ('location', 'project_id', 'product_id', 'gcp_conn_id')
    # [END vision_product_create_template_fields]

    @apply_defaults
    def __init__(
        self,
        location,
        product,
        project_id=None,
        product_id=None,
        retry=None,
        timeout=None,
        metadata=None,
        gcp_conn_id='google_cloud_default',
        *args,
        **kwargs
    ):
        super(CloudVisionProductCreateOperator, self).__init__(*args, **kwargs)
        self.location = location
        self.product = product
        self.project_id = project_id
        self.product_id = product_id
        self.retry = retry
        self.timeout = timeout
        self.metadata = metadata
        self.gcp_conn_id = gcp_conn_id
        self._hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)

    def execute(self, context):
        try:
            return self._hook.create_product(
                location=self.location,
                product=self.product,
                project_id=self.project_id,
                product_id=self.product_id,
                retry=self.retry,
                timeout=self.timeout,
                metadata=self.metadata,
            )
        except AlreadyExists:
            self.log.info(
                'Product with id %s already exists. Exiting from the create operation.', self.product_id
            )
            return self.product_id

class CloudVisionProductSetGetOperator(BaseOperator):
    """
    Gets information associated with a ProductSet.

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:CloudVisionProductSetGetOperator`

    :param location: (Required) The region where the ProductSet is located. Valid regions (as of 2019-02-05)
        are: us-east1, us-west1, europe-west1, asia-east1
    :type location: str
    :param product_set_id: (Required) The resource id of this ProductSet.
    :type product_set_id: str
    :param project_id: (Optional) The project in which the ProductSet 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: The connection ID used to connect to Google Cloud Platform.
    :type gcp_conn_id: str
    """

    # [START vision_productset_get_template_fields]
    template_fields = ('location', 'project_id', 'product_set_id', 'gcp_conn_id')
    # [END vision_productset_get_template_fields]

    @apply_defaults
    def __init__(
        self,
        location,
        product_set_id,
        project_id=None,
        retry=None,
        timeout=None,
        metadata=None,
        gcp_conn_id='google_cloud_default',
        *args,
        **kwargs
    ):
        super(CloudVisionProductSetGetOperator, self).__init__(*args, **kwargs)
        self.location = location
        self.project_id = project_id
        self.product_set_id = product_set_id
        self.retry = retry
        self.timeout = timeout
        self.metadata = metadata
        self.gcp_conn_id = gcp_conn_id
        self._hook = CloudVisionHook(gcp_conn_id=self.gcp_conn_id)

    def execute(self, context):
        return self._hook.get_product_set(
            location=self.location,
            product_set_id=self.product_set_id,
            project_id=self.project_id,
            retry=self.retry,
            timeout=self.timeout,
            metadata=self.metadata,
        )