From f6ec3298a7fc153da4b78946ec254cef61c48cf6 Mon Sep 17 00:00:00 2001 From: qizhicheng Date: Fri, 3 Nov 2023 21:10:02 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=E5=88=9D=E5=A7=8B=E5=8C=96=20notify=20?= =?UTF-8?q?=E6=97=B6=E8=87=AA=E5=8A=A8=E6=A0=B9=E6=8D=AE=20workflow=20audi?= =?UTF-8?q?t=20=E5=8F=96=20workflow?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- sql/models.py | 12 +++++++++ sql/notify.py | 54 +++++++++++++++---------------------- sql/test_notify.py | 13 ++++++++- sql/utils/workflow_audit.py | 10 ++----- 4 files changed, 48 insertions(+), 41 deletions(-) diff --git a/sql/models.py b/sql/models.py index fc01328c27..02b77a9db6 100755 --- a/sql/models.py +++ b/sql/models.py @@ -324,6 +324,18 @@ class WorkflowAudit(models.Model): create_time = models.DateTimeField("申请时间", auto_now_add=True) sys_time = models.DateTimeField("系统时间", auto_now=True) + def get_workflow(self): + """尝试从 audit 中取出 workflow""" + if self.workflow_type == WorkflowType.QUERY: + return QueryPrivilegesApply.objects.get( + apply_id=self.workflow_id + ) + elif self.workflow_type == WorkflowType.SQL_REVIEW: + return SqlWorkflow.objects.get(id=self.workflow_id) + elif self.workflow_type == WorkflowType.ARCHIVE: + return ArchiveConfig.objects.get(id=self.workflow_id) + raise ValueError("无法获取到关联工单") + def __int__(self): return self.audit_id diff --git a/sql/notify.py b/sql/notify.py index b478bba9bc..c8e703aa12 100755 --- a/sql/notify.py +++ b/sql/notify.py @@ -52,23 +52,25 @@ class My2SqlResult: error: str = "" +@dataclass class Notifier: - name = "base" - sys_config_key: str = "" - - def __init__( - self, - workflow: Union[SqlWorkflow, ArchiveConfig, QueryPrivilegesApply, My2SqlResult], - sys_config: SysConfig, - audit: WorkflowAudit = None, - audit_detail: WorkflowAuditDetail = None, - event_type: EventType = EventType.AUDIT, - ): - self.workflow = workflow - self.audit = audit - self.audit_detail = audit_detail - self.event_type = event_type - self.sys_config = sys_config + workflow: Union[SqlWorkflow, ArchiveConfig, QueryPrivilegesApply, My2SqlResult] + sys_config: SysConfig = None + # init false, class property, 不是 instance property + name: str = field(init=False, default="base") + sys_config_key: str = field(init=False, default="") + event_type: EventType = EventType.AUDIT + audit: WorkflowAudit = None + audit_detail: WorkflowAuditDetail = None + + def __post_init__(self): + if not self.workflow: + if not self.audit: + raise ValueError("需要提供 WorkflowAudit 或 workflow") + self.workflow = self.audit.get_workflow() + # 防止 get_auditor 显式的传了个 None + if not self.sys_config: + self.sys_config = SysConfig() def render(self): raise NotImplementedError @@ -91,12 +93,9 @@ def run(self): class GenericWebhookNotifier(Notifier): - name = "generic_webhook" + name: str = "generic_webhook" sys_config_key: str = "generic_webhook_url" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.request_data = None + request_data: dict = None def render(self): self.request_data = {} @@ -133,13 +132,9 @@ class LegacyMessage: msg_cc: List[Users] = field(default_factory=list) +@dataclass class LegacyRender(Notifier): - messages: List[LegacyMessage] - sys_config_key: str = "" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.messages = [] + messages: List[LegacyMessage] = field(default_factory=list) def render_audit(self): # 获取审核信息 @@ -476,11 +471,6 @@ def auto_notify( 加载所有的 notifier, 调用 notifier 的 render 和 send 方法 内部方法, 有数据库查询, 为了方便测试, 请勿使用 async_task 调用, 防止 patch 后调用失败 """ - if not workflow and event_type == EventType.AUDIT: - if audit.workflow_type == 1: - workflow = QueryPrivilegesApply.objects.get(apply_id=audit.workflow_id) - if audit.workflow_type == 2: - workflow = SqlWorkflow.objects.get(id=audit.workflow_id) for notifier in settings.ENABLED_NOTIFIERS: file, _class = notifier.split(":") try: diff --git a/sql/test_notify.py b/sql/test_notify.py index 2c175faccb..b48f696bc8 100644 --- a/sql/test_notify.py +++ b/sql/test_notify.py @@ -199,7 +199,7 @@ def test_base_notifier(self): @patch("sql.notify.FeishuWebhookNotifier.run") def test_auto_notify(self, mock_run): with self.settings(ENABLED_NOTIFIERS=("sql.notify:FeishuWebhookNotifier",)): - auto_notify(self.sys_config, event_type=EventType.EXECUTE) + auto_notify(self.sys_config, event_type=EventType.EXECUTE, workflow=self.wf) mock_run.assert_called_once() @patch("sql.notify.auto_notify") @@ -494,10 +494,13 @@ def tearDownClass(cls): def setUp(self): self.patcher = patch("sql.notify.MsgSender") self.mock_msg_sender = self.patcher.start() + self.get_workflow_patcher = patch("sql.models.WorkflowAudit.get_workflow") + self.mock_get_workflow = self.get_workflow_patcher.start() self.sys_config = SysConfig() def tearDown(self): self.patcher.stop() + self.get_workflow_patcher.stop() def generate_notifier(self, module) -> Notifier: return module(workflow=None, audit=self.audit_wf, sys_config=self.sys_config) @@ -561,3 +564,11 @@ def test_mail(self): ] notifier.send() mocker.assert_called_once() + + +def test_override_sys_key(): + """dataclass 的继承有时候让人有点困惑, 在这里补一个测试确认可以正常覆盖一些值""" + class OverrideNotifier(Notifier): + sys_config_key = "test" + n = OverrideNotifier(workflow="test") + assert n.sys_config_key == "test" diff --git a/sql/utils/workflow_audit.py b/sql/utils/workflow_audit.py index 4f6578b102..509138ebb0 100644 --- a/sql/utils/workflow_audit.py +++ b/sql/utils/workflow_audit.py @@ -136,14 +136,8 @@ def review_info(self) -> (str, str): def get_workflow(self): """尝试从 audit 中取出 workflow""" - if self.audit.workflow_type == WorkflowType.QUERY: - self.workflow = QueryPrivilegesApply.objects.get( - apply_id=self.audit.workflow_id - ) - elif self.audit.workflow_type == WorkflowType.SQL_REVIEW: - self.workflow = SqlWorkflow.objects.get(id=self.audit.workflow_id) - elif self.audit.workflow_type == WorkflowType.ARCHIVE: - self.workflow = ArchiveConfig.objects.get(id=self.audit.workflow_id) + self.workflow = self.audit.get_workflow() + if self.audit.workflow_type == WorkflowType.ARCHIVE: self.resource_group = self.audit.group_name self.resource_group_id = self.audit.group_id From 060ae6a6d691554eb2d5fcade9df29a2674e882a Mon Sep 17 00:00:00 2001 From: qizhicheng Date: Mon, 6 Nov 2023 15:38:24 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E5=88=9D=E5=A7=8B?= =?UTF-8?q?=E5=8C=96=E6=97=B6=E4=B8=8D=E4=BC=A0=20workflow=20=E7=9A=84?= =?UTF-8?q?=E5=88=9D=E5=A7=8B=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- sql/models.py | 4 +--- sql/test_notify.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/sql/models.py b/sql/models.py index 02b77a9db6..0da4dcf1b8 100755 --- a/sql/models.py +++ b/sql/models.py @@ -327,9 +327,7 @@ class WorkflowAudit(models.Model): def get_workflow(self): """尝试从 audit 中取出 workflow""" if self.workflow_type == WorkflowType.QUERY: - return QueryPrivilegesApply.objects.get( - apply_id=self.workflow_id - ) + return QueryPrivilegesApply.objects.get(apply_id=self.workflow_id) elif self.workflow_type == WorkflowType.SQL_REVIEW: return SqlWorkflow.objects.get(id=self.workflow_id) elif self.workflow_type == WorkflowType.ARCHIVE: diff --git a/sql/test_notify.py b/sql/test_notify.py index b48f696bc8..ecb598939c 100644 --- a/sql/test_notify.py +++ b/sql/test_notify.py @@ -280,6 +280,17 @@ def test_legacy_render_audit(self): notifier.render() self.assertEqual(len(notifier.messages), 1) self.assertIn("新的工单申请", notifier.messages[0].msg_title) + # 测试一下不传 workflow + notifier = LegacyRender( + event_type=EventType.AUDIT, + workflow=None, + audit=self.audit_wf, + audit_detail=self.audit_wf_detail, + sys_config=self.sys_config, + ) + notifier.render() + self.assertEqual(len(notifier.messages), 1) + self.assertIn("新的工单申请", notifier.messages[0].msg_title) def test_legacy_render_query_audit(self): # 默认是库权限的 @@ -568,7 +579,9 @@ def test_mail(self): def test_override_sys_key(): """dataclass 的继承有时候让人有点困惑, 在这里补一个测试确认可以正常覆盖一些值""" + class OverrideNotifier(Notifier): sys_config_key = "test" + n = OverrideNotifier(workflow="test") assert n.sys_config_key == "test"