From e6ecfdc04712928ba9d5b55313e45ba630ffaced Mon Sep 17 00:00:00 2001 From: subham611 Date: Mon, 22 Apr 2024 23:12:12 +0530 Subject: [PATCH 1/5] Adds notification setting parameter in create job config --- airflow/providers/databricks/operators/databricks.py | 3 +++ tests/providers/databricks/operators/test_databricks.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py index eab772d233b13..4514355579996 100644 --- a/airflow/providers/databricks/operators/databricks.py +++ b/airflow/providers/databricks/operators/databricks.py @@ -254,6 +254,7 @@ def __init__( job_clusters: list[dict] | None = None, email_notifications: dict | None = None, webhook_notifications: dict | None = None, + notification_settings: dict | None = None, timeout_seconds: int | None = None, schedule: dict | None = None, max_concurrent_runs: int | None = None, @@ -286,6 +287,8 @@ def __init__( self.json["email_notifications"] = email_notifications if webhook_notifications is not None: self.json["webhook_notifications"] = webhook_notifications + if notification_settings: + self.json["notification_settings"] = notification_settings if timeout_seconds is not None: self.json["timeout_seconds"] = timeout_seconds if schedule is not None: diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py index 46e14a917ab4e..885d5c293d055 100644 --- a/tests/providers/databricks/operators/test_databricks.py +++ b/tests/providers/databricks/operators/test_databricks.py @@ -202,6 +202,10 @@ } ], } +NOTIFICATION_SETTINGS = { + "no_alert_for_canceled_runs": True, + "no_alert_for_skipped_runs": True +} TIMEOUT_SECONDS = 86400 SCHEDULE = { "quartz_cron_expression": "20 30 * * * ?", @@ -414,6 +418,7 @@ def test_exec_create(self, db_mock_class): "job_clusters": JOB_CLUSTERS, "email_notifications": EMAIL_NOTIFICATIONS, "webhook_notifications": WEBHOOK_NOTIFICATIONS, + "notification_settings": NOTIFICATION_SETTINGS, "timeout_seconds": TIMEOUT_SECONDS, "schedule": SCHEDULE, "max_concurrent_runs": MAX_CONCURRENT_RUNS, @@ -436,6 +441,7 @@ def test_exec_create(self, db_mock_class): "job_clusters": JOB_CLUSTERS, "email_notifications": EMAIL_NOTIFICATIONS, "webhook_notifications": WEBHOOK_NOTIFICATIONS, + "notification_settings": NOTIFICATION_SETTINGS, "timeout_seconds": TIMEOUT_SECONDS, "schedule": SCHEDULE, "max_concurrent_runs": MAX_CONCURRENT_RUNS, @@ -466,6 +472,7 @@ def test_exec_reset(self, db_mock_class): "job_clusters": JOB_CLUSTERS, "email_notifications": EMAIL_NOTIFICATIONS, "webhook_notifications": WEBHOOK_NOTIFICATIONS, + "notification_settings": NOTIFICATION_SETTINGS, "timeout_seconds": TIMEOUT_SECONDS, "schedule": SCHEDULE, "max_concurrent_runs": MAX_CONCURRENT_RUNS, @@ -486,6 +493,7 @@ def test_exec_reset(self, db_mock_class): "job_clusters": JOB_CLUSTERS, "email_notifications": EMAIL_NOTIFICATIONS, "webhook_notifications": WEBHOOK_NOTIFICATIONS, + "notification_settings": NOTIFICATION_SETTINGS, "timeout_seconds": TIMEOUT_SECONDS, "schedule": SCHEDULE, "max_concurrent_runs": MAX_CONCURRENT_RUNS, From 5710000e464429655c57f44ef622981451ac9de6 Mon Sep 17 00:00:00 2001 From: subham611 Date: Mon, 22 Apr 2024 23:16:11 +0530 Subject: [PATCH 2/5] Adds is not None check --- airflow/providers/databricks/operators/databricks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py index 4514355579996..57abd958b8cc8 100644 --- a/airflow/providers/databricks/operators/databricks.py +++ b/airflow/providers/databricks/operators/databricks.py @@ -287,7 +287,7 @@ def __init__( self.json["email_notifications"] = email_notifications if webhook_notifications is not None: self.json["webhook_notifications"] = webhook_notifications - if notification_settings: + if notification_settings is not None: self.json["notification_settings"] = notification_settings if timeout_seconds is not None: self.json["timeout_seconds"] = timeout_seconds From a1326629fc9a125aef3787be938b70b3a9c917c6 Mon Sep 17 00:00:00 2001 From: subham611 Date: Mon, 22 Apr 2024 23:19:28 +0530 Subject: [PATCH 3/5] Updates doc --- airflow/providers/databricks/operators/databricks.py | 1 + .../operators/jobs_create.rst | 1 + 2 files changed, 2 insertions(+) diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py index 57abd958b8cc8..f1bdc1d45ebd3 100644 --- a/airflow/providers/databricks/operators/databricks.py +++ b/airflow/providers/databricks/operators/databricks.py @@ -214,6 +214,7 @@ class DatabricksCreateJobsOperator(BaseOperator): tasks of this job. Array of objects (JobCluster). :param email_notifications: Object (JobEmailNotifications). :param webhook_notifications: Object (WebhookNotifications). + :param notification_settings: Optional notification settings :param timeout_seconds: An optional timeout applied to each run of this job. :param schedule: Object (CronSchedule). :param max_concurrent_runs: An optional maximum allowed number of concurrent runs of the job. diff --git a/docs/apache-airflow-providers-databricks/operators/jobs_create.rst b/docs/apache-airflow-providers-databricks/operators/jobs_create.rst index 779095e92cd6b..6f93ce68579f1 100644 --- a/docs/apache-airflow-providers-databricks/operators/jobs_create.rst +++ b/docs/apache-airflow-providers-databricks/operators/jobs_create.rst @@ -49,6 +49,7 @@ Currently the named parameters that ``DatabricksCreateJobsOperator`` supports ar - ``job_clusters`` - ``email_notifications`` - ``webhook_notifications`` + - ``notification_settings`` - ``timeout_seconds`` - ``schedule`` - ``max_concurrent_runs`` From 0899589253abf37e5b53c096190eb5a6946c999a Mon Sep 17 00:00:00 2001 From: subham611 Date: Mon, 22 Apr 2024 23:35:49 +0530 Subject: [PATCH 4/5] Adds description param --- airflow/providers/databricks/operators/databricks.py | 6 +++++- .../operators/jobs_create.rst | 1 + tests/providers/databricks/operators/test_databricks.py | 5 +++++ 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py index f1bdc1d45ebd3..22f37e95555d6 100644 --- a/airflow/providers/databricks/operators/databricks.py +++ b/airflow/providers/databricks/operators/databricks.py @@ -207,6 +207,7 @@ class DatabricksCreateJobsOperator(BaseOperator): .. seealso:: For more information about templating see :ref:`concepts:jinja-templating`. :param name: An optional name for the job. + :param description: An optional description for the job. :param tags: A map of tags associated with the job. :param tasks: A list of task specifications to be executed by this job. Array of objects (JobTaskSettings). @@ -214,7 +215,7 @@ class DatabricksCreateJobsOperator(BaseOperator): tasks of this job. Array of objects (JobCluster). :param email_notifications: Object (JobEmailNotifications). :param webhook_notifications: Object (WebhookNotifications). - :param notification_settings: Optional notification settings + :param notification_settings: Optional notification settings. :param timeout_seconds: An optional timeout applied to each run of this job. :param schedule: Object (CronSchedule). :param max_concurrent_runs: An optional maximum allowed number of concurrent runs of the job. @@ -250,6 +251,7 @@ def __init__( *, json: Any | None = None, name: str | None = None, + description: str | None = None, tags: dict[str, str] | None = None, tasks: list[dict] | None = None, job_clusters: list[dict] | None = None, @@ -278,6 +280,8 @@ def __init__( self.databricks_retry_args = databricks_retry_args if name is not None: self.json["name"] = name + if description is not None: + self.json["description"] = description if tags is not None: self.json["tags"] = tags if tasks is not None: diff --git a/docs/apache-airflow-providers-databricks/operators/jobs_create.rst b/docs/apache-airflow-providers-databricks/operators/jobs_create.rst index 6f93ce68579f1..7e6765eba420a 100644 --- a/docs/apache-airflow-providers-databricks/operators/jobs_create.rst +++ b/docs/apache-airflow-providers-databricks/operators/jobs_create.rst @@ -44,6 +44,7 @@ override the top level ``json`` keys. Currently the named parameters that ``DatabricksCreateJobsOperator`` supports are: - ``name`` + - ``description`` - ``tags`` - ``tasks`` - ``job_clusters`` diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py index 885d5c293d055..66f7bf878142a 100644 --- a/tests/providers/databricks/operators/test_databricks.py +++ b/tests/providers/databricks/operators/test_databricks.py @@ -63,6 +63,7 @@ RUN_PAGE_URL = "run-page-url" JOB_ID = "42" JOB_NAME = "job-name" +JOB_DESCRIPTION = "job-description" NOTEBOOK_PARAMS = {"dry-run": "true", "oldest-time-to-consider": "1457570074236"} JAR_PARAMS = ["param1", "param2"] RENDERED_TEMPLATED_JAR_PARAMS = [f"/test-{DATE}"] @@ -413,6 +414,7 @@ def test_exec_create(self, db_mock_class): """ json = { "name": JOB_NAME, + "description": JOB_DESCRIPTION, "tags": TAGS, "tasks": TASKS, "job_clusters": JOB_CLUSTERS, @@ -436,6 +438,7 @@ def test_exec_create(self, db_mock_class): expected = utils.normalise_json_content( { "name": JOB_NAME, + "description": JOB_DESCRIPTION, "tags": TAGS, "tasks": TASKS, "job_clusters": JOB_CLUSTERS, @@ -467,6 +470,7 @@ def test_exec_reset(self, db_mock_class): """ json = { "name": JOB_NAME, + "description": JOB_DESCRIPTION, "tags": TAGS, "tasks": TASKS, "job_clusters": JOB_CLUSTERS, @@ -488,6 +492,7 @@ def test_exec_reset(self, db_mock_class): expected = utils.normalise_json_content( { "name": JOB_NAME, + "description": JOB_DESCRIPTION, "tags": TAGS, "tasks": TASKS, "job_clusters": JOB_CLUSTERS, From 7927c5e84edec07b480543de323a70af2d916f51 Mon Sep 17 00:00:00 2001 From: subham611 Date: Tue, 23 Apr 2024 00:02:57 +0530 Subject: [PATCH 5/5] Fix static checks --- tests/providers/databricks/operators/test_databricks.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py index 66f7bf878142a..26d59baa61610 100644 --- a/tests/providers/databricks/operators/test_databricks.py +++ b/tests/providers/databricks/operators/test_databricks.py @@ -203,10 +203,7 @@ } ], } -NOTIFICATION_SETTINGS = { - "no_alert_for_canceled_runs": True, - "no_alert_for_skipped_runs": True -} +NOTIFICATION_SETTINGS = {"no_alert_for_canceled_runs": True, "no_alert_for_skipped_runs": True} TIMEOUT_SECONDS = 86400 SCHEDULE = { "quartz_cron_expression": "20 30 * * * ?",