From f33e29ec38375da3d79fc3df5ce5e2b2a8a0bb7d Mon Sep 17 00:00:00 2001 From: Joshua Carp Date: Sat, 4 Apr 2020 17:52:42 -0400 Subject: [PATCH] Add mypy plugin for decorators. * Preserve types for decorators that don't change signatures * Augment types for decorators that change signatures * Drop redundant type checks --- UPDATING.md | 29 ++ airflow/mypy/__init__.py | 17 ++ airflow/mypy/plugin/__init__.py | 17 ++ airflow/mypy/plugin/decorators.py | 79 +++++ .../example_dags/example_automl_tables.py | 1 - .../providers/google/cloud/hooks/automl.py | 53 +--- .../google/cloud/hooks/bigquery_dts.py | 16 +- .../providers/google/cloud/hooks/bigtable.py | 17 +- .../google/cloud/hooks/cloud_build.py | 4 +- .../google/cloud/hooks/cloud_memorystore.py | 16 +- .../providers/google/cloud/hooks/cloud_sql.py | 22 +- .../hooks/cloud_storage_transfer_service.py | 8 +- .../providers/google/cloud/hooks/compute.py | 28 +- .../google/cloud/hooks/datacatalog.py | 82 +++--- .../providers/google/cloud/hooks/dataflow.py | 25 +- .../google/cloud/hooks/datafusion.py | 25 +- .../providers/google/cloud/hooks/dataproc.py | 26 +- .../providers/google/cloud/hooks/datastore.py | 52 ++-- airflow/providers/google/cloud/hooks/dlp.py | 28 +- .../providers/google/cloud/hooks/functions.py | 10 +- .../google/cloud/hooks/kubernetes_engine.py | 8 +- .../providers/google/cloud/hooks/mlengine.py | 31 +- .../providers/google/cloud/hooks/pubsub.py | 30 +- .../providers/google/cloud/hooks/spanner.py | 41 +-- .../google/cloud/hooks/stackdriver.py | 24 +- airflow/providers/google/cloud/hooks/tasks.py | 26 +- .../providers/google/cloud/hooks/vision.py | 48 +-- .../google/cloud/operators/automl.py | 6 - .../google/cloud/operators/dataflow.py | 25 +- .../google/cloud/operators/datafusion.py | 2 +- .../google/firebase/hooks/firestore.py | 2 - scripts/ci/ci_pylint_main.sh | 2 +- setup.cfg | 2 + setup.py | 2 +- .../google/cloud/hooks/test_automl.py | 2 +- .../google/cloud/hooks/test_bigtable.py | 103 ++----- .../google/cloud/hooks/test_cloud_build.py | 10 +- .../cloud/hooks/test_cloud_memorystore.py | 14 +- .../google/cloud/hooks/test_cloud_sql.py | 274 ++---------------- .../test_cloud_storage_transfer_service.py | 6 +- .../google/cloud/hooks/test_compute.py | 206 ++----------- .../google/cloud/hooks/test_datacatalog.py | 84 +++--- .../google/cloud/hooks/test_dataflow.py | 28 +- .../google/cloud/hooks/test_datastore.py | 37 +-- .../providers/google/cloud/hooks/test_dlp.py | 28 +- .../google/cloud/hooks/test_functions.py | 62 +--- .../cloud/hooks/test_kubernetes_engine.py | 4 +- .../google/cloud/hooks/test_mlengine.py | 163 +---------- .../google/cloud/hooks/test_spanner.py | 240 ++------------- .../google/cloud/hooks/test_stackdriver.py | 16 +- .../google/cloud/operators/test_automl.py | 2 - .../google/cloud/operators/test_dataflow.py | 24 +- .../cloud/operators/test_mlengine_utils.py | 3 +- 53 files changed, 624 insertions(+), 1486 deletions(-) create mode 100644 airflow/mypy/__init__.py create mode 100644 airflow/mypy/plugin/__init__.py create mode 100644 airflow/mypy/plugin/decorators.py diff --git a/UPDATING.md b/UPDATING.md index 96c75ab82e0e9..ec5aab4d8ea18 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -62,6 +62,35 @@ https://developers.google.com/style/inclusive-documentation --> +### Added mypy plugin to preserve types of decorated functions + +Mypy currently doesn't support precise type information for decorated +functions; see https://github.com/python/mypy/issues/3157 for details. +To preserve precise type definitions for decorated functions, we now +include a mypy plugin to preserve precise type definitions for decorated +functions. To use the plugin, update your setup.cfg: + +``` +[mypy] +plugins = + airflow.mypy.plugin.decorators +``` + +### Use project_id argument consistently across GCP hooks and operators + +- Changed order of arguments in DataflowHook.start_python_dataflow. Uses + with positional arguments may break. +- Changed order of arguments in DataflowHook.is_job_dataflow_running. Uses + with positional arguments may break. +- Changed order of arguments in DataflowHook.cancel_job. Uses + with positional arguments may break. +- Added optional project_id argument to DataflowCreateJavaJobOperator + constructor. +- Added optional project_id argument to DataflowTemplatedJobStartOperator + constructor. +- Added optional project_id argument to DataflowCreatePythonJobOperator + constructor. + ### Rename pool statsd metrics Used slot has been renamed to running slot to make the name self-explanatory diff --git a/airflow/mypy/__init__.py b/airflow/mypy/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/mypy/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/mypy/plugin/__init__.py b/airflow/mypy/plugin/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/mypy/plugin/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/mypy/plugin/decorators.py b/airflow/mypy/plugin/decorators.py new file mode 100644 index 0000000000000..a53aa9c1def9b --- /dev/null +++ b/airflow/mypy/plugin/decorators.py @@ -0,0 +1,79 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import copy +import functools +from typing import List + +from mypy.nodes import ARG_NAMED_OPT # pylint: disable=no-name-in-module +from mypy.plugin import FunctionContext, Plugin # pylint: disable=no-name-in-module +from mypy.types import CallableType, NoneType, UnionType # pylint: disable=no-name-in-module + +TYPED_DECORATORS = { + "fallback_to_default_project_id of GoogleBaseHook": ["project_id"], + "airflow.providers.google.cloud.hooks.dataflow._fallback_to_project_id_from_variables": ["project_id"], + "provide_gcp_credential_file of GoogleBaseHook": [], +} + + +class TypedDecoratorPlugin(Plugin): + """Mypy plugin for typed decorators.""" + def get_function_hook(self, fullname: str): + """Check for known typed decorators by name.""" + if fullname in TYPED_DECORATORS: + return functools.partial( + _analyze_decorator, + provided_arguments=TYPED_DECORATORS[fullname], + ) + return None + + +def _analyze_decorator(function_ctx: FunctionContext, provided_arguments: List[str]): + if not isinstance(function_ctx.arg_types[0][0], CallableType): + return function_ctx.default_return_type + if not isinstance(function_ctx.default_return_type, CallableType): + return function_ctx.default_return_type + return _change_decorator_function_type( + function_ctx.arg_types[0][0], + function_ctx.default_return_type, + provided_arguments, + ) + + +def _change_decorator_function_type( + decorated: CallableType, + decorator: CallableType, + provided_arguments: List[str], +) -> CallableType: + decorator.arg_kinds = decorated.arg_kinds + decorator.arg_names = decorated.arg_names + + # Mark provided arguments as optional + decorator.arg_types = copy.copy(decorated.arg_types) + for argument in provided_arguments: + index = decorated.arg_names.index(argument) + decorated_type = decorated.arg_types[index] + decorator.arg_types[index] = UnionType.make_union([decorated_type, NoneType()]) + decorated.arg_kinds[index] = ARG_NAMED_OPT + + return decorator + + +def plugin(version: str): # pylint: disable=unused-argument + """Mypy plugin entrypoint.""" + return TypedDecoratorPlugin diff --git a/airflow/providers/google/cloud/example_dags/example_automl_tables.py b/airflow/providers/google/cloud/example_dags/example_automl_tables.py index 47ad310c3caaf..78ecd10a18721 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_tables.py +++ b/airflow/providers/google/cloud/example_dags/example_automl_tables.py @@ -139,7 +139,6 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str: task_id="update_dataset_task", dataset=update, location=GCP_AUTOML_LOCATION, - project_id=GCP_PROJECT_ID, ) # [END howto_operator_automl_update_dataset] diff --git a/airflow/providers/google/cloud/hooks/automl.py b/airflow/providers/google/cloud/hooks/automl.py index f1b7e593be59b..5c2ea890bee65 100644 --- a/airflow/providers/google/cloud/hooks/automl.py +++ b/airflow/providers/google/cloud/hooks/automl.py @@ -83,7 +83,7 @@ def create_model( self, model: Union[dict, Model], location: str, - project_id: Optional[str] = None, + project_id: str, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, retry: Optional[Retry] = None, @@ -113,8 +113,6 @@ def create_model( :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() parent = client.location_path(project_id, location) return client.create_model( @@ -128,7 +126,7 @@ def batch_predict( input_config: Union[dict, BatchPredictInputConfig], output_config: Union[dict, BatchPredictOutputConfig], location: str, - project_id: Optional[str] = None, + project_id: str, params: Optional[Dict[str, str]] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -168,8 +166,6 @@ def batch_predict( :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance """ - if not project_id: - raise ValueError("The project_id should be set") client = self.prediction_client name = client.model_path(project=project_id, location=location, model=model_id) result = client.batch_predict( @@ -189,7 +185,7 @@ def predict( model_id: str, payload: Union[dict, ExamplePayload], location: str, - project_id: Optional[str] = None, + project_id: str, params: Optional[Dict[str, str]] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -223,8 +219,6 @@ def predict( :return: `google.cloud.automl_v1beta1.types.PredictResponse` instance """ - if not project_id: - raise ValueError("The project_id should be set") client = self.prediction_client name = client.model_path(project=project_id, location=location, model=model_id) result = client.predict( @@ -242,7 +236,7 @@ def create_dataset( self, dataset: Union[dict, Dataset], location: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -269,8 +263,6 @@ def create_dataset( :return: `google.cloud.automl_v1beta1.types.Dataset` instance. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() parent = client.location_path(project=project_id, location=location) result = client.create_dataset( @@ -288,7 +280,7 @@ def import_data( dataset_id: str, location: str, input_config: Union[dict, InputConfig], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -317,8 +309,6 @@ def import_data( :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() name = client.dataset_path( project=project_id, location=location, dataset=dataset_id @@ -338,10 +328,10 @@ def list_column_specs( # pylint: disable=too-many-arguments dataset_id: str, table_spec_id: str, location: str, + project_id: str, field_mask: Union[dict, FieldMask] = None, filter_: Optional[str] = None, page_size: Optional[int] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -380,8 +370,6 @@ def list_column_specs( # pylint: disable=too-many-arguments :return: `google.cloud.automl_v1beta1.types.ColumnSpec` instance. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() parent = client.table_spec_path( project=project_id, @@ -405,7 +393,7 @@ def get_model( self, model_id: str, location: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -431,8 +419,6 @@ def get_model( :return: `google.cloud.automl_v1beta1.types.Model` instance. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() name = client.model_path(project=project_id, location=location, model=model_id) result = client.get_model( @@ -445,7 +431,7 @@ def delete_model( self, model_id: str, location: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -471,8 +457,6 @@ def delete_model( :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() name = client.model_path(project=project_id, location=location, model=model_id) result = client.delete_model( @@ -480,12 +464,10 @@ def delete_model( ) return result - @GoogleBaseHook.fallback_to_default_project_id def update_dataset( self, dataset: Union[dict, Dataset], update_mask: Union[dict, FieldMask] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -499,9 +481,6 @@ def update_dataset( :param update_mask: The update mask applies to the resource. If a dict is provided, it must be of the same form as the protobuf message FieldMask. :type update_mask: Union[dict, FieldMask] - :param project_id: ID of the Google Cloud project where dataset is located if None then - default project_id is used. - :type project_id: str :param retry: A retry object used to retry requests. If `None` is specified, requests will not be retried. :type retry: Optional[google.api_core.retry.Retry] @@ -513,8 +492,6 @@ def update_dataset( :return: `google.cloud.automl_v1beta1.types.Dataset` instance.. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() result = client.update_dataset( dataset=dataset, @@ -530,7 +507,7 @@ def deploy_model( self, model_id: str, location: str, - project_id: Optional[str] = None, + project_id: str, image_detection_metadata: Union[ ImageObjectDetectionModelDeploymentMetadata, dict ] = None, @@ -568,8 +545,6 @@ def deploy_model( :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() name = client.model_path(project=project_id, location=location, model=model_id) result = client.deploy_model( @@ -624,8 +599,6 @@ def list_table_specs( This object can also be configured to iterate over the pages of the response through the `options` parameter. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() parent = client.dataset_path( project=project_id, location=location, dataset=dataset_id @@ -644,7 +617,7 @@ def list_table_specs( def list_datasets( self, location: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -671,8 +644,6 @@ def list_datasets( This object can also be configured to iterate over the pages of the response through the `options` parameter. """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() parent = client.location_path(project=project_id, location=location) result = client.list_datasets( @@ -685,7 +656,7 @@ def delete_dataset( self, dataset_id: str, location: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -711,8 +682,6 @@ def delete_dataset( :return: `google.cloud.automl_v1beta1.types._OperationFuture` instance """ - if not project_id: - raise ValueError("The project_id should be set") client = self.get_conn() name = client.dataset_path( project=project_id, location=location, dataset=dataset_id diff --git a/airflow/providers/google/cloud/hooks/bigquery_dts.py b/airflow/providers/google/cloud/hooks/bigquery_dts.py index f03ef10652ae3..6196a0d29398f 100644 --- a/airflow/providers/google/cloud/hooks/bigquery_dts.py +++ b/airflow/providers/google/cloud/hooks/bigquery_dts.py @@ -97,7 +97,7 @@ def get_conn(self) -> DataTransferServiceClient: def create_transfer_config( self, transfer_config: Union[dict, TransferConfig], - project_id: Optional[str] = None, + project_id: str, authorization_code: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -125,8 +125,6 @@ def create_transfer_config( :type metadata: Optional[Sequence[Tuple[str, str]]] :return: A ``google.cloud.bigquery_datatransfer_v1.types.TransferConfig`` instance. """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() parent = client.project_path(project_id) return client.create_transfer_config( @@ -142,7 +140,7 @@ def create_transfer_config( def delete_transfer_config( self, transfer_config_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -166,8 +164,6 @@ def delete_transfer_config( :type metadata: Optional[Sequence[Tuple[str, str]]] :return: None """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() name = client.project_transfer_config_path( project=project_id, transfer_config=transfer_config_id @@ -180,7 +176,7 @@ def delete_transfer_config( def start_manual_transfer_runs( self, transfer_config_id: str, - project_id: Optional[str] = None, + project_id: str, requested_time_range: Optional[dict] = None, requested_run_time: Optional[dict] = None, retry: Optional[Retry] = None, @@ -218,8 +214,6 @@ def start_manual_transfer_runs( :type metadata: Optional[Sequence[Tuple[str, str]]] :return: An ``google.cloud.bigquery_datatransfer_v1.types.StartManualTransferRunsResponse`` instance. """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() parent = client.project_transfer_config_path( project=project_id, transfer_config=transfer_config_id @@ -238,7 +232,7 @@ def get_transfer_run( self, run_id: str, transfer_config_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -264,8 +258,6 @@ def get_transfer_run( :type metadata: Optional[Sequence[Tuple[str, str]]] :return: An ``google.cloud.bigquery_datatransfer_v1.types.TransferRun`` instance. """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() name = client.project_run_path( project=project_id, transfer_config=transfer_config_id, run=run_id diff --git a/airflow/providers/google/cloud/hooks/bigtable.py b/airflow/providers/google/cloud/hooks/bigtable.py index c69c65588737e..49188a161fca2 100644 --- a/airflow/providers/google/cloud/hooks/bigtable.py +++ b/airflow/providers/google/cloud/hooks/bigtable.py @@ -54,7 +54,7 @@ def _get_client(self, project_id: str): return self._client @GoogleBaseHook.fallback_to_default_project_id - def get_instance(self, instance_id: str, project_id: Optional[str] = None) -> Instance: + def get_instance(self, instance_id: str, project_id: str) -> Instance: """ Retrieves and returns the specified Cloud Bigtable instance if it exists. Otherwise, returns None. @@ -66,16 +66,13 @@ def get_instance(self, instance_id: str, project_id: Optional[str] = None) -> In the default project_id from the GCP connection is used. :type project_id: str """ - if not project_id: - raise ValueError("Project ID should be set.") - instance = self._get_client(project_id=project_id).instance(instance_id) if not instance.exists(): return None return instance @GoogleBaseHook.fallback_to_default_project_id - def delete_instance(self, instance_id: str, project_id: Optional[str] = None) -> None: + def delete_instance(self, instance_id: str, project_id: str) -> None: """ Deletes the specified Cloud Bigtable instance. Raises google.api_core.exceptions.NotFound if the Cloud Bigtable instance does @@ -88,8 +85,6 @@ def delete_instance(self, instance_id: str, project_id: Optional[str] = None) -> :param instance_id: The ID of the Cloud Bigtable instance. :type instance_id: str """ - if not project_id: - raise ValueError("Project ID should be set.") instance = self.get_instance(instance_id=instance_id, project_id=project_id) if instance: instance.delete() @@ -103,7 +98,7 @@ def create_instance( instance_id: str, main_cluster_id: str, main_cluster_zone: str, - project_id: Optional[str] = None, + project_id: str, replica_cluster_id: Optional[str] = None, replica_cluster_zone: Optional[str] = None, instance_display_name: Optional[str] = None, @@ -148,8 +143,6 @@ def create_instance( :param timeout: (optional) timeout (in seconds) for instance creation. If None is not specified, Operator will wait indefinitely. """ - if not project_id: - raise ValueError("Project ID should be set.") cluster_storage_type = enums.StorageType(cluster_storage_type) instance_type = enums.Instance.Type(instance_type) @@ -213,7 +206,7 @@ def create_table( table.create(initial_split_keys, column_families) @GoogleBaseHook.fallback_to_default_project_id - def delete_table(self, instance_id: str, table_id: str, project_id: Optional[str] = None) -> None: + def delete_table(self, instance_id: str, table_id: str, project_id: str) -> None: """ Deletes the specified table in Cloud Bigtable. Raises google.api_core.exceptions.NotFound if the table does not exist. @@ -227,8 +220,6 @@ def delete_table(self, instance_id: str, table_id: str, project_id: Optional[str BigTable exists. If set to None or missing, the default project_id from the GCP connection is used. """ - if not project_id: - raise ValueError("Project ID should be set.") table = self.get_instance(instance_id=instance_id, project_id=project_id).table(table_id=table_id) table.delete() diff --git a/airflow/providers/google/cloud/hooks/cloud_build.py b/airflow/providers/google/cloud/hooks/cloud_build.py index d78b2e73c5651..76f3b15913828 100644 --- a/airflow/providers/google/cloud/hooks/cloud_build.py +++ b/airflow/providers/google/cloud/hooks/cloud_build.py @@ -70,7 +70,7 @@ def get_conn(self): return self._conn @GoogleBaseHook.fallback_to_default_project_id - def create_build(self, body: Dict, project_id: Optional[str] = None) -> Dict: + def create_build(self, body: Dict, project_id: str) -> Dict: """ Starts a build with the specified configuration. @@ -82,8 +82,6 @@ def create_build(self, body: Dict, project_id: Optional[str] = None) -> Dict: :type project_id: str :return: Dict """ - if not project_id: - raise ValueError("The project_id should be set") service = self.get_conn() # Create build diff --git a/airflow/providers/google/cloud/hooks/cloud_memorystore.py b/airflow/providers/google/cloud/hooks/cloud_memorystore.py index 16572e27b4510..64ef45a2a0573 100644 --- a/airflow/providers/google/cloud/hooks/cloud_memorystore.py +++ b/airflow/providers/google/cloud/hooks/cloud_memorystore.py @@ -87,7 +87,7 @@ def create_instance( location: str, instance_id: str, instance: Union[Dict, Instance], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -162,7 +162,7 @@ def delete_instance( self, location: str, instance: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -205,7 +205,7 @@ def export_instance( location: str, instance: str, output_config: Union[Dict, OutputConfig], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -251,7 +251,7 @@ def failover_instance( location: str, instance: str, data_protection_mode: FailoverInstanceRequest.DataProtectionMode, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -299,7 +299,7 @@ def get_instance( self, location: str, instance: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -335,7 +335,7 @@ def import_instance( location: str, instance: str, input_config: Union[Dict, InputConfig], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -381,7 +381,7 @@ def list_instances( self, location: str, page_size: int, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -424,9 +424,9 @@ def update_instance( self, update_mask: Union[Dict, FieldMask], instance: Union[Dict, Instance], + project_id: str, location: Optional[str] = None, instance_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, diff --git a/airflow/providers/google/cloud/hooks/cloud_sql.py b/airflow/providers/google/cloud/hooks/cloud_sql.py index de4ec5674a044..44a99da3744f4 100644 --- a/airflow/providers/google/cloud/hooks/cloud_sql.py +++ b/airflow/providers/google/cloud/hooks/cloud_sql.py @@ -101,7 +101,7 @@ def get_conn(self): return self._conn @GoogleBaseHook.fallback_to_default_project_id - def get_instance(self, instance: str, project_id: Optional[str] = None) -> Dict: + def get_instance(self, instance: str, project_id: str) -> Dict: """ Retrieves a resource containing information about a Cloud SQL instance. @@ -120,7 +120,7 @@ def get_instance(self, instance: str, project_id: Optional[str] = None) -> Dict: @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def create_instance(self, body: Dict, project_id: Optional[str] = None) -> None: + def create_instance(self, body: Dict, project_id: str) -> None: """ Creates a new Cloud SQL instance. @@ -142,7 +142,7 @@ def create_instance(self, body: Dict, project_id: Optional[str] = None) -> None: @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def patch_instance(self, body: Dict, instance: str, project_id: Optional[str] = None) -> None: + def patch_instance(self, body: Dict, instance: str, project_id: str) -> None: """ Updates settings of a Cloud SQL instance. @@ -170,7 +170,7 @@ def patch_instance(self, body: Dict, instance: str, project_id: Optional[str] = @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def delete_instance(self, instance: str, project_id: Optional[str] = None) -> None: + def delete_instance(self, instance: str, project_id: str) -> None: """ Deletes a Cloud SQL instance. @@ -190,7 +190,7 @@ def delete_instance(self, instance: str, project_id: Optional[str] = None) -> No operation_name=operation_name) @GoogleBaseHook.fallback_to_default_project_id - def get_database(self, instance: str, database: str, project_id: Optional[str] = None) -> Dict: + def get_database(self, instance: str, database: str, project_id: str) -> Dict: """ Retrieves a database resource from a Cloud SQL instance. @@ -213,7 +213,7 @@ def get_database(self, instance: str, database: str, project_id: Optional[str] = @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def create_database(self, instance: str, body: Dict, project_id: Optional[str] = None) -> None: + def create_database(self, instance: str, body: Dict, project_id: str) -> None: """ Creates a new database inside a Cloud SQL instance. @@ -243,7 +243,7 @@ def patch_database( instance: str, database: str, body: Dict, - project_id: Optional[str] = None + project_id: str, ) -> None: """ Updates a database resource inside a Cloud SQL instance. @@ -275,7 +275,7 @@ def patch_database( @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def delete_database(self, instance: str, database: str, project_id: Optional[str] = None) -> None: + def delete_database(self, instance: str, database: str, project_id: str) -> None: """ Deletes a database from a Cloud SQL instance. @@ -299,7 +299,7 @@ def delete_database(self, instance: str, database: str, project_id: Optional[str @GoogleBaseHook.fallback_to_default_project_id @GoogleBaseHook.operation_in_progress_retry() - def export_instance(self, instance: str, body: Dict, project_id: Optional[str] = None) -> None: + def export_instance(self, instance: str, body: Dict, project_id: str) -> None: """ Exports data from a Cloud SQL instance to a Cloud Storage bucket as a SQL dump or CSV file. @@ -325,7 +325,7 @@ def export_instance(self, instance: str, body: Dict, project_id: Optional[str] = operation_name=operation_name) @GoogleBaseHook.fallback_to_default_project_id - def import_instance(self, instance: str, body: Dict, project_id: Optional[str] = None) -> None: + def import_instance(self, instance: str, body: Dict, project_id: str) -> None: """ Imports data into a Cloud SQL instance from a SQL dump or CSV file in Cloud Storage. @@ -366,8 +366,6 @@ def _wait_for_operation_to_complete(self, project_id: str, operation_name: str) :type operation_name: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") service = self.get_conn() while True: operation_response = service.operations().get( # noqa # pylint: disable=no-member diff --git a/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py b/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py index cc73bab2f9bd8..2f5f9453f83e6 100644 --- a/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +++ b/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py @@ -147,7 +147,7 @@ def create_transfer_job(self, body: Dict) -> Dict: num_retries=self.num_retries) @GoogleBaseHook.fallback_to_default_project_id - def get_transfer_job(self, job_name: str, project_id: Optional[str] = None) -> Dict: + def get_transfer_job(self, job_name: str, project_id: str) -> Dict: """ Gets the latest state of a long-running operation in Google Storage Transfer Service. @@ -161,8 +161,6 @@ def get_transfer_job(self, job_name: str, project_id: Optional[str] = None) -> D :return: Transfer Job :rtype: dict """ - if not project_id: - raise ValueError("Project ID should be set.") return ( self.get_conn() # pylint: disable=no-member .transferJobs() @@ -227,7 +225,7 @@ def update_transfer_job(self, job_name: str, body: Dict) -> Dict: ) @GoogleBaseHook.fallback_to_default_project_id - def delete_transfer_job(self, job_name: str, project_id: Optional[str] = None) -> None: + def delete_transfer_job(self, job_name: str, project_id: str) -> None: """ Deletes a transfer job. This is a soft delete. After a transfer job is deleted, the job and all the transfer executions are subject to garbage @@ -242,8 +240,6 @@ def delete_transfer_job(self, job_name: str, project_id: Optional[str] = None) - :type project_id: str :rtype: None """ - if not project_id: - raise ValueError("Project ID should be set.") ( self.get_conn() # pylint: disable=no-member .transferJobs() diff --git a/airflow/providers/google/cloud/hooks/compute.py b/airflow/providers/google/cloud/hooks/compute.py index 4b550abbc432c..e3075fa2a4de9 100644 --- a/airflow/providers/google/cloud/hooks/compute.py +++ b/airflow/providers/google/cloud/hooks/compute.py @@ -73,7 +73,7 @@ def get_conn(self): return self._conn @GoogleBaseHook.fallback_to_default_project_id - def start_instance(self, zone: str, resource_id: str, project_id: Optional[str] = None) -> None: + def start_instance(self, zone: str, resource_id: str, project_id: str) -> None: """ Starts an existing instance defined by project_id, zone and resource_id. Must be called with keyword arguments rather than positional. @@ -88,8 +88,6 @@ def start_instance(self, zone: str, resource_id: str, project_id: Optional[str] :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().instances().start( # pylint: disable=no-member project=project_id, zone=zone, @@ -106,7 +104,7 @@ def start_instance(self, zone: str, resource_id: str, project_id: Optional[str] zone=zone) @GoogleBaseHook.fallback_to_default_project_id - def stop_instance(self, zone: str, resource_id: str, project_id: Optional[str] = None) -> None: + def stop_instance(self, zone: str, resource_id: str, project_id: str) -> None: """ Stops an instance defined by project_id, zone and resource_id Must be called with keyword arguments rather than positional. @@ -121,8 +119,6 @@ def stop_instance(self, zone: str, resource_id: str, project_id: Optional[str] = :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().instances().stop( # pylint: disable=no-member project=project_id, zone=zone, @@ -144,7 +140,7 @@ def set_machine_type( zone: str, resource_id: str, body: Dict, - project_id: Optional[str] = None + project_id: str ) -> None: """ Sets machine type of an instance defined by project_id, zone and resource_id. @@ -164,8 +160,6 @@ def set_machine_type( :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") response = self._execute_set_machine_type(zone, resource_id, body, project_id) try: operation_name = response["name"] @@ -189,7 +183,7 @@ def _execute_set_machine_type( .execute(num_retries=self.num_retries) @GoogleBaseHook.fallback_to_default_project_id - def get_instance_template(self, resource_id: str, project_id: Optional[str] = None) -> Dict: + def get_instance_template(self, resource_id: str, project_id: str) -> Dict: """ Retrieves instance template by project_id and resource_id. Must be called with keyword arguments rather than positional. @@ -204,8 +198,6 @@ def get_instance_template(self, resource_id: str, project_id: Optional[str] = No https://cloud.google.com/compute/docs/reference/rest/v1/instanceTemplates :rtype: dict """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().instanceTemplates().get( # pylint: disable=no-member project=project_id, instanceTemplate=resource_id @@ -216,8 +208,8 @@ def get_instance_template(self, resource_id: str, project_id: Optional[str] = No def insert_instance_template( self, body: Dict, + project_id: str, request_id: Optional[str] = None, - project_id: Optional[str] = None ) -> None: """ Inserts instance template using body specified @@ -237,8 +229,6 @@ def insert_instance_template( :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().instanceTemplates().insert( # pylint: disable=no-member project=project_id, body=body, @@ -258,7 +248,7 @@ def get_instance_group_manager( self, zone: str, resource_id: str, - project_id: Optional[str] = None + project_id: str, ) -> Dict: """ Retrieves Instance Group Manager by project_id, zone and resource_id. @@ -276,8 +266,6 @@ def get_instance_group_manager( https://cloud.google.com/compute/docs/reference/rest/beta/instanceGroupManagers :rtype: dict """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().instanceGroupManagers().get( # pylint: disable=no-member project=project_id, zone=zone, @@ -291,8 +279,8 @@ def patch_instance_group_manager( zone: str, resource_id: str, body: Dict, + project_id: str, request_id: Optional[str] = None, - project_id: Optional[str] = None ) -> None: """ Patches Instance Group Manager with the specified body. @@ -317,8 +305,6 @@ def patch_instance_group_manager( :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().instanceGroupManagers().patch( # pylint: disable=no-member project=project_id, zone=zone, diff --git a/airflow/providers/google/cloud/hooks/datacatalog.py b/airflow/providers/google/cloud/hooks/datacatalog.py index c7abd717c18ef..49b4c1603fd09 100644 --- a/airflow/providers/google/cloud/hooks/datacatalog.py +++ b/airflow/providers/google/cloud/hooks/datacatalog.py @@ -61,7 +61,7 @@ def create_entry( entry_group: str, entry_id: str, entry: Union[Dict, Entry], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -83,7 +83,7 @@ def create_entry( :type entry: Union[Dict, google.cloud.datacatalog_v1beta1.types.Entry] :param project_id: The ID of the GCP project that owns the entry. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If set to ``None`` or missing, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -108,7 +108,7 @@ def create_entry_group( location: str, entry_group_id: str, entry_group: Union[Dict, EntryGroup], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -129,7 +129,7 @@ def create_entry_group( :type entry_group: Union[Dict, google.cloud.datacatalog_v1beta1.types.EntryGroup] :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -162,8 +162,8 @@ def create_tag( entry_group: str, entry: str, tag: Union[Dict, Tag], + project_id: str, template_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -185,7 +185,7 @@ def create_tag( :type template_id: Optional[str] :param project_id: The ID of the GCP project that owns the tag. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -216,7 +216,7 @@ def create_tag_template( location, tag_template_id: str, tag_template: Union[Dict, TagTemplate], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -235,7 +235,7 @@ def create_tag_template( :type tag_template: Union[Dict, google.cloud.datacatalog_v1beta1.types.TagTemplate] :param project_id: The ID of the GCP project that owns the tag template. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -269,7 +269,7 @@ def create_tag_template_field( tag_template: str, tag_template_field_id: str, tag_template_field: Union[Dict, TagTemplateField], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -293,7 +293,7 @@ def create_tag_template_field( :type tag_template_field: Union[Dict, google.cloud.datacatalog_v1beta1.types.TagTemplateField] :param project_id: The ID of the GCP project that owns the tag template field. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -327,7 +327,7 @@ def delete_entry( location: str, entry_group: str, entry: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -343,7 +343,7 @@ def delete_entry( :type entry: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -364,7 +364,7 @@ def delete_entry_group( self, location, entry_group, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -380,7 +380,7 @@ def delete_entry_group( :type entry_group: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -404,7 +404,7 @@ def delete_tag( entry_group: str, entry: str, tag: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -422,7 +422,7 @@ def delete_tag( :type tag: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -445,7 +445,7 @@ def delete_tag_template( location, tag_template, force: bool, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -459,7 +459,7 @@ def delete_tag_template( :type tag_template: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param force: Required. Currently, this field must always be set to ``true``. This confirms the deletion of any possible tags using this template. ``force = false`` will be supported in the future. @@ -487,7 +487,7 @@ def delete_tag_template_field( tag_template: str, field: str, force: bool, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -505,7 +505,7 @@ def delete_tag_template_field( :type force: bool :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -530,7 +530,7 @@ def get_entry( location: str, entry_group: str, entry: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -546,7 +546,7 @@ def get_entry( :type entry: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -570,8 +570,8 @@ def get_entry_group( self, location: str, entry_group: str, + project_id: str, read_mask: Union[Dict, FieldMask] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -590,7 +590,7 @@ def get_entry_group( :type read_mask: Union[Dict, google.cloud.datacatalog_v1beta1.types.FieldMask] :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -618,7 +618,7 @@ def get_tag_template( self, location: str, tag_template: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -632,7 +632,7 @@ def get_tag_template( :type tag_template: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -659,8 +659,8 @@ def list_tags( location: str, entry_group: str, entry: str, - page_size: int = 100, - project_id: Optional[str] = None, + project_id: str, + page_size: Optional[int] = 100, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -680,7 +680,7 @@ def list_tags( :type page_size: int :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -710,7 +710,7 @@ def get_tag_for_template_name( entry_group: str, entry: str, template_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -728,7 +728,7 @@ def get_tag_for_template_name( :type template_name: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -808,7 +808,7 @@ def rename_tag_template_field( tag_template: str, field: str, new_tag_template_field_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -828,7 +828,7 @@ def rename_tag_template_field( :type new_tag_template_field_id: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -940,10 +940,10 @@ def update_entry( self, entry: Union[Dict, Entry], update_mask: Union[Dict, FieldMask], + project_id: str, location: Optional[str] = None, entry_group: Optional[str] = None, entry_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -970,7 +970,7 @@ def update_entry( :type entry_id: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -1011,11 +1011,11 @@ def update_tag( # pylint: disable=too-many-arguments self, tag: Union[Dict, Tag], update_mask: Union[Dict, FieldMask], + project_id: str, location: Optional[str] = None, entry_group: Optional[str] = None, entry: Optional[str] = None, tag_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -1044,7 +1044,7 @@ def update_tag( # pylint: disable=too-many-arguments :type tag_id: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -1084,9 +1084,9 @@ def update_tag_template( self, tag_template: Union[Dict, TagTemplate], update_mask: Union[Dict, FieldMask], + project_id: str, location: Optional[str] = None, tag_template_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -1116,7 +1116,7 @@ def update_tag_template( :type tag_template_id: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry @@ -1162,11 +1162,11 @@ def update_tag_template_field( # pylint: disable=too-many-arguments self, tag_template_field: Union[Dict, TagTemplateField], update_mask: Union[Dict, FieldMask], + project_id: str, tag_template_field_name: Optional[str] = None, location: Optional[str] = None, tag_template: Optional[str] = None, tag_template_field_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -1202,7 +1202,7 @@ def update_tag_template_field( # pylint: disable=too-many-arguments :type tag_template_field_id: str :param project_id: The ID of the GCP project that owns the entry group. If set to ``None`` or missing, the default project_id from the GCP connection is used. - :type project_id: Optional[str] + :type project_id: str :param retry: A retry object used to retry requests. If ``None`` is specified, requests will be retried using a default configuration. :type retry: google.api_core.retry.Retry diff --git a/airflow/providers/google/cloud/hooks/dataflow.py b/airflow/providers/google/cloud/hooks/dataflow.py index 50626ec7c1f07..17d8a9e1e8098 100644 --- a/airflow/providers/google/cloud/hooks/dataflow.py +++ b/airflow/providers/google/cloud/hooks/dataflow.py @@ -458,7 +458,7 @@ def start_java_dataflow( job_name: str, variables: Dict, jar: str, - project_id: Optional[str] = None, + project_id: str, job_class: Optional[str] = None, append_job_name: bool = True, multiple_jobs: bool = False, @@ -484,9 +484,6 @@ def start_java_dataflow( :param on_new_job_id_callback: Callback called when the job ID is known. :type on_new_job_id_callback: callable """ - if not project_id: - raise ValueError("The project_id should be set") - name = self._build_dataflow_job_name(job_name, append_job_name) variables['jobName'] = name @@ -514,7 +511,7 @@ def start_template_dataflow( variables: Dict, parameters: Dict, dataflow_template: str, - project_id: Optional[str] = None, + project_id: str, append_job_name: bool = True, on_new_job_id_callback: Optional[Callable[[str], None]] = None ) -> Dict: @@ -536,9 +533,6 @@ def start_template_dataflow( :param on_new_job_id_callback: Callback called when the job ID is known. :type on_new_job_id_callback: callable """ - if not project_id: - raise ValueError("The project_id should be set") - variables = self._set_variables(variables) name = self._build_dataflow_job_name(job_name, append_job_name) # Builds RuntimeEnvironment from variables dictionary @@ -585,10 +579,10 @@ def start_python_dataflow( # pylint: disable=too-many-arguments variables: Dict, dataflow: str, py_options: List[str], + project_id: str, py_interpreter: str = "python3", py_requirements: Optional[List[str]] = None, py_system_site_packages: bool = False, - project_id: Optional[str] = None, append_job_name: bool = True, on_new_job_id_callback: Optional[Callable[[str], None]] = None ): @@ -626,9 +620,6 @@ def start_python_dataflow( # pylint: disable=too-many-arguments :param on_new_job_id_callback: Callback called when the job ID is known. :type on_new_job_id_callback: callable """ - if not project_id: - raise ValueError("The project_id should be set") - name = self._build_dataflow_job_name(job_name, append_job_name) variables['job_name'] = name @@ -701,7 +692,7 @@ def _build_cmd(variables: Dict, label_formatter: Callable, project_id: str) -> L @_fallback_to_project_id_from_variables @GoogleBaseHook.fallback_to_default_project_id - def is_job_dataflow_running(self, name: str, variables: Dict, project_id: Optional[str] = None) -> bool: + def is_job_dataflow_running(self, name: str, variables: Dict, project_id: str) -> bool: """ Helper method to check if jos is still running in dataflow @@ -714,9 +705,6 @@ def is_job_dataflow_running(self, name: str, variables: Dict, project_id: Option :return: True if job is running. :rtype: bool """ - if not project_id: - raise ValueError("The project_id should be set") - variables = self._set_variables(variables) jobs_controller = _DataflowJobsController( dataflow=self.get_conn(), @@ -730,10 +718,10 @@ def is_job_dataflow_running(self, name: str, variables: Dict, project_id: Option @GoogleBaseHook.fallback_to_default_project_id def cancel_job( self, + project_id: str, job_name: Optional[str] = None, job_id: Optional[str] = None, location: Optional[str] = None, - project_id: Optional[str] = None ) -> None: """ Cancels the job with the specified name prefix or Job ID. @@ -750,9 +738,6 @@ def cancel_job( If set to None or missing, the default project_id from the GCP connection is used. :type project_id: """ - if not project_id: - raise ValueError("The project_id should be set") - jobs_controller = _DataflowJobsController( dataflow=self.get_conn(), project_number=project_id, diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py index 53c4185a57abf..6360e3fabe9ae 100644 --- a/airflow/providers/google/cloud/hooks/datafusion.py +++ b/airflow/providers/google/cloud/hooks/datafusion.py @@ -107,7 +107,7 @@ def get_conn(self) -> Resource: @GoogleBaseHook.fallback_to_default_project_id def restart_instance( - self, instance_name: str, location: str, project_id: Optional[str] = None + self, instance_name: str, location: str, project_id: str ) -> Operation: """ Restart a single Data Fusion instance. @@ -120,9 +120,6 @@ def restart_instance( :param project_id: The ID of the Google Cloud Platform project that the instance belongs to. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") - operation = ( self.get_conn() # pylint: disable=no-member .projects() @@ -135,7 +132,7 @@ def restart_instance( @GoogleBaseHook.fallback_to_default_project_id def delete_instance( - self, instance_name: str, location: str, project_id: Optional[str] = None + self, instance_name: str, location: str, project_id: str ) -> Operation: """ Deletes a single Date Fusion instance. @@ -147,9 +144,6 @@ def delete_instance( :param project_id: The ID of the Google Cloud Platform project that the instance belongs to. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") - operation = ( self.get_conn() # pylint: disable=no-member .projects() @@ -166,7 +160,7 @@ def create_instance( instance_name: str, instance: Dict[str, Any], location: str, - project_id: Optional[str] = None, + project_id: str, ) -> Operation: """ Creates a new Data Fusion instance in the specified project and location. @@ -181,9 +175,6 @@ def create_instance( :param project_id: The ID of the Google Cloud Platform project that the instance belongs to. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") - operation = ( self.get_conn() # pylint: disable=no-member .projects() @@ -200,7 +191,7 @@ def create_instance( @GoogleBaseHook.fallback_to_default_project_id def get_instance( - self, instance_name: str, location: str, project_id: Optional[str] = None + self, instance_name: str, location: str, project_id: str ) -> Dict[str, Any]: """ Gets details of a single Data Fusion instance. @@ -212,9 +203,6 @@ def get_instance( :param project_id: The ID of the Google Cloud Platform project that the instance belongs to. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") - instance = ( self.get_conn() # pylint: disable=no-member .projects() @@ -232,7 +220,7 @@ def patch_instance( instance: Dict[str, Any], update_mask: str, location: str, - project_id: Optional[str] = None, + project_id: str, ) -> Operation: """ Updates a single Data Fusion instance. @@ -254,9 +242,6 @@ def patch_instance( :param project_id: The ID of the Google Cloud Platform project that the instance belongs to. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") - operation = ( self.get_conn() # pylint: disable=no-member .projects() diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index 97b7404749079..b42a8d92a3187 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -254,7 +254,7 @@ def create_cluster( self, region: str, cluster: Union[Dict, Cluster], - project_id: Optional[str] = None, + project_id: str, request_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -301,8 +301,8 @@ def delete_cluster( self, region: str, cluster_name: str, + project_id: str, cluster_uuid: Optional[str] = None, - project_id: Optional[str] = None, request_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -351,7 +351,7 @@ def diagnose_cluster( self, region: str, cluster_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -391,7 +391,7 @@ def get_cluster( self, region: str, cluster_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -430,8 +430,8 @@ def list_clusters( self, region: str, filter_: str, + project_id: str, page_size: Optional[int] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -477,7 +477,7 @@ def update_cluster( # pylint: disable=too-many-arguments cluster_name: str, cluster: Union[Dict, Cluster], update_mask: Union[Dict, FieldMask], - project_id: Optional[str] = None, + project_id: str, graceful_decommission_timeout: Optional[Union[Dict, Duration]] = None, request_id: Optional[str] = None, retry: Optional[Retry] = None, @@ -562,7 +562,7 @@ def create_workflow_template( self, location: str, template: Union[Dict, WorkflowTemplate], - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -601,7 +601,7 @@ def instantiate_workflow_template( self, location: str, template_name: str, - project_id: Optional[str] = None, + project_id: str, version: Optional[int] = None, request_id: Optional[str] = None, parameters: Optional[Dict[str, str]] = None, @@ -658,7 +658,7 @@ def instantiate_inline_workflow_template( self, location: str, template: Union[Dict, WorkflowTemplate], - project_id: Optional[str] = None, + project_id: str, request_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -704,7 +704,7 @@ def wait_for_job( self, job_id: str, location: str, - project_id: Optional[str] = None, + project_id: str, wait_time: int = 10 ): """ @@ -738,7 +738,7 @@ def get_job( self, location: str, job_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -777,7 +777,7 @@ def submit_job( self, location: str, job: Union[Dict, Job], - project_id: Optional[str] = None, + project_id: str, request_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -858,7 +858,7 @@ def submit( def cancel_job( self, job_id: str, - project_id: Optional[str] = None, + project_id: str, location: str = 'global', retry: Optional[Retry] = None, timeout: Optional[float] = None, diff --git a/airflow/providers/google/cloud/hooks/datastore.py b/airflow/providers/google/cloud/hooks/datastore.py index 8926577521333..7da30af0c638a 100644 --- a/airflow/providers/google/cloud/hooks/datastore.py +++ b/airflow/providers/google/cloud/hooks/datastore.py @@ -71,7 +71,7 @@ def get_conn(self) -> Any: return self.connection @GoogleBaseHook.fallback_to_default_project_id - def allocate_ids(self, partial_keys: List, project_id: Optional[str] = None) -> List: + def allocate_ids(self, partial_keys: List, project_id: str) -> List: """ Allocate IDs for incomplete keys. @@ -95,7 +95,7 @@ def allocate_ids(self, partial_keys: List, project_id: Optional[str] = None) -> return resp['keys'] @GoogleBaseHook.fallback_to_default_project_id - def begin_transaction(self, project_id: Optional[str] = None) -> str: + def begin_transaction(self, project_id: str) -> str: """ Begins a new transaction. @@ -117,7 +117,7 @@ def begin_transaction(self, project_id: Optional[str] = None) -> str: return resp['transaction'] @GoogleBaseHook.fallback_to_default_project_id - def commit(self, body: Dict, project_id: Optional[str] = None) -> Dict: + def commit(self, body: Dict, project_id: str) -> Dict: """ Commit a transaction, optionally creating, deleting or modifying some entities. @@ -141,11 +141,13 @@ def commit(self, body: Dict, project_id: Optional[str] = None) -> Dict: return resp @GoogleBaseHook.fallback_to_default_project_id - def lookup(self, - keys: List, - read_consistency: Optional[str] = None, - transaction: Optional[str] = None, - project_id: Optional[str] = None) -> Dict: + def lookup( + self, + keys: List, + project_id: str, + read_consistency: Optional[str] = None, + transaction: Optional[str] = None, + ) -> Dict: """ Lookup some entities by key. @@ -179,7 +181,7 @@ def lookup(self, return resp @GoogleBaseHook.fallback_to_default_project_id - def rollback(self, transaction: str, project_id: Optional[str] = None) -> Any: + def rollback(self, transaction: str, project_id: str) -> Any: """ Roll back a transaction. @@ -198,7 +200,7 @@ def rollback(self, transaction: str, project_id: Optional[str] = None) -> Any: ).execute(num_retries=self.num_retries) @GoogleBaseHook.fallback_to_default_project_id - def run_query(self, body: Dict, project_id: Optional[str] = None) -> Dict: + def run_query(self, body: Dict, project_id: str) -> Dict: """ Run a query for entities. @@ -289,12 +291,14 @@ def poll_operation_until_done(self, name: str, polling_interval_in_seconds: int) return result @GoogleBaseHook.fallback_to_default_project_id - def export_to_storage_bucket(self, - bucket: str, - namespace: Optional[str] = None, - entity_filter: Optional[Dict] = None, - labels: Optional[Dict[str, str]] = None, - project_id: Optional[str] = None) -> Dict: + def export_to_storage_bucket( + self, + bucket: str, + project_id: str, + namespace: Optional[str] = None, + entity_filter: Optional[Dict] = None, + labels: Optional[Dict[str, str]] = None, + ) -> Dict: """ Export entities from Cloud Datastore to Cloud Storage for backup. @@ -337,13 +341,15 @@ def export_to_storage_bucket(self, return resp @GoogleBaseHook.fallback_to_default_project_id - def import_from_storage_bucket(self, - bucket: str, - file: str, - namespace: Optional[str] = None, - entity_filter: Optional[Dict] = None, - labels: Optional[Union[Dict, str]] = None, - project_id: Optional[str] = None) -> Dict: + def import_from_storage_bucket( + self, + bucket: str, + file: str, + project_id: str, + namespace: Optional[str] = None, + entity_filter: Optional[Dict] = None, + labels: Optional[Union[Dict, str]] = None, + ) -> Dict: """ Import a backup from Cloud Storage to Cloud Datastore. diff --git a/airflow/providers/google/cloud/hooks/dlp.py b/airflow/providers/google/cloud/hooks/dlp.py index 8a40fbd3a90cf..71187d832fd1d 100644 --- a/airflow/providers/google/cloud/hooks/dlp.py +++ b/airflow/providers/google/cloud/hooks/dlp.py @@ -79,7 +79,7 @@ def get_conn(self) -> DlpServiceClient: def cancel_dlp_job( self, dlp_job_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -172,7 +172,7 @@ def create_deidentify_template( @GoogleBaseHook.fallback_to_default_project_id def create_dlp_job( self, - project_id: Optional[str] = None, + project_id: str, inspect_job: Optional[Union[dict, InspectJobConfig]] = None, risk_job: Optional[Union[dict, RiskAnalysisJobConfig]] = None, job_id: Optional[str] = None, @@ -312,7 +312,7 @@ def create_inspect_template( @GoogleBaseHook.fallback_to_default_project_id def create_job_trigger( self, - project_id: Optional[str] = None, + project_id: str, job_trigger: Optional[Union[dict, JobTrigger]] = None, trigger_id: Optional[str] = None, retry: Optional[Retry] = None, @@ -415,7 +415,7 @@ def create_stored_info_type( @GoogleBaseHook.fallback_to_default_project_id def deidentify_content( self, - project_id: Optional[str] = None, + project_id: str, deidentify_config: Optional[Union[dict, DeidentifyConfig]] = None, inspect_config: Optional[Union[dict, InspectConfig]] = None, item: Optional[Union[dict, ContentItem]] = None, @@ -523,7 +523,7 @@ def delete_deidentify_template( def delete_dlp_job( self, dlp_job_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -610,7 +610,7 @@ def delete_inspect_template( def delete_job_trigger( self, job_trigger_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -746,7 +746,7 @@ def get_deidentify_template( def get_dlp_job( self, dlp_job_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -834,7 +834,7 @@ def get_inspect_template( def get_job_trigger( self, job_trigger_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -921,7 +921,7 @@ def get_stored_info_type( @GoogleBaseHook.fallback_to_default_project_id def inspect_content( self, - project_id: Optional[str] = None, + project_id: str, inspect_config: Optional[Union[dict, InspectConfig]] = None, item: Optional[Union[dict, ContentItem]] = None, inspect_template_name: Optional[str] = None, @@ -1034,7 +1034,7 @@ def list_deidentify_templates( @GoogleBaseHook.fallback_to_default_project_id def list_dlp_jobs( self, - project_id: Optional[str] = None, + project_id: str, results_filter: Optional[str] = None, page_size: Optional[int] = None, job_type: Optional[str] = None, @@ -1189,7 +1189,7 @@ def list_inspect_templates( @GoogleBaseHook.fallback_to_default_project_id def list_job_triggers( self, - project_id: Optional[str] = None, + project_id: str, page_size: Optional[int] = None, order_by: Optional[str] = None, results_filter: Optional[str] = None, @@ -1301,7 +1301,7 @@ def list_stored_info_types( @GoogleBaseHook.fallback_to_default_project_id def redact_image( self, - project_id: Optional[str] = None, + project_id: str, inspect_config: Optional[Union[dict, InspectConfig]] = None, image_redaction_configs: Optional[ Union[List[dict], List[RedactImageRequest.ImageRedactionConfig]] @@ -1361,7 +1361,7 @@ def redact_image( @GoogleBaseHook.fallback_to_default_project_id def reidentify_content( self, - project_id: Optional[str] = None, + project_id: str, reidentify_config: Optional[Union[dict, DeidentifyConfig]] = None, inspect_config: Optional[Union[dict, InspectConfig]] = None, item: Optional[Union[dict, ContentItem]] = None, @@ -1548,7 +1548,7 @@ def update_inspect_template( def update_job_trigger( self, job_trigger_id: str, - project_id: Optional[str] = None, + project_id: str, job_trigger: Optional[Union[dict, JobTrigger]] = None, update_mask: Optional[Union[dict, FieldMask]] = None, retry: Optional[Retry] = None, diff --git a/airflow/providers/google/cloud/hooks/functions.py b/airflow/providers/google/cloud/hooks/functions.py index ff4257d827669..dc3de06cd350f 100644 --- a/airflow/providers/google/cloud/hooks/functions.py +++ b/airflow/providers/google/cloud/hooks/functions.py @@ -90,7 +90,7 @@ def get_function(self, name: str) -> Dict: name=name).execute(num_retries=self.num_retries) @GoogleBaseHook.fallback_to_default_project_id - def create_new_function(self, location: str, body: Dict, project_id: Optional[str] = None) -> None: + def create_new_function(self, location: str, body: Dict, project_id: str) -> None: """ Creates a new function in Cloud Function in the location specified in the body. @@ -103,8 +103,6 @@ def create_new_function(self, location: str, body: Dict, project_id: Optional[st :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") response = self.get_conn().projects().locations().functions().create( # pylint: disable=no-member location=self._full_location(project_id, location), body=body @@ -133,7 +131,7 @@ def update_function(self, name: str, body: Dict, update_mask: List[str]) -> None self._wait_for_operation_to_complete(operation_name=operation_name) @GoogleBaseHook.fallback_to_default_project_id - def upload_function_zip(self, location: str, zip_path: str, project_id: Optional[str] = None) -> str: + def upload_function_zip(self, location: str, zip_path: str, project_id: str) -> str: """ Uploads zip file with sources. @@ -147,8 +145,6 @@ def upload_function_zip(self, location: str, zip_path: str, project_id: Optional :return: The upload URL that was returned by generateUploadUrl method. :rtype: str """ - if not project_id: - raise ValueError("The project_id should be set") response = \ self.get_conn().projects().locations().functions().generateUploadUrl( # pylint: disable=no-member # noqa parent=self._full_location(project_id, location) @@ -188,7 +184,7 @@ def call_function( function_id: str, input_data: Dict, location: str, - project_id: Optional[str] = None + project_id: str, ) -> Dict: """ Synchronously invokes a deployed Cloud Function. To be used for testing diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 94ea13832bf01..e2a3c1478d47c 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -142,7 +142,7 @@ def _append_label(cluster_proto: Cluster, key: str, val: str) -> Cluster: def delete_cluster( self, name: str, - project_id: Optional[str] = None, + project_id: str, retry: Retry = DEFAULT, timeout: float = DEFAULT ) -> Optional[str]: @@ -189,7 +189,7 @@ def delete_cluster( def create_cluster( self, cluster: Union[Dict, Cluster], - project_id: Optional[str] = None, + project_id: str, retry: Retry = DEFAULT, timeout: float = DEFAULT ) -> str: @@ -241,13 +241,13 @@ def create_cluster( return resource.target_link except AlreadyExists as error: self.log.info('Assuming Success: %s', error.message) - return self.get_cluster(name=cluster.name) + return self.get_cluster(name=cluster.name, project_id=project_id) @GoogleBaseHook.fallback_to_default_project_id def get_cluster( self, name: str, - project_id: Optional[str] = None, + project_id: str, retry: Retry = DEFAULT, timeout: float = DEFAULT ) -> Cluster: diff --git a/airflow/providers/google/cloud/hooks/mlengine.py b/airflow/providers/google/cloud/hooks/mlengine.py index dfa0bbed3826f..ef118b6f3d4a4 100644 --- a/airflow/providers/google/cloud/hooks/mlengine.py +++ b/airflow/providers/google/cloud/hooks/mlengine.py @@ -78,7 +78,7 @@ def get_conn(self): def create_job( self, job: Dict, - project_id: Optional[str] = None, + project_id: str, use_existing_job_fn: Optional[Callable] = None ) -> Dict: """ @@ -112,9 +112,6 @@ def create_job( terminal state (which might be FAILED or CANCELLED state). :rtype: dict """ - if not project_id: - raise ValueError("The project_id should be set") - hook = self.get_conn() self._append_label(job) @@ -152,7 +149,7 @@ def create_job( def cancel_job( self, job_id: str, - project_id: Optional[str] = None + project_id: str, ) -> Dict: """ @@ -169,10 +166,6 @@ def cancel_job( :rtype: dict :raises: googleapiclient.errors.HttpError """ - - if not project_id: - raise ValueError("The project_id should be set") - hook = self.get_conn() request = hook.projects().jobs().cancel( # pylint: disable=no-member @@ -251,7 +244,7 @@ def create_version( self, model_name: str, version_spec: Dict, - project_id: Optional[str] = None, + project_id: str, ) -> Dict: """ Creates the Version on Google Cloud ML Engine. @@ -291,7 +284,7 @@ def set_default_version( self, model_name: str, version_name: str, - project_id: Optional[str] = None, + project_id: str, ) -> Dict: """ Sets a version to be the default. Blocks until finished. @@ -327,7 +320,7 @@ def set_default_version( def list_versions( self, model_name: str, - project_id: Optional[str] = None, + project_id: str, ) -> List[Dict]: """ Lists all available versions of a model. Blocks until finished. @@ -364,7 +357,7 @@ def delete_version( self, model_name: str, version_name: str, - project_id: Optional[str] = None, + project_id: str, ) -> Dict: """ Deletes the given version of a model. Blocks until finished. @@ -379,9 +372,6 @@ def delete_version( Otherwise raises an error. :rtype: Dict """ - if not project_id: - raise ValueError("The project_id should be set") - hook = self.get_conn() full_name = 'projects/{}/models/{}/versions/{}'.format( project_id, model_name, version_name) @@ -401,7 +391,7 @@ def delete_version( def create_model( self, model: Dict, - project_id: Optional[str] = None + project_id: str, ) -> Dict: """ Create a Model. Blocks until finished. @@ -458,7 +448,7 @@ def create_model( def get_model( self, model_name: str, - project_id: Optional[str] = None, + project_id: str, ) -> Optional[Dict]: """ Gets a Model. Blocks until finished. @@ -492,8 +482,8 @@ def get_model( def delete_model( self, model_name: str, + project_id: str, delete_contents: bool = False, - project_id: Optional[str] = None, ) -> None: """ Delete a Model. Blocks until finished. @@ -509,9 +499,6 @@ def delete_model( :type project_id: str :raises: googleapiclient.errors.HttpError """ - if not project_id: - raise ValueError("The project_id should be set") - hook = self.get_conn() if not model_name: diff --git a/airflow/providers/google/cloud/hooks/pubsub.py b/airflow/providers/google/cloud/hooks/pubsub.py index 5e949b8734d3f..7c8d6d33a75fa 100644 --- a/airflow/providers/google/cloud/hooks/pubsub.py +++ b/airflow/providers/google/cloud/hooks/pubsub.py @@ -86,7 +86,7 @@ def publish( self, topic: str, messages: List[Dict], - project_id: Optional[str] = None, + project_id: str, ) -> None: """ Publishes messages to a Pub/Sub topic. @@ -102,8 +102,6 @@ def publish( If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str """ - if not project_id: - raise ValueError("Project ID should be set.") self._validate_messages(messages) publisher = self.get_conn() @@ -154,7 +152,7 @@ def _validate_messages(messages) -> None: def create_topic( self, topic: str, - project_id: Optional[str] = None, + project_id: str, fail_if_exists: bool = False, labels: Optional[Dict[str, str]] = None, message_storage_policy: Union[Dict, MessageStoragePolicy] = None, @@ -199,8 +197,6 @@ def create_topic( :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]]] """ - if not project_id: - raise ValueError("Project ID should be set.") publisher = self.get_conn() topic_path = PublisherClient.topic_path(project_id, topic) # pylint: disable=no-member @@ -233,7 +229,7 @@ def create_topic( def delete_topic( self, topic: str, - project_id: Optional[str] = None, + project_id: str, fail_if_not_exists: bool = False, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -261,8 +257,6 @@ def delete_topic( :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]]] """ - if not project_id: - raise ValueError("Project ID should be set.") publisher = self.get_conn() topic_path = PublisherClient.topic_path(project_id, topic) # pylint: disable=no-member @@ -288,7 +282,7 @@ def delete_topic( def create_subscription( self, topic: str, - project_id: Optional[str] = None, + project_id: str, subscription: Optional[str] = None, subscription_project_id: Optional[str] = None, ack_deadline_secs: int = 10, @@ -355,8 +349,6 @@ def create_subscription( the ``subscription`` parameter is not supplied :rtype: str """ - if not project_id: - raise ValueError("Project ID should be set.") subscriber = self.subscriber_client if not subscription: @@ -400,7 +392,7 @@ def create_subscription( def delete_subscription( self, subscription: str, - project_id: Optional[str] = None, + project_id: str, fail_if_not_exists: bool = False, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -427,8 +419,6 @@ def delete_subscription( :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]]] """ - if not project_id: - raise ValueError("Project ID should be set.") subscriber = self.subscriber_client subscription_path = SubscriberClient.subscription_path(project_id, subscription) # noqa E501 # pylint: disable=no-member,line-too-long @@ -456,7 +446,7 @@ def pull( self, subscription: str, max_messages: int, - project_id: Optional[str] = None, + project_id: str, return_immediately: bool = False, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -492,8 +482,6 @@ def pull( the base64-encoded message content. See https://cloud.google.com/pubsub/docs/reference/rest/v1/projects.subscriptions/pull#ReceivedMessage """ - if not project_id: - raise ValueError("Project ID should be set.") subscriber = self.subscriber_client subscription_path = SubscriberClient.subscription_path(project_id, subscription) # noqa E501 # pylint: disable=no-member,line-too-long @@ -518,9 +506,9 @@ def pull( def acknowledge( self, subscription: str, + project_id: str, ack_ids: Optional[List[str]] = None, messages: Optional[List[ReceivedMessage]] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -550,10 +538,6 @@ def acknowledge( :param metadata: (Optional) Additional metadata that is provided to the method. :type metadata: Sequence[Tuple[str, str]]] """ - - if not project_id: - raise ValueError("Project ID should be set.") - if ack_ids is not None and messages is None: pass elif ack_ids is None and messages is not None: diff --git a/airflow/providers/google/cloud/hooks/spanner.py b/airflow/providers/google/cloud/hooks/spanner.py index aebbfa2372608..0a03d19e84a8c 100644 --- a/airflow/providers/google/cloud/hooks/spanner.py +++ b/airflow/providers/google/cloud/hooks/spanner.py @@ -64,7 +64,7 @@ def _get_client(self, project_id: str) -> Client: def get_instance( self, instance_id: str, - project_id: Optional[str] = None + project_id: str, ) -> Instance: """ Gets information about a particular instance. @@ -77,8 +77,6 @@ def get_instance( :return: Spanner instance :rtype: google.cloud.spanner_v1.instance.Instance """ - if not project_id: - raise ValueError("The project_id should be set") instance = self._get_client(project_id=project_id).instance(instance_id=instance_id) if not instance.exists(): return None @@ -133,7 +131,7 @@ def create_instance( configuration_name: str, node_count: int, display_name: str, - project_id: Optional[str] = None + project_id: str, ) -> None: """ Creates a new Cloud Spanner instance. @@ -156,8 +154,6 @@ def create_instance( :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") self._apply_to_instance(project_id, instance_id, configuration_name, node_count, display_name, lambda x: x.create()) @@ -168,7 +164,7 @@ def update_instance( configuration_name: str, node_count: int, display_name: str, - project_id: Optional[str] = None + project_id: str, ) -> None: """ Updates an existing Cloud Spanner instance. @@ -191,13 +187,11 @@ def update_instance( :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") self._apply_to_instance(project_id, instance_id, configuration_name, node_count, display_name, lambda x: x.update()) @GoogleBaseHook.fallback_to_default_project_id - def delete_instance(self, instance_id: str, project_id: Optional[str] = None) -> None: + def delete_instance(self, instance_id: str, project_id: str) -> None: """ Deletes an existing Cloud Spanner instance. @@ -208,8 +202,6 @@ def delete_instance(self, instance_id: str, project_id: Optional[str] = None) -> :type project_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") instance = self._get_client(project_id=project_id).instance(instance_id) try: instance.delete() @@ -223,7 +215,7 @@ def get_database( self, instance_id: str, database_id: str, - project_id: Optional[str] = None + project_id: str, ) -> Optional[Database]: """ Retrieves a database in Cloud Spanner. If the database does not exist @@ -239,9 +231,6 @@ def get_database( :return: Database object or None if database does not exist :rtype: google.cloud.spanner_v1.database.Database or None """ - if not project_id: - raise ValueError("The project_id should be set") - instance = self._get_client(project_id=project_id).instance( instance_id=instance_id) if not instance.exists(): @@ -259,7 +248,7 @@ def create_database( instance_id: str, database_id: str, ddl_statements: List[str], - project_id: Optional[str] = None + project_id: str, ) -> None: """ Creates a new database in Cloud Spanner. @@ -275,9 +264,6 @@ def create_database( database. If set to None or missing, the default project_id from the GCP connection is used. :return: None """ - if not project_id: - raise ValueError("The project_id should be set") - instance = self._get_client(project_id=project_id).instance( instance_id=instance_id) if not instance.exists(): @@ -301,7 +287,7 @@ def update_database( instance_id: str, database_id: str, ddl_statements: List[str], - project_id: Optional[str] = None, + project_id: str, operation_id: Optional[str] = None ) -> None: """ @@ -321,9 +307,6 @@ def update_database( :type operation_id: str :return: None """ - if not project_id: - raise ValueError("The project_id should be set") - instance = self._get_client(project_id=project_id).instance( instance_id=instance_id) if not instance.exists(): @@ -347,7 +330,7 @@ def update_database( raise e @GoogleBaseHook.fallback_to_default_project_id - def delete_database(self, instance_id: str, database_id, project_id: Optional[str] = None) -> bool: + def delete_database(self, instance_id: str, database_id, project_id: str) -> bool: """ Drops a database in Cloud Spanner. @@ -361,9 +344,6 @@ def delete_database(self, instance_id: str, database_id, project_id: Optional[st :return: True if everything succeeded :rtype: bool """ - if not project_id: - raise ValueError("The project_id should be set") - instance = self._get_client(project_id=project_id).\ instance(instance_id=instance_id) if not instance.exists(): @@ -390,7 +370,7 @@ def execute_dml( instance_id: str, database_id: str, queries: List[str], - project_id: Optional[str] = None, + project_id: str, ) -> None: """ Executes an arbitrary DML query (INSERT, UPDATE, DELETE). @@ -405,9 +385,6 @@ def execute_dml( database. If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") - self._get_client(project_id=project_id).instance(instance_id=instance_id).\ database(database_id=database_id).run_in_transaction( lambda transaction: self._execute_sql_in_transaction(transaction, queries)) diff --git a/airflow/providers/google/cloud/hooks/stackdriver.py b/airflow/providers/google/cloud/hooks/stackdriver.py index 51b0bd4f47343..e0ebd7b9c7979 100644 --- a/airflow/providers/google/cloud/hooks/stackdriver.py +++ b/airflow/providers/google/cloud/hooks/stackdriver.py @@ -56,7 +56,7 @@ def _get_channel_client(self): @GoogleBaseHook.fallback_to_default_project_id def list_alert_policies( self, - project_id: Optional[str] = None, + project_id: str, format_: Optional[str] = None, filter_: Optional[str] = None, order_by: Optional[str] = None, @@ -124,7 +124,7 @@ def list_alert_policies( def _toggle_policy_status( self, new_state: bool, - project_id: Optional[str] = None, + project_id: str, filter_: Optional[str] = None, retry: Optional[str] = DEFAULT, timeout: Optional[float] = DEFAULT, @@ -148,7 +148,7 @@ def _toggle_policy_status( @GoogleBaseHook.fallback_to_default_project_id def enable_alert_policies( self, - project_id: Optional[str] = None, + project_id: str, filter_: Optional[str] = None, retry: Optional[str] = DEFAULT, timeout: Optional[float] = DEFAULT, @@ -186,10 +186,10 @@ def enable_alert_policies( @GoogleBaseHook.fallback_to_default_project_id def disable_alert_policies( self, - project_id: Optional[str] = None, + project_id: str, filter_: Optional[str] = None, retry: Optional[str] = DEFAULT, - timeout: Optional[str] = DEFAULT, + timeout: Optional[float] = DEFAULT, metadata: Optional[str] = None ) -> None: """ @@ -225,7 +225,7 @@ def disable_alert_policies( def upsert_alert( self, alerts: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[str] = DEFAULT, timeout: Optional[float] = DEFAULT, metadata: Optional[str] = None @@ -365,6 +365,7 @@ def delete_alert_policy( @GoogleBaseHook.fallback_to_default_project_id def list_notification_channels( self, + project_id: str, format_: Optional[str] = None, filter_: Optional[str] = None, order_by: Optional[str] = None, @@ -372,7 +373,6 @@ def list_notification_channels( retry: Optional[str] = DEFAULT, timeout: Optional[str] = DEFAULT, metadata: Optional[str] = None, - project_id: Optional[str] = None ) -> Any: """ Fetches all the Notification Channels identified by the filter passed as @@ -433,8 +433,8 @@ def list_notification_channels( @GoogleBaseHook.fallback_to_default_project_id def _toggle_channel_status( self, - new_state: str, - project_id: Optional[str] = None, + new_state: bool, + project_id: str, filter_: Optional[str] = None, retry: Optional[str] = DEFAULT, timeout: Optional[str] = DEFAULT, @@ -461,7 +461,7 @@ def _toggle_channel_status( @GoogleBaseHook.fallback_to_default_project_id def enable_notification_channels( self, - project_id: Optional[str] = None, + project_id: str, filter_: Optional[str] = None, retry: Optional[str] = DEFAULT, timeout: Optional[str] = DEFAULT, @@ -500,8 +500,8 @@ def enable_notification_channels( @GoogleBaseHook.fallback_to_default_project_id def disable_notification_channels( self, + project_id: str, filter_: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[str] = DEFAULT, timeout: Optional[str] = DEFAULT, metadata: Optional[str] = None @@ -540,7 +540,7 @@ def disable_notification_channels( def upsert_channel( self, channels: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[str] = DEFAULT, timeout: Optional[float] = DEFAULT, metadata: Optional[str] = None diff --git a/airflow/providers/google/cloud/hooks/tasks.py b/airflow/providers/google/cloud/hooks/tasks.py index e44312773f417..a567cc5269d6c 100644 --- a/airflow/providers/google/cloud/hooks/tasks.py +++ b/airflow/providers/google/cloud/hooks/tasks.py @@ -70,7 +70,7 @@ def create_queue( self, location: str, task_queue: Union[Dict, Queue], - project_id: Optional[str] = None, + project_id: str, queue_name: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -126,7 +126,7 @@ def create_queue( def update_queue( self, task_queue: Queue, - project_id: Optional[str] = None, + project_id: str, location: Optional[str] = None, queue_name: Optional[str] = None, update_mask: Optional[FieldMask] = None, @@ -189,7 +189,7 @@ def get_queue( self, location: str, queue_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -227,7 +227,7 @@ def get_queue( def list_queues( self, location: str, - project_id: Optional[str] = None, + project_id: str, results_filter: Optional[str] = None, page_size: Optional[int] = None, retry: Optional[Retry] = None, @@ -277,7 +277,7 @@ def delete_queue( self, location: str, queue_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -315,7 +315,7 @@ def purge_queue( self, location: str, queue_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -354,7 +354,7 @@ def pause_queue( self, location: str, queue_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -393,7 +393,7 @@ def resume_queue( self, location: str, queue_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -433,7 +433,7 @@ def create_task( location: str, queue_name: str, task: Union[Dict, Task], - project_id: Optional[str] = None, + project_id: str, task_name: Optional[str] = None, response_view: Optional[enums.Task.View] = None, retry: Optional[Retry] = None, @@ -499,7 +499,7 @@ def get_task( location: str, queue_name: str, task_name: str, - project_id: Optional[str] = None, + project_id: str, response_view: Optional[enums.Task.View] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -548,7 +548,7 @@ def list_tasks( self, location: str, queue_name: str, - project_id: Optional[str] = None, + project_id: str, response_view: Optional[enums.Task.View] = None, page_size: Optional[int] = None, retry: Optional[Retry] = None, @@ -601,7 +601,7 @@ def delete_task( location: str, queue_name: str, task_name: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -642,7 +642,7 @@ def run_task( location: str, queue_name: str, task_name: str, - project_id: Optional[str] = None, + project_id: str, response_view: Optional[enums.Task.View] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, diff --git a/airflow/providers/google/cloud/hooks/vision.py b/airflow/providers/google/cloud/hooks/vision.py index a8fb1a2a27c87..19a696ef4517f 100644 --- a/airflow/providers/google/cloud/hooks/vision.py +++ b/airflow/providers/google/cloud/hooks/vision.py @@ -167,7 +167,7 @@ def create_product_set( self, location: str, product_set: Union[dict, ProductSet], - project_id: Optional[str] = None, + project_id: str, product_set_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -177,8 +177,6 @@ def create_product_set( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetCreateOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() parent = ProductSearchClient.location_path(project_id, location) self.log.info('Creating a new ProductSet under the parent: %s', parent) @@ -205,7 +203,7 @@ def get_product_set( self, location: str, product_set_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None ) -> Dict: @@ -213,8 +211,6 @@ def get_product_set( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetGetOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Retrieving ProductSet: %s', name) @@ -227,10 +223,10 @@ def get_product_set( def update_product_set( self, product_set: Union[dict, ProductSet], + project_id: str, location: Optional[str] = None, product_set_id: Optional[str] = None, update_mask: Union[dict, FieldMask] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -239,8 +235,6 @@ def update_product_set( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetUpdateOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() product_set = self.product_set_name_determiner.get_entity_with_name( product_set, product_set_id, location, project_id @@ -258,7 +252,7 @@ def delete_product_set( self, location: str, product_set_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -267,8 +261,6 @@ def delete_product_set( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductSetDeleteOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() name = ProductSearchClient.product_set_path(project_id, location, product_set_id) self.log.info('Deleting ProductSet: %s', name) @@ -280,7 +272,7 @@ def create_product( self, location: str, product: Union[dict, Product], - project_id: Optional[str] = None, + project_id: str, product_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, @@ -290,8 +282,6 @@ def create_product( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductCreateOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() parent = ProductSearchClient.location_path(project_id, location) self.log.info('Creating a new Product under the parent: %s', parent) @@ -318,7 +308,7 @@ def get_product( self, location: str, product_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -327,8 +317,6 @@ def get_product( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductGetOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() name = ProductSearchClient.product_path(project_id, location, product_id) self.log.info('Retrieving Product: %s', name) @@ -341,10 +329,10 @@ def get_product( def update_product( self, product: Union[dict, Product], + project_id: str, location: Optional[str] = None, product_id: Optional[str] = None, update_mask: Optional[Dict[str, FieldMask]] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -353,8 +341,6 @@ def update_product( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductUpdateOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() product = self.product_name_determiner.get_entity_with_name(product, product_id, location, project_id) self.log.info('Updating ProductSet: %s', product.name) @@ -370,7 +356,7 @@ def delete_product( self, location: str, product_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None @@ -379,8 +365,6 @@ def delete_product( For the documentation see: :class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionProductDeleteOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() name = ProductSearchClient.product_path(project_id, location, product_id) self.log.info('Deleting ProductSet: %s', name) @@ -393,8 +377,8 @@ def create_reference_image( location: str, product_id: str, reference_image: Union[Dict, ReferenceImage], + project_id: str, reference_image_id: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -403,8 +387,6 @@ def create_reference_image( For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionReferenceImageCreateOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() self.log.info('Creating ReferenceImage') parent = ProductSearchClient.product_path(project=project_id, location=location, product=product_id) @@ -436,7 +418,7 @@ def delete_reference_image( location: str, product_id: str, reference_image_id: str, - project_id: Optional[str] = None, + project_id: str, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -445,8 +427,6 @@ def delete_reference_image( For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionReferenceImageCreateOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() self.log.info('Deleting ReferenceImage') name = ProductSearchClient.reference_image_path( @@ -465,8 +445,8 @@ def add_product_to_product_set( self, product_set_id: str, product_id: str, + project_id: str, location: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -475,8 +455,6 @@ def add_product_to_product_set( For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionAddProductToProductSetOperator` """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() product_name = ProductSearchClient.product_path(project_id, location, product_id) @@ -495,8 +473,8 @@ def remove_product_from_product_set( self, product_set_id: str, product_id: str, + project_id: str, location: Optional[str] = None, - project_id: Optional[str] = None, retry: Optional[Retry] = None, timeout: Optional[float] = None, metadata: Optional[Sequence[Tuple[str, str]]] = None, @@ -505,8 +483,6 @@ def remove_product_from_product_set( For the documentation see: :py:class:`~airflow.contrib.operators.gcp_vision_operator.CloudVisionRemoveProductFromProductSetOperator` # pylint: disable=line-too-long # noqa """ - if not project_id: - raise ValueError("Project ID should be set.") client = self.get_conn() product_name = ProductSearchClient.product_path(project_id, location, product_id) diff --git a/airflow/providers/google/cloud/operators/automl.py b/airflow/providers/google/cloud/operators/automl.py index 4e010006b4697..833b60089c070 100644 --- a/airflow/providers/google/cloud/operators/automl.py +++ b/airflow/providers/google/cloud/operators/automl.py @@ -542,9 +542,6 @@ class AutoMLTablesUpdateDatasetOperator(BaseOperator): :param update_mask: The update mask applies to the resource. If a dict is provided, it must be of the same form as the protobuf message FieldMask. :type update_mask: Union[dict, FieldMask] - :param project_id: ID of the Google Cloud project where dataset is located if None then - default project_id is used. - :type project_id: str :param location: The location of the project. :type location: str :param params: Additional domain-specific parameters for the predictions. @@ -568,7 +565,6 @@ def __init__( self, dataset: dict, location: str, - project_id: Optional[str] = None, update_mask: Optional[dict] = None, metadata: Optional[MetaData] = None, timeout: Optional[float] = None, @@ -582,7 +578,6 @@ def __init__( self.dataset = dataset self.update_mask = update_mask self.location = location - self.project_id = project_id self.metadata = metadata self.timeout = timeout self.retry = retry @@ -594,7 +589,6 @@ def execute(self, context): result = hook.update_dataset( dataset=self.dataset, update_mask=self.update_mask, - project_id=self.project_id, retry=self.retry, timeout=self.timeout, metadata=self.metadata, diff --git a/airflow/providers/google/cloud/operators/dataflow.py b/airflow/providers/google/cloud/operators/dataflow.py index 9826965591d10..ed8a603e1e2d0 100644 --- a/airflow/providers/google/cloud/operators/dataflow.py +++ b/airflow/providers/google/cloud/operators/dataflow.py @@ -172,6 +172,7 @@ def __init__( job_name: str = '{{task.task_id}}', dataflow_default_options: Optional[dict] = None, options: Optional[dict] = None, + project_id: Optional[str] = None, gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None, poll_sleep: int = 10, @@ -186,6 +187,7 @@ def __init__( options = options or {} options.setdefault('labels', {}).update( {'airflow-version': 'v' + version.replace('.', '-').replace('+', '-')}) + self.project_id = project_id self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to self.jar = jar @@ -211,10 +213,12 @@ def execute(self, context): if self.check_if_running != CheckJobRunning.IgnoreJob: is_running = self.hook.is_job_dataflow_running( name=self.job_name, - variables=dataflow_options + variables=dataflow_options, + project_id=self.project_id, ) while is_running and self.check_if_running == CheckJobRunning.WaitForRun: - is_running = self.hook.is_job_dataflow_running(name=self.job_name, variables=dataflow_options) + is_running = self.hook.is_job_dataflow_running( + name=self.job_name, variables=dataflow_options, project_id=self.project_id) if not is_running: bucket_helper = GoogleCloudBucketHelper( @@ -231,13 +235,14 @@ def set_current_job_id(job_id): job_class=self.job_class, append_job_name=True, multiple_jobs=self.multiple_jobs, - on_new_job_id_callback=set_current_job_id + on_new_job_id_callback=set_current_job_id, + project_id=self.project_id, ) def on_kill(self) -> None: self.log.info("On kill.") if self.job_id: - self.hook.cancel_job(job_id=self.job_id) + self.hook.cancel_job(job_id=self.job_id, project_id=self.project_id) class DataflowTemplatedJobStartOperator(BaseOperator): @@ -360,7 +365,8 @@ def set_current_job_id(job_id): variables=self.dataflow_default_options, parameters=self.parameters, dataflow_template=self.template, - on_new_job_id_callback=set_current_job_id + on_new_job_id_callback=set_current_job_id, + project_id=self.project_id, ) return job @@ -368,7 +374,7 @@ def set_current_job_id(job_id): def on_kill(self) -> None: self.log.info("On kill.") if self.job_id: - self.hook.cancel_job(job_id=self.job_id) + self.hook.cancel_job(job_id=self.job_id, project_id=self.project_id) class DataflowCreatePythonJobOperator(BaseOperator): @@ -437,6 +443,7 @@ def __init__( # pylint: disable=too-many-arguments py_options: Optional[List[str]] = None, py_requirements: Optional[List[str]] = None, py_system_site_packages: bool = False, + project_id: Optional[str] = None, gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None, poll_sleep: int = 10, @@ -455,6 +462,7 @@ def __init__( # pylint: disable=too-many-arguments self.py_interpreter = py_interpreter self.py_requirements = py_requirements or [] self.py_system_site_packages = py_system_site_packages + self.project_id = project_id self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to self.poll_sleep = poll_sleep @@ -490,13 +498,14 @@ def set_current_job_id(job_id): py_interpreter=self.py_interpreter, py_requirements=self.py_requirements, py_system_site_packages=self.py_system_site_packages, - on_new_job_id_callback=set_current_job_id + on_new_job_id_callback=set_current_job_id, + project_id=self.project_id, ) def on_kill(self) -> None: self.log.info("On kill.") if self.job_id: - self.hook.cancel_job(job_id=self.job_id) + self.hook.cancel_job(job_id=self.job_id, project_id=self.project_id) class GoogleCloudBucketHelper: diff --git a/airflow/providers/google/cloud/operators/datafusion.py b/airflow/providers/google/cloud/operators/datafusion.py index 4d5b747fff121..0d9b0ebc488b5 100644 --- a/airflow/providers/google/cloud/operators/datafusion.py +++ b/airflow/providers/google/cloud/operators/datafusion.py @@ -221,7 +221,7 @@ def execute(self, context: Dict): raise self.log.info("Instance %s already exists", self.instance_name) instance = hook.get_instance( - instance_name=self.instance_name, location=self.location + instance_name=self.instance_name, location=self.location, project_id=self.project_id ) return instance diff --git a/airflow/providers/google/firebase/hooks/firestore.py b/airflow/providers/google/firebase/hooks/firestore.py index 141e1737484fc..162ca4de6fa93 100644 --- a/airflow/providers/google/firebase/hooks/firestore.py +++ b/airflow/providers/google/firebase/hooks/firestore.py @@ -96,8 +96,6 @@ def export_documents( If set to None or missing, the default project_id from the GCP connection is used. :type project_id: str """ - if not project_id: - raise ValueError("The project_id should be set") service = self.get_conn() name = f"projects/{project_id}/databases/{database_id}" diff --git a/scripts/ci/ci_pylint_main.sh b/scripts/ci/ci_pylint_main.sh index 5796ca039eccd..53a89914ac5e0 100755 --- a/scripts/ci/ci_pylint_main.sh +++ b/scripts/ci/ci_pylint_main.sh @@ -51,7 +51,7 @@ function run_pylint_main() { prepare_ci_build -rebuild_ci_image_if_needed +# rebuild_ci_image_if_needed if [[ "${#@}" != "0" ]]; then filter_out_files_from_pylint_todo_list "$@" diff --git a/setup.cfg b/setup.cfg index bbdd571abe236..b10293a546148 100644 --- a/setup.cfg +++ b/setup.cfg @@ -57,6 +57,8 @@ packages = airflow [mypy] ignore_missing_imports = True no_implicit_optional = True +plugins = + airflow.mypy.plugin.decorators [mypy-airflow.migrations.*] ignore_errors = True diff --git a/setup.py b/setup.py index e3334ceb297d0..c1cacc14491b8 100644 --- a/setup.py +++ b/setup.py @@ -455,7 +455,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version ############################################################################################################ if PY3: - devel += ['mypy==0.740'] + devel += ['mypy==0.770'] else: devel += ['unittest2'] diff --git a/tests/providers/google/cloud/hooks/test_automl.py b/tests/providers/google/cloud/hooks/test_automl.py index fc3828441a86e..07bdd70efd2ae 100644 --- a/tests/providers/google/cloud/hooks/test_automl.py +++ b/tests/providers/google/cloud/hooks/test_automl.py @@ -215,7 +215,7 @@ def test_delete_model(self, mock_delete_model): @mock.patch("airflow.providers.google.cloud.hooks.automl.AutoMlClient.update_dataset") def test_update_dataset(self, mock_update_dataset): self.hook.update_dataset( - dataset=DATASET, update_mask=MASK, project_id=GCP_PROJECT_ID + dataset=DATASET, update_mask=MASK, ) mock_update_dataset.assert_called_once_with( diff --git a/tests/providers/google/cloud/hooks/test_bigtable.py b/tests/providers/google/cloud/hooks/test_bigtable.py index 4a9430093fbc4..69390695a4974 100644 --- a/tests/providers/google/cloud/hooks/test_bigtable.py +++ b/tests/providers/google/cloud/hooks/test_bigtable.py @@ -24,7 +24,6 @@ from google.cloud.bigtable.instance import Instance from mock import PropertyMock -from airflow.exceptions import AirflowException from airflow.providers.google.cloud.hooks.bigtable import BigtableHook from tests.providers.google.cloud.utils.base_gcp_mock import ( GCP_PROJECT_ID_HOOK_UNIT_TEST, mock_base_gcp_hook_default_project_id, @@ -61,23 +60,6 @@ def test_bigtable_client_creation(self, mock_client, mock_get_creds, mock_client self.assertEqual(mock_client.return_value, result) self.assertEqual(self.bigtable_hook_no_default_project_id._client, result) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') - def test_get_instance_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - with self.assertRaises(AirflowException) as cm: - self.bigtable_hook_no_default_project_id.get_instance(instance_id=CBT_INSTANCE) - instance_exists_method.assert_not_called() - instance_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_get_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -91,25 +73,6 @@ def test_get_instance_overridden_project_id(self, get_client): get_client.assert_called_once_with(project_id='example-project') self.assertIsNotNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') - def test_delete_instance_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - delete_method = instance_method.return_value.delete - instance_exists_method.return_value = True - with self.assertRaises(AirflowException) as cm: - self.bigtable_hook_no_default_project_id.delete_instance(instance_id=CBT_INSTANCE) - instance_exists_method.assert_not_called() - instance_method.assert_not_called() - delete_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -124,27 +87,6 @@ def test_delete_instance_overridden_project_id(self, get_client): get_client.assert_called_once_with(project_id='example-project') self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('google.cloud.bigtable.instance.Instance.create') - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') - def test_create_instance_missing_project_id(self, get_client, instance_create, mock_project_id): - operation = mock.Mock() - operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client) - instance_create.return_value = operation - with self.assertRaises(AirflowException) as cm: - self.bigtable_hook_no_default_project_id.create_instance( - instance_id=CBT_INSTANCE, - main_cluster_id=CBT_CLUSTER, - main_cluster_zone=CBT_ZONE) - get_client.assert_not_called() - instance_create.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('google.cloud.bigtable.instance.Instance.create') @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_create_instance_overridden_project_id(self, get_client, instance_create): @@ -160,27 +102,6 @@ def test_create_instance_overridden_project_id(self, get_client, instance_create instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') - def test_delete_table_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - table_delete_method = instance_method.return_value.table.return_value.delete - instance_exists_method.return_value = True - with self.assertRaises(AirflowException) as cm: - self.bigtable_hook_no_default_project_id.delete_table( - instance_id=CBT_INSTANCE, - table_id=CBT_TABLE) - get_client.assert_not_called() - instance_exists_method.assert_not_called() - table_delete_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client') def test_delete_table_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -231,7 +152,9 @@ def test_get_instance(self, get_client, mock_project_id): instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = True res = self.bigtable_hook_default_project_id.get_instance( - instance_id=CBT_INSTANCE) + instance_id=CBT_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') @@ -261,7 +184,9 @@ def test_get_instance_no_instance(self, get_client, mock_project_id): instance_exists_method = instance_method.return_value.exists instance_exists_method.return_value = False res = self.bigtable_hook_default_project_id.get_instance( - instance_id=CBT_INSTANCE) + instance_id=CBT_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() get_client.assert_called_once_with(project_id='example-project') @@ -279,7 +204,9 @@ def test_delete_instance(self, get_client, mock_project_id): instance_exists_method.return_value = True delete_method = instance_method.return_value.delete res = self.bigtable_hook_default_project_id.delete_instance( - instance_id=CBT_INSTANCE) + instance_id=CBT_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_called_once_with() @@ -312,7 +239,9 @@ def test_delete_instance_no_instance(self, get_client, mock_project_id): instance_exists_method.return_value = False delete_method = instance_method.return_value.delete self.bigtable_hook_default_project_id.delete_instance( - instance_id=CBT_INSTANCE) + instance_id=CBT_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) instance_method.assert_called_once_with('instance') instance_exists_method.assert_called_once_with() delete_method.assert_not_called() @@ -332,7 +261,9 @@ def test_create_instance(self, get_client, instance_create, mock_project_id): res = self.bigtable_hook_default_project_id.create_instance( instance_id=CBT_INSTANCE, main_cluster_id=CBT_CLUSTER, - main_cluster_zone=CBT_ZONE) + main_cluster_zone=CBT_ZONE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_create.assert_called_once_with(clusters=mock.ANY) self.assertEqual(res.instance_id, 'instance') @@ -365,7 +296,9 @@ def test_delete_table(self, get_client, mock_project_id): instance_exists_method.return_value = True self.bigtable_hook_default_project_id.delete_table( instance_id=CBT_INSTANCE, - table_id=CBT_TABLE) + table_id=CBT_TABLE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_exists_method.assert_called_once_with() table_delete_method.assert_called_once_with() diff --git a/tests/providers/google/cloud/hooks/test_cloud_build.py b/tests/providers/google/cloud/hooks/test_cloud_build.py index 77576da960966..6a149b4205144 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_build.py +++ b/tests/providers/google/cloud/hooks/test_cloud_build.py @@ -131,7 +131,7 @@ def test_error_operation(self, _, get_conn_mock, mock_project_id): execute_mock = mock.Mock(**{"side_effect": [TEST_WAITING_OPERATION, TEST_ERROR_OPERATION]}) service_mock.operations.return_value.get.return_value.execute = execute_mock with self.assertRaisesRegex(AirflowException, "error"): - self.hook.create_build(body={}) + self.hook.create_build(body={}) # pylint: disable=no-value-for-parameter class TestGcpComputeHookWithDefaultProjectIdFromConnection(unittest.TestCase): @@ -173,7 +173,7 @@ def test_build_immediately_complete(self, get_conn_mock, mock_project_id): service_mock.operations.return_value.get.return_value.execute.return_value = TEST_DONE_OPERATION - result = self.hook.create_build(body={}) + result = self.hook.create_build(body={}) # pylint: disable=no-value-for-parameter service_mock.projects.return_value.builds.return_value.create.assert_called_once_with( body={}, projectId='example-project' @@ -204,7 +204,7 @@ def test_waiting_operation(self, _, get_conn_mock, mock_project_id): ) service_mock.operations.return_value.get.return_value.execute = execute_mock - result = self.hook.create_build(body={}) + result = self.hook.create_build(body={}) # pylint: disable=no-value-for-parameter self.assertEqual(result, TEST_BUILD) @@ -225,7 +225,7 @@ def test_error_operation(self, _, get_conn_mock, mock_project_id): execute_mock = mock.Mock(**{"side_effect": [TEST_WAITING_OPERATION, TEST_ERROR_OPERATION]}) service_mock.operations.return_value.get.return_value.execute = execute_mock with self.assertRaisesRegex(AirflowException, "error"): - self.hook.create_build(body={}) + self.hook.create_build(body={}) # pylint: disable=no-value-for-parameter class TestCloudBuildHookWithoutProjectId(unittest.TestCase): @@ -256,7 +256,7 @@ def test_cloud_build_client_creation(self, mock_build, mock_authorize): @mock.patch("airflow.providers.google.cloud.hooks.cloud_build.CloudBuildHook.get_conn") def test_create_build(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException) as e: - self.hook.create_build(body={}) + self.hook.create_build(body={}) # pylint: disable=no-value-for-parameter self.assertEqual( "The project id must be passed either as keyword project_id parameter or as project_id extra in " diff --git a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py index 8b4f350c5b33c..d49abd71e33df 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py +++ b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py @@ -69,7 +69,7 @@ def setUp(self,): ) def test_create_instance_when_exists(self, mock_get_conn, mock_project_id): mock_get_conn.return_value.get_instance.return_value = Instance(name=TEST_NAME) - result = self.hook.create_instance( + result = self.hook.create_instance( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), @@ -96,7 +96,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): Instance(name=TEST_NAME), ] mock_get_conn.return_value.create_instance.return_value.result.return_value = Instance(name=TEST_NAME) - result = self.hook.create_instance( + result = self.hook.create_instance( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, instance_id=TEST_INSTANCE_ID, instance=Instance(name=TEST_NAME), @@ -132,7 +132,7 @@ def test_create_instance_when_not_exists(self, mock_get_conn, mock_project_id): "airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_delete_instance(self, mock_get_conn, mock_project_id): - self.hook.delete_instance( + self.hook.delete_instance( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, instance=TEST_INSTANCE_ID, retry=TEST_RETRY, @@ -152,7 +152,7 @@ def test_delete_instance(self, mock_get_conn, mock_project_id): "airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_get_instance(self, mock_get_conn, mock_project_id): - self.hook.get_instance( + self.hook.get_instance( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, instance=TEST_INSTANCE_ID, retry=TEST_RETRY, @@ -172,7 +172,7 @@ def test_get_instance(self, mock_get_conn, mock_project_id): "airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_list_instances(self, mock_get_conn, mock_project_id): - self.hook.list_instances( + self.hook.list_instances( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, page_size=TEST_PAGE_SIZE, retry=TEST_RETRY, @@ -196,7 +196,7 @@ def test_list_instances(self, mock_get_conn, mock_project_id): "airflow.providers.google.cloud.hooks.cloud_memorystore.CloudMemorystoreHook.get_conn" ) def test_update_instance(self, mock_get_conn, mock_project_id): - self.hook.update_instance( + self.hook.update_instance( # pylint: disable=no-value-for-parameter update_mask=TEST_UPDATE_MASK, instance=Instance(name=TEST_NAME), location=TEST_LOCATION, @@ -450,7 +450,7 @@ def test_update_instance(self, mock_get_conn): ) def test_update_instance_without_project_id(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.update_instance( + self.hook.update_instance( # pylint: disable=no-value-for-parameter update_mask=TEST_UPDATE_MASK, instance=Instance(name=TEST_NAME), retry=TEST_RETRY, diff --git a/tests/providers/google/cloud/hooks/test_cloud_sql.py b/tests/providers/google/cloud/hooks/test_cloud_sql.py index 1c07752395f8f..e5472b5241742 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_sql.py +++ b/tests/providers/google/cloud/hooks/test_cloud_sql.py @@ -49,7 +49,7 @@ def test_instance_import_exception(self, mock_get_credentials): side_effect=HttpError(resp=httplib2.Response({'status': 400}), content=b'Error content') ) with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook.import_instance( + self.cloudsql_hook.import_instance( # pylint: disable=no-value-for-parameter instance='instance', body={}) err = cm.exception @@ -64,7 +64,7 @@ def test_instance_export_exception(self, mock_get_credentials): content=b'Error content') ) with self.assertRaises(HttpError) as cm: - self.cloudsql_hook.export_instance( + self.cloudsql_hook.export_instance( # pylint: disable=no-value-for-parameter instance='instance', body={}) err = cm.exception @@ -80,7 +80,7 @@ def test_instance_import(self, wait_for_operation_to_complete, get_conn, mock_ge execute_method = import_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.import_instance( + self.cloudsql_hook.import_instance( # pylint: disable=no-value-for-parameter instance='instance', body={}) @@ -99,7 +99,7 @@ def test_instance_export(self, wait_for_operation_to_complete, get_conn, mock_ge execute_method = export_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.export_instance( + self.cloudsql_hook.export_instance( # pylint: disable=no-value-for-parameter instance='instance', body={}) @@ -140,7 +140,7 @@ def test_get_instance(self, wait_for_operation_to_complete, get_conn, mock_get_c execute_method = get_method.return_value.execute execute_method.return_value = {"name": "instance"} wait_for_operation_to_complete.return_value = None - res = self.cloudsql_hook.get_instance( + res = self.cloudsql_hook.get_instance( # pylint: disable=no-value-for-parameter instance='instance') self.assertIsNotNone(res) self.assertEqual('instance', res['name']) @@ -158,7 +158,7 @@ def test_create_instance(self, wait_for_operation_to_complete, get_conn, mock_ge execute_method = insert_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.create_instance( + self.cloudsql_hook.create_instance( # pylint: disable=no-value-for-parameter body={}) insert_method.assert_called_once_with(body={}, project='example-project') @@ -181,7 +181,7 @@ def test_create_instance_with_in_progress_retry(self, wait_for_operation_to_comp {"name": "operation_id"} ] wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.create_instance( + self.cloudsql_hook.create_instance( # pylint: disable=no-value-for-parameter body={}) self.assertEqual(1, mock_get_credentials.call_count) @@ -204,7 +204,7 @@ def test_patch_instance_with_in_progress_retry(self, wait_for_operation_to_compl {"name": "operation_id"} ] wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.patch_instance( + self.cloudsql_hook.patch_instance( # pylint: disable=no-value-for-parameter instance='instance', body={}) @@ -224,7 +224,7 @@ def test_patch_instance(self, wait_for_operation_to_complete, get_conn, mock_get execute_method = patch_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.patch_instance( + self.cloudsql_hook.patch_instance( # pylint: disable=no-value-for-parameter instance='instance', body={}) @@ -244,7 +244,7 @@ def test_delete_instance(self, wait_for_operation_to_complete, get_conn, mock_ge execute_method = delete_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.delete_instance( + self.cloudsql_hook.delete_instance( # pylint: disable=no-value-for-parameter instance='instance') delete_method.assert_called_once_with(instance='instance', project='example-project') @@ -267,7 +267,7 @@ def test_delete_instance_with_in_progress_retry(self, wait_for_operation_to_comp {"name": "operation_id"} ] wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.delete_instance( + self.cloudsql_hook.delete_instance( # pylint: disable=no-value-for-parameter instance='instance') self.assertEqual(1, mock_get_credentials.call_count) @@ -286,7 +286,7 @@ def test_get_database(self, wait_for_operation_to_complete, get_conn, mock_get_c execute_method = get_method.return_value.execute execute_method.return_value = {"name": "database"} wait_for_operation_to_complete.return_value = None - res = self.cloudsql_hook.get_database( + res = self.cloudsql_hook.get_database( # pylint: disable=no-value-for-parameter database='database', instance='instance') self.assertIsNotNone(res) @@ -307,7 +307,7 @@ def test_create_database(self, wait_for_operation_to_complete, get_conn, mock_ge execute_method = insert_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.create_database( + self.cloudsql_hook.create_database( # pylint: disable=no-value-for-parameter instance='instance', body={}) @@ -331,7 +331,7 @@ def test_create_database_with_in_progress_retry(self, wait_for_operation_to_comp {"name": "operation_id"} ] wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.create_database( + self.cloudsql_hook.create_database( # pylint: disable=no-value-for-parameter instance='instance', body={}) @@ -351,7 +351,7 @@ def test_patch_database(self, wait_for_operation_to_complete, get_conn, mock_get execute_method = patch_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.patch_database( + self.cloudsql_hook.patch_database( # pylint: disable=no-value-for-parameter instance='instance', database='database', body={}) @@ -379,7 +379,7 @@ def test_patch_database_with_in_progress_retry(self, wait_for_operation_to_compl {"name": "operation_id"} ] wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.patch_database( + self.cloudsql_hook.patch_database( # pylint: disable=no-value-for-parameter instance='instance', database='database', body={}) @@ -400,7 +400,7 @@ def test_delete_database(self, wait_for_operation_to_complete, get_conn, mock_ge execute_method = delete_method.return_value.execute execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.delete_database( + self.cloudsql_hook.delete_database( # pylint: disable=no-value-for-parameter instance='instance', database='database') @@ -426,7 +426,7 @@ def test_delete_database_with_in_progress_retry(self, wait_for_operation_to_comp {"name": "operation_id"} ] wait_for_operation_to_complete.return_value = None - self.cloudsql_hook.delete_database( + self.cloudsql_hook.delete_database( # pylint: disable=no-value-for-parameter instance='instance', database='database') @@ -467,30 +467,6 @@ def test_instance_import_overridden_project_id( wait_for_operation_to_complete.assert_called_once_with(project_id='example-project', operation_name='operation_id') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_instance_import_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - import_method = get_conn.return_value.instances.return_value.import_ - execute_method = import_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.import_instance( - instance='instance', - body={}) - import_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -514,30 +490,6 @@ def test_instance_export_overridden_project_id( wait_for_operation_to_complete.assert_called_once_with(project_id='example-project', operation_name='operation_id') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_instance_export_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - export_method = get_conn.return_value.instances.return_value.export - execute_method = export_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.export_instance( - instance='instance', - body={}) - export_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -561,29 +513,6 @@ def test_get_instance_overridden_project_id( execute_method.assert_called_once_with(num_retries=5) wait_for_operation_to_complete.assert_not_called() - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_get_instance_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - get_method = get_conn.return_value.instances.return_value.get - execute_method = get_method.return_value.execute - execute_method.return_value = {"name": "instance"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.get_instance( - instance='instance') - get_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -607,29 +536,6 @@ def test_create_instance_overridden_project_id( operation_name='operation_id', project_id='example-project' ) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_create_instance_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - insert_method = get_conn.return_value.instances.return_value.insert - execute_method = insert_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.create_instance( - body={}) - insert_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -654,30 +560,6 @@ def test_patch_instance_overridden_project_id( operation_name='operation_id', project_id='example-project' ) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_patch_instance_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - patch_method = get_conn.return_value.instances.return_value.patch - execute_method = patch_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.patch_instance( - instance='instance', - body={}) - patch_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -701,29 +583,6 @@ def test_delete_instance_overridden_project_id( operation_name='operation_id', project_id='example-project' ) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_delete_instance_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - delete_method = get_conn.return_value.instances.return_value.delete - execute_method = delete_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.delete_instance( - instance='instance') - delete_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -750,30 +609,6 @@ def test_get_database_overridden_project_id( execute_method.assert_called_once_with(num_retries=5) wait_for_operation_to_complete.assert_not_called() - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_get_database_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - get_method = get_conn.return_value.databases.return_value.get - execute_method = get_method.return_value.execute - execute_method.return_value = {"name": "database"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.get_database( - database='database', - instance='instance') - get_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -798,30 +633,6 @@ def test_create_database_overridden_project_id( operation_name='operation_id', project_id='example-project' ) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_create_database_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - insert_method = get_conn.return_value.databases.return_value.insert - execute_method = insert_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.create_database( - instance='instance', - body={}) - insert_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -850,31 +661,6 @@ def test_patch_database_overridden_project_id( operation_name='operation_id', project_id='example-project' ) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_patch_database_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - patch_method = get_conn.return_value.databases.return_value.patch - execute_method = patch_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.patch_database( - instance='instance', - database='database', - body={}) - patch_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, @@ -901,30 +687,6 @@ def test_delete_database_overridden_project_id( operation_name='operation_id', project_id='example-project' ) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook.get_conn') - @mock.patch('airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLHook._wait_for_operation_to_complete') - def test_delete_database_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - delete_method = get_conn.return_value.databases.return_value.delete - execute_method = delete_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.cloudsql_hook_no_default_project_id.delete_database( - instance='instance', - database='database') - delete_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - class TestCloudSqlDatabaseHook(unittest.TestCase): diff --git a/tests/providers/google/cloud/hooks/test_cloud_storage_transfer_service.py b/tests/providers/google/cloud/hooks/test_cloud_storage_transfer_service.py index b143148235ed4..dce2cb46fe228 100644 --- a/tests/providers/google/cloud/hooks/test_cloud_storage_transfer_service.py +++ b/tests/providers/google/cloud/hooks/test_cloud_storage_transfer_service.py @@ -464,7 +464,8 @@ def test_get_transfer_job(self, get_conn, mock_project_id): get_method = get_conn.return_value.transferJobs.return_value.get execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB - res = self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) + res = self.gct_hook.get_transfer_job( # pylint: disable=no-value-for-parameter + job_name=TEST_TRANSFER_JOB_NAME) self.assertIsNotNone(res) self.assertEqual(TEST_TRANSFER_JOB_NAME, res[NAME]) get_method.assert_called_once_with(jobName=TEST_TRANSFER_JOB_NAME, projectId='example-project') @@ -670,7 +671,8 @@ def test_get_transfer_job(self, get_conn, mock_project_id): execute_method = get_method.return_value.execute execute_method.return_value = TEST_TRANSFER_JOB with self.assertRaises(AirflowException) as e: - self.gct_hook.get_transfer_job(job_name=TEST_TRANSFER_JOB_NAME) + self.gct_hook.get_transfer_job( # pylint: disable=no-value-for-parameter + job_name=TEST_TRANSFER_JOB_NAME) self.assertEqual( 'The project id must be passed either as keyword project_id ' 'parameter or as project_id extra in GCP connection definition. ' diff --git a/tests/providers/google/cloud/hooks/test_compute.py b/tests/providers/google/cloud/hooks/test_compute.py index 61874a07263e2..7d6ae29dcb4a5 100644 --- a/tests/providers/google/cloud/hooks/test_compute.py +++ b/tests/providers/google/cloud/hooks/test_compute.py @@ -73,30 +73,6 @@ def test_start_instance_overridden_project_id(self, wait_for_operation_to_comple operation_name='operation_id', zone='zone') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_start_instance_no_project_id(self, wait_for_operation_to_complete, get_conn, mock_project_id): - start_method = get_conn.return_value.instances.return_value.start - execute_method = start_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.start_instance( - zone=GCE_ZONE, - resource_id=GCE_INSTANCE) - start_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' @@ -117,30 +93,6 @@ def test_stop_instance_overridden_project_id(self, wait_for_operation_to_complet operation_name='operation_id', zone='zone') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_stop_instance_no_project_id(self, wait_for_operation_to_complete, get_conn, mock_project_id): - stop_method = get_conn.return_value.instances.return_value.stop - execute_method = stop_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.stop_instance( - zone=GCE_ZONE, - resource_id=GCE_INSTANCE) - stop_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' @@ -163,31 +115,6 @@ def test_set_machine_type_overridden_project_id(self, wait_for_operation_to_comp operation_name='operation_id', zone='zone') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_set_machine_type_no_project_id(self, wait_for_operation_to_complete, get_conn, mock_project_id): - set_machine_type_method = get_conn.return_value.instances.return_value.setMachineType - execute_method = set_machine_type_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.set_machine_type( - body={}, - zone=GCE_ZONE, - resource_id=GCE_INSTANCE) - set_machine_type_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' @@ -206,32 +133,6 @@ def test_get_instance_template_overridden_project_id(self, wait_for_operation_to execute_method.assert_called_once_with(num_retries=5) wait_for_operation_to_complete.assert_not_called() - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_get_instance_template_no_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - get_method = get_conn.return_value.instanceTemplates.return_value.get - execute_method = get_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.get_instance_template( - resource_id=GCE_INSTANCE_TEMPLATE - ) - get_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' @@ -252,33 +153,6 @@ def test_insert_instance_template_overridden_project_id(self, wait_for_operation wait_for_operation_to_complete.assert_called_once_with(project_id='example-project', operation_name='operation_id') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_insert_instance_template_no_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - insert_method = get_conn.return_value.instanceTemplates.return_value.insert - execute_method = insert_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.insert_instance_template( - body={}, - request_id=GCE_REQUEST_ID - ) - insert_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' @@ -300,33 +174,6 @@ def test_get_instance_group_manager_overridden_project_id(self, wait_for_operati execute_method.assert_called_once_with(num_retries=5) wait_for_operation_to_complete.assert_not_called() - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_get_instance_group_manager_no_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - get_method = get_conn.return_value.instanceGroupManagers.return_value.get - execute_method = get_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.get_instance_group_manager( - zone=GCE_ZONE, - resource_id=GCE_INSTANCE_GROUP_MANAGER - ) - get_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' @@ -357,35 +204,6 @@ def test_patch_instance_group_manager_overridden_project_id(self, project_id='example-project', zone='zone') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.compute.ComputeEngineHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.compute.ComputeEngineHook._wait_for_operation_to_complete' - ) - def test_patch_instance_group_manager_no_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - patch_method = get_conn.return_value.instanceGroupManagers.return_value.patch - execute_method = patch_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gce_hook_no_project_id.patch_instance_group_manager( - zone=GCE_ZONE, - resource_id=GCE_INSTANCE_GROUP_MANAGER, - body={}, - request_id=GCE_REQUEST_ID - ) - patch_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - class TestGcpComputeHookDefaultProjectId(unittest.TestCase): def setUp(self): @@ -409,7 +227,9 @@ def test_start_instance(self, wait_for_operation_to_complete, get_conn, mock_pro wait_for_operation_to_complete.return_value = None res = self.gce_hook.start_instance( zone=GCE_ZONE, - resource_id=GCE_INSTANCE) + resource_id=GCE_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) self.assertIsNone(res) start_method.assert_called_once_with(instance='instance', project='example-project', zone='zone') execute_method.assert_called_once_with(num_retries=5) @@ -453,7 +273,9 @@ def test_stop_instance(self, wait_for_operation_to_complete, get_conn, mock_proj wait_for_operation_to_complete.return_value = None res = self.gce_hook.stop_instance( zone=GCE_ZONE, - resource_id=GCE_INSTANCE) + resource_id=GCE_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) self.assertIsNone(res) stop_method.assert_called_once_with(instance='instance', project='example-project', zone='zone') execute_method.assert_called_once_with(num_retries=5) @@ -497,7 +319,9 @@ def test_set_machine_type_instance(self, wait_for_operation_to_complete, get_con res = self.gce_hook.set_machine_type( body={}, zone=GCE_ZONE, - resource_id=GCE_INSTANCE) + resource_id=GCE_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) self.assertIsNone(res) execute_method.assert_called_once_with(num_retries=5) wait_for_operation_to_complete.assert_called_once_with(project_id='example-project', @@ -538,7 +362,8 @@ def test_get_instance_template(self, wait_for_operation_to_complete, get_conn, m execute_method.return_value = {"name": "operation_id"} wait_for_operation_to_complete.return_value = None res = self.gce_hook.get_instance_template( - resource_id=GCE_INSTANCE_TEMPLATE) + resource_id=GCE_INSTANCE_TEMPLATE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertIsNotNone(res) get_method.assert_called_once_with(instanceTemplate='instance-template', project='example-project') execute_method.assert_called_once_with(num_retries=5) @@ -577,7 +402,8 @@ def test_insert_instance_template(self, wait_for_operation_to_complete, get_conn wait_for_operation_to_complete.return_value = None res = self.gce_hook.insert_instance_template( body={}, - request_id=GCE_REQUEST_ID + request_id=GCE_REQUEST_ID, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) self.assertIsNone(res) insert_method.assert_called_once_with(body={}, project='example-project', requestId='request_id') @@ -621,7 +447,8 @@ def test_get_instance_group_manager(self, wait_for_operation_to_complete, get_co wait_for_operation_to_complete.return_value = None res = self.gce_hook.get_instance_group_manager( zone=GCE_ZONE, - resource_id=GCE_INSTANCE_GROUP_MANAGER + resource_id=GCE_INSTANCE_GROUP_MANAGER, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) self.assertIsNotNone(res) get_method.assert_called_once_with(instanceGroupManager='instance_group_manager', @@ -669,7 +496,8 @@ def test_patch_instance_group_manager(self, wait_for_operation_to_complete, get_ zone=GCE_ZONE, resource_id=GCE_INSTANCE_GROUP_MANAGER, body={}, - request_id=GCE_REQUEST_ID + request_id=GCE_REQUEST_ID, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) self.assertIsNone(res) patch_method.assert_called_once_with( diff --git a/tests/providers/google/cloud/hooks/test_datacatalog.py b/tests/providers/google/cloud/hooks/test_datacatalog.py index e4d66a2775bef..735f56af8caaa 100644 --- a/tests/providers/google/cloud/hooks/test_datacatalog.py +++ b/tests/providers/google/cloud/hooks/test_datacatalog.py @@ -178,7 +178,7 @@ def setUp(self,) -> None: "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_create_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.create_entry( + self.hook.create_entry( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry_id=TEST_ENTRY_ID, @@ -204,7 +204,7 @@ def test_create_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> Non "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_create_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.create_entry_group( + self.hook.create_entry_group( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group_id=TEST_ENTRY_GROUP_ID, entry_group=TEST_ENTRY_GROUP, @@ -229,7 +229,7 @@ def test_create_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_create_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.create_tag( + self.hook.create_tag( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -255,7 +255,7 @@ def test_create_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_create_tag_protobuff(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.create_tag( + self.hook.create_tag( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -281,7 +281,7 @@ def test_create_tag_protobuff(self, mock_get_conn, mock_get_creds_and_project_id "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_create_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.create_tag_template( + self.hook.create_tag_template( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template_id=TEST_TAG_TEMPLATE_ID, tag_template=TEST_TAG_TEMPLATE, @@ -306,7 +306,7 @@ def test_create_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_create_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.create_tag_template_field( + self.hook.create_tag_template_field( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, tag_template_field_id=TEST_TAG_TEMPLATE_FIELD_ID, @@ -332,7 +332,7 @@ def test_create_tag_template_field(self, mock_get_conn, mock_get_creds_and_proje "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_delete_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.delete_entry( + self.hook.delete_entry( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -355,7 +355,7 @@ def test_delete_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> Non "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_delete_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.delete_entry_group( + self.hook.delete_entry_group( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, retry=TEST_RETRY, @@ -377,7 +377,7 @@ def test_delete_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_delete_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.delete_tag( + self.hook.delete_tag( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -401,7 +401,7 @@ def test_delete_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_delete_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.delete_tag_template( + self.hook.delete_tag_template( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, force=TEST_FORCE, @@ -425,7 +425,7 @@ def test_delete_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_delete_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.delete_tag_template_field( + self.hook.delete_tag_template_field( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, field=TEST_TAG_TEMPLATE_FIELD_ID, @@ -450,7 +450,7 @@ def test_delete_tag_template_field(self, mock_get_conn, mock_get_creds_and_proje "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_get_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.get_entry( + self.hook.get_entry( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -473,7 +473,7 @@ def test_get_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_get_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.get_entry_group( + self.hook.get_entry_group( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, read_mask=TEST_READ_MASK, @@ -497,7 +497,7 @@ def test_get_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_get_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.get_tag_template( + self.hook.get_tag_template( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, retry=TEST_RETRY, @@ -519,7 +519,7 @@ def test_get_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_list_tags(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.list_tags( + self.hook.list_tags( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -548,7 +548,7 @@ def test_get_tag_for_template_name(self, mock_get_conn, mock_get_creds_and_proje tag_2 = mock.MagicMock(template=TEST_TAG_TEMPLATE_PATH.format(TEST_PROJECT_ID_1)) mock_get_conn.return_value.list_tags.return_value = [tag_1, tag_2] - result = self.hook.get_tag_for_template_name( + result = self.hook.get_tag_for_template_name( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -574,7 +574,7 @@ def test_get_tag_for_template_name(self, mock_get_conn, mock_get_creds_and_proje "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_rename_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.rename_tag_template_field( + self.hook.rename_tag_template_field( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, field=TEST_TAG_TEMPLATE_FIELD_ID, @@ -599,7 +599,7 @@ def test_rename_tag_template_field(self, mock_get_conn, mock_get_creds_and_proje "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_update_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.update_entry( + self.hook.update_entry( # pylint: disable=no-value-for-parameter entry=TEST_ENTRY, update_mask=TEST_UPDATE_MASK, location=TEST_LOCATION, @@ -625,7 +625,7 @@ def test_update_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> Non "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_update_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.update_tag( + self.hook.update_tag( # pylint: disable=no-value-for-parameter tag=deepcopy(TEST_TAG), update_mask=TEST_UPDATE_MASK, location=TEST_LOCATION, @@ -652,7 +652,7 @@ def test_update_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_update_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.update_tag_template( + self.hook.update_tag_template( # pylint: disable=no-value-for-parameter tag_template=TEST_TAG_TEMPLATE, update_mask=TEST_UPDATE_MASK, location=TEST_LOCATION, @@ -677,7 +677,7 @@ def test_update_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) "airflow.providers.google.cloud.hooks.datacatalog.CloudDataCatalogHook.get_conn" ) def test_update_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: - self.hook.update_tag_template_field( + self.hook.update_tag_template_field( # pylint: disable=no-value-for-parameter tag_template_field=TEST_TAG_TEMPLATE_FIELD, update_mask=TEST_UPDATE_MASK, tag_template=TEST_TAG_TEMPLATE_ID, @@ -1276,7 +1276,7 @@ def setUp(self,) -> None: ) def test_create_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.create_entry( + self.hook.create_entry( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry_id=TEST_ENTRY_ID, @@ -1295,7 +1295,7 @@ def test_create_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> Non ) def test_create_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.create_entry_group( + self.hook.create_entry_group( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group_id=TEST_ENTRY_GROUP_ID, entry_group=TEST_ENTRY_GROUP, @@ -1314,7 +1314,7 @@ def test_create_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) def test_create_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.create_tag( + self.hook.create_tag( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1335,7 +1335,7 @@ def test_create_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: def test_create_tag_protobuff(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.create_tag( + self.hook.create_tag( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1356,7 +1356,7 @@ def test_create_tag_protobuff(self, mock_get_conn, mock_get_creds_and_project_id def test_create_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.create_tag_template( + self.hook.create_tag_template( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template_id=TEST_TAG_TEMPLATE_ID, tag_template=TEST_TAG_TEMPLATE, @@ -1375,7 +1375,7 @@ def test_create_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) def test_create_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.create_tag_template_field( + self.hook.create_tag_template_field( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, tag_template_field_id=TEST_TAG_TEMPLATE_FIELD_ID, @@ -1395,7 +1395,7 @@ def test_create_tag_template_field(self, mock_get_conn, mock_get_creds_and_proje def test_delete_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.delete_entry( + self.hook.delete_entry( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1413,7 +1413,7 @@ def test_delete_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> Non ) def test_delete_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.delete_entry_group( + self.hook.delete_entry_group( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, retry=TEST_RETRY, @@ -1430,7 +1430,7 @@ def test_delete_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) ) def test_delete_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.delete_tag( + self.hook.delete_tag( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1449,7 +1449,7 @@ def test_delete_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: ) def test_delete_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.delete_tag_template( + self.hook.delete_tag_template( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, force=TEST_FORCE, @@ -1467,7 +1467,7 @@ def test_delete_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) ) def test_delete_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.delete_tag_template_field( + self.hook.delete_tag_template_field( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, field=TEST_TAG_TEMPLATE_FIELD_ID, @@ -1486,7 +1486,7 @@ def test_delete_tag_template_field(self, mock_get_conn, mock_get_creds_and_proje ) def test_get_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.get_entry( + self.hook.get_entry( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1504,7 +1504,7 @@ def test_get_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: ) def test_get_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.get_entry_group( + self.hook.get_entry_group( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, read_mask=TEST_READ_MASK, @@ -1522,7 +1522,7 @@ def test_get_entry_group(self, mock_get_conn, mock_get_creds_and_project_id) -> ) def test_get_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.get_tag_template( + self.hook.get_tag_template( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, retry=TEST_RETRY, @@ -1539,7 +1539,7 @@ def test_get_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> ) def test_list_tags(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.list_tags( + self.hook.list_tags( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1562,7 +1562,7 @@ def test_get_tag_for_template_name(self, mock_get_conn, mock_get_creds_and_proje mock_get_conn.return_value.list_tags.return_value = [tag_1, tag_2] with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.get_tag_for_template_name( + self.hook.get_tag_for_template_name( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, entry_group=TEST_ENTRY_GROUP_ID, entry=TEST_ENTRY_ID, @@ -1581,7 +1581,7 @@ def test_get_tag_for_template_name(self, mock_get_conn, mock_get_creds_and_proje ) def test_rename_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.rename_tag_template_field( + self.hook.rename_tag_template_field( # pylint: disable=no-value-for-parameter location=TEST_LOCATION, tag_template=TEST_TAG_TEMPLATE_ID, field=TEST_TAG_TEMPLATE_FIELD_ID, @@ -1600,7 +1600,7 @@ def test_rename_tag_template_field(self, mock_get_conn, mock_get_creds_and_proje ) def test_update_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.update_entry( + self.hook.update_entry( # pylint: disable=no-value-for-parameter entry=TEST_ENTRY, update_mask=TEST_UPDATE_MASK, location=TEST_LOCATION, @@ -1620,7 +1620,7 @@ def test_update_entry(self, mock_get_conn, mock_get_creds_and_project_id) -> Non ) def test_update_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.update_tag( + self.hook.update_tag( # pylint: disable=no-value-for-parameter tag=deepcopy(TEST_TAG), update_mask=TEST_UPDATE_MASK, location=TEST_LOCATION, @@ -1641,7 +1641,7 @@ def test_update_tag(self, mock_get_conn, mock_get_creds_and_project_id) -> None: ) def test_update_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.update_tag_template( + self.hook.update_tag_template( # pylint: disable=no-value-for-parameter tag_template=TEST_TAG_TEMPLATE, update_mask=TEST_UPDATE_MASK, location=TEST_LOCATION, @@ -1660,7 +1660,7 @@ def test_update_tag_template(self, mock_get_conn, mock_get_creds_and_project_id) ) def test_update_tag_template_field(self, mock_get_conn, mock_get_creds_and_project_id) -> None: with self.assertRaisesRegex(AirflowException, TEST_MESSAGE): - self.hook.update_tag_template_field( + self.hook.update_tag_template_field( # pylint: disable=no-value-for-parameter tag_template_field=TEST_TAG_TEMPLATE_FIELD, update_mask=TEST_UPDATE_MASK, tag_template=TEST_TAG_TEMPLATE_ID, diff --git a/tests/providers/google/cloud/hooks/test_dataflow.py b/tests/providers/google/cloud/hooks/test_dataflow.py index 18225b2271197..451cc22c5ad8b 100644 --- a/tests/providers/google/cloud/hooks/test_dataflow.py +++ b/tests/providers/google/cloud/hooks/test_dataflow.py @@ -171,9 +171,10 @@ def test_start_python_dataflow( dataflow_instance.wait_for_done.return_value = None dataflowjob_instance = mock_dataflowjob.return_value dataflowjob_instance.wait_for_done.return_value = None - self.dataflow_hook.start_python_dataflow( + self.dataflow_hook.start_python_dataflow( # pylint: disable=no-value-for-parameter job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_PY, - dataflow=PY_FILE, py_options=PY_OPTIONS) + dataflow=PY_FILE, py_options=PY_OPTIONS, + ) expected_cmd = ["python3", '-m', PY_FILE, '--region=us-central1', '--runner=DataflowRunner', '--project=test', @@ -194,7 +195,7 @@ def test_start_python_dataflow( @mock.patch(DATAFLOW_STRING.format('_DataflowRunner')) @mock.patch(DATAFLOW_STRING.format('DataflowHook.get_conn')) def test_start_python_dataflow_with_custom_interpreter( - self, name, py_interpreter, mock_conn, mock_dataflow, mock_dataflowjob, mock_uuid + self, name, py_interpreter, mock_conn, mock_dataflow, mock_dataflowjob, mock_uuid, ): del name # unused variable mock_uuid.return_value = MOCK_UUID @@ -203,10 +204,11 @@ def test_start_python_dataflow_with_custom_interpreter( dataflow_instance.wait_for_done.return_value = None dataflowjob_instance = mock_dataflowjob.return_value dataflowjob_instance.wait_for_done.return_value = None - self.dataflow_hook.start_python_dataflow( + self.dataflow_hook.start_python_dataflow( # pylint: disable=no-value-for-parameter job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_PY, dataflow=PY_FILE, py_options=PY_OPTIONS, - py_interpreter=py_interpreter) + py_interpreter=py_interpreter, + ) expected_cmd = [py_interpreter, '-m', PY_FILE, '--region=us-central1', '--runner=DataflowRunner', '--project=test', @@ -228,7 +230,7 @@ def test_start_java_dataflow(self, mock_conn, dataflow_instance.wait_for_done.return_value = None dataflowjob_instance = mock_dataflowjob.return_value dataflowjob_instance.wait_for_done.return_value = None - self.dataflow_hook.start_java_dataflow( + self.dataflow_hook.start_java_dataflow( # pylint: disable=no-value-for-parameter job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_JAVA, jar=JAR_FILE) expected_cmd = ['java', '-jar', JAR_FILE, @@ -244,14 +246,15 @@ def test_start_java_dataflow(self, mock_conn, @mock.patch(DATAFLOW_STRING.format('_DataflowJobsController')) @mock.patch(DATAFLOW_STRING.format('_DataflowRunner')) @mock.patch(DATAFLOW_STRING.format('DataflowHook.get_conn')) - def test_start_java_dataflow_with_job_class(self, mock_conn, mock_dataflow, mock_dataflowjob, mock_uuid): + def test_start_java_dataflow_with_job_class( + self, mock_conn, mock_dataflow, mock_dataflowjob, mock_uuid): mock_uuid.return_value = MOCK_UUID mock_conn.return_value = None dataflow_instance = mock_dataflow.return_value dataflow_instance.wait_for_done.return_value = None dataflowjob_instance = mock_dataflowjob.return_value dataflowjob_instance.wait_for_done.return_value = None - self.dataflow_hook.start_java_dataflow( + self.dataflow_hook.start_java_dataflow( # pylint: disable=no-value-for-parameter job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_JAVA, jar=JAR_FILE, job_class=JOB_CLASS) expected_cmd = ['java', '-cp', JAR_FILE, JOB_CLASS, @@ -314,9 +317,10 @@ def test_start_template_dataflow(self, mock_conn, mock_controller, mock_uuid): launch ) launch_method.return_value.execute.return_value = {"job": {"id": TEST_JOB_ID}} - self.dataflow_hook.start_template_dataflow( + self.dataflow_hook.start_template_dataflow( # pylint: disable=no-value-for-parameter job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_TEMPLATE, parameters=PARAMETERS, - dataflow_template=TEMPLATE) + dataflow_template=TEMPLATE, + ) options_with_region = {'region': 'us-central1'} options_with_region.update(DATAFLOW_OPTIONS_TEMPLATE) options_with_region_without_project = copy.deepcopy(options_with_region) @@ -364,11 +368,11 @@ def test_start_template_dataflow_with_runtime_env(self, mock_conn, mock_dataflow .launch) method.return_value.execute.return_value = {'job': {'id': TEST_JOB_ID}} - self.dataflow_hook.start_template_dataflow( + self.dataflow_hook.start_template_dataflow( # pylint: disable=no-value-for-parameter job_name=JOB_NAME, variables=options_with_runtime_env, parameters=PARAMETERS, - dataflow_template=TEMPLATE + dataflow_template=TEMPLATE, ) body = {"jobName": mock.ANY, "parameters": PARAMETERS, diff --git a/tests/providers/google/cloud/hooks/test_datastore.py b/tests/providers/google/cloud/hooks/test_datastore.py index 073332c28a398..b1a0cdeb89662 100644 --- a/tests/providers/google/cloud/hooks/test_datastore.py +++ b/tests/providers/google/cloud/hooks/test_datastore.py @@ -73,7 +73,8 @@ def test_allocate_ids_no_project_id(self, mock_get_conn, mock_project_id): partial_keys = [] with self.assertRaises(AirflowException) as err: - self.datastore_hook.allocate_ids(partial_keys=partial_keys) + self.datastore_hook.allocate_ids( # pylint: disable=no-value-for-parameter + partial_keys=partial_keys) self.assertIn("project_id", str(err.exception)) @patch('airflow.providers.google.cloud.hooks.datastore.DatastoreHook.get_conn') @@ -96,7 +97,7 @@ def test_begin_transaction(self, mock_get_conn): def test_begin_transaction_no_project_id(self, mock_get_conn, mock_project_id): self.datastore_hook.connection = mock_get_conn.return_value with self.assertRaises(AirflowException) as err: - self.datastore_hook.begin_transaction() + self.datastore_hook.begin_transaction() # pylint: disable=no-value-for-parameter self.assertIn("project_id", str(err.exception)) @patch('airflow.providers.google.cloud.hooks.datastore.DatastoreHook.get_conn') @@ -122,7 +123,7 @@ def test_commit_no_project_id(self, mock_get_conn, mock_project_id): body = {'item': 'a'} with self.assertRaises(AirflowException) as err: - self.datastore_hook.commit(body=body) + self.datastore_hook.commit(body=body) # pylint: disable=no-value-for-parameter self.assertIn("project_id", str(err.exception)) @patch('airflow.providers.google.cloud.hooks.datastore.DatastoreHook.get_conn') @@ -161,7 +162,7 @@ def test_lookup_no_project_id(self, mock_get_conn, mock_project_id): transaction = 'transaction' with self.assertRaises(AirflowException) as err: - self.datastore_hook.lookup(keys=keys, + self.datastore_hook.lookup(keys=keys, # pylint: disable=no-value-for-parameter read_consistency=read_consistency, transaction=transaction, ) @@ -190,7 +191,7 @@ def test_rollback_no_project_id(self, mock_get_conn, mock_project_id): transaction = 'transaction' with self.assertRaises(AirflowException) as err: - self.datastore_hook.rollback(transaction=transaction) + self.datastore_hook.rollback(transaction=transaction) # pylint: disable=no-value-for-parameter self.assertIn("project_id", str(err.exception)) @patch('airflow.providers.google.cloud.hooks.datastore.DatastoreHook.get_conn') @@ -216,7 +217,7 @@ def test_run_query_no_project_id(self, mock_get_conn, mock_project_id): body = {'item': 'a'} with self.assertRaises(AirflowException) as err: - self.datastore_hook.run_query(body=body) + self.datastore_hook.run_query(body=body) # pylint: disable=no-value-for-parameter self.assertIn("project_id", str(err.exception)) @patch('airflow.providers.google.cloud.hooks.datastore.DatastoreHook.get_conn') @@ -310,11 +311,12 @@ def test_export_to_storage_bucket_no_project_id(self, mock_get_conn, mock_projec labels = {} with self.assertRaises(AirflowException) as err: - self.datastore_hook.export_to_storage_bucket(bucket=bucket, - namespace=namespace, - entity_filter=entity_filter, - labels=labels, - ) + self.datastore_hook.export_to_storage_bucket( # pylint: disable=no-value-for-parameter + bucket=bucket, + namespace=namespace, + entity_filter=entity_filter, + labels=labels, + ) self.assertIn("project_id", str(err.exception)) @patch('airflow.providers.google.cloud.hooks.datastore.DatastoreHook.get_conn') @@ -361,10 +363,11 @@ def test_import_from_storage_bucket_no_project_id(self, mock_get_conn, mock_proj labels = {} with self.assertRaises(AirflowException) as err: - self.datastore_hook.import_from_storage_bucket(bucket=bucket, - file=file, - namespace=namespace, - entity_filter=entity_filter, - labels=labels, - ) + self.datastore_hook.import_from_storage_bucket( # pylint: disable=no-value-for-parameter + bucket=bucket, + file=file, + namespace=namespace, + entity_filter=entity_filter, + labels=labels, + ) self.assertIn("project_id", str(err.exception)) diff --git a/tests/providers/google/cloud/hooks/test_dlp.py b/tests/providers/google/cloud/hooks/test_dlp.py index b8e6fa5c265f3..b58a8e03e11f5 100644 --- a/tests/providers/google/cloud/hooks/test_dlp.py +++ b/tests/providers/google/cloud/hooks/test_dlp.py @@ -108,7 +108,7 @@ def test_cancel_dlp_job_without_dlp_job_id(self, _): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_cancel_dlp_job_without_parent(self, _, mock_project_id): with self.assertRaises(AirflowException): - self.hook.cancel_dlp_job(dlp_job_id=DLP_JOB_ID) + self.hook.cancel_dlp_job(dlp_job_id=DLP_JOB_ID) # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -181,7 +181,7 @@ def test_create_dlp_job(self, get_conn): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_create_dlp_job_without_project_id(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.create_dlp_job() + self.hook.create_dlp_job() # pylint: disable=no-value-for-parameter @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_create_dlp_job_with_wait_until_finished(self, get_conn): @@ -266,7 +266,7 @@ def test_create_job_trigger(self, get_conn): ) def test_create_job_trigger_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.create_job_trigger() + self.hook.create_job_trigger() # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -343,7 +343,7 @@ def test_deidentify_content(self, get_conn): ) def test_deidentify_content_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.deidentify_content() + self.hook.deidentify_content() # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -414,7 +414,7 @@ def test_delete_dlp_job_without_dlp_job_id(self, _): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_delete_dlp_job_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.delete_dlp_job(dlp_job_id=DLP_JOB_ID) + self.hook.delete_dlp_job(dlp_job_id=DLP_JOB_ID) # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -483,7 +483,7 @@ def test_delete_job_trigger_without_trigger_id(self, _): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_delete_job_trigger_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.delete_job_trigger(job_trigger_id=TRIGGER_ID) + self.hook.delete_job_trigger(job_trigger_id=TRIGGER_ID) # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -604,7 +604,7 @@ def test_get_dlp_job_without_dlp_job_id(self, _): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_get_dlp_job_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.get_dlp_job(dlp_job_id=DLP_JOB_ID) + self.hook.get_dlp_job(dlp_job_id=DLP_JOB_ID) # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -681,7 +681,7 @@ def test_get_job_trigger_without_trigger_id(self, _): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_get_job_trigger_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.get_job_trigger(job_trigger_id=TRIGGER_ID) + self.hook.get_job_trigger(job_trigger_id=TRIGGER_ID) # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -757,7 +757,7 @@ def test_inspect_content(self, get_conn): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_inspect_content_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.inspect_content() + self.hook.inspect_content() # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -826,7 +826,7 @@ def test_list_dlp_jobs(self, get_conn): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_list_dlp_jobs_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.list_dlp_jobs() + self.hook.list_dlp_jobs() # pylint: disable=no-value-for-parameter @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_list_info_types(self, get_conn): @@ -904,7 +904,7 @@ def test_list_job_triggers(self, get_conn): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_list_job_triggers_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.list_job_triggers() + self.hook.list_job_triggers() # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -974,7 +974,7 @@ def test_redact_image(self, get_conn): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_redact_image_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.redact_image() + self.hook.redact_image() # pylint: disable=no-value-for-parameter @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_reidentify_content(self, get_conn): @@ -1004,7 +1004,7 @@ def test_reidentify_content(self, get_conn): ) def test_reidentify_content_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.reidentify_content() + self.hook.reidentify_content() # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', @@ -1148,7 +1148,7 @@ def test_update_job_trigger_without_job_trigger_id(self, _): @mock.patch("airflow.providers.google.cloud.hooks.dlp.CloudDLPHook.get_conn") def test_update_job_trigger_without_parent(self, mock_get_conn, mock_project_id): with self.assertRaises(AirflowException): - self.hook.update_job_trigger(job_trigger_id=TRIGGER_ID) + self.hook.update_job_trigger(job_trigger_id=TRIGGER_ID) # pylint: disable=no-value-for-parameter @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', diff --git a/tests/providers/google/cloud/hooks/test_functions.py b/tests/providers/google/cloud/hooks/test_functions.py index bdcf52fcd4749..e656aaeee9e0c 100644 --- a/tests/providers/google/cloud/hooks/test_functions.py +++ b/tests/providers/google/cloud/hooks/test_functions.py @@ -49,34 +49,6 @@ def test_gcf_client_creation(self, mock_build, mock_authorize): self.assertEqual(mock_build.return_value, result) self.assertEqual(self.gcf_function_hook_no_project_id._conn, result) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn') - @mock.patch( - 'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete' - ) - def test_create_new_function_missing_project_id( - self, wait_for_operation_to_complete, get_conn, mock_project_id - ): - create_method = get_conn.return_value.projects.return_value.locations. \ - return_value.functions.return_value.create - execute_method = create_method.return_value.execute - execute_method.return_value = {"name": "operation_id"} - wait_for_operation_to_complete.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gcf_function_hook_no_project_id.create_new_function( - location=GCF_LOCATION, - body={} - ) - create_method.assert_not_called() - execute_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - wait_for_operation_to_complete.assert_not_called() - @mock.patch('airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn') @mock.patch( 'airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook._wait_for_operation_to_complete' @@ -98,34 +70,6 @@ def test_create_new_function_overridden_project_id(self, wait_for_operation_to_c execute_method.assert_called_once_with(num_retries=5) wait_for_operation_to_complete.assert_called_once_with(operation_name='operation_id') - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('requests.put') - @mock.patch('airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn') - def test_upload_function_zip_missing_project_id( - self, get_conn, requests_put, mock_project_id - ): - mck = mock.mock_open() - with mock.patch('builtins.open', mck): - generate_upload_url_method = get_conn.return_value.projects.return_value.locations. \ - return_value.functions.return_value.generateUploadUrl - execute_method = generate_upload_url_method.return_value.execute - execute_method.return_value = {"uploadUrl": "http://uploadHere"} - requests_put.return_value = None - with self.assertRaises(AirflowException) as cm: - self.gcf_function_hook_no_project_id.upload_function_zip( - location=GCF_LOCATION, - zip_path="/tmp/path.zip" - ) - generate_upload_url_method.assert_not_called() - execute_method.assert_not_called() - mck.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('requests.put') @mock.patch('airflow.providers.google.cloud.hooks.functions.CloudFunctionsHook.get_conn') def test_upload_function_zip_overridden_project_id(self, get_conn, requests_put): @@ -186,7 +130,8 @@ def test_create_new_function(self, wait_for_operation_to_complete, get_conn, moc wait_for_operation_to_complete.return_value = None res = self.gcf_function_hook.create_new_function( location=GCF_LOCATION, - body={} + body={}, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) self.assertIsNone(res) create_method.assert_called_once_with(body={}, @@ -287,7 +232,8 @@ def test_upload_function_zip(self, get_conn, requests_put, mock_project_id): requests_put.return_value = None res = self.gcf_function_hook.upload_function_zip( location=GCF_LOCATION, - zip_path="/tmp/path.zip" + zip_path="/tmp/path.zip", + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) self.assertEqual("http://uploadHere", res) generate_upload_url_method.assert_called_once_with( diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py index 59d52f0ff4d3c..c0fa397dd583d 100644 --- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py @@ -116,7 +116,7 @@ def test_delete_cluster_error(self, wait_mock, convert_mock, mock_project_id): self.gke_hook._client.delete_cluster.side_effect = AirflowException('400') with self.assertRaises(AirflowException): - self.gke_hook.delete_cluster(name='a-cluster') + self.gke_hook.delete_cluster(name='a-cluster') # pylint: disable=no-value-for-parameter wait_mock.assert_not_called() convert_mock.assert_not_called() @@ -192,7 +192,7 @@ def test_create_cluster_error(self, wait_mock, convert_mock): mock_cluster_proto = None with self.assertRaises(AirflowException): - self.gke_hook.create_cluster(mock_cluster_proto) + self.gke_hook.create_cluster(mock_cluster_proto) # pylint: disable=no-value-for-parameter wait_mock.assert_not_called() convert_mock.assert_not_called() diff --git a/tests/providers/google/cloud/hooks/test_mlengine.py b/tests/providers/google/cloud/hooks/test_mlengine.py index 27323baeaa995..e32e5f636c173 100644 --- a/tests/providers/google/cloud/hooks/test_mlengine.py +++ b/tests/providers/google/cloud/hooks/test_mlengine.py @@ -18,16 +18,14 @@ import unittest from copy import deepcopy from unittest import mock +from unittest.mock import PropertyMock import httplib2 from googleapiclient.errors import HttpError -from mock import PropertyMock -from airflow.exceptions import AirflowException from airflow.providers.google.cloud.hooks import mlengine as hook from tests.providers.google.cloud.utils.base_gcp_mock import ( GCP_PROJECT_ID_HOOK_UNIT_TEST, mock_base_gcp_hook_default_project_id, - mock_base_gcp_hook_no_default_project_id, ) @@ -874,7 +872,8 @@ def test_create_version(self, mock_get_conn, mock_project_id): create_version_response = self.hook.create_version( model_name=model_name, - version_spec=version + version_spec=version, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST ) self.assertEqual(create_version_response, operation_done) @@ -910,7 +909,8 @@ def test_set_default_version(self, mock_get_conn, mock_project_id): set_default_version_response = self.hook.set_default_version( model_name=model_name, - version_name=version_name + version_name=version_name, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, ) self.assertEqual(set_default_version_response, operation_done) @@ -950,7 +950,8 @@ def test_list_versions(self, mock_get_conn, mock_sleep, mock_project_id): versions.return_value ) = versions_mock - list_versions_response = self.hook.list_versions(model_name=model_name) + list_versions_response = self.hook.list_versions( + model_name=model_name, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertEqual(list_versions_response, version_names) mock_get_conn.assert_has_calls([ @@ -998,7 +999,8 @@ def test_delete_version(self, mock_get_conn, mock_project_id): execute.return_value ) = version - delete_version_response = self.hook.delete_version(model_name=model_name, version_name=version_name) + delete_version_response = self.hook.delete_version( + model_name=model_name, version_name=version_name, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertEqual(delete_version_response, operation_done) mock_get_conn.assert_has_calls([ @@ -1029,7 +1031,7 @@ def test_create_model(self, mock_get_conn, mock_project_id): execute.return_value ) = model - create_model_response = self.hook.create_model(model=model) + create_model_response = self.hook.create_model(model=model, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertEqual(create_model_response, model) mock_get_conn.assert_has_calls([ @@ -1056,7 +1058,8 @@ def test_get_model(self, mock_get_conn, mock_project_id): execute.return_value ) = model - get_model_response = self.hook.get_model(model_name=model_name) + get_model_response = self.hook.get_model( + model_name=model_name, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertEqual(get_model_response, model) mock_get_conn.assert_has_calls([ @@ -1082,7 +1085,7 @@ def test_delete_model(self, mock_get_conn, mock_project_id): execute.return_value ) = model - self.hook.delete_model(model_name=model_name) + self.hook.delete_model(model_name=model_name, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) mock_get_conn.assert_has_calls([ mock.call().projects().models().delete(name=model_path), @@ -1128,7 +1131,7 @@ def test_create_mlengine_job(self, mock_get_conn, mock_sleep, mock_project_id): execute.side_effect ) = [job_queued, job_succeeded] - create_job_response = self.hook.create_job(job=new_job) + create_job_response = self.hook.create_job(job=new_job, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertEqual(create_job_response, job_succeeded) mock_get_conn.assert_has_calls([ @@ -1157,145 +1160,9 @@ def test_cancel_mlengine_job(self, mock_get_conn, mock_project_id): execute.return_value ) = job_cancelled - cancel_job_response = self.hook.cancel_job(job_id=job_id) + cancel_job_response = self.hook.cancel_job(job_id=job_id, project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST) self.assertEqual(cancel_job_response, job_cancelled) mock_get_conn.assert_has_calls([ mock.call().projects().jobs().cancel(name=job_path), ], any_order=True) - - -class TestMLEngineHookWithoutProjectId(unittest.TestCase): - def setUp(self) -> None: - super().setUp() - with mock.patch( - 'airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.__init__', - new=mock_base_gcp_hook_no_default_project_id, - ): - self.hook = hook.MLEngineHook() - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_create_version(self, mock_get_conn, mock_project_id): - model_name = 'test-model' - version_name = 'test-version' - version = {'name': version_name} - - with self.assertRaises(AirflowException): - self.hook.create_version( - model_name=model_name, - version_spec=version - ) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_set_default_version(self, mock_get_conn, mock_project_id): - model_name = 'test-model' - version_name = 'test-version' - - with self.assertRaises(AirflowException): - self.hook.set_default_version( - model_name=model_name, - version_name=version_name - ) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.time.sleep") - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_list_versions(self, mock_get_conn, mock_sleep, mock_project_id): - model_name = 'test-model' - - with self.assertRaises(AirflowException): - self.hook.list_versions(model_name=model_name) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_delete_version(self, mock_get_conn, mock_project_id): - model_name = 'test-model' - version_name = 'test-version' - - with self.assertRaises(AirflowException): - self.hook.delete_version(model_name=model_name, version_name=version_name) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_create_model(self, mock_get_conn, mock_project_id): - model_name = 'test-model' - model = { - 'name': model_name, - } - - with self.assertRaises(AirflowException): - self.hook.create_model(model=model) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_get_model(self, mock_get_conn, mock_project_id): - model_name = 'test-model' - with self.assertRaises(AirflowException): - self.hook.get_model(model_name=model_name) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_delete_model(self, mock_get_conn, mock_project_id): - model_name = 'test-model' - - with self.assertRaises(AirflowException): - self.hook.delete_model(model_name=model_name) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.time.sleep") - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_create_mlengine_job(self, mock_get_conn, mock_sleep, mock_project_id): - job_id = 'test-job-id' - new_job = { - 'jobId': job_id, - 'foo': 4815162342, - } - - with self.assertRaises(AirflowException): - self.hook.create_job(job=new_job) - - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch("airflow.providers.google.cloud.hooks.mlengine.MLEngineHook.get_conn") - def test_cancel_mlengine_job(self, mock_get_conn, mock_project_id): - job_id = 'test-job-id' - - with self.assertRaises(AirflowException): - self.hook.cancel_job(job_id=job_id) diff --git a/tests/providers/google/cloud/hooks/test_spanner.py b/tests/providers/google/cloud/hooks/test_spanner.py index eb5f642059a22..71d8b9bce7d15 100644 --- a/tests/providers/google/cloud/hooks/test_spanner.py +++ b/tests/providers/google/cloud/hooks/test_spanner.py @@ -21,7 +21,6 @@ import mock from mock import PropertyMock -from airflow.exceptions import AirflowException from airflow.providers.google.cloud.hooks.spanner import SpannerHook from tests.providers.google.cloud.utils.base_gcp_mock import ( GCP_PROJECT_ID_HOOK_UNIT_TEST, mock_base_gcp_hook_default_project_id, @@ -92,7 +91,9 @@ def test_create_instance(self, get_client, mock_project_id): instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=1, - display_name=SPANNER_DATABASE) + display_name=SPANNER_DATABASE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with( instance_id='instance', @@ -136,7 +137,9 @@ def test_update_instance(self, get_client, mock_project_id): instance_id=SPANNER_INSTANCE, configuration_name=SPANNER_CONFIGURATION, node_count=2, - display_name=SPANNER_DATABASE) + display_name=SPANNER_DATABASE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with( instance_id='instance', configuration_name='configuration', display_name='database-name', @@ -177,7 +180,9 @@ def test_delete_instance(self, get_client, mock_project_id): delete_method = instance_method.return_value.delete delete_method.return_value = False res = self.spanner_hook_default_project_id.delete_instance( - instance_id=SPANNER_INSTANCE) + instance_id=SPANNER_INSTANCE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with( 'instance') @@ -215,7 +220,9 @@ def test_get_database(self, get_client, mock_project_id): database_exists_method.return_value = True res = self.spanner_hook_default_project_id.get_database( instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE) + database_id=SPANNER_DATABASE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') @@ -255,7 +262,9 @@ def test_create_database(self, get_client, mock_project_id): res = self.spanner_hook_default_project_id.create_database( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, - ddl_statements=[]) + ddl_statements=[], + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name', ddl_statements=[]) @@ -295,7 +304,9 @@ def test_update_database(self, get_client, mock_project_id): res = self.spanner_hook_default_project_id.update_database( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, - ddl_statements=[]) + ddl_statements=[], + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') @@ -336,7 +347,9 @@ def test_delete_database(self, get_client, mock_project_id): database_exists_method.return_value = True res = self.spanner_hook_default_project_id.delete_database( instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE) + database_id=SPANNER_DATABASE, + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') @@ -379,7 +392,9 @@ def test_execute_dml(self, get_client, mock_project_id): res = self.spanner_hook_default_project_id.execute_dml( instance_id=SPANNER_INSTANCE, database_id=SPANNER_DATABASE, - queries='') + queries='', + project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, + ) get_client.assert_called_once_with(project_id='example-project') instance_method.assert_called_once_with(instance_id='instance') database_method.assert_called_once_with(database_id='database-name') @@ -431,24 +446,6 @@ def test_spanner_client_creation(self, mock_client, mock_get_creds, mock_client_ self.assertEqual(mock_client.return_value, result) self.assertEqual(self.spanner_hook_no_default_project_id._client, result) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_get_existing_instance_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.get_instance(instance_id=SPANNER_INSTANCE) - get_client.assert_not_called() - instance_method.assert_not_called() - instance_exists_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_get_existing_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -471,29 +468,6 @@ def test_get_non_existing_instance(self, get_client): instance_method.assert_called_once_with(instance_id='instance') self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_create_instance_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - create_method = instance_method.return_value.create - create_method.return_value = False - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.create_instance( - instance_id=SPANNER_INSTANCE, - configuration_name=SPANNER_CONFIGURATION, - node_count=1, - display_name=SPANNER_DATABASE) - get_client.assert_called_once_with(project_id='example-project') - get_client.assert_not_called() - instance_method.assert_not_called() - create_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_create_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -513,30 +487,6 @@ def test_create_instance_overridden_project_id(self, get_client): node_count=1) self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_update_instance_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - update_method = instance_method.return_value.update - update_method.return_value = False - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.update_instance( - instance_id=SPANNER_INSTANCE, - configuration_name=SPANNER_CONFIGURATION, - node_count=2, - display_name=SPANNER_DATABASE) - get_client.assert_not_called() - instance_method.assert_not_called() - update_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_update_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -557,27 +507,6 @@ def test_update_instance_overridden_project_id(self, get_client): update_method.assert_called_once_with() self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_delete_instance_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - delete_method = instance_method.return_value.delete - delete_method.return_value = False - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.delete_instance( - instance_id=SPANNER_INSTANCE) - get_client.assert_not_called() - instance_method.assert_not_called() - delete_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_delete_instance_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -594,29 +523,6 @@ def test_delete_instance_overridden_project_id(self, get_client): delete_method.assert_called_once_with() self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_get_database_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - database_method = instance_method.return_value.database - database_exists_method = instance_method.return_value.exists - database_exists_method.return_value = True - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.get_database( - instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE) - get_client.assert_not_called() - instance_method.assert_not_called() - database_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_get_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -635,30 +541,6 @@ def test_get_database_overridden_project_id(self, get_client): database_exists_method.assert_called_once_with() self.assertIsNotNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_create_database_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - database_method = instance_method.return_value.database - database_create_method = database_method.return_value.create - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.create_database( - instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE, - ddl_statements=[]) - get_client.assert_not_called() - instance_method.assert_not_called() - database_method.assert_not_called() - database_create_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_create_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -677,30 +559,6 @@ def test_create_database_overridden_project_id(self, get_client): database_create_method.assert_called_once_with() self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_update_database_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - database_method = instance_method.return_value.database - database_update_method = database_method.return_value.update - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.update_database( - instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE, - ddl_statements=[]) - get_client.assert_not_called() - instance_method.assert_not_called() - database_method.assert_not_called() - database_update_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_update_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -738,32 +596,6 @@ def test_update_database_overridden_project_id_and_operation(self, get_client): database_update_ddl_method.assert_called_once_with(ddl_statements=[], operation_id="operation") self.assertIsNone(res) - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_delete_database_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - database_method = instance_method.return_value.database - database_drop_method = database_method.return_value.drop - database_exists_method = database_method.return_value.exists - database_exists_method.return_value = True - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.delete_database( - instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE) - get_client.assert_not_called() - instance_method.assert_not_called() - database_method.assert_not_called() - database_exists_method.assert_not_called() - database_drop_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_delete_database_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance @@ -803,30 +635,6 @@ def test_delete_database_missing_database(self, get_client): database_exists_method.assert_called_once_with() database_drop_method.assert_not_called() - @mock.patch( - 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', - new_callable=PropertyMock, - return_value=None - ) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') - def test_execute_dml_missing_project_id(self, get_client, mock_project_id): - instance_method = get_client.return_value.instance - instance_exists_method = instance_method.return_value.exists - instance_exists_method.return_value = True - database_method = instance_method.return_value.database - run_in_transaction_method = database_method.return_value.run_in_transaction - with self.assertRaises(AirflowException) as cm: - self.spanner_hook_no_default_project_id.execute_dml( - instance_id=SPANNER_INSTANCE, - database_id=SPANNER_DATABASE, - queries='') - get_client.assert_not_called() - instance_method.assert_not_called() - database_method.assert_not_called() - run_in_transaction_method.assert_not_called() - err = cm.exception - self.assertIn("The project id must be passed", str(err)) - @mock.patch('airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client') def test_execute_dml_overridden_project_id(self, get_client): instance_method = get_client.return_value.instance diff --git a/tests/providers/google/cloud/hooks/test_stackdriver.py b/tests/providers/google/cloud/hooks/test_stackdriver.py index 8408b9191f421..7340b243bd25a 100644 --- a/tests/providers/google/cloud/hooks/test_stackdriver.py +++ b/tests/providers/google/cloud/hooks/test_stackdriver.py @@ -114,7 +114,8 @@ def test_stackdriver_list_alert_policies(self, mock_policy_client, mock_get_cred method = mock_policy_client.return_value.list_alert_policies hook = stackdriver.StackdriverHook() hook.list_alert_policies( - filter_=TEST_FILTER + filter_=TEST_FILTER, + project_id=PROJECT_ID, ) method.assert_called_once_with( name='projects/{project}'.format(project=PROJECT_ID), @@ -142,6 +143,7 @@ def test_stackdriver_enable_alert_policy(self, mock_policy_client, mock_get_cred mock_policy_client.return_value.list_alert_policies.return_value = alert_policies hook.enable_alert_policies( filter_=TEST_FILTER, + project_id=PROJECT_ID, ) mock_policy_client.return_value.list_alert_policies.assert_called_once_with( name='projects/{project}'.format(project=PROJECT_ID), @@ -179,6 +181,7 @@ def test_stackdriver_disable_alert_policy(self, mock_policy_client, mock_get_cre ] hook.disable_alert_policies( filter_=TEST_FILTER, + project_id=PROJECT_ID, ) mock_policy_client.return_value.list_alert_policies.assert_called_once_with( name='projects/{project}'.format(project=PROJECT_ID), @@ -216,7 +219,8 @@ def test_stackdriver_upsert_alert_policy(self, mock_channel_client, mock_policy_ mock_channel_client.return_value.list_notification_channels.return_value = [] hook.upsert_alert( - alerts=json.dumps({"policies": [TEST_ALERT_POLICY_1, TEST_ALERT_POLICY_2], "channels": []}) + alerts=json.dumps({"policies": [TEST_ALERT_POLICY_1, TEST_ALERT_POLICY_2], "channels": []}), + project_id=PROJECT_ID, ) mock_channel_client.return_value.list_notification_channels.assert_called_once_with( name='projects/{project}'.format(project=PROJECT_ID), @@ -281,7 +285,8 @@ def test_stackdriver_delete_alert_policy(self, mock_policy_client, mock_get_cred def test_stackdriver_list_notification_channel(self, mock_channel_client, mock_get_creds_and_project_id): hook = stackdriver.StackdriverHook() hook.list_notification_channels( - filter_=TEST_FILTER + filter_=TEST_FILTER, + project_id=PROJECT_ID, ) mock_channel_client.return_value.list_notification_channels.assert_called_once_with( name='projects/{project}'.format(project=PROJECT_ID), @@ -312,6 +317,7 @@ def test_stackdriver_enable_notification_channel(self, mock_channel_client, hook.enable_notification_channels( filter_=TEST_FILTER, + project_id=PROJECT_ID, ) notification_channel_disabled.enabled.value = True # pylint: disable=no-member @@ -344,6 +350,7 @@ def test_stackdriver_disable_notification_channel(self, mock_channel_client, hook.disable_notification_channels( filter_=TEST_FILTER, + project_id=PROJECT_ID, ) notification_channel_enabled.enabled.value = False # pylint: disable=no-member @@ -374,7 +381,8 @@ def test_stackdriver_upsert_channel(self, mock_channel_client, mock_get_creds_an existing_notification_channel ] hook.upsert_channel( - channels=json.dumps({"channels": [TEST_NOTIFICATION_CHANNEL_1, TEST_NOTIFICATION_CHANNEL_2]}) + channels=json.dumps({"channels": [TEST_NOTIFICATION_CHANNEL_1, TEST_NOTIFICATION_CHANNEL_2]}), + project_id=PROJECT_ID, ) mock_channel_client.return_value.list_notification_channels.assert_called_once_with( name='projects/{project}'.format(project=PROJECT_ID), diff --git a/tests/providers/google/cloud/operators/test_automl.py b/tests/providers/google/cloud/operators/test_automl.py index 04392042f8fa7..73b081b5f9983 100644 --- a/tests/providers/google/cloud/operators/test_automl.py +++ b/tests/providers/google/cloud/operators/test_automl.py @@ -188,14 +188,12 @@ def test_execute(self, mock_hook): dataset=dataset, update_mask=MASK, location=GCP_LOCATION, - project_id=GCP_PROJECT_ID, task_id=TASK_ID, ) op.execute(context=None) mock_hook.return_value.update_dataset.assert_called_once_with( dataset=dataset, metadata=None, - project_id=GCP_PROJECT_ID, retry=None, timeout=None, update_mask=MASK, diff --git a/tests/providers/google/cloud/operators/test_dataflow.py b/tests/providers/google/cloud/operators/test_dataflow.py index 483839dc12bf2..246d53d5859ce 100644 --- a/tests/providers/google/cloud/operators/test_dataflow.py +++ b/tests/providers/google/cloud/operators/test_dataflow.py @@ -113,7 +113,8 @@ def test_exec(self, gcs_hook, dataflow_mock): py_interpreter=PY_INTERPRETER, py_requirements=[], py_system_site_packages=False, - on_new_job_id_callback=mock.ANY + on_new_job_id_callback=mock.ANY, + project_id=None, ) self.assertTrue(self.dataflow.py_file.startswith('/tmp/dataflow')) @@ -163,7 +164,8 @@ def test_exec(self, gcs_hook, dataflow_mock): job_class=JOB_CLASS, append_job_name=True, multiple_jobs=None, - on_new_job_id_callback=mock.ANY + on_new_job_id_callback=mock.ANY, + project_id=None, ) @mock.patch('airflow.providers.google.cloud.operators.dataflow.DataflowHook') @@ -182,7 +184,8 @@ def test_check_job_running_exec(self, gcs_hook, dataflow_mock): self.assertTrue(dataflow_mock.called) gcs_download_hook.assert_not_called() start_java_hook.assert_not_called() - dataflow_running.assert_called_once_with(name=JOB_NAME, variables=mock.ANY) + dataflow_running.assert_called_once_with( + name=JOB_NAME, variables=mock.ANY, project_id=None) @mock.patch('airflow.providers.google.cloud.operators.dataflow.DataflowHook') @mock.patch(GCS_HOOK_STRING.format('GoogleCloudBucketHelper')) @@ -206,9 +209,11 @@ def test_check_job_not_running_exec(self, gcs_hook, dataflow_mock): job_class=JOB_CLASS, append_job_name=True, multiple_jobs=None, - on_new_job_id_callback=mock.ANY + on_new_job_id_callback=mock.ANY, + project_id=None, ) - dataflow_running.assert_called_once_with(name=JOB_NAME, variables=mock.ANY) + dataflow_running.assert_called_once_with( + name=JOB_NAME, variables=mock.ANY, project_id=None) @mock.patch('airflow.providers.google.cloud.operators.dataflow.DataflowHook') @mock.patch(GCS_HOOK_STRING.format('GoogleCloudBucketHelper')) @@ -233,9 +238,11 @@ def test_check_multiple_job_exec(self, gcs_hook, dataflow_mock): job_class=JOB_CLASS, append_job_name=True, multiple_jobs=True, - on_new_job_id_callback=mock.ANY + on_new_job_id_callback=mock.ANY, + project_id=None, ) - dataflow_running.assert_called_once_with(name=JOB_NAME, variables=mock.ANY) + dataflow_running.assert_called_once_with( + name=JOB_NAME, variables=mock.ANY, project_id=None) class TestDataflowTemplateOperator(unittest.TestCase): @@ -279,7 +286,8 @@ def test_exec(self, dataflow_mock): variables=expected_options, parameters=PARAMETERS, dataflow_template=TEMPLATE, - on_new_job_id_callback=mock.ANY + on_new_job_id_callback=mock.ANY, + project_id=None, ) diff --git a/tests/providers/google/cloud/operators/test_mlengine_utils.py b/tests/providers/google/cloud/operators/test_mlengine_utils.py index 0dfbc9c24f7d7..fab26b5b2eb38 100644 --- a/tests/providers/google/cloud/operators/test_mlengine_utils.py +++ b/tests/providers/google/cloud/operators/test_mlengine_utils.py @@ -117,7 +117,8 @@ def test_successful_run(self): py_requirements=['apache-beam[gcp]>=2.14.0'], py_interpreter='python3', py_system_site_packages=False, - on_new_job_id_callback=ANY + on_new_job_id_callback=ANY, + project_id='test-project', ) with patch('airflow.providers.google.cloud.utils.mlengine_operator_utils.GCSHook') as mock_gcs_hook: