apache-airflow-providers-google 10.20.0rc1__py3-none-any.whl → 10.21.0__py3-none-any.whl

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

@@ -908,7 +908,8 @@ class CloudSQLDatabaseHook(BaseHook):
     def _set_temporary_ssl_file(
         self, cert_name: str, cert_path: str | None = None, cert_value: str | None = None
     ) -> str | None:
-        """Save the certificate as a temporary file.
+        """
+        Save the certificate as a temporary file.
 
         This method was implemented in order to overcome psql connection error caused by excessive file
         permissions: "private key file "..." has group or world access; file must have permissions
@@ -1019,7 +1020,8 @@ class CloudSQLDatabaseHook(BaseHook):
 
     @staticmethod
     def _generate_unique_path() -> str:
-        """Generate a unique path.
+        """
+        Generate a unique path.
 
         We don't using mkdtemp here since it can generate paths close to 60
         characters. We append project/location/instance to the path, Postgres
@@ -1106,7 +1108,8 @@ class CloudSQLDatabaseHook(BaseHook):
         return instance_specification
 
     def create_connection(self) -> Connection:
-        """Create a connection.
+        """
+        Create a connection.
 
         Connection ID will be randomly generated according to whether it uses
         proxy, TCP, UNIX sockets, SSL.
@@ -1117,7 +1120,8 @@ class CloudSQLDatabaseHook(BaseHook):
         return connection
 
     def get_sqlproxy_runner(self) -> CloudSqlProxyRunner:
-        """Retrieve Cloud SQL Proxy runner.
+        """
+        Retrieve Cloud SQL Proxy runner.
 
         It is used to manage the proxy lifecycle per task.
 
@@ -1137,7 +1141,8 @@ class CloudSQLDatabaseHook(BaseHook):
         )
 
     def get_database_hook(self, connection: Connection) -> PostgresHook | MySqlHook:
-        """Retrieve database hook.
+        """
+        Retrieve database hook.
 
         This is the actual Postgres or MySQL database hook that uses proxy or
         connects directly to the Google Cloud SQL database.
@@ -1168,7 +1173,8 @@ class CloudSQLDatabaseHook(BaseHook):
         self.sql_proxy_tcp_port = self.reserved_tcp_socket.getsockname()[1]
 
     def free_reserved_port(self) -> None:
-        """Free TCP port.
+        """
+        Free TCP port.
 
         Makes it immediately ready to be used by Cloud SQL Proxy.
         """

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

@@ -134,7 +134,8 @@ NEGATIVE_STATUSES = {GcpTransferOperationStatus.FAILED, GcpTransferOperationStat
 
 
 def gen_job_name(job_name: str) -> str:
-    """Add a unique suffix to the job name.
+    """
+    Add a unique suffix to the job name.
 
     :param job_name:
     :return: job_name with suffix
@@ -144,7 +145,8 @@ def gen_job_name(job_name: str) -> str:
 
 
 class CloudDataTransferServiceHook(GoogleBaseHook):
-    """Google Storage Transfer Service functionalities.
+    """
+    Google Storage Transfer Service functionalities.
 
     All methods in the hook with *project_id* in the signature must be called
     with keyword arguments rather than positional.
@@ -171,7 +173,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         self._conn = None
 
     def get_conn(self) -> Resource:
-        """Retrieve connection to Google Storage Transfer service.
+        """
+        Retrieve connection to Google Storage Transfer service.
 
         :return: Google Storage Transfer service object
         """
@@ -183,7 +186,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         return self._conn
 
     def create_transfer_job(self, body: dict) -> dict:
-        """Create a transfer job that runs periodically.
+        """
+        Create a transfer job that runs periodically.
 
         :param body: (Required) The request body, as described in
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/patch#request-body
@@ -225,7 +229,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def get_transfer_job(self, job_name: str, project_id: str) -> dict:
-        """Get latest state of a long-running Google Storage Transfer Service job.
+        """
+        Get latest state of a long-running Google Storage Transfer Service job.
 
         :param job_name: (Required) Name of the job to be fetched
         :param project_id: (Optional) the ID of the project that owns the Transfer
@@ -241,7 +246,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def list_transfer_job(self, request_filter: dict | None = None, **kwargs) -> list[dict]:
-        """List long-running operations in Google Storage Transfer Service.
+        """
+        List long-running operations in Google Storage Transfer Service.
 
         A filter can be specified to match only certain entries.
 
@@ -279,7 +285,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def enable_transfer_job(self, job_name: str, project_id: str) -> dict:
-        """Make new transfers be performed based on the schedule.
+        """
+        Make new transfers be performed based on the schedule.
 
         :param job_name: (Required) Name of the job to be updated
         :param project_id: (Optional) the ID of the project that owns the Transfer
@@ -302,7 +309,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def update_transfer_job(self, job_name: str, body: dict) -> dict:
-        """Update a transfer job that runs periodically.
+        """
+        Update a transfer job that runs periodically.
 
         :param job_name: (Required) Name of the job to be updated
         :param body: A request body, as described in
@@ -319,7 +327,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def delete_transfer_job(self, job_name: str, project_id: str) -> None:
-        """Delete a transfer job.
+        """
+        Delete 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 collection. Transfer jobs
@@ -346,7 +355,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def run_transfer_job(self, job_name: str, project_id: str) -> dict:
-        """Run Google Storage Transfer Service job.
+        """
+        Run Google Storage Transfer Service job.
 
         :param job_name: (Required) Name of the job to be fetched
         :param project_id: (Optional) the ID of the project that owns the Transfer
@@ -371,14 +381,16 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def cancel_transfer_operation(self, operation_name: str) -> None:
-        """Cancel a transfer operation in Google Storage Transfer Service.
+        """
+        Cancel a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: Name of the transfer operation.
         """
         self.get_conn().transferOperations().cancel(name=operation_name).execute(num_retries=self.num_retries)
 
     def get_transfer_operation(self, operation_name: str) -> dict:
-        """Get a transfer operation in Google Storage Transfer Service.
+        """
+        Get a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: (Required) Name of the transfer operation.
         :return: transfer operation
@@ -393,7 +405,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         )
 
     def list_transfer_operations(self, request_filter: dict | None = None, **kwargs) -> list[dict]:
-        """Get a transfer operation in Google Storage Transfer Service.
+        """
+        Get a transfer operation in Google Storage Transfer Service.
 
         :param request_filter: (Required) A request filter, as described in
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/list#body.QUERY_PARAMETERS.filter
@@ -440,14 +453,16 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         return operations
 
     def pause_transfer_operation(self, operation_name: str) -> None:
-        """Pause a transfer operation in Google Storage Transfer Service.
+        """
+        Pause a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: (Required) Name of the transfer operation.
         """
         self.get_conn().transferOperations().pause(name=operation_name).execute(num_retries=self.num_retries)
 
     def resume_transfer_operation(self, operation_name: str) -> None:
-        """Resume a transfer operation in Google Storage Transfer Service.
+        """
+        Resume a transfer operation in Google Storage Transfer Service.
 
         :param operation_name: (Required) Name of the transfer operation.
         """
@@ -459,7 +474,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
         expected_statuses: set[str] | None = None,
         timeout: float | timedelta | None = None,
     ) -> None:
-        """Wait until the job reaches the expected state.
+        """
+        Wait until the job reaches the expected state.
 
         :param job: The transfer job to wait for. See:
             https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs#TransferJob
@@ -503,7 +519,8 @@ class CloudDataTransferServiceHook(GoogleBaseHook):
     def operations_contain_expected_statuses(
         operations: list[dict], expected_statuses: set[str] | str
     ) -> bool:
-        """Check whether an operation exists with the expected status.
+        """
+        Check whether an operation exists with the expected status.
 
         :param operations: (Required) List of transfer operations to check.
         :param expected_statuses: (Required) The expected status. See:

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

@@ -80,7 +80,8 @@ T = TypeVar("T", bound=Callable)
 def process_line_and_extract_dataflow_job_id_callback(
     on_new_job_id_callback: Callable[[str], None] | None,
 ) -> Callable[[str], None]:
-    """Build callback that triggers the specified function.
+    """
+    Build callback that triggers the specified function.
 
     The returned callback is intended to be used as ``process_line_callback`` in
     :py:class:`~airflow.providers.apache.beam.hooks.beam.BeamCommandRunner`.
@@ -418,32 +419,34 @@ class _DataflowJobsController(LoggingMixin):
         current_state = job["currentState"]
         is_streaming = job.get("type") == DataflowJobType.JOB_TYPE_STREAMING
 
-        if self._expected_terminal_state is None:
+        current_expected_state = self._expected_terminal_state
+
+        if current_expected_state is None:
             if is_streaming:
-                self._expected_terminal_state = DataflowJobStatus.JOB_STATE_RUNNING
+                current_expected_state = DataflowJobStatus.JOB_STATE_RUNNING
             else:
-                self._expected_terminal_state = DataflowJobStatus.JOB_STATE_DONE
-        else:
-            terminal_states = DataflowJobStatus.TERMINAL_STATES | {DataflowJobStatus.JOB_STATE_RUNNING}
-            if self._expected_terminal_state not in terminal_states:
-                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 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 AirflowException(
-                    "Google Cloud Dataflow job's expected terminal state cannot be "
-                    "JOB_STATE_DRAINED while it is a batch job"
-                )
+                current_expected_state = DataflowJobStatus.JOB_STATE_DONE
 
-        if current_state == self._expected_terminal_state:
-            if self._expected_terminal_state == DataflowJobStatus.JOB_STATE_RUNNING:
+        terminal_states = DataflowJobStatus.TERMINAL_STATES | {DataflowJobStatus.JOB_STATE_RUNNING}
+        if current_expected_state not in terminal_states:
+            raise AirflowException(
+                f"Google Cloud Dataflow job's expected terminal state "
+                f"'{current_expected_state}' is invalid."
+                f" The value should be any of the following: {terminal_states}"
+            )
+        elif is_streaming and current_expected_state == DataflowJobStatus.JOB_STATE_DONE:
+            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 current_expected_state == DataflowJobStatus.JOB_STATE_DRAINED:
+            raise AirflowException(
+                "Google Cloud Dataflow job's expected terminal state cannot be "
+                "JOB_STATE_DRAINED while it is a batch job"
+            )
+
+        if current_state == current_expected_state:
+            if current_expected_state == DataflowJobStatus.JOB_STATE_RUNNING:
                 return not self._wait_until_finished
             return True
 
@@ -453,7 +456,7 @@ class _DataflowJobsController(LoggingMixin):
         self.log.debug("Current job: %s", job)
         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}"
+            f"expected terminal state: {current_expected_state}"
         )
 
     def wait_for_done(self) -> None:
@@ -1583,7 +1586,8 @@ class AsyncDataflowHook(GoogleBaseAsyncHook):
         page_size: int | None = None,
         page_token: str | None = None,
     ) -> ListJobsAsyncPager:
-        """List jobs.
+        """
+        List jobs.
 
         For detail see:
         https://cloud.google.com/python/docs/reference/dataflow/latest/google.cloud.dataflow_v1beta3.types.ListJobsRequest

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

@@ -663,7 +663,8 @@ class DataformHook(GoogleBaseHook):
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
     ) -> InstallNpmPackagesResponse:
-        """Install NPM dependencies in the provided workspace.
+        """
+        Install NPM dependencies in the provided workspace.
 
         Requires "package.json" to be created in the workspace.
 

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

@@ -523,7 +523,8 @@ class DataFusionHook(GoogleBaseHook):
 
     @staticmethod
     def cdap_program_type(pipeline_type: DataFusionPipelineType) -> str:
-        """Retrieve CDAP Program type depending on the pipeline type.
+        """
+        Retrieve CDAP Program type depending on the pipeline type.
 
         :param pipeline_type: Pipeline type.
         """
@@ -535,7 +536,8 @@ class DataFusionHook(GoogleBaseHook):
 
     @staticmethod
     def cdap_program_id(pipeline_type: DataFusionPipelineType) -> str:
-        """Retrieve CDAP Program id depending on the pipeline type.
+        """
+        Retrieve CDAP Program id depending on the pipeline type.
 
         :param pipeline_type: Pipeline type.
         """