From dc3ceec33d97e5076c723f267ad44c6d688d6622 Mon Sep 17 00:00:00 2001 From: Maksim Zinal Date: Tue, 28 Feb 2023 19:14:01 +0300 Subject: [PATCH] add support for Yandex Dataproc cluster labels --- airflow/providers/yandex/operators/yandexcloud_dataproc.py | 5 +++++ airflow/providers/yandex/provider.yaml | 2 +- .../providers/yandex/operators/test_yandexcloud_dataproc.py | 1 + 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/airflow/providers/yandex/operators/yandexcloud_dataproc.py b/airflow/providers/yandex/operators/yandexcloud_dataproc.py index 625827d109963..dfd3a07fe4f94 100644 --- a/airflow/providers/yandex/operators/yandexcloud_dataproc.py +++ b/airflow/providers/yandex/operators/yandexcloud_dataproc.py @@ -95,6 +95,8 @@ class DataprocCreateClusterOperator(BaseOperator): Docs: https://cloud.yandex.com/docs/data-proc/concepts/logs :param initialization_actions: Set of init-actions to run when cluster starts. Docs: https://cloud.yandex.com/docs/data-proc/concepts/init-action + :param labels: Cluster labels as key:value pairs. No more than 64 per resource. + Docs: https://cloud.yandex.ru/docs/resource-manager/concepts/labels """ def __init__( @@ -135,6 +137,7 @@ def __init__( security_group_ids: Iterable[str] | None = None, log_group_id: str | None = None, initialization_actions: Iterable[InitializationAction] | None = None, + labels: dict[str, str] | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -173,6 +176,7 @@ def __init__( self.security_group_ids = security_group_ids self.log_group_id = log_group_id self.initialization_actions = initialization_actions + self.labels = labels self.hook: DataprocHook | None = None @@ -214,6 +218,7 @@ def execute(self, context: Context) -> dict: host_group_ids=self.host_group_ids, security_group_ids=self.security_group_ids, log_group_id=self.log_group_id, + labels=self.labels, initialization_actions=self.initialization_actions and [ self.hook.sdk.wrappers.InitializationAction( diff --git a/airflow/providers/yandex/provider.yaml b/airflow/providers/yandex/provider.yaml index 5a23ce60af16e..798294315bbac 100644 --- a/airflow/providers/yandex/provider.yaml +++ b/airflow/providers/yandex/provider.yaml @@ -37,7 +37,7 @@ versions: dependencies: - apache-airflow>=2.3.0 - - yandexcloud>=0.173.0 + - yandexcloud>=0.202.0 integrations: - integration-name: Yandex.Cloud diff --git a/tests/providers/yandex/operators/test_yandexcloud_dataproc.py b/tests/providers/yandex/operators/test_yandexcloud_dataproc.py index 8ffe2050a037c..879645daf60cb 100644 --- a/tests/providers/yandex/operators/test_yandexcloud_dataproc.py +++ b/tests/providers/yandex/operators/test_yandexcloud_dataproc.py @@ -130,6 +130,7 @@ def test_create_cluster(self, create_cluster_mock, *_): enable_ui_proxy=False, host_group_ids=None, security_group_ids=None, + labels=None, initialization_actions=None, ) context["task_instance"].xcom_push.assert_has_calls(