fix: try002 for provider google (#38803)

Author: dondaum

airflow/providers/google/cloud/hooks/bigquery.py

@@ -2295,7 +2295,7 @@ def var_print(var_name):
                 return f"Format exception for {var_name}: "
 
         if table_input.count(".") + table_input.count(":") > 3:
-            raise Exception(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")
+            raise ValueError(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")
         cmpt = table_input.rsplit(":", 1)
         project_id = None
         rest = table_input
@@ -2307,7 +2307,7 @@ def var_print(var_name):
                 project_id = cmpt[0]
                 rest = cmpt[1]
         else:
-            raise Exception(
+            raise ValueError(
                 f"{var_print(var_name)}Expect format of (<project:)<dataset>.<table>, got {table_input}"
             )
 
@@ -2323,7 +2323,7 @@ def var_print(var_name):
             dataset_id = cmpt[0]
             table_id = cmpt[1]
         else:
-            raise Exception(
+            raise ValueError(
                 f"{var_print(var_name)} Expect format of (<project.|<project:)<dataset>.<table>, "
                 f"got {table_input}"
             )
@@ -3107,7 +3107,7 @@ def var_print(var_name):
             return f"Format exception for {var_name}: "
 
     if table_input.count(".") + table_input.count(":") > 3:
-        raise Exception(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")
+        raise ValueError(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")
     cmpt = table_input.rsplit(":", 1)
     project_id = None
     rest = table_input
@@ -3119,7 +3119,7 @@ def var_print(var_name):
             project_id = cmpt[0]
             rest = cmpt[1]
     else:
-        raise Exception(
+        raise ValueError(
             f"{var_print(var_name)}Expect format of (<project:)<dataset>.<table>, got {table_input}"
         )
 
@@ -3135,7 +3135,7 @@ def var_print(var_name):
         dataset_id = cmpt[0]
         table_id = cmpt[1]
     else:
-        raise Exception(
+        raise ValueError(
             f"{var_print(var_name)}Expect format of (<project.|<project:)<dataset>.<table>, got {table_input}"
         )
 

airflow/providers/google/cloud/hooks/dataflow.py

@@ -252,7 +252,7 @@ def _get_current_jobs(self) -> list[dict]:
                 self._job_id = jobs[0]["id"]
             return jobs
         else:
-            raise Exception("Missing both dataflow job ID and name.")
+            raise ValueError("Missing both dataflow job ID and name.")
 
     def fetch_job_by_id(self, job_id: str) -> dict:
         """
@@ -410,18 +410,18 @@ def _check_dataflow_job_state(self, job) -> bool:
         else:
             terminal_states = DataflowJobStatus.TERMINAL_STATES | {DataflowJobStatus.JOB_STATE_RUNNING}
             if self._expected_terminal_state not in terminal_states:
-                raise Exception(
+                raise AirflowException(
                     f"Google Cloud Dataflow job's expected terminal state "
                     f"'{self._expected_terminal_state}' is invalid."
                     f" The value should be any of the following: {terminal_states}"
                 )
             elif is_streaming and self._expected_terminal_state == DataflowJobStatus.JOB_STATE_DONE:
-                raise Exception(
+                raise AirflowException(
                     "Google Cloud Dataflow job's expected terminal state cannot be "
                     "JOB_STATE_DONE while it is a streaming job"
                 )
             elif not is_streaming and self._expected_terminal_state == DataflowJobStatus.JOB_STATE_DRAINED:
-                raise Exception(
+                raise AirflowException(
                     "Google Cloud Dataflow job's expected terminal state cannot be "
                     "JOB_STATE_DRAINED while it is a batch job"
                 )
@@ -435,7 +435,7 @@ def _check_dataflow_job_state(self, job) -> bool:
             return self._wait_until_finished is False
 
         self.log.debug("Current job: %s", job)
-        raise Exception(
+        raise AirflowException(
             f"Google Cloud Dataflow job {job['name']} is in an unexpected terminal state: {current_state}, "
             f"expected terminal state: {self._expected_terminal_state}"
         )

airflow/providers/google/cloud/hooks/dataproc.py

@@ -158,10 +158,10 @@ def set_main(self, main_jar: str | None = None, main_class: str | None = None) -
 
         :param main_jar: URI for the main file.
         :param main_class: Name of the main class.
-        :raises: Exception
+        :raises: ValueError
         """
         if main_class is not None and main_jar is not None:
-            raise Exception("Set either main_jar or main_class")
+            raise ValueError("Set either main_jar or main_class")
         if main_jar:
             self.job["job"][self.job_type]["main_jar_file_uri"] = main_jar
         else:

airflow/providers/google/cloud/operators/dataproc.py

@@ -1430,13 +1430,13 @@ def _generate_job_template(self) -> str:
         if self.job_template:
             job = self.job_template.build()
             return job["job"]
-        raise Exception("Create a job template before")
+        raise AirflowException("Create a job template before")
 
     def execute(self, context: Context):
         if self.job_template:
             self.job = self.job_template.build()
             if self.job is None:
-                raise Exception("The job should be set here.")
+                raise AirflowException("The job should be set here.")
             self.dataproc_job_id = self.job["job"]["reference"]["job_id"]
             self.log.info("Submitting %s job %s", self.job_type, self.dataproc_job_id)
             job_object = self.hook.submit_job(

airflow/providers/google/leveldb/hooks/leveldb.py

@@ -94,17 +94,17 @@ def run(
         """
         if command == "put":
             if not value:
-                raise Exception("Please provide `value`!")
+                raise ValueError("Please provide `value`!")
             return self.put(key, value)
         elif command == "get":
             return self.get(key)
         elif command == "delete":
             return self.delete(key)
         elif command == "write_batch":
             if not keys:
-                raise Exception("Please provide `keys`!")
+                raise ValueError("Please provide `keys`!")
             if not values:
-                raise Exception("Please provide `values`!")
+                raise ValueError("Please provide `values`!")
             return self.write_batch(keys, values)
         else:
             raise LevelDBHookException("Unknown command for LevelDB hook")
@@ -117,7 +117,7 @@ def put(self, key: bytes, value: bytes):
         :param value: value for put execution e.g. ``b'value'``, ``b'another-value'``
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         self.db.put(key, value)
 
     def get(self, key: bytes) -> bytes:
@@ -128,7 +128,7 @@ def get(self, key: bytes) -> bytes:
         :returns: value of key from db.get
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         return self.db.get(key)
 
     def delete(self, key: bytes):
@@ -138,7 +138,7 @@ def delete(self, key: bytes):
         :param key: key for delete execution, e.g. ``b'key'``, ``b'another-key'``
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         self.db.delete(key)
 
     def write_batch(self, keys: list[bytes], values: list[bytes]):
@@ -149,7 +149,7 @@ def write_batch(self, keys: list[bytes], values: list[bytes]):
         :param values: values for write_batch execution e.g. ``[b'value', b'another-value']``
         """
         if not self.db:
-            raise Exception(DB_NOT_INITIALIZED_BEFORE)
+            raise AirflowException(DB_NOT_INITIALIZED_BEFORE)
         with self.db.write_batch() as batch:
             for i, key in enumerate(keys):
                 batch.put(key, values[i])

pyproject.toml

@@ -381,12 +381,6 @@ combine-as-imports = true
 # All the providers modules which do not follow TRY002 yet
 # cncf.kubernetes
 "airflow/providers/cncf/kubernetes/operators/pod.py" = ["TRY002"]
-# google
-"airflow/providers/google/cloud/hooks/bigquery.py" = ["TRY002"]
-"airflow/providers/google/cloud/hooks/dataflow.py" = ["TRY002"]
-"airflow/providers/google/cloud/hooks/dataproc.py" = ["TRY002"]
-"airflow/providers/google/cloud/operators/dataproc.py" = ["TRY002"]
-"airflow/providers/google/leveldb/hooks/leveldb.py" = ["TRY002"]
 # imap
 "airflow/providers/imap/hooks/imap.py" = ["TRY002"]
 

tests/providers/google/cloud/hooks/test_bigquery.py

@@ -155,7 +155,7 @@ def test_get_pandas_df(self, mock_read_gbq):
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     def test_invalid_schema_update_options(self, mock_get_service):
         with pytest.raises(
-            Exception,
+            ValueError,
             match=(
                 r"\['THIS IS NOT VALID'\] contains invalid schema update options. "
                 r"Please only use one or more of the following options: "
@@ -172,7 +172,7 @@ def test_invalid_schema_update_options(self, mock_get_service):
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     def test_invalid_schema_update_and_write_disposition(self, mock_get_service):
         with pytest.raises(
-            Exception,
+            ValueError,
             match="schema_update_options is only allowed if"
             " write_disposition is 'WRITE_APPEND' or 'WRITE_TRUNCATE'.",
         ):
@@ -846,7 +846,7 @@ def test_update_table_schema_without_policy_tags(self, mock_update, mock_get_sch
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     def test_invalid_source_format(self, mock_get_service):
         with pytest.raises(
-            Exception,
+            ValueError,
             match=r"JSON is not a valid source format. Please use one of the following types: \['CSV', "
             r"'NEWLINE_DELIMITED_JSON', 'AVRO', 'GOOGLE_SHEETS', 'DATASTORE_BACKUP', 'PARQUET'\]",
         ):
@@ -959,7 +959,7 @@ def test_job_id_validity(self, mock_md5, test_dag_id, expected_job_id):
 
 class TestBigQueryTableSplitter:
     def test_internal_need_default_project(self):
-        with pytest.raises(Exception, match="INTERNAL: No default project is specified"):
+        with pytest.raises(ValueError, match="INTERNAL: No default project is specified"):
             split_tablename("dataset.table", None)
 
     @pytest.mark.parametrize(
@@ -1008,7 +1008,7 @@ def test_split_tablename(self, project_expected, dataset_expected, table_expecte
     )
     def test_invalid_syntax(self, table_input, var_name, exception_message):
         default_project_id = "project"
-        with pytest.raises(Exception, match=exception_message.format(table_input)):
+        with pytest.raises(ValueError, match=exception_message.format(table_input)):
             split_tablename(table_input, default_project_id, var_name)
 
 

tests/providers/google/cloud/hooks/test_dataflow.py

@@ -1284,7 +1284,7 @@ def test_dataflow_job_wait_for_multiple_jobs_and_one_in_terminal_state(self, sta
             num_retries=20,
             multiple_jobs=True,
         )
-        with pytest.raises(Exception, match=exception_regex):
+        with pytest.raises(AirflowException, match=exception_regex):
             dataflow_job.wait_for_done()
 
     def test_dataflow_job_wait_for_multiple_jobs_and_streaming_jobs(self):
@@ -1517,7 +1517,7 @@ def test_check_dataflow_job_state_terminal_state(self, job_type, job_state, exce
             num_retries=20,
             multiple_jobs=True,
         )
-        with pytest.raises(Exception, match=exception_regex):
+        with pytest.raises(AirflowException, match=exception_regex):
             dataflow_job._check_dataflow_job_state(job)
 
     @pytest.mark.parametrize(
@@ -1558,7 +1558,7 @@ def test_check_dataflow_job_state__invalid_expected_state(self, job_type, expect
             multiple_jobs=False,
             expected_terminal_state=expected_terminal_state,
         )
-        with pytest.raises(Exception, match=match):
+        with pytest.raises(AirflowException, match=match):
             dataflow_job._check_dataflow_job_state(job)
 
     def test_dataflow_job_cancel_job(self):

tests/providers/google/cloud/hooks/test_dataproc.py

@@ -1084,7 +1084,7 @@ def test_add_python_file_uris(self):
         assert python_file_uris == self.builder.job["job"][self.job_type]["python_file_uris"]
 
     def test_set_main_error(self):
-        with pytest.raises(Exception):
+        with pytest.raises(ValueError):
             self.builder.set_main("test", "test")
 
     def test_set_main_class(self):