Skip to content

Commit 892522f

Browse files
turbaszekturbaszek
authored
Change signature of GSheetsHook methods (#7853)
1 parent 2d0db94 commit 892522f

File tree

2 files changed

+43
-22
lines changed

2 files changed

+43
-22
lines changed

β€Žairflow/providers/google/suite/hooks/sheets.py

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,6 @@ class GSheetsHook(CloudBaseHook):
3636
3737
:param gcp_conn_id: The connection ID to use when fetching connection info.
3838
:type gcp_conn_id: str
39-
:param spreadsheet_id: The Google Sheet ID to interact with
40-
:type spreadsheet_id: str
4139
:param api_version: API Version
4240
:type api_version: str
4341
:param delegate_to: The account to impersonate, if any.
@@ -48,13 +46,11 @@ class GSheetsHook(CloudBaseHook):
4846

4947
def __init__(
5048
self,
51-
spreadsheet_id: str,
5249
gcp_conn_id: str = 'google_cloud_default',
5350
api_version: str = 'v4',
5451
delegate_to: Optional[str] = None
5552
) -> None:
5653
super().__init__(gcp_conn_id, delegate_to)
57-
self.spreadsheet_id = spreadsheet_id
5854
self.gcp_conn_id = gcp_conn_id
5955
self.api_version = api_version
6056
self.delegate_to = delegate_to
@@ -75,15 +71,18 @@ def get_conn(self) -> Any:
7571

7672
def get_values(
7773
self,
74+
spreadsheet_id: str,
7875
range_: str,
7976
major_dimension: str = 'DIMENSION_UNSPECIFIED',
8077
value_render_option: str = 'FORMATTED_VALUE',
8178
date_time_render_option: str = 'SERIAL_NUMBER'
82-
) -> Dict:
79+
) -> List:
8380
"""
8481
Gets values from Google Sheet from a single range
8582
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/get
8683
84+
:param spreadsheet_id: The Google Sheet ID to interact with
85+
:type spreadsheet_id: str
8786
:param range_: The A1 notation of the values to retrieve.
8887
:type range_: str
8988
:param major_dimension: Indicates which dimension an operation should apply to.
@@ -95,22 +94,23 @@ def get_values(
9594
:param date_time_render_option: Determines how dates should be rendered in the output.
9695
SERIAL_NUMBER or FORMATTED_STRING
9796
:type date_time_render_option: str
98-
:return: Google Sheets API response.
99-
:rtype: Dict
97+
:return: An array of sheet values from the specified sheet.
98+
:rtype: List
10099
"""
101100
service = self.get_conn()
102101
response = service.spreadsheets().values().get( # pylint: disable=no-member
103-
spreadsheetId=self.spreadsheet_id,
102+
spreadsheetId=spreadsheet_id,
104103
range=range_,
105104
majorDimension=major_dimension,
106105
valueRenderOption=value_render_option,
107106
dateTimeRenderOption=date_time_render_option
108107
).execute(num_retries=self.num_retries)
109108

110-
return response
109+
return response['values']
111110

112111
def batch_get_values(
113112
self,
113+
spreadsheet_id: str,
114114
ranges: List,
115115
major_dimension: str = 'DIMENSION_UNSPECIFIED',
116116
value_render_option: str = 'FORMATTED_VALUE',
@@ -120,6 +120,8 @@ def batch_get_values(
120120
Gets values from Google Sheet from a list of ranges
121121
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/batchGet
122122
123+
:param spreadsheet_id: The Google Sheet ID to interact with
124+
:type spreadsheet_id: str
123125
:param ranges: The A1 notation of the values to retrieve.
124126
:type ranges: List
125127
:param major_dimension: Indicates which dimension an operation should apply to.
@@ -136,7 +138,7 @@ def batch_get_values(
136138
"""
137139
service = self.get_conn()
138140
response = service.spreadsheets().values().batchGet( # pylint: disable=no-member
139-
spreadsheetId=self.spreadsheet_id,
141+
spreadsheetId=spreadsheet_id,
140142
ranges=ranges,
141143
majorDimension=major_dimension,
142144
valueRenderOption=value_render_option,
@@ -147,6 +149,7 @@ def batch_get_values(
147149

148150
def update_values(
149151
self,
152+
spreadsheet_id: str,
150153
range_: str,
151154
values: List,
152155
major_dimension: str = 'ROWS',
@@ -159,6 +162,8 @@ def update_values(
159162
Updates values from Google Sheet from a single range
160163
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/update
161164
165+
:param spreadsheet_id: The Google Sheet ID to interact with
166+
:type spreadsheet_id: str
162167
:param range_: The A1 notation of the values to retrieve.
163168
:type range_: str
164169
:param values: Data within a range of the spreadsheet.
@@ -188,7 +193,7 @@ def update_values(
188193
"values": values
189194
}
190195
response = service.spreadsheets().values().update( # pylint: disable=no-member
191-
spreadsheetId=self.spreadsheet_id,
196+
spreadsheetId=spreadsheet_id,
192197
range=range_,
193198
valueInputOption=value_input_option,
194199
includeValuesInResponse=include_values_in_response,
@@ -201,6 +206,7 @@ def update_values(
201206

202207
def batch_update_values(
203208
self,
209+
spreadsheet_id: str,
204210
ranges: List,
205211
values: List,
206212
major_dimension: str = 'ROWS',
@@ -213,6 +219,8 @@ def batch_update_values(
213219
Updates values from Google Sheet for multiple ranges
214220
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/batchUpdate
215221
222+
:param spreadsheet_id: The Google Sheet ID to interact with
223+
:type spreadsheet_id: str
216224
:param ranges: The A1 notation of the values to retrieve.
217225
:type ranges: List
218226
:param values: Data within a range of the spreadsheet.
@@ -257,14 +265,15 @@ def batch_update_values(
257265
"responseDateTimeRenderOption": date_time_render_option
258266
}
259267
response = service.spreadsheets().values().batchUpdate( # pylint: disable=no-member
260-
spreadsheetId=self.spreadsheet_id,
268+
spreadsheetId=spreadsheet_id,
261269
body=body
262270
).execute(num_retries=self.num_retries)
263271

264272
return response
265273

266274
def append_values(
267275
self,
276+
spreadsheet_id: str,
268277
range_: str,
269278
values: List,
270279
major_dimension: str = 'ROWS',
@@ -278,6 +287,8 @@ def append_values(
278287
Append values from Google Sheet from a single range
279288
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/append
280289
290+
:param spreadsheet_id: The Google Sheet ID to interact with
291+
:type spreadsheet_id: str
281292
:param range_: The A1 notation of the values to retrieve.
282293
:type range_: str
283294
:param values: Data within a range of the spreadsheet.
@@ -310,7 +321,7 @@ def append_values(
310321
"values": values
311322
}
312323
response = service.spreadsheets().values().append( # pylint: disable=no-member
313-
spreadsheetId=self.spreadsheet_id,
324+
spreadsheetId=spreadsheet_id,
314325
range=range_,
315326
valueInputOption=value_input_option,
316327
insertDataOption=insert_data_option,
@@ -322,29 +333,33 @@ def append_values(
322333

323334
return response
324335

325-
def clear(self, range_: str) -> Dict:
336+
def clear(self, spreadsheet_id: str, range_: str) -> Dict:
326337
"""
327338
Clear values from Google Sheet from a single range
328339
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/clear
329340
341+
:param spreadsheet_id: The Google Sheet ID to interact with
342+
:type spreadsheet_id: str
330343
:param range_: The A1 notation of the values to retrieve.
331344
:type range_: str
332345
:return: Google Sheets API response.
333346
:rtype: Dict
334347
"""
335348
service = self.get_conn()
336349
response = service.spreadsheets().values().clear( # pylint: disable=no-member
337-
spreadsheetId=self.spreadsheet_id,
350+
spreadsheetId=spreadsheet_id,
338351
range=range_
339352
).execute(num_retries=self.num_retries)
340353

341354
return response
342355

343-
def batch_clear(self, ranges: List) -> Dict:
356+
def batch_clear(self, spreadsheet_id: str, ranges: List) -> Dict:
344357
"""
345358
Clear values from Google Sheet from a list of ranges
346359
https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/batchClear
347360
361+
:param spreadsheet_id: The Google Sheet ID to interact with
362+
:type spreadsheet_id: str
348363
:param ranges: The A1 notation of the values to retrieve.
349364
:type ranges: List
350365
:return: Google Sheets API response.
@@ -355,7 +370,7 @@ def batch_clear(self, ranges: List) -> Dict:
355370
"ranges": ranges
356371
}
357372
response = service.spreadsheets().values().batchClear( # pylint: disable=no-member
358-
spreadsheetId=self.spreadsheet_id,
373+
spreadsheetId=spreadsheet_id,
359374
body=body
360375
).execute(num_retries=self.num_retries)
361376

β€Žtests/providers/google/suite/hooks/test_sheets.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ class TestGSheetsHook(unittest.TestCase):
4848
def setUp(self):
4949
with mock.patch('airflow.providers.google.cloud.hooks.base.CloudBaseHook.__init__',
5050
new=mock_base_gcp_hook_default_project_id):
51-
self.hook = GSheetsHook(gcp_conn_id=GCP_CONN_ID, spreadsheet_id=SPREADHSEET_ID)
51+
self.hook = GSheetsHook(gcp_conn_id=GCP_CONN_ID)
5252

5353
@mock.patch("airflow.providers.google.suite.hooks.sheets.GSheetsHook._authorize")
5454
@mock.patch("airflow.providers.google.suite.hooks.sheets.build")
@@ -63,13 +63,14 @@ def test_gsheets_client_creation(self, mock_build, mock_authorize):
6363
def test_get_values(self, get_conn):
6464
get_method = get_conn.return_value.spreadsheets.return_value.values.return_value.get
6565
execute_method = get_method.return_value.execute
66-
execute_method.return_value = API_RESPONSE
66+
execute_method.return_value = {"values": VALUES}
6767
result = self.hook.get_values(
68+
spreadsheet_id=SPREADHSEET_ID,
6869
range_=RANGE_,
6970
major_dimension=MAJOR_DIMENSION,
7071
value_render_option=VALUE_RENDER_OPTION,
7172
date_time_render_option=DATE_TIME_RENDER_OPTION)
72-
self.assertIs(result, API_RESPONSE)
73+
self.assertIs(result, VALUES)
7374
execute_method.assert_called_once_with(num_retries=NUM_RETRIES)
7475
get_method.assert_called_once_with(
7576
spreadsheetId=SPREADHSEET_ID,
@@ -85,6 +86,7 @@ def test_batch_get_values(self, get_conn):
8586
execute_method = batch_get_method.return_value.execute
8687
execute_method.return_value = API_RESPONSE
8788
result = self.hook.batch_get_values(
89+
spreadsheet_id=SPREADHSEET_ID,
8890
ranges=RANGES,
8991
major_dimension=MAJOR_DIMENSION,
9092
value_render_option=VALUE_RENDER_OPTION,
@@ -105,6 +107,7 @@ def test_update_values(self, get_conn):
105107
execute_method = update_method.return_value.execute
106108
execute_method.return_value = API_RESPONSE
107109
result = self.hook.update_values(
110+
spreadsheet_id=SPREADHSEET_ID,
108111
range_=RANGE_,
109112
values=VALUES,
110113
major_dimension=MAJOR_DIMENSION,
@@ -135,6 +138,7 @@ def test_batch_update_values(self, get_conn):
135138
execute_method = batch_update_method.return_value.execute
136139
execute_method.return_value = API_RESPONSE
137140
result = self.hook.batch_update_values(
141+
spreadsheet_id=SPREADHSEET_ID,
138142
ranges=RANGES,
139143
values=VALUES_BATCH,
140144
major_dimension=MAJOR_DIMENSION,
@@ -171,6 +175,7 @@ def test_batch_update_values_with_bad_data(self, get_conn):
171175
execute_method.return_value = API_RESPONSE
172176
with self.assertRaises(AirflowException) as cm:
173177
self.hook.batch_update_values(
178+
spreadsheet_id=SPREADHSEET_ID,
174179
ranges=['test!A1:B2', 'test!C1:C2'],
175180
values=[[1, 2, 3]], # bad data
176181
major_dimension=MAJOR_DIMENSION,
@@ -189,6 +194,7 @@ def test_append_values(self, get_conn):
189194
execute_method = append_method.return_value.execute
190195
execute_method.return_value = API_RESPONSE
191196
result = self.hook.append_values(
197+
spreadsheet_id=SPREADHSEET_ID,
192198
range_=RANGE_,
193199
values=VALUES,
194200
major_dimension=MAJOR_DIMENSION,
@@ -220,7 +226,7 @@ def test_clear_values(self, get_conn):
220226
clear_method = get_conn.return_value.spreadsheets.return_value.values.return_value.clear
221227
execute_method = clear_method.return_value.execute
222228
execute_method.return_value = API_RESPONSE
223-
result = self.hook.clear(range_=RANGE_)
229+
result = self.hook.clear(spreadsheet_id=SPREADHSEET_ID, range_=RANGE_)
224230

225231
self.assertIs(result, API_RESPONSE)
226232
execute_method.assert_called_once_with(num_retries=NUM_RETRIES)
@@ -234,7 +240,7 @@ def test_batch_clear_values(self, get_conn):
234240
batch_clear_method = get_conn.return_value.spreadsheets.return_value.values.return_value.batchClear
235241
execute_method = batch_clear_method.return_value.execute
236242
execute_method.return_value = API_RESPONSE
237-
result = self.hook.batch_clear(ranges=RANGES)
243+
result = self.hook.batch_clear(spreadsheet_id=SPREADHSEET_ID, ranges=RANGES)
238244
body = {"ranges": RANGES}
239245
self.assertIs(result, API_RESPONSE)
240246
execute_method.assert_called_once_with(num_retries=NUM_RETRIES)

0 commit comments

Comments
 (0)