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 26f94c5370 Cloud Data Loss Prevention Operators assets (#26618)
26f94c5370 is described below
commit 26f94c5370587f73ebd935cecf208c6a36bdf9b6
Author: Beata Kossakowska <[email protected]>
AuthorDate: Mon Sep 26 16:21:54 2022 +0200
Cloud Data Loss Prevention Operators assets (#26618)
---
.../google/cloud/links/data_loss_prevention.py | 318 +++++++++++++++++++++
airflow/providers/google/cloud/operators/dlp.py | 307 +++++++++++++++++++-
airflow/providers/google/provider.yaml | 11 +
tests/providers/google/cloud/operators/test_dlp.py | 78 ++---
.../data_loss_prevention/example_dlp_info_types.py | 2 +-
5 files changed, 676 insertions(+), 40 deletions(-)
diff --git a/airflow/providers/google/cloud/links/data_loss_prevention.py
b/airflow/providers/google/cloud/links/data_loss_prevention.py
new file mode 100644
index 0000000000..46c824f907
--- /dev/null
+++ b/airflow/providers/google/cloud/links/data_loss_prevention.py
@@ -0,0 +1,318 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.google.cloud.links.base import BaseGoogleLink
+
+if TYPE_CHECKING:
+ from airflow.utils.context import Context
+
+BASE_LINK = "https://console.cloud.google.com";
+
+DLP_BASE_LINK = BASE_LINK + "/security/dlp"
+
+DLP_DEIDENTIFY_TEMPLATES_LIST_LINK = (
+ DLP_BASE_LINK +
"/landing/configuration/templates/deidentify?project={project_id}"
+)
+DLP_DEIDENTIFY_TEMPLATE_DETAILS_LINK = (
+ DLP_BASE_LINK
+ +
"/projects/{project_id}/locations/global/deidentifyTemplates/{template_name}?project={project_id}"
+)
+
+DLP_JOB_TRIGGER_LIST_LINK = DLP_BASE_LINK +
"/landing/inspection/triggers?project={project_id}"
+DLP_JOB_TRIGGER_DETAILS_LINK = (
+ DLP_BASE_LINK +
"/projects/{project_id}/locations/global/jobTriggers/{trigger_name}?project={project_id}"
+)
+
+DLP_JOBS_LIST_LINK = DLP_BASE_LINK +
"/landing/inspection/jobs?project={project_id}"
+DLP_JOB_DETAILS_LINK = (
+ DLP_BASE_LINK +
"/projects/{project_id}/locations/global/dlpJobs/{job_name}?project={project_id}"
+)
+
+DLP_INSPECT_TEMPLATES_LIST_LINK = (
+ DLP_BASE_LINK +
"/landing/configuration/templates/inspect?project={project_id}"
+)
+DLP_INSPECT_TEMPLATE_DETAILS_LINK = (
+ DLP_BASE_LINK
+ +
"/projects/{project_id}/locations/global/inspectTemplates/{template_name}?project={project_id}"
+)
+
+DLP_INFO_TYPES_LIST_LINK = (
+ DLP_BASE_LINK +
"/landing/configuration/infoTypes/stored?cloudshell=false&project={project_id}"
+)
+DLP_INFO_TYPE_DETAILS_LINK = (
+ DLP_BASE_LINK
+ +
"/projects/{project_id}/locations/global/storedInfoTypes/{info_type_name}?project={project_id}"
+)
+DLP_POSSIBLE_INFO_TYPES_LIST_LINK = (
+ DLP_BASE_LINK +
"/landing/configuration/infoTypes/built-in?project={project_id}"
+)
+
+
+class CloudDLPDeidentifyTemplatesListLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Deidentify Templates List"
+ key = "cloud_dlp_deidentify_templates_list_key"
+ format_str = DLP_DEIDENTIFY_TEMPLATES_LIST_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPDeidentifyTemplatesListLink.key,
+ value={
+ "project_id": project_id,
+ },
+ )
+
+
+class CloudDLPDeidentifyTemplateDetailsLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Deidentify Template Details"
+ key = "cloud_dlp_deidentify_template_details_key"
+ format_str = DLP_DEIDENTIFY_TEMPLATE_DETAILS_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ template_name: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPDeidentifyTemplateDetailsLink.key,
+ value={
+ "project_id": project_id,
+ "template_name": template_name,
+ },
+ )
+
+
+class CloudDLPJobTriggersListLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Job Triggers List"
+ key = "cloud_dlp_job_triggers_list_key"
+ format_str = DLP_JOB_TRIGGER_LIST_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPJobTriggersListLink.key,
+ value={
+ "project_id": project_id,
+ },
+ )
+
+
+class CloudDLPJobTriggerDetailsLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Job Triggers Details"
+ key = "cloud_dlp_job_trigger_details_key"
+ format_str = DLP_JOB_TRIGGER_DETAILS_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ trigger_name: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPJobTriggerDetailsLink.key,
+ value={
+ "project_id": project_id,
+ "trigger_name": trigger_name,
+ },
+ )
+
+
+class CloudDLPJobsListLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Jobs List"
+ key = "cloud_dlp_jobs_list_key"
+ format_str = DLP_JOBS_LIST_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPJobsListLink.key,
+ value={
+ "project_id": project_id,
+ },
+ )
+
+
+class CloudDLPJobDetailsLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Job Details"
+ key = "cloud_dlp_job_details_key"
+ format_str = DLP_JOB_DETAILS_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ job_name: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPJobDetailsLink.key,
+ value={
+ "project_id": project_id,
+ "job_name": job_name,
+ },
+ )
+
+
+class CloudDLPInspectTemplatesListLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Inspect Templates List"
+ key = "cloud_dlp_inspect_templates_list_key"
+ format_str = DLP_INSPECT_TEMPLATES_LIST_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPInspectTemplatesListLink.key,
+ value={
+ "project_id": project_id,
+ },
+ )
+
+
+class CloudDLPInspectTemplateDetailsLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Inspect Template Details"
+ key = "cloud_dlp_inspect_template_details_key"
+ format_str = DLP_INSPECT_TEMPLATE_DETAILS_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ template_name: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPInspectTemplateDetailsLink.key,
+ value={
+ "project_id": project_id,
+ "template_name": template_name,
+ },
+ )
+
+
+class CloudDLPInfoTypesListLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Info Types List"
+ key = "cloud_dlp_info_types_list_key"
+ format_str = DLP_INFO_TYPES_LIST_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPInfoTypesListLink.key,
+ value={
+ "project_id": project_id,
+ },
+ )
+
+
+class CloudDLPInfoTypeDetailsLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Info Type Details"
+ key = "cloud_dlp_info_type_details_key"
+ format_str = DLP_INFO_TYPE_DETAILS_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ info_type_name: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPInfoTypeDetailsLink.key,
+ value={
+ "project_id": project_id,
+ "info_type_name": info_type_name,
+ },
+ )
+
+
+class CloudDLPPossibleInfoTypesListLink(BaseGoogleLink):
+ """Helper class for constructing Cloud Data Loss Prevention link"""
+
+ name = "Cloud DLP Possible Info Types List"
+ key = "cloud_dlp_possible_info_types_list_key"
+ format_str = DLP_POSSIBLE_INFO_TYPES_LIST_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=CloudDLPPossibleInfoTypesListLink.key,
+ value={
+ "project_id": project_id,
+ },
+ )
diff --git a/airflow/providers/google/cloud/operators/dlp.py
b/airflow/providers/google/cloud/operators/dlp.py
index 1d7d09467b..bae212e9fa 100644
--- a/airflow/providers/google/cloud/operators/dlp.py
+++ b/airflow/providers/google/cloud/operators/dlp.py
@@ -45,6 +45,19 @@ from google.protobuf.json_format import MessageToDict
from airflow.models import BaseOperator
from airflow.providers.google.cloud.hooks.dlp import CloudDLPHook
+from airflow.providers.google.cloud.links.data_loss_prevention import (
+ CloudDLPDeidentifyTemplateDetailsLink,
+ CloudDLPDeidentifyTemplatesListLink,
+ CloudDLPInfoTypeDetailsLink,
+ CloudDLPInfoTypesListLink,
+ CloudDLPInspectTemplateDetailsLink,
+ CloudDLPInspectTemplatesListLink,
+ CloudDLPJobDetailsLink,
+ CloudDLPJobsListLink,
+ CloudDLPJobTriggerDetailsLink,
+ CloudDLPJobTriggersListLink,
+ CloudDLPPossibleInfoTypesListLink,
+)
if TYPE_CHECKING:
from airflow.utils.context import Context
@@ -85,6 +98,7 @@ class CloudDLPCancelDLPJobOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobDetailsLink(),)
def __init__(
self,
@@ -120,6 +134,15 @@ class CloudDLPCancelDLPJobOperator(BaseOperator):
metadata=self.metadata,
)
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ job_name=self.dlp_job_id,
+ )
+
class CloudDLPCreateDeidentifyTemplateOperator(BaseOperator):
"""
@@ -164,6 +187,7 @@ class
CloudDLPCreateDeidentifyTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPDeidentifyTemplateDetailsLink(),)
def __init__(
self,
@@ -216,8 +240,19 @@ class
CloudDLPCreateDeidentifyTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+ result = MessageToDict(template)
+
+ project_id = self.project_id or hook.project_id
+ template_id = self.template_id or result['name'].split("/")[-1] if
result['name'] else None
+ if project_id and template_id:
+ CloudDLPDeidentifyTemplateDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ template_name=template_id,
+ )
- return MessageToDict(template)
+ return result
class CloudDLPCreateDLPJobOperator(BaseOperator):
@@ -263,6 +298,7 @@ class CloudDLPCreateDLPJobOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobDetailsLink(),)
def __init__(
self,
@@ -317,7 +353,19 @@ class CloudDLPCreateDLPJobOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
- return MessageToDict(job)
+
+ result = MessageToDict(job)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ job_name=result['name'].split("/")[-1] if result['name'] else
None,
+ )
+
+ return result
class CloudDLPCreateInspectTemplateOperator(BaseOperator):
@@ -363,6 +411,7 @@ class CloudDLPCreateInspectTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInspectTemplateDetailsLink(),)
def __init__(
self,
@@ -415,7 +464,20 @@ class CloudDLPCreateInspectTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
- return MessageToDict(template)
+
+ result = MessageToDict(template)
+
+ template_id = self.template_id or result['name'].split("/")[-1] if
result['name'] else None
+ project_id = self.project_id or hook.project_id
+ if project_id and template_id:
+ CloudDLPInspectTemplateDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ template_name=template_id,
+ )
+
+ return result
class CloudDLPCreateJobTriggerOperator(BaseOperator):
@@ -458,6 +520,7 @@ class CloudDLPCreateJobTriggerOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobTriggerDetailsLink(),)
def __init__(
self,
@@ -508,7 +571,20 @@ class CloudDLPCreateJobTriggerOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
- return MessageToDict(trigger)
+
+ result = MessageToDict(trigger)
+
+ project_id = self.project_id or hook.project_id
+ trigger_name = result['name'].split("/")[-1] if result['name'] else
None
+ if project_id:
+ CloudDLPJobTriggerDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ trigger_name=trigger_name,
+ )
+
+ return result
class CloudDLPCreateStoredInfoTypeOperator(BaseOperator):
@@ -553,6 +629,7 @@ class CloudDLPCreateStoredInfoTypeOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInfoTypeDetailsLink(),)
def __init__(
self,
@@ -607,7 +684,22 @@ class CloudDLPCreateStoredInfoTypeOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
- return MessageToDict(info)
+
+ result = MessageToDict(info)
+
+ project_id = self.project_id or hook.project_id
+ stored_info_type_id = (
+ self.stored_info_type_id or result['name'].split("/")[-1] if
result['name'] else None
+ )
+ if project_id and stored_info_type_id:
+ CloudDLPInfoTypeDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ info_type_name=stored_info_type_id,
+ )
+
+ return result
class CloudDLPDeidentifyContentOperator(BaseOperator):
@@ -749,6 +841,7 @@ class
CloudDLPDeleteDeidentifyTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPDeidentifyTemplatesListLink(),)
def __init__(
self,
@@ -787,6 +880,13 @@ class
CloudDLPDeleteDeidentifyTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPDeidentifyTemplatesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
except NotFound:
self.log.error("Template %s not found.", self.template_id)
@@ -827,6 +927,7 @@ class CloudDLPDeleteDLPJobOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobsListLink(),)
def __init__(
self,
@@ -862,6 +963,15 @@ class CloudDLPDeleteDLPJobOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobsListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
except NotFound:
self.log.error("Job %s id not found.", self.dlp_job_id)
@@ -904,6 +1014,7 @@ class CloudDLPDeleteInspectTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInspectTemplatesListLink(),)
def __init__(
self,
@@ -942,6 +1053,15 @@ class CloudDLPDeleteInspectTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInspectTemplatesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
except NotFound:
self.log.error("Template %s not found", self.template_id)
@@ -981,6 +1101,7 @@ class CloudDLPDeleteJobTriggerOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobTriggersListLink(),)
def __init__(
self,
@@ -1016,6 +1137,15 @@ class CloudDLPDeleteJobTriggerOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobTriggersListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
except NotFound:
self.log.error("Trigger %s not found", self.job_trigger_id)
@@ -1058,6 +1188,7 @@ class CloudDLPDeleteStoredInfoTypeOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInfoTypesListLink(),)
def __init__(
self,
@@ -1099,6 +1230,14 @@ class CloudDLPDeleteStoredInfoTypeOperator(BaseOperator):
except NotFound:
self.log.error("Stored info %s not found",
self.stored_info_type_id)
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInfoTypesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
class CloudDLPGetDeidentifyTemplateOperator(BaseOperator):
"""
@@ -1140,6 +1279,7 @@ class CloudDLPGetDeidentifyTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPDeidentifyTemplateDetailsLink(),)
def __init__(
self,
@@ -1177,6 +1317,13 @@ class
CloudDLPGetDeidentifyTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPDeidentifyTemplateDetailsLink.persist(
+ context=context, task_instance=self, project_id=project_id,
template_name=self.template_id
+ )
+
return MessageToDict(template)
@@ -1217,6 +1364,7 @@ class CloudDLPGetDLPJobOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobDetailsLink(),)
def __init__(
self,
@@ -1251,6 +1399,16 @@ class CloudDLPGetDLPJobOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ job_name=self.dlp_job_id,
+ )
+
return MessageToDict(job)
@@ -1294,6 +1452,7 @@ class CloudDLPGetInspectTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInspectTemplateDetailsLink(),)
def __init__(
self,
@@ -1331,6 +1490,16 @@ class CloudDLPGetInspectTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInspectTemplateDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ template_name=self.template_id,
+ )
+
return MessageToDict(template)
@@ -1371,6 +1540,7 @@ class CloudDLPGetDLPJobTriggerOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobTriggerDetailsLink(),)
def __init__(
self,
@@ -1405,6 +1575,16 @@ class CloudDLPGetDLPJobTriggerOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobTriggerDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ trigger_name=self.job_trigger_id,
+ )
+
return MessageToDict(trigger)
@@ -1448,6 +1628,7 @@ class CloudDLPGetStoredInfoTypeOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInfoTypeDetailsLink(),)
def __init__(
self,
@@ -1485,6 +1666,16 @@ class CloudDLPGetStoredInfoTypeOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInfoTypeDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ info_type_name=self.stored_info_type_id,
+ )
+
return MessageToDict(info)
@@ -1617,6 +1808,7 @@ class
CloudDLPListDeidentifyTemplatesOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPDeidentifyTemplatesListLink(),)
def __init__(
self,
@@ -1658,6 +1850,15 @@ class
CloudDLPListDeidentifyTemplatesOperator(BaseOperator):
metadata=self.metadata,
)
# the MessageToDict does not have the right type defined as possible
to pass in constructor
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPDeidentifyTemplatesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
return [MessageToDict(template) for template in templates] # type:
ignore[arg-type]
@@ -1702,6 +1903,7 @@ class CloudDLPListDLPJobsOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobsListLink(),)
def __init__(
self,
@@ -1745,6 +1947,15 @@ class CloudDLPListDLPJobsOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobsListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
# the MessageToDict does not have the right type defined as possible
to pass in constructor
return [MessageToDict(job) for job in jobs] # type: ignore[arg-type]
@@ -1785,10 +1996,12 @@ class CloudDLPListInfoTypesOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPPossibleInfoTypesListLink(),)
def __init__(
self,
*,
+ project_id: str | None = None,
language_code: str | None = None,
results_filter: str | None = None,
retry: Retry | _MethodDefault = DEFAULT,
@@ -1799,6 +2012,7 @@ class CloudDLPListInfoTypesOperator(BaseOperator):
**kwargs,
) -> None:
super().__init__(**kwargs)
+ self.project_id = project_id
self.language_code = language_code
self.results_filter = results_filter
self.retry = retry
@@ -1819,6 +2033,15 @@ class CloudDLPListInfoTypesOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPPossibleInfoTypesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
return MessageToDict(response)
@@ -1864,6 +2087,7 @@ class CloudDLPListInspectTemplatesOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInspectTemplatesListLink(),)
def __init__(
self,
@@ -1904,6 +2128,15 @@ class CloudDLPListInspectTemplatesOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInspectTemplatesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
return [MessageToDict(t) for t in templates]
@@ -1947,6 +2180,7 @@ class CloudDLPListJobTriggersOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobTriggersListLink(),)
def __init__(
self,
@@ -1987,6 +2221,15 @@ class CloudDLPListJobTriggersOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobTriggersListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
return [MessageToDict(j) for j in jobs]
@@ -2032,6 +2275,7 @@ class CloudDLPListStoredInfoTypesOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInfoTypesListLink(),)
def __init__(
self,
@@ -2072,6 +2316,15 @@ class CloudDLPListStoredInfoTypesOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInfoTypesListLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ )
+
return [MessageToDict(i) for i in infos]
@@ -2310,6 +2563,7 @@ class
CloudDLPUpdateDeidentifyTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPDeidentifyTemplateDetailsLink(),)
def __init__(
self,
@@ -2353,6 +2607,16 @@ class
CloudDLPUpdateDeidentifyTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPDeidentifyTemplateDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ template_name=self.template_id,
+ )
+
return MessageToDict(template)
@@ -2400,6 +2664,7 @@ class CloudDLPUpdateInspectTemplateOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInspectTemplateDetailsLink(),)
def __init__(
self,
@@ -2443,6 +2708,16 @@ class
CloudDLPUpdateInspectTemplateOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInspectTemplateDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ template_name=self.template_id,
+ )
+
return MessageToDict(template)
@@ -2487,6 +2762,7 @@ class CloudDLPUpdateJobTriggerOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPJobTriggerDetailsLink(),)
def __init__(
self,
@@ -2527,6 +2803,16 @@ class CloudDLPUpdateJobTriggerOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPJobTriggerDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ trigger_name=self.job_trigger_id,
+ )
+
return MessageToDict(trigger)
@@ -2575,6 +2861,7 @@ class CloudDLPUpdateStoredInfoTypeOperator(BaseOperator):
"gcp_conn_id",
"impersonation_chain",
)
+ operator_extra_links = (CloudDLPInfoTypeDetailsLink(),)
def __init__(
self,
@@ -2618,4 +2905,14 @@ class CloudDLPUpdateStoredInfoTypeOperator(BaseOperator):
timeout=self.timeout,
metadata=self.metadata,
)
+
+ project_id = self.project_id or hook.project_id
+ if project_id:
+ CloudDLPInfoTypeDetailsLink.persist(
+ context=context,
+ task_instance=self,
+ project_id=project_id,
+ info_type_name=self.stored_info_type_id,
+ )
+
return MessageToDict(info)
diff --git a/airflow/providers/google/provider.yaml
b/airflow/providers/google/provider.yaml
index 702b73392d..27d20d47bf 100644
--- a/airflow/providers/google/provider.yaml
+++ b/airflow/providers/google/provider.yaml
@@ -1020,6 +1020,17 @@ extra-links:
-
airflow.providers.google.cloud.links.cloud_storage_transfer.CloudStorageTransferListLink
-
airflow.providers.google.cloud.links.cloud_storage_transfer.CloudStorageTransferJobLink
-
airflow.providers.google.cloud.links.cloud_storage_transfer.CloudStorageTransferDetailsLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPDeidentifyTemplatesListLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPDeidentifyTemplateDetailsLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPJobTriggersListLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPJobTriggerDetailsLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPJobsListLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPJobDetailsLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPInspectTemplatesListLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPInspectTemplateDetailsLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPInfoTypesListLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPInfoTypeDetailsLink
+ -
airflow.providers.google.cloud.links.data_loss_prevention.CloudDLPPossibleInfoTypesListLink
- airflow.providers.google.common.links.storage.StorageLink
- airflow.providers.google.common.links.storage.FileDetailsLink
diff --git a/tests/providers/google/cloud/operators/test_dlp.py
b/tests/providers/google/cloud/operators/test_dlp.py
index 4c66039c7b..e33638caab 100644
--- a/tests/providers/google/cloud/operators/test_dlp.py
+++ b/tests/providers/google/cloud/operators/test_dlp.py
@@ -24,6 +24,7 @@ import unittest
from unittest import mock
from google.api_core.gapic_v1.method import DEFAULT
+from google.cloud.dlp_v2.types import DeidentifyTemplate, DlpJob,
InspectTemplate, JobTrigger, StoredInfoType
from airflow.providers.google.cloud.operators.dlp import (
CloudDLPCancelDLPJobOperator,
@@ -65,6 +66,8 @@ DLP_JOB_ID = "job123"
TEMPLATE_ID = "template123"
STORED_INFO_TYPE_ID = "type123"
TRIGGER_ID = "trigger123"
+DLP_JOB_PATH = f"projects/{PROJECT_ID}/dlpJobs/{DLP_JOB_ID}"
+DLP_JOB_TRIGGER_PATH = f"projects/{PROJECT_ID}/jobTriggers/{TRIGGER_ID}"
class TestCloudDLPCancelDLPJobOperator(unittest.TestCase):
@@ -72,7 +75,7 @@ class TestCloudDLPCancelDLPJobOperator(unittest.TestCase):
def test_cancel_dlp_job(self, mock_hook):
mock_hook.return_value.cancel_dlp_job.return_value = mock.MagicMock()
operator = CloudDLPCancelDLPJobOperator(dlp_job_id=DLP_JOB_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -90,8 +93,9 @@ class
TestCloudDLPCreateDeidentifyTemplateOperator(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.dlp.CloudDLPHook")
def test_create_deidentify_template(self, mock_hook):
mock_hook.return_value.create_deidentify_template.return_value =
mock.MagicMock()
+ mock_hook.return_value.create_deidentify_template.return_value =
DeidentifyTemplate(name=DLP_JOB_PATH)
operator =
CloudDLPCreateDeidentifyTemplateOperator(organization_id=ORGANIZATION_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -110,9 +114,12 @@ class
TestCloudDLPCreateDeidentifyTemplateOperator(unittest.TestCase):
class TestCloudDLPCreateDLPJobOperator(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.dlp.CloudDLPHook")
def test_create_dlp_job(self, mock_hook):
- mock_hook.return_value.create_dlp_job.return_value = mock.MagicMock()
+ mock_hook.return_value.create_dlp_job.return_value = DlpJob(
+ name=DLP_JOB_PATH, state=DlpJob.JobState.PENDING
+ )
operator = CloudDLPCreateDLPJobOperator(project_id=PROJECT_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
+
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -132,9 +139,9 @@ class TestCloudDLPCreateDLPJobOperator(unittest.TestCase):
class TestCloudDLPCreateInspectTemplateOperator(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.dlp.CloudDLPHook")
def test_create_inspect_template(self, mock_hook):
- mock_hook.return_value.create_inspect_template.return_value =
mock.MagicMock()
+ mock_hook.return_value.create_inspect_template.return_value =
InspectTemplate(name=DLP_JOB_PATH)
operator =
CloudDLPCreateInspectTemplateOperator(organization_id=ORGANIZATION_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -153,9 +160,9 @@ class
TestCloudDLPCreateInspectTemplateOperator(unittest.TestCase):
class TestCloudDLPCreateJobTriggerOperator(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.dlp.CloudDLPHook")
def test_create_job_trigger(self, mock_hook):
- mock_hook.return_value.create_job_trigger.return_value =
mock.MagicMock()
+ mock_hook.return_value.create_job_trigger.return_value =
JobTrigger(name=DLP_JOB_TRIGGER_PATH)
operator = CloudDLPCreateJobTriggerOperator(project_id=PROJECT_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -173,9 +180,9 @@ class
TestCloudDLPCreateJobTriggerOperator(unittest.TestCase):
class TestCloudDLPCreateStoredInfoTypeOperator(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.dlp.CloudDLPHook")
def test_create_stored_info_type(self, mock_hook):
- mock_hook.return_value.create_stored_info_type.return_value =
mock.MagicMock()
+ mock_hook.return_value.create_stored_info_type.return_value =
StoredInfoType(name=DLP_JOB_PATH)
operator =
CloudDLPCreateStoredInfoTypeOperator(organization_id=ORGANIZATION_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -196,7 +203,7 @@ class
TestCloudDLPDeidentifyContentOperator(unittest.TestCase):
def test_deidentify_content(self, mock_hook):
mock_hook.return_value.deidentify_content.return_value =
mock.MagicMock()
operator = CloudDLPDeidentifyContentOperator(project_id=PROJECT_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -221,7 +228,7 @@ class
TestCloudDLPDeleteDeidentifyTemplateOperator(unittest.TestCase):
operator = CloudDLPDeleteDeidentifyTemplateOperator(
template_id=TEMPLATE_ID, organization_id=ORGANIZATION_ID,
task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -241,7 +248,7 @@ class TestCloudDLPDeleteDlpJobOperator(unittest.TestCase):
def test_delete_dlp_job(self, mock_hook):
mock_hook.return_value.delete_dlp_job.return_value = mock.MagicMock()
operator = CloudDLPDeleteDLPJobOperator(dlp_job_id=DLP_JOB_ID,
project_id=PROJECT_ID, task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -262,7 +269,7 @@ class
TestCloudDLPDeleteInspectTemplateOperator(unittest.TestCase):
operator = CloudDLPDeleteInspectTemplateOperator(
template_id=TEMPLATE_ID, organization_id=ORGANIZATION_ID,
task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -284,7 +291,7 @@ class
TestCloudDLPDeleteJobTriggerOperator(unittest.TestCase):
operator = CloudDLPDeleteJobTriggerOperator(
job_trigger_id=TRIGGER_ID, project_id=PROJECT_ID, task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -307,7 +314,7 @@ class
TestCloudDLPDeleteStoredInfoTypeOperator(unittest.TestCase):
organization_id=ORGANIZATION_ID,
task_id="id",
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -329,7 +336,7 @@ class
TestCloudDLPGetDeidentifyTemplateOperator(unittest.TestCase):
operator = CloudDLPGetDeidentifyTemplateOperator(
template_id=TEMPLATE_ID, organization_id=ORGANIZATION_ID,
task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -349,7 +356,7 @@ class TestCloudDLPGetDlpJobOperator(unittest.TestCase):
def test_get_dlp_job(self, mock_hook):
mock_hook.return_value.get_dlp_job.return_value = mock.MagicMock()
operator = CloudDLPGetDLPJobOperator(dlp_job_id=DLP_JOB_ID,
project_id=PROJECT_ID, task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -370,7 +377,7 @@ class
TestCloudDLPGetInspectTemplateOperator(unittest.TestCase):
operator = CloudDLPGetInspectTemplateOperator(
template_id=TEMPLATE_ID, organization_id=ORGANIZATION_ID,
task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -392,7 +399,7 @@ class TestCloudDLPGetJobTripperOperator(unittest.TestCase):
operator = CloudDLPGetDLPJobTriggerOperator(
job_trigger_id=TRIGGER_ID, project_id=PROJECT_ID, task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -415,7 +422,7 @@ class
TestCloudDLPGetStoredInfoTypeOperator(unittest.TestCase):
organization_id=ORGANIZATION_ID,
task_id="id",
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -433,9 +440,12 @@ class
TestCloudDLPGetStoredInfoTypeOperator(unittest.TestCase):
class TestCloudDLPInspectContentOperator(unittest.TestCase):
@mock.patch("airflow.providers.google.cloud.operators.dlp.CloudDLPHook")
def test_inspect_content(self, mock_hook):
+ inspect_template_name = "inspect_template_name/name"
mock_hook.return_value.inspect_content.return_value = mock.MagicMock()
- operator = CloudDLPInspectContentOperator(project_id=PROJECT_ID,
task_id="id")
- operator.execute(context=None)
+ operator = CloudDLPInspectContentOperator(
+ project_id=PROJECT_ID, task_id="id",
inspect_template_name=inspect_template_name
+ )
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -444,7 +454,7 @@ class TestCloudDLPInspectContentOperator(unittest.TestCase):
project_id=PROJECT_ID,
inspect_config=None,
item=None,
- inspect_template_name=None,
+ inspect_template_name=inspect_template_name,
retry=DEFAULT,
timeout=None,
metadata=(),
@@ -456,7 +466,7 @@ class
TestCloudDLPListDeidentifyTemplatesOperator(unittest.TestCase):
def test_list_deidentify_templates(self, mock_hook):
mock_hook.return_value.list_deidentify_templates.return_value =
mock.MagicMock()
operator =
CloudDLPListDeidentifyTemplatesOperator(organization_id=ORGANIZATION_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -477,7 +487,7 @@ class TestCloudDLPListDlpJobsOperator(unittest.TestCase):
def test_list_dlp_jobs(self, mock_hook):
mock_hook.return_value.list_dlp_jobs.return_value = mock.MagicMock()
operator = CloudDLPListDLPJobsOperator(project_id=PROJECT_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -499,7 +509,7 @@ class TestCloudDLPListInfoTypesOperator(unittest.TestCase):
def test_list_info_types(self, mock_hook):
mock_hook.return_value.list_info_types.return_value = mock.MagicMock()
operator = CloudDLPListInfoTypesOperator(task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -518,7 +528,7 @@ class
TestCloudDLPListInspectTemplatesOperator(unittest.TestCase):
def test_list_inspect_templates(self, mock_hook):
mock_hook.return_value.list_inspect_templates.return_value =
mock.MagicMock()
operator =
CloudDLPListInspectTemplatesOperator(organization_id=ORGANIZATION_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -539,7 +549,7 @@ class
TestCloudDLPListJobTriggersOperator(unittest.TestCase):
def test_list_job_triggers(self, mock_hook):
mock_hook.return_value.list_job_triggers.return_value =
mock.MagicMock()
operator = CloudDLPListJobTriggersOperator(project_id=PROJECT_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -560,7 +570,7 @@ class
TestCloudDLPListStoredInfoTypesOperator(unittest.TestCase):
def test_list_stored_info_types(self, mock_hook):
mock_hook.return_value.list_stored_info_types.return_value =
mock.MagicMock()
operator =
CloudDLPListStoredInfoTypesOperator(organization_id=ORGANIZATION_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -628,7 +638,7 @@ class
TestCloudDLPUpdateDeidentifyTemplateOperator(unittest.TestCase):
operator = CloudDLPUpdateDeidentifyTemplateOperator(
template_id=TEMPLATE_ID, organization_id=ORGANIZATION_ID,
task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -652,7 +662,7 @@ class
TestCloudDLPUpdateInspectTemplateOperator(unittest.TestCase):
operator = CloudDLPUpdateInspectTemplateOperator(
template_id=TEMPLATE_ID, organization_id=ORGANIZATION_ID,
task_id="id"
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -674,7 +684,7 @@ class
TestCloudDLPUpdateJobTriggerOperator(unittest.TestCase):
def test_update_job_trigger(self, mock_hook):
mock_hook.return_value.update_job_trigger.return_value =
mock.MagicMock()
operator = CloudDLPUpdateJobTriggerOperator(job_trigger_id=TRIGGER_ID,
task_id="id")
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
@@ -699,7 +709,7 @@ class
TestCloudDLPUpdateStoredInfoTypeOperator(unittest.TestCase):
organization_id=ORGANIZATION_ID,
task_id="id",
)
- operator.execute(context=None)
+ operator.execute(context=mock.MagicMock())
mock_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=None,
diff --git
a/tests/system/providers/google/cloud/data_loss_prevention/example_dlp_info_types.py
b/tests/system/providers/google/cloud/data_loss_prevention/example_dlp_info_types.py
index e39566dc28..003810b957 100644
---
a/tests/system/providers/google/cloud/data_loss_prevention/example_dlp_info_types.py
+++
b/tests/system/providers/google/cloud/data_loss_prevention/example_dlp_info_types.py
@@ -112,7 +112,7 @@ with models.DAG(
)
get_stored_info_type = CloudDLPGetStoredInfoTypeOperator(
- task_id="list_stored_info_type", project_id=PROJECT_ID,
stored_info_type_id=CUSTOM_INFO_TYPE_ID
+ task_id="get_stored_info_type", project_id=PROJECT_ID,
stored_info_type_id=CUSTOM_INFO_TYPE_ID
)
# [START howto_operator_dlp_update_info_type]
