[BEAM-13388] Cherry-pick more changes for google cloud dlp update. (#16268)

Co-authored-by: Brian Hulette <hulettbh@gmail.com>
Co-authored-by: Yichi Zhang <zyichi@google.com>
Co-authored-by: Kyle Weaver <kcweaver@google.com>
Co-authored-by: Andy Ye <andyye333@gmail.com>

Author: 5 people

sdks/python/apache_beam/io/gcp/pubsub_integration_test.py

@@ -210,6 +210,8 @@ def test_streaming_data_only(self):
 
   @pytest.mark.it_postcommit
   def test_streaming_with_attributes(self):
+    if self.runner_name == 'TestDataflowRunner':
+      pytest.skip("BEAM-13218")
     self._test_streaming(with_attributes=True)
 
 

sdks/python/apache_beam/ml/gcp/cloud_dlp.py

@@ -20,9 +20,11 @@
 """
 
 import logging
+from typing import List
 
 from google.cloud import dlp_v2
 
+from apache_beam import typehints
 from apache_beam.options.pipeline_options import GoogleCloudOptions
 from apache_beam.transforms import DoFn
 from apache_beam.transforms import ParDo
@@ -35,6 +37,8 @@
 
 
 @experimental()
+@typehints.with_input_types(str)
+@typehints.with_output_types(str)
 class MaskDetectedDetails(PTransform):
   """Scrubs sensitive information detected in text.
   The ``PTransform`` returns a ``PCollection`` of ``str``
@@ -126,6 +130,8 @@ def expand(self, pcoll):
 
 
 @experimental()
+@typehints.with_input_types(str)
+@typehints.with_output_types(List[dlp_v2.types.dlp.Finding])
 class InspectForDetails(PTransform):
   """Inspects input text for sensitive information.
   the ``PTransform`` returns a ``PCollection`` of
@@ -190,13 +196,13 @@ def setup(self):
       self.client = dlp_v2.DlpServiceClient()
     self.params = {
         'timeout': self.timeout,
-        'parent': self.client.project_path(self.project)
     }
-    self.params.update(self.config)
+    self.parent = self.client.common_project_path(self.project)
 
   def process(self, element, **kwargs):
-    operation = self.client.deidentify_content(
-        item={"value": element}, **self.params)
+    request = {'item': {'value': element}, 'parent': self.parent}
+    request.update(self.config)
+    operation = self.client.deidentify_content(request=request, **self.params)
     yield operation.item.value
 
 
@@ -213,12 +219,12 @@ def setup(self):
       self.client = dlp_v2.DlpServiceClient()
     self.params = {
         'timeout': self.timeout,
-        "parent": self.client.project_path(self.project)
     }
-    self.params.update(self.config)
+    self.parent = self.client.common_project_path(self.project)
 
   def process(self, element, **kwargs):
-    operation = self.client.inspect_content(
-        item={"value": element}, **self.params)
+    request = {'item': {'value': element}, 'parent': self.parent}
+    request.update(self.config)
+    operation = self.client.inspect_content(request=request, **self.params)
     hits = [x for x in operation.result.findings]
     yield hits

sdks/python/apache_beam/ml/gcp/cloud_dlp_test.py

@@ -37,6 +37,7 @@
   from apache_beam.ml.gcp.cloud_dlp import MaskDetectedDetails
   from apache_beam.ml.gcp.cloud_dlp import _DeidentifyFn
   from apache_beam.ml.gcp.cloud_dlp import _InspectFn
+  from google.cloud.dlp_v2.types import dlp
 # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
 
 _LOGGER = logging.getLogger(__name__)
@@ -56,6 +57,9 @@ class TestDeidentifyFn(unittest.TestCase):
   def test_deidentify_called(self):
     class ClientMock(object):
       def deidentify_content(self, *args, **kwargs):
+        # Check that we can marshal a valid request.
+        dlp.DeidentifyContentRequest(kwargs['request'])
+
         called = Metrics.counter('test_deidentify_text', 'called')
         called.inc()
         operation = mock.Mock()
@@ -64,27 +68,29 @@ def deidentify_content(self, *args, **kwargs):
         operation.item = item
         return operation
 
-      def project_path(self, *args):
+      def common_project_path(self, *args):
         return 'test'
 
     with mock.patch('google.cloud.dlp_v2.DlpServiceClient', ClientMock):
       p = TestPipeline()
-      deidentify_config = {
-          "info_type_transformations": {
-              "transformations": [{
-                  "primitive_transformation": {
-                      "character_mask_config": {
-                          "masking_character": '#'
+      config = {
+          "deidentify_config": {
+              "info_type_transformations": {
+                  "transformations": [{
+                      "primitive_transformation": {
+                          "character_mask_config": {
+                              "masking_character": '#'
+                          }
                       }
-                  }
-              }]
+                  }]
+              }
           }
       }
       # pylint: disable=expression-not-assigned
       (
           p
           | beam.Create(['mary.sue@example.com', 'john.doe@example.com'])
-          | beam.ParDo(_DeidentifyFn(config=deidentify_config)))
+          | beam.ParDo(_DeidentifyFn(config=config)))
       result = p.run()
       result.wait_until_finish()
     called = result.metrics().query()['counters'][0]
@@ -101,28 +107,31 @@ def test_exception_raised_then_no_config_provided(self):
 
 
 @unittest.skipIf(dlp_v2 is None, 'GCP dependencies are not installed')
-class TestDeidentifyFn(unittest.TestCase):
+class TestInspectFn(unittest.TestCase):
   def test_inspect_called(self):
     class ClientMock(object):
       def inspect_content(self, *args, **kwargs):
+        # Check that we can marshal a valid request.
+        dlp.InspectContentRequest(kwargs['request'])
+
         called = Metrics.counter('test_inspect_text', 'called')
         called.inc()
         operation = mock.Mock()
         operation.result = mock.Mock()
         operation.result.findings = [None]
         return operation
 
-      def project_path(self, *args):
+      def common_project_path(self, *args):
         return 'test'
 
     with mock.patch('google.cloud.dlp_v2.DlpServiceClient', ClientMock):
       p = TestPipeline()
-      inspect_config = {"info_types": [{"name": "EMAIL_ADDRESS"}]}
+      config = {"inspect_config": {"info_types": [{"name": "EMAIL_ADDRESS"}]}}
       # pylint: disable=expression-not-assigned
       (
           p
           | beam.Create(['mary.sue@example.com', 'john.doe@example.com'])
-          | beam.ParDo(_InspectFn(config=inspect_config)))
+          | beam.ParDo(_InspectFn(config=config)))
       result = p.run()
       result.wait_until_finish()
       called = result.metrics().query()['counters'][0]