Improve memory usage in Dataproc deferrable operators (#28117)

Łukasz Wyszomirski · web-flow · commit bdf3175e9616 · 2022-12-06T00:58:33.000+01:00
diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py
@@ -992,6 +992,7 @@ def __init__(
         impersonation_chain: str | Sequence[str] | None = None,
     ) -> None:
         super().__init__(gcp_conn_id, delegate_to, impersonation_chain)
+        self._cached_client: JobControllerAsyncClient | None = None
 
     def get_cluster_client(self, region: str | None = None) -> ClusterControllerAsyncClient:
         """Returns ClusterControllerAsyncClient."""
@@ -1015,15 +1016,17 @@ def get_template_client(self, region: str | None = None) -> WorkflowTemplateServ
 
     def get_job_client(self, region: str | None = None) -> JobControllerAsyncClient:
         """Returns JobControllerAsyncClient."""
-        client_options = None
-        if region and region != "global":
-            client_options = ClientOptions(api_endpoint=f"{region}-dataproc.googleapis.com:443")
-
-        return JobControllerAsyncClient(
-            credentials=self.get_credentials(),
-            client_info=CLIENT_INFO,
-            client_options=client_options,
-        )
+        if self._cached_client is None:
+            client_options = None
+            if region and region != "global":
+                client_options = ClientOptions(api_endpoint=f"{region}-dataproc.googleapis.com:443")
+
+            self._cached_client = JobControllerAsyncClient(
+                credentials=self.get_credentials(),
+                client_info=CLIENT_INFO,
+                client_options=client_options,
+            )
+        return self._cached_client
 
     def get_batch_client(self, region: str | None = None) -> BatchControllerAsyncClient:
         """Returns BatchControllerAsyncClient"""
