Improve memory usage in Dataproc deferrable operators (#28117)

apache · Dec 5, 2022 · bdf3175 · bdf3175
1 parent 3fef462
commit bdf3175
Showing 1 changed file with 12 additions and 9 deletions.
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"""