diff --git a/providers/src/airflow/providers/google/cloud/hooks/dataproc.py b/providers/src/airflow/providers/google/cloud/hooks/dataproc.py index b704c23107e71..6cc57128bab4b 100644 --- a/providers/src/airflow/providers/google/cloud/hooks/dataproc.py +++ b/providers/src/airflow/providers/google/cloud/hooks/dataproc.py @@ -116,13 +116,17 @@ def add_args(self, args: list[str] | None = None) -> None: if args is not None: self.job["job"][self.job_type]["args"] = args - def add_query(self, query: str) -> None: + def add_query(self, query: str | list[str]) -> None: """ - Set query for Dataproc job. + Add query for Dataproc job. :param query: query for the job. """ - self.job["job"][self.job_type]["query_list"] = {"queries": [query]} + queries = self.job["job"][self.job_type].setdefault("query_list", {"queries": []})["queries"] + if isinstance(query, str): + queries.append(query) + elif isinstance(query, list): + queries.extend(query) def add_query_uri(self, query_uri: str) -> None: """ diff --git a/providers/tests/google/cloud/hooks/test_dataproc.py b/providers/tests/google/cloud/hooks/test_dataproc.py index b15d7855454e9..1e0349d8bfdcd 100644 --- a/providers/tests/google/cloud/hooks/test_dataproc.py +++ b/providers/tests/google/cloud/hooks/test_dataproc.py @@ -1097,9 +1097,16 @@ def test_add_args(self): assert args == self.builder.job["job"][self.job_type]["args"] def test_add_query(self): - query = ["query"] - self.builder.add_query(query) - assert self.builder.job["job"][self.job_type]["query_list"] == {"queries": [query]} + query1 = "query1" + self.builder.add_query(query1) + query2 = "query2" + self.builder.add_query(query2) + assert self.builder.job["job"][self.job_type]["query_list"] == {"queries": [query1, query2]} + new_queries = ["query3", "query4"] + self.builder.add_query(new_queries) + assert self.builder.job["job"][self.job_type]["query_list"] == { + "queries": [query1, query2] + new_queries + } def test_add_query_uri(self): query_uri = "query_uri"