fix: call DLP API in the region of DataFlow job (GoogleCloudPlatform#9144)

minherz · web-flow · commit 5a3675e48c2a · 2023-02-21T23:28:53.000Z
* fix: call DLP API in the region of DataFlow job

modify final code to call DLP API for the region where DataFlow job is deployed.
add placeholder for this modification to boilerplate code.
populate README with information about the folder and link to Neos tutorial.
diff --git a/logging/redaction/README.md b/logging/redaction/README.md
@@ -1,16 +1,25 @@
 # Logging redaction tutorial code samples
 
-> **Warning**
-> This section is still **W**ork **I**n **P**rogress.
-> Cloud Shell button now opens this README. It will open the tutorial _AFTER_ its official launch.
-> Tests to validate the code samples will be added.
+This is a sample of the [DataFlow] pipeline that detects and masks US Social Security Numbers from the log payloads at ingestion time.
+The folder contains the following files:
 
-This section contains code that is used in the "Redact confidential information in logs" tutorial.
-You can open the tutorial in Cloud Shell:
+* [Dockerfile] to use for a customized a DataFlow job container in order to save initialization time
+* [Boilerplate code][boilerplate] that implements a pipeline for streaming log entries from [PubSub] to a destination Log bucket
+* [Final version][final] of the pipeline that includes all modifications to the [boilerplate] code that are required to implement log redaction
+* [Requirements] file to be install the DataFlow job's environment with missing component(s)
+
+If you have a Google Cloud account and an access to a GCP project you can launch an interactive tutorial in Cloud Console and see how the sample works.
+To run the tutorial press the button below.
 
 [![Open in Cloud Shell][shell_img]][shell_link]
 
-_NOTE:_ You will need a valid Google Cloud credentials to open the tutorial.
+NOTE: To run this tutorial you will need to have permissions to enable Google APIs, provision PubSub, DataFlow and Cloud Storage resources as well as permissions to call DLP API
 
+[dataflow]: https://cloud.google.com/dataflow
+[pubsub]: https://cloud.google.com/pubsub
+[dockerfile]: https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/logging/redaction/Dockerfile
+[boilerplate]: https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/logging/redaction/log_redaction.py
+[final]: https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/logging/redaction/log_redaction_final.py
+[requirements]: https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/logging/redaction/requirements.txt
 [shell_img]: http://gstatic.com/cloudssh/images/open-btn.png
-[shell_link]: https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/GoogleCloudPlatform/python-docs-samples&page=editor&open_in_editor=logging/redaction/README.md
+[shell_link]: https://console.cloud.google.com/?walkthrough_id=cloud-ops-log-redacting-on-ingestion
diff --git a/logging/redaction/log_redaction.py b/logging/redaction/log_redaction.py
@@ -105,6 +105,9 @@ def run(
         streaming=True,
         save_main_session=True
     )
+
+    # TODO: Read job's deployment region
+
     pipeline = Pipeline(options=pipeline_options)
     _ = (
         pipeline
diff --git a/logging/redaction/log_redaction_final.py b/logging/redaction/log_redaction_final.py
@@ -20,7 +20,7 @@
 
 from apache_beam import CombineFn, CombineGlobally, DoFn, io, ParDo, Pipeline, WindowInto
 from apache_beam.error import PipelineError
-from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import GoogleCloudOptions, PipelineOptions
 from apache_beam.transforms.window import FixedWindows
 
 from google.cloud import dlp_v2, logging_v2
@@ -82,8 +82,9 @@ def extract_output(self, accumulator):
 class LogRedaction(DoFn):
     '''Apply inspection and redaction to textPayload field of log entries'''
 
-    def __init__(self, project_id: str):
+    def __init__(self, region, project_id: str):
         self.project_id = project_id
+        self.region = region
         self.dlp_client = None
 
     def _log_to_row(self, entry):
@@ -113,7 +114,7 @@ def process(self, logs):
 
         response = self.dlp_client.deidentify_content(
             request={
-                'parent': f'projects/{self.project_id}',
+                'parent': f'projects/{self.project_id}/locations/{self.region}',
                 'inspect_config': INSPECT_CFG,
                 'deidentify_config': REDACTION_CFG,
                 'item': table,
@@ -177,6 +178,13 @@ def run(
         streaming=True,
         save_main_session=True
     )
+
+    region = "us-central1"
+    try:
+        region = pipeline_options.view_as(GoogleCloudOptions).region
+    except AttributeError:
+        pass
+
     pipeline = Pipeline(options=pipeline_options)
     _ = (
         pipeline
@@ -186,7 +194,7 @@ def run(
         # Optimize Google API consumption and avoid possible throttling
         # by calling APIs for batched data and not per each element
         | 'Batch aggregated payloads' >> CombineGlobally(BatchPayloads()).without_defaults()
-        | 'Redact SSN info from logs' >> ParDo(LogRedaction(destination_log_name.split('/')[1]))
+        | 'Redact SSN info from logs' >> ParDo(LogRedaction(region, destination_log_name.split('/')[1]))
         | 'Ingest to output log' >> ParDo(IngestLogs(destination_log_name))
     )
     pipeline.run()
