[airflow] branch main updated: Don't push secret in XCOM in BigQueryCreateDataTransferOperator (#29348)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new f51742d20b Don't push secret in XCOM in BigQueryCreateDataTransferOperator (#29348)
f51742d20b is described below

commit f51742d20b2e53bcd90a19db21e4e12d2a287677
Author: Pankaj Singh <98...@users.noreply.github.com>
AuthorDate: Tue Feb 21 04:36:50 2023 +0530

    Don't push secret in XCOM in BigQueryCreateDataTransferOperator (#29348)
    
    * Don't push secret in xcom in BigQueryCreateDataTransferOperator
---
 airflow/providers/google/cloud/operators/bigquery_dts.py    |  3 +++
 tests/providers/google/cloud/operators/test_bigquery_dts.py | 10 ++++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/google/cloud/operators/bigquery_dts.py b/airflow/providers/google/cloud/operators/bigquery_dts.py
index ee1a3b548b..d786c903e7 100644
--- a/airflow/providers/google/cloud/operators/bigquery_dts.py
+++ b/airflow/providers/google/cloud/operators/bigquery_dts.py
@@ -138,6 +138,9 @@ class BigQueryCreateDataTransferOperator(BaseOperator):
         result = TransferConfig.to_dict(response)
         self.log.info("Created DTS transfer config %s", get_object_id(result))
         self.xcom_push(context, key="transfer_config_id", value=get_object_id(result))
+        # don't push AWS secret in XCOM
+        result.get("params", {}).pop("secret_access_key", None)
+        result.get("params", {}).pop("access_key_id", None)
         return result
 
 
diff --git a/tests/providers/google/cloud/operators/test_bigquery_dts.py b/tests/providers/google/cloud/operators/test_bigquery_dts.py
index 78c92d52ed..aa52169a77 100644
--- a/tests/providers/google/cloud/operators/test_bigquery_dts.py
+++ b/tests/providers/google/cloud/operators/test_bigquery_dts.py
@@ -46,12 +46,15 @@ TRANSFER_CONFIG_ID = "id1234"
 
 TRANSFER_CONFIG_NAME = "projects/123abc/locations/321cba/transferConfig/1a2b3c"
 RUN_NAME = "projects/123abc/locations/321cba/transferConfig/1a2b3c/runs/123"
+transfer_config = TransferConfig(
+    name=TRANSFER_CONFIG_NAME, params={"secret_access_key": "AIRFLOW_KEY", "access_key_id": "AIRFLOW_KEY_ID"}
+)
 
 
 class BigQueryCreateDataTransferOperatorTestCase(unittest.TestCase):
     @mock.patch(
         "airflow.providers.google.cloud.operators.bigquery_dts.BiqQueryDataTransferServiceHook",
-        **{"return_value.create_transfer_config.return_value": TransferConfig(name=TRANSFER_CONFIG_NAME)},
+        **{"return_value.create_transfer_config.return_value": transfer_config},
     )
     def test_execute(self, mock_hook):
         op = BigQueryCreateDataTransferOperator(
@@ -59,7 +62,7 @@ class BigQueryCreateDataTransferOperatorTestCase(unittest.TestCase):
         )
         ti = mock.MagicMock()
 
-        op.execute({"ti": ti})
+        return_value = op.execute({"ti": ti})
 
         mock_hook.return_value.create_transfer_config.assert_called_once_with(
             authorization_code=None,
@@ -71,6 +74,9 @@ class BigQueryCreateDataTransferOperatorTestCase(unittest.TestCase):
         )
         ti.xcom_push.assert_called_with(execution_date=None, key="transfer_config_id", value="1a2b3c")
 
+        assert "secret_access_key" not in return_value.get("params", {})
+        assert "access_key_id" not in return_value.get("params", {})
+
 
 class BigQueryDeleteDataTransferConfigOperatorTestCase(unittest.TestCase):
     @mock.patch("airflow.providers.google.cloud.operators.bigquery_dts.BiqQueryDataTransferServiceHook")