apache-airflow-providers-google 10.7.0rc1__py3-none-any.whl → 10.8.0__py3-none-any.whl

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

@@ -20,10 +20,9 @@ from __future__ import annotations
 
 from copy import deepcopy
 from functools import cached_property
-from typing import Any, Callable, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.vision_v1 import (
     AnnotateImageRequest,
     Image,
@@ -33,13 +32,16 @@ from google.cloud.vision_v1 import (
     ProductSet,
     ReferenceImage,
 )
-from google.protobuf import field_mask_pb2
 from google.protobuf.json_format import MessageToDict
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+    from google.protobuf import field_mask_pb2
+
 ERR_DIFF_NAMES = """The {label} name provided in the object ({explicit_name}) is different
     than the name created from the input parameters ({constructed_name}). Please either:
     1) Remove the {label} name,

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

@@ -16,20 +16,22 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.operation import Operation
-from google.api_core.retry import Retry
 from google.cloud.workflows.executions_v1beta import Execution, ExecutionsClient
-from google.cloud.workflows.executions_v1beta.services.executions.pagers import ListExecutionsPager
 from google.cloud.workflows_v1beta import Workflow, WorkflowsClient
-from google.cloud.workflows_v1beta.services.workflows.pagers import ListWorkflowsPager
-from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 
+if TYPE_CHECKING:
+    from google.api_core.operation import Operation
+    from google.api_core.retry import Retry
+    from google.cloud.workflows.executions_v1beta.services.executions.pagers import ListExecutionsPager
+    from google.cloud.workflows_v1beta.services.workflows.pagers import ListWorkflowsPager
+    from google.protobuf.field_mask_pb2 import FieldMask
+
 
 class WorkflowsHook(GoogleBaseHook):
     """

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

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 BIGQUERY_BASE_LINK = "/bigquery"

airflow/providers/google/cloud/links/bigquery_dts.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 BIGQUERY_BASE_LINK = "/bigquery/transfers"

airflow/providers/google/cloud/links/cloud_functions.py

@@ -21,10 +21,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 

airflow/providers/google/cloud/links/cloud_memorystore.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 BASE_LINK = "/memorystore"

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

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 

airflow/providers/google/cloud/links/cloud_tasks.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 CLOUD_TASKS_BASE_LINK = "/cloudtasks"

airflow/providers/google/cloud/links/compute.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 COMPUTE_BASE_LINK = "https://console.cloud.google.com/compute"

airflow/providers/google/cloud/links/datacatalog.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 DATACATALOG_BASE_LINK = "/datacatalog"

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

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 DATAFLOW_BASE_LINK = "/dataflow/jobs"

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

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 DATAFORM_BASE_LINK = "/bigquery/dataform"

airflow/providers/google/cloud/links/pubsub.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 PUBSUB_BASE_LINK = "/cloudpubsub"

airflow/providers/google/cloud/links/spanner.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 SPANNER_BASE_LINK = "/spanner/instances"

airflow/providers/google/cloud/links/stackdriver.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
 STACKDRIVER_BASE_LINK = "/monitoring/alerting"

airflow/providers/google/cloud/links/workflows.py

@@ -20,13 +20,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
+    from airflow.models import BaseOperator
     from airflow.utils.context import Context
 
-WORKFLOWS_BASE_LINK = "workflows"
+WORKFLOWS_BASE_LINK = "/workflows"
 WORKFLOW_LINK = WORKFLOWS_BASE_LINK + "/workflow/{location_id}/{workflow_id}/executions?project={project_id}"
 WORKFLOWS_LINK = WORKFLOWS_BASE_LINK + "?project={project_id}"
 EXECUTION_LINK = (

airflow/providers/google/cloud/log/gcs_task_handler.py

@@ -241,11 +241,9 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         try:
             blob = storage.Blob.from_string(remote_log_location, self.client)
             old_log = blob.download_as_bytes().decode()
-            log = "\n".join([old_log, log]) if old_log else log
+            log = f"{old_log}\n{log}" if old_log else log
         except Exception as e:
-            if self.no_log_found(e):
-                pass
-            else:
+            if not self.no_log_found(e):
                 log += self._add_message(
                     f"Error checking for previous log; if exists, may be overwritten: {e}"
                 )
@@ -265,8 +263,8 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
 
         :meta private:
         """
-        if exc.args and isinstance(exc.args[0], str) and "No such object" in exc.args[0]:
-            return True
-        elif getattr(exc, "resp", {}).get("status") == "404":
+        if (exc.args and isinstance(exc.args[0], str) and "No such object" in exc.args[0]) or getattr(
+            exc, "resp", {}
+        ).get("status") == "404":
             return True
         return False

airflow/providers/google/cloud/log/stackdriver_task_handler.py

@@ -18,22 +18,26 @@
 from __future__ import annotations
 
 import logging
-from contextvars import ContextVar
 from functools import cached_property
-from typing import Collection
+from typing import TYPE_CHECKING, Collection
 from urllib.parse import urlencode
 
-from google.auth.credentials import Credentials
 from google.cloud import logging as gcp_logging
 from google.cloud.logging import Resource
 from google.cloud.logging.handlers.transports import BackgroundThreadTransport, Transport
 from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client
 from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse
 
-from airflow.models import TaskInstance
 from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
 from airflow.providers.google.common.consts import CLIENT_INFO
 
+if TYPE_CHECKING:
+    from contextvars import ContextVar
+
+    from google.auth.credentials import Credentials
+
+    from airflow.models import TaskInstance
+
 try:
     # todo: remove this conditional import when min airflow version >= 2.6
     ctx_indiv_trigger: ContextVar | None

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

@@ -22,7 +22,6 @@ import ast
 from typing import TYPE_CHECKING, Sequence, Tuple
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.automl_v1beta1 import (
     BatchPredictResult,
     ColumnSpec,
@@ -43,6 +42,8 @@ from airflow.providers.google.cloud.links.automl import (
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
     from airflow.utils.context import Context
 
 MetaData = Sequence[Tuple[str, str]]

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

@@ -26,7 +26,6 @@ from typing import TYPE_CHECKING, Any, Iterable, Sequence, SupportsAbs
 
 import attr
 from google.api_core.exceptions import Conflict
-from google.api_core.retry import Retry
 from google.cloud.bigquery import DEFAULT_RETRY, CopyJob, ExtractJob, LoadJob, QueryJob
 from google.cloud.bigquery.table import RowIterator
 
@@ -56,6 +55,7 @@ from airflow.providers.google.cloud.triggers.bigquery import (
 from airflow.providers.google.cloud.utils.bigquery import convert_job_id
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
     from google.cloud.bigquery import UnknownJob
 
     from airflow.models.taskinstancekey import TaskInstanceKey
@@ -443,6 +443,10 @@ class BigQueryValueCheckOperator(_BigQueryDbHookMixin, SQLValueCheckOperator):
                     method_name="execute_complete",
                 )
             self._handle_job_error(job)
+            # job.result() returns a RowIterator. Mypy expects an instance of SupportsNext[Any] for
+            # the next() call which the RowIterator does not resemble to. Hence, ignore the arg-type error.
+            records = next(job.result())  # type: ignore[arg-type]
+            self.check_value(records)
             self.log.info("Current state of job %s is %s", job.job_id, job.state)
 
     @staticmethod
@@ -1039,7 +1043,7 @@ class BigQueryGetDataOperator(GoogleCloudBaseOperator):
             self.log.info("Total extracted rows: %s", len(rows))
 
             if self.as_dict:
-                table_data = [{k: v for k, v in row.items()} for row in rows]
+                table_data = [dict(row) for row in rows]
             else:
                 table_data = [row.values() for row in rows]
 

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

@@ -23,7 +23,6 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.bigquery_datatransfer_v1 import (
     StartManualTransferRunsResponse,
     TransferConfig,
@@ -39,6 +38,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.cloud.triggers.bigquery_dts import BigQueryDataTransferRunTrigger
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
     from airflow.utils.context import Context
 
 

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

@@ -18,12 +18,9 @@
 """This module contains Google Cloud Bigtable operators."""
 from __future__ import annotations
 
-import enum
 from typing import TYPE_CHECKING, Iterable, Sequence
 
 import google.api_core.exceptions
-from google.cloud.bigtable import enums
-from google.cloud.bigtable.column_family import GarbageCollectionRule
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
@@ -35,6 +32,11 @@ from airflow.providers.google.cloud.links.bigtable import (
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 if TYPE_CHECKING:
+    import enum
+
+    from google.cloud.bigtable import enums
+    from google.cloud.bigtable.column_family import GarbageCollectionRule
+
     from airflow.utils.context import Context
 
 

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

@@ -17,9 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
-from google.api_core import operation  # type: ignore
 from google.cloud.batch_v1 import Job, Task
 
 from airflow.configuration import conf
@@ -27,7 +26,11 @@ from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_batch import CloudBatchJobFinishedTrigger
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from google.api_core import operation
+
+    from airflow.utils.context import Context
 
 
 class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator):

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

@@ -25,7 +25,6 @@ from typing import TYPE_CHECKING, Any, Sequence
 from urllib.parse import unquote, urlsplit
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
 
 from airflow.configuration import conf
@@ -44,6 +43,8 @@ from airflow.utils import yaml
 from airflow.utils.helpers import exactly_one
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
     from airflow.utils.context import Context
 
 

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

@@ -21,10 +21,8 @@ from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.orchestration.airflow.service_v1 import ImageVersion
 from google.cloud.orchestration.airflow.service_v1.types import Environment
-from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow import AirflowException
 from airflow.configuration import conf
@@ -35,6 +33,9 @@ from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposer
 from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+    from google.protobuf.field_mask_pb2 import FieldMask
+
     from airflow.utils.context import Context
 
 CLOUD_COMPOSER_BASE_LINK = "https://console.cloud.google.com/composer/environments"

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

@@ -28,10 +28,8 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.memcache_v1beta2.types import cloud_memcache
 from google.cloud.redis_v1 import FailoverInstanceRequest, InputConfig, Instance, OutputConfig
-from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow.providers.google.cloud.hooks.cloud_memorystore import (
     CloudMemorystoreHook,
@@ -46,6 +44,9 @@ from airflow.providers.google.cloud.links.cloud_memorystore import (
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+    from google.protobuf.field_mask_pb2 import FieldMask
+
     from airflow.utils.context import Context
 
 

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

@@ -19,9 +19,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
 
-from google.api_core import operation
 from google.cloud.run_v2 import Job
-from google.cloud.run_v2.types import Execution
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -30,6 +28,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.cloud.triggers.cloud_run import CloudRunJobFinishedTrigger, RunJobStatus
 
 if TYPE_CHECKING:
+    from google.api_core import operation
+    from google.cloud.run_v2.types import Execution
+
     from airflow.utils.context import Context