Skip to content

Commit

Permalink
Append query to the queries sent to dataproc instead of replacing
Browse files Browse the repository at this point in the history
Requested by reviewer, instead of adding a function that sets a list of queries, just fix the original function add_query that will actually append query to the queries list send to dataproc job.
  • Loading branch information
Amir Mor authored and amirmor1 committed Dec 14, 2024
1 parent 3ba9ba4 commit 8155625
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 18 deletions.
12 changes: 2 additions & 10 deletions providers/src/airflow/providers/google/cloud/hooks/dataproc.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,19 +118,11 @@ def add_args(self, args: list[str] | None = None) -> None:

def add_query(self, query: 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]}

def set_queries(self, queries: list[str]) -> None:
"""
Set queries for Dataproc job.
:param queries: queries for the job.
"""
self.job["job"][self.job_type]["query_list"] = {"queries": queries}
self.job["job"][self.job_type].setdefault("query_list", {"queries": []})["queries"].append(query)

def add_query_uri(self, query_uri: str) -> None:
"""
Expand Down
13 changes: 5 additions & 8 deletions providers/tests/google/cloud/hooks/test_dataproc.py
Original file line number Diff line number Diff line change
Expand Up @@ -1097,14 +1097,11 @@ 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]}

def test_set_queries(self):
queries = ["query1", "query2"]
self.builder.set_queries(queries)
assert self.builder.job["job"][self.job_type]["query_list"] == {"queries": queries}
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]}

def test_add_query_uri(self):
query_uri = "query_uri"
Expand Down

0 comments on commit 8155625

Please sign in to comment.