Edit on GitHub

sqlmesh.integrations.github.notification_operator_provider

 1from __future__ import annotations
 2
 3import typing as t
 4
 5from sqlmesh.core.notification_target import NotificationStatus
 6from sqlmesh.core.plan import PlanStatus
 7from sqlmesh.integrations.github.notification_target import GithubNotificationTarget
 8from sqlmesh.integrations.github.shared import add_comment_to_pr
 9from sqlmesh.schedulers.airflow import common
10from sqlmesh.schedulers.airflow.operators.notification import (
11    BaseNotificationOperatorProvider,
12)
13
14if t.TYPE_CHECKING:
15    from airflow.providers.github.operators.github import GithubOperator
16
17
18class GithubNotificationOperatorProvider(
19    BaseNotificationOperatorProvider[GithubNotificationTarget]
20):
21    """
22    Github Notification Operator for Airflow.
23    """
24
25    def operator(
26        self,
27        target: GithubNotificationTarget,
28        plan_status: PlanStatus,
29        plan_dag_spec: common.PlanDagSpec,
30        **dag_kwargs: t.Any,
31    ) -> GithubOperator:
32        from airflow.providers.github.operators.github import GithubOperator
33
34        if plan_status.is_started:
35            notification_status = NotificationStatus.PROGRESS
36            msg = f"Updating Environment `{plan_dag_spec.environment_name}`"
37        elif plan_status.is_finished:
38            notification_status = NotificationStatus.SUCCESS
39            msg = f"Updated Environment `{plan_dag_spec.environment_name}`"
40        else:
41            notification_status = NotificationStatus.FAILURE
42            msg = f"Failed to Update Environment `{plan_dag_spec.environment_name}`"
43
44        bot_users = [user for user in plan_dag_spec.users if user.is_bot]
45        bot_user_to_append_to = bot_users[0] if bot_users else None
46
47        return GithubOperator(
48            task_id=self.get_task_id(target, plan_status),
49            trigger_rule=self.get_trigger_rule(plan_status),
50            github_method="get_repo",
51            github_method_args={"full_name_or_id": target.pull_request_info.full_repo_path},
52            result_processor=lambda repo: add_comment_to_pr(
53                repo,
54                target.pull_request_info,
55                notification_status,
56                msg,
57                user_to_append_to=bot_user_to_append_to,
58            ),
59        )
19class GithubNotificationOperatorProvider(
20    BaseNotificationOperatorProvider[GithubNotificationTarget]
21):
22    """
23    Github Notification Operator for Airflow.
24    """
25
26    def operator(
27        self,
28        target: GithubNotificationTarget,
29        plan_status: PlanStatus,
30        plan_dag_spec: common.PlanDagSpec,
31        **dag_kwargs: t.Any,
32    ) -> GithubOperator:
33        from airflow.providers.github.operators.github import GithubOperator
34
35        if plan_status.is_started:
36            notification_status = NotificationStatus.PROGRESS
37            msg = f"Updating Environment `{plan_dag_spec.environment_name}`"
38        elif plan_status.is_finished:
39            notification_status = NotificationStatus.SUCCESS
40            msg = f"Updated Environment `{plan_dag_spec.environment_name}`"
41        else:
42            notification_status = NotificationStatus.FAILURE
43            msg = f"Failed to Update Environment `{plan_dag_spec.environment_name}`"
44
45        bot_users = [user for user in plan_dag_spec.users if user.is_bot]
46        bot_user_to_append_to = bot_users[0] if bot_users else None
47
48        return GithubOperator(
49            task_id=self.get_task_id(target, plan_status),
50            trigger_rule=self.get_trigger_rule(plan_status),
51            github_method="get_repo",
52            github_method_args={"full_name_or_id": target.pull_request_info.full_repo_path},
53            result_processor=lambda repo: add_comment_to_pr(
54                repo,
55                target.pull_request_info,
56                notification_status,
57                msg,
58                user_to_append_to=bot_user_to_append_to,
59            ),
60        )

Github Notification Operator for Airflow.

GithubNotificationOperatorProvider()
def operator( self, target: sqlmesh.integrations.github.notification_target.GithubNotificationTarget, plan_status: sqlmesh.core.plan.definition.PlanStatus, plan_dag_spec: sqlmesh.schedulers.airflow.common.PlanDagSpec, **dag_kwargs: Any) -> <MagicMock id='6049944224'>:
26    def operator(
27        self,
28        target: GithubNotificationTarget,
29        plan_status: PlanStatus,
30        plan_dag_spec: common.PlanDagSpec,
31        **dag_kwargs: t.Any,
32    ) -> GithubOperator:
33        from airflow.providers.github.operators.github import GithubOperator
34
35        if plan_status.is_started:
36            notification_status = NotificationStatus.PROGRESS
37            msg = f"Updating Environment `{plan_dag_spec.environment_name}`"
38        elif plan_status.is_finished:
39            notification_status = NotificationStatus.SUCCESS
40            msg = f"Updated Environment `{plan_dag_spec.environment_name}`"
41        else:
42            notification_status = NotificationStatus.FAILURE
43            msg = f"Failed to Update Environment `{plan_dag_spec.environment_name}`"
44
45        bot_users = [user for user in plan_dag_spec.users if user.is_bot]
46        bot_user_to_append_to = bot_users[0] if bot_users else None
47
48        return GithubOperator(
49            task_id=self.get_task_id(target, plan_status),
50            trigger_rule=self.get_trigger_rule(plan_status),
51            github_method="get_repo",
52            github_method_args={"full_name_or_id": target.pull_request_info.full_repo_path},
53            result_processor=lambda repo: add_comment_to_pr(
54                repo,
55                target.pull_request_info,
56                notification_status,
57                msg,
58                user_to_append_to=bot_user_to_append_to,
59            ),
60        )