Skip to content

Commit a22d5bd

Browse files
potiukpotiuk
authored
Fix mypy errors in Google Cloud provider (#20611)
Part of #19891 Another attempt to clean-up all MyPy errors in Google Provider.
1 parent 2d09202 commit a22d5bd

28 files changed

Lines changed: 220 additions & 168 deletions

‎airflow/providers/google/cloud/example_dags/example_cloud_build.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
from pathlib import Path
3434
from typing import Any, Dict
3535

36+
import yaml
3637
from future.backports.urllib.parse import urlparse
3738

3839
from airflow import models
@@ -190,7 +191,7 @@
190191
create_build_from_file = CloudBuildCreateBuildOperator(
191192
task_id="create_build_from_file",
192193
project_id=GCP_PROJECT_ID,
193-
build=str(CURRENT_FOLDER.joinpath('example_cloud_build.yaml')),
194+
build=yaml.safe_load((Path(CURRENT_FOLDER) / 'example_cloud_build.yaml').read_text()),
194195
params={'name': 'Airflow'},
195196
)
196197
# [END howto_operator_gcp_create_build_from_yaml_body]

‎airflow/providers/google/cloud/example_dags/example_tasks.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from google.api_core.retry import Retry
2828
from google.cloud.tasks_v2.types import Queue
2929
from google.protobuf import timestamp_pb2
30+
from google.protobuf.field_mask_pb2 import FieldMask
3031

3132
from airflow import models
3233
from airflow.models.baseoperator import chain
@@ -136,7 +137,7 @@
136137
task_queue=Queue(stackdriver_logging_config=dict(sampling_ratio=1)),
137138
location=LOCATION,
138139
queue_name=QUEUE_ID,
139-
update_mask={"paths": ["stackdriver_logging_config.sampling_ratio"]},
140+
update_mask=FieldMask(paths=["stackdriver_logging_config.sampling_ratio"]),
140141
task_id="update_queue",
141142
)
142143
# [END update_queue]

‎airflow/providers/google/cloud/example_dags/example_workflows.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import os
1919
from datetime import datetime
2020

21+
from google.protobuf.field_mask_pb2 import FieldMask
22+
2123
from airflow import DAG
2224
from airflow.providers.google.cloud.operators.workflows import (
2325
WorkflowsCancelExecutionOperator,
@@ -102,7 +104,7 @@
102104
location=LOCATION,
103105
project_id=PROJECT_ID,
104106
workflow_id=WORKFLOW_ID,
105-
update_mask={"paths": ["name", "description"]},
107+
update_mask=FieldMask(paths=["name", "description"]),
106108
)
107109
# [END how_to_update_workflow]
108110

‎airflow/providers/google/cloud/hooks/datastore.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -262,21 +262,21 @@ def delete_operation(self, name: str) -> dict:
262262

263263
return resp
264264

265-
def poll_operation_until_done(self, name: str, polling_interval_in_seconds: int) -> Dict:
265+
def poll_operation_until_done(self, name: str, polling_interval_in_seconds: float) -> Dict:
266266
"""
267267
Poll backup operation state until it's completed.
268268
269269
:param name: the name of the operation resource
270270
:type name: str
271271
:param polling_interval_in_seconds: The number of seconds to wait before calling another request.
272-
:type polling_interval_in_seconds: int
272+
:type polling_interval_in_seconds: float
273273
:return: a resource operation instance.
274274
:rtype: dict
275275
"""
276276
while True:
277-
result = self.get_operation(name) # type: Dict
277+
result: Dict = self.get_operation(name)
278278

279-
state = result['metadata']['common']['state'] # type: str
279+
state: str = result['metadata']['common']['state']
280280
if state == 'PROCESSING':
281281
self.log.info(
282282
'Operation is processing. Re-polling state in %s seconds', polling_interval_in_seconds

‎airflow/providers/google/cloud/hooks/dlp.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ def create_inspect_template(
285285
self,
286286
organization_id: Optional[str] = None,
287287
project_id: Optional[str] = None,
288-
inspect_template: Optional[Union[dict, InspectTemplate]] = None,
288+
inspect_template: Optional[InspectTemplate] = None,
289289
template_id: Optional[str] = None,
290290
retry: Optional[Retry] = None,
291291
timeout: Optional[float] = None,

‎airflow/providers/google/cloud/hooks/gcs.py‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
from io import BytesIO
3030
from os import path
3131
from tempfile import NamedTemporaryFile
32-
from typing import Callable, List, Optional, Sequence, Set, Tuple, TypeVar, Union, cast
32+
from typing import Callable, List, Optional, Sequence, Set, Tuple, TypeVar, Union, cast, overload
3333
from urllib.parse import urlparse
3434

3535
from google.api_core.exceptions import NotFound
@@ -273,6 +273,30 @@ def rewrite(
273273
destination_bucket.name, # type: ignore[attr-defined]
274274
)
275275

276+
@overload
277+
def download(
278+
self,
279+
bucket_name: str,
280+
object_name: str,
281+
filename: None = None,
282+
chunk_size: Optional[int] = None,
283+
timeout: Optional[int] = DEFAULT_TIMEOUT,
284+
num_max_attempts: Optional[int] = 1,
285+
) -> bytes:
286+
...
287+
288+
@overload
289+
def download(
290+
self,
291+
bucket_name: str,
292+
object_name: str,
293+
filename: str,
294+
chunk_size: Optional[int] = None,
295+
timeout: Optional[int] = DEFAULT_TIMEOUT,
296+
num_max_attempts: Optional[int] = 1,
297+
) -> str:
298+
...
299+
276300
def download(
277301
self,
278302
bucket_name: str,
@@ -366,15 +390,12 @@ def download_as_byte_array(
366390
:type num_max_attempts: int
367391
"""
368392
# We do not pass filename, so will never receive string as response
369-
return cast(
370-
bytes,
371-
self.download(
372-
bucket_name=bucket_name,
373-
object_name=object_name,
374-
chunk_size=chunk_size,
375-
timeout=timeout,
376-
num_max_attempts=num_max_attempts,
377-
),
393+
return self.download(
394+
bucket_name=bucket_name,
395+
object_name=object_name,
396+
chunk_size=chunk_size,
397+
timeout=timeout,
398+
num_max_attempts=num_max_attempts,
378399
)
379400

380401
@_fallback_object_url_to_object_name_and_bucket_name()

‎airflow/providers/google/cloud/hooks/stackdriver.py‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def list_alert_policies(
6868
order_by: Optional[str] = None,
6969
page_size: Optional[int] = None,
7070
retry: Optional[str] = DEFAULT,
71-
timeout: Optional[float] = DEFAULT,
71+
timeout: Optional[float] = None,
7272
metadata: Sequence[Tuple[str, str]] = (),
7373
) -> Any:
7474
"""
@@ -135,7 +135,7 @@ def _toggle_policy_status(
135135
project_id: str = PROVIDE_PROJECT_ID,
136136
filter_: Optional[str] = None,
137137
retry: Optional[str] = DEFAULT,
138-
timeout: Optional[float] = DEFAULT,
138+
timeout: Optional[float] = None,
139139
metadata: Sequence[Tuple[str, str]] = (),
140140
):
141141
client = self._get_policy_client()
@@ -157,7 +157,7 @@ def enable_alert_policies(
157157
project_id: str = PROVIDE_PROJECT_ID,
158158
filter_: Optional[str] = None,
159159
retry: Optional[str] = DEFAULT,
160-
timeout: Optional[float] = DEFAULT,
160+
timeout: Optional[float] = None,
161161
metadata: Sequence[Tuple[str, str]] = (),
162162
) -> None:
163163
"""
@@ -195,7 +195,7 @@ def disable_alert_policies(
195195
project_id: str = PROVIDE_PROJECT_ID,
196196
filter_: Optional[str] = None,
197197
retry: Optional[str] = DEFAULT,
198-
timeout: Optional[float] = DEFAULT,
198+
timeout: Optional[float] = None,
199199
metadata: Sequence[Tuple[str, str]] = (),
200200
) -> None:
201201
"""
@@ -233,7 +233,7 @@ def upsert_alert(
233233
alerts: str,
234234
project_id: str = PROVIDE_PROJECT_ID,
235235
retry: Optional[str] = DEFAULT,
236-
timeout: Optional[float] = DEFAULT,
236+
timeout: Optional[float] = None,
237237
metadata: Sequence[Tuple[str, str]] = (),
238238
) -> None:
239239
"""
@@ -334,7 +334,7 @@ def delete_alert_policy(
334334
self,
335335
name: str,
336336
retry: Optional[str] = DEFAULT,
337-
timeout: Optional[float] = DEFAULT,
337+
timeout: Optional[float] = None,
338338
metadata: Sequence[Tuple[str, str]] = (),
339339
) -> None:
340340
"""
@@ -370,7 +370,7 @@ def list_notification_channels(
370370
order_by: Optional[str] = None,
371371
page_size: Optional[int] = None,
372372
retry: Optional[str] = DEFAULT,
373-
timeout: Optional[str] = DEFAULT,
373+
timeout: Optional[float] = None,
374374
metadata: Sequence[Tuple[str, str]] = (),
375375
) -> Any:
376376
"""
@@ -437,7 +437,7 @@ def _toggle_channel_status(
437437
project_id: str = PROVIDE_PROJECT_ID,
438438
filter_: Optional[str] = None,
439439
retry: Optional[str] = DEFAULT,
440-
timeout: Optional[str] = DEFAULT,
440+
timeout: Optional[float] = None,
441441
metadata: Sequence[Tuple[str, str]] = (),
442442
) -> None:
443443
client = self._get_channel_client()
@@ -461,7 +461,7 @@ def enable_notification_channels(
461461
project_id: str = PROVIDE_PROJECT_ID,
462462
filter_: Optional[str] = None,
463463
retry: Optional[str] = DEFAULT,
464-
timeout: Optional[str] = DEFAULT,
464+
timeout: Optional[float] = None,
465465
metadata: Sequence[Tuple[str, str]] = (),
466466
) -> None:
467467
"""
@@ -499,7 +499,7 @@ def disable_notification_channels(
499499
project_id: str,
500500
filter_: Optional[str] = None,
501501
retry: Optional[str] = DEFAULT,
502-
timeout: Optional[str] = DEFAULT,
502+
timeout: Optional[float] = None,
503503
metadata: Sequence[Tuple[str, str]] = (),
504504
) -> None:
505505
"""
@@ -537,7 +537,7 @@ def upsert_channel(
537537
channels: str,
538538
project_id: str,
539539
retry: Optional[str] = DEFAULT,
540-
timeout: Optional[float] = DEFAULT,
540+
timeout: Optional[float] = None,
541541
metadata: Sequence[Tuple[str, str]] = (),
542542
) -> dict:
543543
"""
@@ -603,7 +603,7 @@ def delete_notification_channel(
603603
self,
604604
name: str,
605605
retry: Optional[str] = DEFAULT,
606-
timeout: Optional[str] = DEFAULT,
606+
timeout: Optional[float] = None,
607607
metadata: Sequence[Tuple[str, str]] = (),
608608
) -> None:
609609
"""

‎airflow/providers/google/cloud/operators/bigquery.py‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
import uuid
2626
import warnings
2727
from datetime import datetime
28-
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Set, SupportsAbs, Union, cast
28+
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Set, SupportsAbs, Union
2929

3030
import attr
3131
from google.api_core.exceptions import Conflict
@@ -620,7 +620,7 @@ def __init__(
620620
sql: Union[str, Iterable],
621621
destination_dataset_table: Optional[str] = None,
622622
write_disposition: str = 'WRITE_EMPTY',
623-
allow_large_results: Optional[bool] = False,
623+
allow_large_results: bool = False,
624624
flatten_results: Optional[bool] = None,
625625
gcp_conn_id: str = 'google_cloud_default',
626626
bigquery_conn_id: Optional[str] = None,
@@ -694,7 +694,7 @@ def execute(self, context: 'Context'):
694694
impersonation_chain=self.impersonation_chain,
695695
)
696696
if isinstance(self.sql, str):
697-
job_id = self.hook.run_query(
697+
job_id: Union[str, List[str]] = self.hook.run_query(
698698
sql=self.sql,
699699
destination_dataset_table=self.destination_dataset_table,
700700
write_disposition=self.write_disposition,
@@ -1211,10 +1211,7 @@ def execute(self, context: 'Context') -> None:
12111211
delegate_to=self.delegate_to,
12121212
impersonation_chain=self.impersonation_chain,
12131213
)
1214-
schema_fields_bytes_or_string = gcs_hook.download(self.bucket, self.schema_object)
1215-
if hasattr(schema_fields_bytes_or_string, 'decode'):
1216-
schema_fields_bytes_or_string = cast(bytes, schema_fields_bytes_or_string).decode("utf-8")
1217-
schema_fields = json.loads(schema_fields_bytes_or_string)
1214+
schema_fields = json.loads(gcs_hook.download(self.bucket, self.schema_object).decode("utf-8"))
12181215
else:
12191216
schema_fields = self.schema_fields
12201217

@@ -2114,9 +2111,9 @@ def __init__(
21142111
self,
21152112
*,
21162113
schema_fields_updates: List[Dict[str, Any]],
2117-
include_policy_tags: Optional[bool] = False,
2118-
dataset_id: Optional[str] = None,
2119-
table_id: Optional[str] = None,
2114+
dataset_id: str,
2115+
table_id: str,
2116+
include_policy_tags: bool = False,
21202117
project_id: Optional[str] = None,
21212118
gcp_conn_id: str = 'google_cloud_default',
21222119
delegate_to: Optional[str] = None,

‎airflow/providers/google/cloud/operators/cloud_build.py‎

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ class CloudBuildCreateBuildOperator(BaseOperator):
159159
def __init__(
160160
self,
161161
*,
162-
build: Optional[Union[Dict, Build, str]] = None,
162+
build: Optional[Union[Dict, Build]] = None,
163163
body: Optional[Dict] = None,
164164
project_id: Optional[str] = None,
165165
wait: bool = True,
@@ -171,28 +171,32 @@ def __init__(
171171
**kwargs,
172172
) -> None:
173173
super().__init__(**kwargs)
174-
self.build = build
175-
# Not template fields to keep original value
176-
self.build_raw = build
177-
self.body = body
178174
self.project_id = project_id
179175
self.wait = wait
180176
self.retry = retry
181177
self.timeout = timeout
182178
self.metadata = metadata
183179
self.gcp_conn_id = gcp_conn_id
184180
self.impersonation_chain = impersonation_chain
181+
self.body = body
185182

186-
if self.body and self.build:
187-
raise AirflowException("Either build or body should be passed.")
188-
189-
if self.body:
183+
if body and build:
184+
raise AirflowException("You should not pass both build or body parameters. Both are set.")
185+
if body is not None:
190186
warnings.warn(
191187
"The body parameter has been deprecated. You should pass body using the build parameter.",
192188
DeprecationWarning,
193189
stacklevel=4,
194190
)
195-
self.build = self.build_raw = self.body
191+
actual_build = body
192+
else:
193+
if build is None:
194+
raise AirflowException("You should pass one of the build or body parameters. Both are None")
195+
actual_build = build
196+
197+
self.build = actual_build
198+
# Not template fields to keep original value
199+
self.build_raw = actual_build
196200

197201
def prepare_template(self) -> None:
198202
# if no file is specified, skip

‎airflow/providers/google/cloud/operators/cloud_sql.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222

2323
from airflow.exceptions import AirflowException
2424
from airflow.hooks.base import BaseHook
25-
from airflow.models import BaseOperator
25+
from airflow.models import BaseOperator, Connection
2626
from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook
2727
from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
2828
from airflow.providers.mysql.hooks.mysql import MySqlHook
@@ -1044,7 +1044,7 @@ def __init__(
10441044
self.gcp_cloudsql_conn_id = gcp_cloudsql_conn_id
10451045
self.autocommit = autocommit
10461046
self.parameters = parameters
1047-
self.gcp_connection = None
1047+
self.gcp_connection: Optional[Connection] = None
10481048

10491049
def _execute_query(
10501050
self, hook: CloudSQLDatabaseHook, database_hook: Union[PostgresHook, MySqlHook]

0 commit comments

Comments
 (0)