Allow users to specify custom audit entries in pipeline options. (#34062)

* Allow users to specify custom audit entries in pipeline options.

* Add sanity check on custom audit entries during pipeline option parsing.

* Fix lints and typos.

* Add docstring

* Change the flag to snake case.

Author: shunping

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

@@ -125,14 +125,21 @@ def create_storage_client(pipeline_options, use_credentials=True):
     from google.api_core import client_info
     beam_client_info = client_info.ClientInfo(
         user_agent="apache-beam/%s (GPN:Beam)" % beam_version.__version__)
+
+    extra_headers = {
+        "x-goog-custom-audit-job": google_cloud_options.job_name
+        if google_cloud_options.job_name else "UNKNOWN"
+    }
+
+    # Note: Custom audit entries with "job" key will overwrite the default above
+    if google_cloud_options.gcs_custom_audit_entries is not None:
+      extra_headers.update(google_cloud_options.gcs_custom_audit_entries)
+
     return storage.Client(
         credentials=credentials.get_google_auth_credentials(),
         project=google_cloud_options.project,
         client_info=beam_client_info,
-        extra_headers={
-            "x-goog-custom-audit-job": google_cloud_options.job_name
-            if google_cloud_options.job_name else "UNKNOWN"
-        })
+        extra_headers=extra_headers)
   else:
     return storage.Client.create_anonymous_client()
 

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

@@ -30,6 +30,7 @@
 from apache_beam.metrics.execution import MetricsContainer
 from apache_beam.metrics.execution import MetricsEnvironment
 from apache_beam.metrics.metricbase import MetricName
+from apache_beam.pipeline import PipelineOptions
 from apache_beam.runners.worker import statesampler
 from apache_beam.utils import counters
 
@@ -712,7 +713,13 @@ def test_headers(self, mock_get_service_credentials, mock_do_request):
     mock_get_service_credentials.return_value = _ApitoolsCredentialsAdapter(
         _make_credentials("test-project"))
 
-    gcs = gcsio.GcsIO(pipeline_options={"job_name": "test-job-name"})
+    options = PipelineOptions([
+        "--job_name=test-job-name",
+        "--gcs_custom_audit_entry=user=test-user-id",
+        "--gcs_custom_audit_entries={\"id\": \"1234\", \"status\": \"ok\"}"
+    ])
+
+    gcs = gcsio.GcsIO(pipeline_options=options)
     # no HTTP request when initializing GcsIO
     mock_do_request.assert_not_called()
 
@@ -732,6 +739,9 @@ def test_headers(self, mock_get_service_credentials, mock_do_request):
     beam_user_agent = "apache-beam/%s (GPN:Beam)" % beam_version.__version__
     self.assertIn(beam_user_agent, actual_headers['User-Agent'])
     self.assertEqual(actual_headers['x-goog-custom-audit-job'], 'test-job-name')
+    self.assertEqual(actual_headers['x-goog-custom-audit-user'], 'test-user-id')
+    self.assertEqual(actual_headers['x-goog-custom-audit-id'], '1234')
+    self.assertEqual(actual_headers['x-goog-custom-audit-status'], 'ok')
 
   @mock.patch('google.cloud._http.JSONConnection._do_request')
   @mock.patch('apache_beam.internal.gcp.auth.get_service_credentials')

sdks/python/apache_beam/options/pipeline_options.py

@@ -149,6 +149,71 @@ def __call__(self, parser, namespace, values, option_string=None):
     getattr(namespace, self.dest).update(values)
 
 
+class _GcsCustomAuditEntriesAction(argparse.Action):
+  """
+  Argparse Action for GCS custom audit entries.
+
+  According to Google Cloud Storage audit logging documentation
+  (https://cloud.google.com/storage/docs/audit-logging#add-custom-metadata), the
+  following limitations apply:
+  - keys must be 64 characters or less,
+  - values 1,200 characters or less, and
+  - a maximum of four custom metadata entries are permitted per request.
+  """
+  MAX_KEY_LENGTH = 64
+  MAX_VALUE_LENGTH = 1200
+  MAX_ENTRIES = 4
+
+  def _exceed_entry_limit(self):
+    if 'x-goog-custom-audit-job' in self._custom_audit_entries:
+      return len(
+          self._custom_audit_entries) > _GcsCustomAuditEntriesAction.MAX_ENTRIES
+    else:
+      return len(self._custom_audit_entries) > (
+          _GcsCustomAuditEntriesAction.MAX_ENTRIES - 1)
+
+  def _add_entry(self, key, value):
+    if len(key) > _GcsCustomAuditEntriesAction.MAX_KEY_LENGTH:
+      raise argparse.ArgumentError(
+          None,
+          "The key '%s' in GCS custom audit entries exceeds the %d-character limit."  # pylint: disable=line-too-long
+          % (key, _GcsCustomAuditEntriesAction.MAX_KEY_LENGTH))
+
+    if len(value) > _GcsCustomAuditEntriesAction.MAX_VALUE_LENGTH:
+      raise argparse.ArgumentError(
+          None,
+          "The value '%s' in GCS custom audit entries exceeds the %d-character limit."  # pylint: disable=line-too-long
+          % (value, _GcsCustomAuditEntriesAction.MAX_VALUE_LENGTH))
+
+    self._custom_audit_entries[f"x-goog-custom-audit-{key}"] = value
+
+  def __call__(self, parser, namespace, values, option_string=None):
+    if not hasattr(namespace,
+                   self.dest) or getattr(namespace, self.dest) is None:
+      setattr(namespace, self.dest, {})
+      self._custom_audit_entries = getattr(namespace, self.dest)
+
+    if option_string == '--gcs_custom_audit_entries':
+      # in the format of {"key": "value"}
+      assert (isinstance(values, str))
+      sub_entries = json.loads(values)
+      for key, value in sub_entries.items():
+        self._add_entry(key, value)
+    else:  # option_string == '--gcs_custom_audit_entry'
+      # in the format of 'key=value'
+      assert (isinstance(values, str))
+      parts = values.split('=', 1)
+      key = parts[0]
+      value = parts[1] if len(parts) > 1 else ''
+      self._add_entry(key, value)
+
+    if self._exceed_entry_limit():
+      raise argparse.ArgumentError(
+          None,
+          "The maximum allowed number of GCS custom audit entries (including the default x-goo-custom-audit-job) is %d."  # pylint: disable=line-too-long
+          % _GcsCustomAuditEntriesAction.MAX_ENTRIES)
+
+
 class PipelineOptions(HasDisplayData):
   """This class and subclasses are used as containers for command line options.
 
@@ -953,6 +1018,16 @@ def _add_argparse_args(cls, parser):
         action='store_true',
         help='Use blob generation when mutating blobs in GCSIO to '
         'mitigate race conditions at the cost of more HTTP requests.')
+    parser.add_argument(
+        '--gcs_custom_audit_entry',
+        '--gcs_custom_audit_entries',
+        dest='gcs_custom_audit_entries',
+        action=_GcsCustomAuditEntriesAction,
+        default=None,
+        help='Custom information to be attached to audit logs. '
+        'Entries are key value pairs separated by = '
+        '(e.g. --gcs_custom_audit_entry key=value) or a JSON string '
+        '(e.g. --gcs_custom_audit_entries=\'{ "user": "test", "id": "12" }\').')
 
   def _create_default_gcs_bucket(self):
     try:

sdks/python/apache_beam/options/pipeline_options_test.py

@@ -723,6 +723,67 @@ def test_options_store_false_with_different_dest(self):
         "the dest and the flag name to the map "
         "_FLAG_THAT_SETS_FALSE_VALUE in PipelineOptions.py")
 
+  def test_gcs_custom_audit_entries(self):
+    options = PipelineOptions([
+        '--gcs_custom_audit_entry=user=test-user',
+        '--gcs_custom_audit_entry=work=test-work',
+        '--gcs_custom_audit_entries={"job":"test-job", "id":"1234"}'
+    ])
+    entries = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries
+    self.assertDictEqual(
+        entries,
+        {
+            'x-goog-custom-audit-user': 'test-user',
+            'x-goog-custom-audit-work': 'test-work',
+            'x-goog-custom-audit-job': 'test-job',
+            'x-goog-custom-audit-id': '1234'
+        })
+
+  @mock.patch('apache_beam.options.pipeline_options._BeamArgumentParser.error')
+  def test_gcs_custom_audit_entries_with_errors(self, mock_error):
+    long_key = 'a' * 65
+    options = PipelineOptions([f'--gcs_custom_audit_entry={long_key}=1'])
+    _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries
+    self.assertRegex(
+        mock_error.call_args[0][0],
+        'The key .* exceeds the 64-character limit.')
+
+    mock_error.reset_mock()
+
+    long_value = 'b' * 1201
+    options = PipelineOptions([f'--gcs_custom_audit_entry=key={long_value}'])
+    _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries
+    self.assertRegex(
+        mock_error.call_args[0][0],
+        'The value .* exceeds the 1200-character limit.')
+
+    mock_error.reset_mock()
+
+    options = PipelineOptions([
+        '--gcs_custom_audit_entry=a=1',
+        '--gcs_custom_audit_entry=b=2',
+        '--gcs_custom_audit_entry=c=3',
+        '--gcs_custom_audit_entry=d=4',
+        '--gcs_custom_audit_entry=job=test-job'
+    ])
+    _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries
+    self.assertRegex(
+        mock_error.call_args[0][0],
+        'The maximum allowed number of GCS custom audit entries .*')
+
+    mock_error.reset_mock()
+
+    options = PipelineOptions([
+        '--gcs_custom_audit_entry=a=1',
+        '--gcs_custom_audit_entry=b=2',
+        '--gcs_custom_audit_entry=c=3',
+        '--gcs_custom_audit_entry=d=4'
+    ])
+    _ = options.view_as(GoogleCloudOptions).gcs_custom_audit_entries
+    self.assertRegex(
+        mock_error.call_args[0][0],
+        'The maximum allowed number of GCS custom audit entries .*')
+
   def _check_errors(self, options, validator, expected):
     if has_gcsio:
       with mock.patch('apache_beam.io.gcp.gcsio.GcsIO.is_soft_delete_enabled',