diff --git a/providers/google/src/airflow/providers/google/ads/operators/ads.py b/providers/google/src/airflow/providers/google/ads/operators/ads.py index f6aaf1df3fda2..67ffb09b9d4ee 100644 --- a/providers/google/src/airflow/providers/google/ads/operators/ads.py +++ b/providers/google/src/airflow/providers/google/ads/operators/ads.py @@ -29,7 +29,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleAdsListAccountsOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py b/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py index baa99c5b4aa5d..8f7fbd26cf151 100644 --- a/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py +++ b/providers/google/src/airflow/providers/google/ads/transfers/ads_to_gcs.py @@ -27,7 +27,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleAdsToGcsOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py index e05aab2bc9293..cdd628cd44e4b 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py @@ -45,6 +45,7 @@ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector +from airflow.providers.common.compat.sdk import timezone from airflow.providers.google.cloud.utils.helpers import normalize_directory_path from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import ( @@ -52,11 +53,6 @@ GoogleBaseAsyncHook, GoogleBaseHook, ) - -try: - from airflow.sdk import timezone -except ImportError: - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] from airflow.version import version if TYPE_CHECKING: diff --git a/providers/google/src/airflow/providers/google/cloud/links/base.py b/providers/google/src/airflow/providers/google/cloud/links/base.py index abd1e5e4e0b95..d9c220bae1fc6 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/base.py +++ b/providers/google/src/airflow/providers/google/cloud/links/base.py @@ -25,8 +25,8 @@ if TYPE_CHECKING: from airflow.models.taskinstancekey import TaskInstanceKey + from airflow.providers.common.compat.sdk import Context from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator - from airflow.utils.context import Context BASE_LINK = "https://console.cloud.google.com" diff --git a/providers/google/src/airflow/providers/google/cloud/links/cloud_storage_transfer.py b/providers/google/src/airflow/providers/google/cloud/links/cloud_storage_transfer.py index 0e30cb7c8dfc4..e8ccffe7b1cf9 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/cloud_storage_transfer.py +++ b/providers/google/src/airflow/providers/google/cloud/links/cloud_storage_transfer.py @@ -24,7 +24,7 @@ from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context CLOUD_STORAGE_TRANSFER_BASE_LINK = "https://console.cloud.google.com/transfer" diff --git a/providers/google/src/airflow/providers/google/cloud/links/cloud_tasks.py b/providers/google/src/airflow/providers/google/cloud/links/cloud_tasks.py index 44f8330425381..bd885e7fd04f7 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/cloud_tasks.py +++ b/providers/google/src/airflow/providers/google/cloud/links/cloud_tasks.py @@ -24,7 +24,7 @@ from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context CLOUD_TASKS_BASE_LINK = "/cloudtasks" CLOUD_TASKS_QUEUE_LINK = CLOUD_TASKS_BASE_LINK + "/queue/{location}/{queue_id}/tasks?project={project_id}" diff --git a/providers/google/src/airflow/providers/google/cloud/links/dataproc.py b/providers/google/src/airflow/providers/google/cloud/links/dataproc.py index a78157c21d2fb..695c86a8e885d 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/dataproc.py +++ b/providers/google/src/airflow/providers/google/cloud/links/dataproc.py @@ -30,8 +30,8 @@ if TYPE_CHECKING: from airflow.models.taskinstancekey import TaskInstanceKey + from airflow.providers.common.compat.sdk import Context from airflow.providers.google.version_compat import BaseOperator - from airflow.utils.context import Context def __getattr__(name: str) -> Any: diff --git a/providers/google/src/airflow/providers/google/cloud/links/kubernetes_engine.py b/providers/google/src/airflow/providers/google/cloud/links/kubernetes_engine.py index 095bb0474a9a2..671341ee0b510 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/kubernetes_engine.py +++ b/providers/google/src/airflow/providers/google/cloud/links/kubernetes_engine.py @@ -24,7 +24,7 @@ from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context KUBERNETES_BASE_LINK = "/kubernetes" KUBERNETES_CLUSTER_LINK = ( diff --git a/providers/google/src/airflow/providers/google/cloud/links/translate.py b/providers/google/src/airflow/providers/google/cloud/links/translate.py index 29ed09df4e5e8..ba7cd1537dbe9 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/translate.py +++ b/providers/google/src/airflow/providers/google/cloud/links/translate.py @@ -23,7 +23,7 @@ from airflow.providers.google.cloud.links.base import BASE_LINK, BaseGoogleLink if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context TRANSLATION_BASE_LINK = BASE_LINK + "/translation" diff --git a/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py b/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py index f735e07f58fe3..d749dda1e1848 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py +++ b/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py @@ -21,7 +21,7 @@ from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context VERTEX_AI_BASE_LINK = "/vertex-ai" VERTEX_AI_MODEL_LINK = ( diff --git a/providers/google/src/airflow/providers/google/cloud/openlineage/utils.py b/providers/google/src/airflow/providers/google/cloud/openlineage/utils.py index d0012f2301bde..e0461aaa5b09b 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/utils.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/utils.py @@ -49,7 +49,7 @@ from google.cloud.bigquery.table import Table from airflow.providers.common.compat.openlineage.facet import Dataset - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context log = logging.getLogger(__name__) diff --git a/providers/google/src/airflow/providers/google/cloud/operators/alloy_db.py b/providers/google/src/airflow/providers/google/cloud/operators/alloy_db.py index 4c157b5245564..102974b8effff 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/alloy_db.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/alloy_db.py @@ -42,7 +42,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class AlloyDBBaseOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/automl.py b/providers/google/src/airflow/providers/google/cloud/operators/automl.py index b965af7c44c80..905b0acf2c074 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/automl.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/automl.py @@ -50,7 +50,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context MetaData = Sequence[tuple[str, str]] diff --git a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py index a928145db8d57..8dfa3910fa2de 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py @@ -66,7 +66,7 @@ from google.api_core.retry import Retry from google.cloud.bigquery import UnknownJob - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context BIGQUERY_JOB_DETAILS_LINK_FMT = "https://console.cloud.google.com/bigquery?j={job_id}" diff --git a/providers/google/src/airflow/providers/google/cloud/operators/bigquery_dts.py b/providers/google/src/airflow/providers/google/cloud/operators/bigquery_dts.py index b06dea7bc0283..daf7de8e4164b 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/bigquery_dts.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/bigquery_dts.py @@ -43,7 +43,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context def _get_transfer_config_details(config_transfer_name: str): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/bigtable.py b/providers/google/src/airflow/providers/google/cloud/operators/bigtable.py index 8c202f84b778f..be38ea18897af 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/bigtable.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/bigtable.py @@ -40,7 +40,7 @@ from google.cloud.bigtable import enums from google.cloud.bigtable.column_family import GarbageCollectionRule - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigtableValidationMixin: diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py index bc6da9ab502ec..f3b858c63a47f 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_batch.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: from google.api_core import operation - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py index 26f1444df36c2..c599889385ea9 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_build.py @@ -48,7 +48,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context REGEX_REPO_PATH = re.compile(r"^/(?P[^/]+)/(?P[^/]+)[\+/]*(?P[^:]+)?") diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py index 5783dd816ad80..e7053fef8c91f 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_composer.py @@ -41,7 +41,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context CLOUD_COMPOSER_BASE_LINK = "https://console.cloud.google.com/composer/environments" CLOUD_COMPOSER_DETAILS_LINK = ( diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_logging_sink.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_logging_sink.py index b9255c32cbe07..e6d4ec1d9f37a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_logging_sink.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_logging_sink.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context def _validate_inputs(obj, required_fields: list[str]) -> None: diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_memorystore.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_memorystore.py index 4b0e5d44d0b95..52d79fb06569c 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_memorystore.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_memorystore.py @@ -50,7 +50,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudMemorystoreCreateInstanceOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_run.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_run.py index c60d638a2d78a..3793b00905179 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_run.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_run.py @@ -35,7 +35,7 @@ from google.api_core import operation from google.cloud.run_v2.types import Execution - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudRunCreateJobOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_sql.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_sql.py index c7392e45dd879..5ec325e909cec 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_sql.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_sql.py @@ -39,8 +39,8 @@ if TYPE_CHECKING: from airflow.models import Connection + from airflow.providers.common.compat.sdk import Context from airflow.providers.openlineage.extractors import OperatorLineage - from airflow.utils.context import Context SETTINGS = "settings" diff --git a/providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py b/providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py index a30b6cda608c8..73ed9f1e9bf56 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py @@ -72,8 +72,8 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context from airflow.providers.openlineage.extractors import OperatorLineage - from airflow.utils.context import Context class TransferJobPreprocessor: diff --git a/providers/google/src/airflow/providers/google/cloud/operators/compute.py b/providers/google/src/airflow/providers/google/cloud/operators/compute.py index 2ba8874cc3773..a35824b350de4 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/compute.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/compute.py @@ -42,7 +42,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class ComputeEngineBaseOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/datacatalog.py b/providers/google/src/airflow/providers/google/cloud/operators/datacatalog.py index ba6692cb1e802..86dced30cc8d7 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/datacatalog.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/datacatalog.py @@ -47,7 +47,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context @deprecated( diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py b/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py index 88542a8f1e401..467ef6964016e 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataflow.py @@ -43,7 +43,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CheckJobRunning(Enum): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataform.py b/providers/google/src/airflow/providers/google/cloud/operators/dataform.py index df21e1b72dd41..f637bff2418b2 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataform.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataform.py @@ -28,7 +28,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.cloud.dataform_v1beta1.types import ( diff --git a/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py b/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py index cfe5e7965fbf1..6617271637d79 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/datafusion.py @@ -40,8 +40,8 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context from airflow.providers.openlineage.extractors import OperatorLineage - from airflow.utils.context import Context class CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataplex.py b/providers/google/src/airflow/providers/google/cloud/operators/dataplex.py index b72c5f094c3dd..ec979ff08f798 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataplex.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataplex.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context from google.api_core.exceptions import AlreadyExists, GoogleAPICallError, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataprep.py b/providers/google/src/airflow/providers/google/cloud/operators/dataprep.py index 1d501a0b6630a..e1af59d389b02 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataprep.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataprep.py @@ -28,7 +28,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DataprepGetJobsForJobGroupOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py index 3dcf64dbeee2b..a51d2ed31ceb5 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataproc.py @@ -72,7 +72,7 @@ from google.protobuf.field_mask_pb2 import FieldMask from google.type.interval_pb2 import Interval - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class PreemptibilityType(Enum): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dataproc_metastore.py b/providers/google/src/airflow/providers/google/cloud/operators/dataproc_metastore.py index 25743a13a701f..20fb483116bea 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dataproc_metastore.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dataproc_metastore.py @@ -40,8 +40,8 @@ from google.protobuf.field_mask_pb2 import FieldMask from airflow.models.taskinstancekey import TaskInstanceKey + from airflow.providers.common.compat.sdk import Context from airflow.providers.google.version_compat import BaseOperator - from airflow.utils.context import Context BASE_LINK = "https://console.cloud.google.com" METASTORE_BASE_LINK = BASE_LINK + "/dataproc/metastore/services/{region}/{service_id}" diff --git a/providers/google/src/airflow/providers/google/cloud/operators/datastore.py b/providers/google/src/airflow/providers/google/cloud/operators/datastore.py index e598b1de2498c..f60a5e91b6f63 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/datastore.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/datastore.py @@ -34,7 +34,7 @@ from airflow.providers.google.common.links.storage import StorageLink if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudDatastoreExportEntitiesOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/dlp.py b/providers/google/src/airflow/providers/google/cloud/operators/dlp.py index 2b050af59e309..cc612612f0f15 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/dlp.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/dlp.py @@ -66,7 +66,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudDLPCancelDLPJobOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/functions.py b/providers/google/src/airflow/providers/google/cloud/operators/functions.py index f31f93c27dbc4..bea1c8c002805 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/functions.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/functions.py @@ -40,7 +40,7 @@ from airflow.version import version if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context def _validate_available_memory_in_mb(value): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py index 54c01b0d470f0..b36c32dcdd730 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py @@ -33,7 +33,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context from google.api_core.exceptions import Conflict from google.cloud.exceptions import GoogleCloudError diff --git a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py index 762665ccdbdac..66328d9b6d34a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py @@ -39,7 +39,7 @@ TuningDatasetOrDict, ) - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GenAIGenerateEmbeddingsOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py b/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py index b2303374a5256..72dfb5788eac3 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -76,7 +76,7 @@ from kubernetes.client.models import V1Job from pendulum import DateTime - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context KUBE_CONFIG_ENV_VAR = "KUBECONFIG" diff --git a/providers/google/src/airflow/providers/google/cloud/operators/looker.py b/providers/google/src/airflow/providers/google/cloud/operators/looker.py index ad18f2c2e9cb7..7e56ec2a85b61 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/looker.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/looker.py @@ -26,7 +26,7 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LookerStartPdtBuildOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/managed_kafka.py b/providers/google/src/airflow/providers/google/cloud/operators/managed_kafka.py index c4f7d09e5789c..3309792638cd8 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/managed_kafka.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/managed_kafka.py @@ -41,7 +41,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class ManagedKafkaBaseOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py b/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py index 96c15be645873..2fa2c75fed899 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/mlengine.py @@ -31,7 +31,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context log = logging.getLogger(__name__) diff --git a/providers/google/src/airflow/providers/google/cloud/operators/natural_language.py b/providers/google/src/airflow/providers/google/cloud/operators/natural_language.py index edd9a4c2ac820..53607acb0a83c 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/natural_language.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/natural_language.py @@ -32,7 +32,7 @@ from google.api_core.retry import Retry from google.cloud.language_v1.types import Document, EncodingType - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context MetaData = Sequence[tuple[str, str]] diff --git a/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py b/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py index c8c91dfb5575b..250bcfde4ef13 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py @@ -53,8 +53,8 @@ if TYPE_CHECKING: from google.api_core.retry import Retry + from airflow.providers.common.compat.sdk import Context from airflow.providers.openlineage.extractors import OperatorLineage - from airflow.utils.context import Context class PubSubCreateTopicOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/spanner.py b/providers/google/src/airflow/providers/google/cloud/operators/spanner.py index 97d9fe797982d..071a068940765 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/spanner.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/spanner.py @@ -30,8 +30,8 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context from airflow.providers.openlineage.extractors import OperatorLineage - from airflow.utils.context import Context class SpannerDeployInstanceOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/speech_to_text.py b/providers/google/src/airflow/providers/google/cloud/operators/speech_to_text.py index 5f21db6806b64..0527d32a0cfdd 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/speech_to_text.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/speech_to_text.py @@ -35,7 +35,7 @@ from google.api_core.retry import Retry from google.cloud.speech_v1.types import RecognitionConfig - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudSpeechToTextRecognizeSpeechOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/stackdriver.py b/providers/google/src/airflow/providers/google/cloud/operators/stackdriver.py index d61d8d40f6d6c..cdbf71d217f45 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/stackdriver.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/stackdriver.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class StackdriverListAlertPoliciesOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/tasks.py b/providers/google/src/airflow/providers/google/cloud/operators/tasks.py index 4cb1f9ff05c0a..4021264e320bc 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/tasks.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/tasks.py @@ -35,7 +35,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context MetaData = Sequence[tuple[str, str]] diff --git a/providers/google/src/airflow/providers/google/cloud/operators/text_to_speech.py b/providers/google/src/airflow/providers/google/cloud/operators/text_to_speech.py index 35ed1f5d34f76..34aa15b7dd3b7 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/text_to_speech.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/text_to_speech.py @@ -36,7 +36,7 @@ from google.api_core.retry import Retry from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudTextToSpeechSynthesizeOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/translate.py b/providers/google/src/airflow/providers/google/cloud/operators/translate.py index 1a112f71636c3..e107bcd820269 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/translate.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/translate.py @@ -56,7 +56,7 @@ ) from google.cloud.translate_v3.types.translation_service import Glossary, GlossaryInputConfig - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudTranslateTextOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/translate_speech.py b/providers/google/src/airflow/providers/google/cloud/operators/translate_speech.py index 01003e070b9e4..acecca5f2b68a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/translate_speech.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/translate_speech.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudTranslateSpeechOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py index ffbb3dc834ed4..07780a2754034 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py @@ -42,7 +42,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py index 4aad729a9dff4..ffe28d2400853 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py @@ -42,7 +42,7 @@ from google.api_core.retry import Retry from google.cloud.aiplatform import BatchPredictionJob as BatchPredictionJobObject, Model, explain - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py index d63c8ed0e992b..6ab4ceedbfc93 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py @@ -53,7 +53,7 @@ ) from google.cloud.aiplatform_v1.types import PscInterfaceConfig - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/dataset.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/dataset.py index c3e2bc369807c..e92b79a36d53a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/dataset.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/dataset.py @@ -35,7 +35,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CreateDatasetOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py index 9871cdccdc49a..cf73596a9d0d8 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py @@ -39,7 +39,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CreateEndpointOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py index 1d60a3f1b5a5b..dd06cf18ced92 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py @@ -30,7 +30,7 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CreateExperimentOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py index 2fe98b75221d4..f00a7ad12650b 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py @@ -39,7 +39,7 @@ FeatureViewDataKey, ) - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class SyncFeatureViewOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py index 31a4bc0adb826..4b783184320ea 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py @@ -33,7 +33,7 @@ from airflow.providers.google.common.deprecated import deprecated if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context @deprecated( diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py index a667778965bc7..a7dc575669809 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py @@ -43,7 +43,7 @@ from google.api_core.retry import Retry from google.cloud.aiplatform import HyperparameterTuningJob, gapic, hyperparameter_tuning - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py index 2b6459d434986..878629fe3106f 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py @@ -37,7 +37,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DeleteModelOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py index 875d2eff00e50..c4e87311e3b9a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py @@ -42,7 +42,7 @@ from google.cloud.aiplatform import PipelineJob from google.cloud.aiplatform.metadata import experiment_resources - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class RunPipelineJobOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py index 95a284c1defc8..40bc50daeb292 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/ray.py @@ -35,7 +35,7 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class RayBaseOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/video_intelligence.py b/providers/google/src/airflow/providers/google/cloud/operators/video_intelligence.py index 0d4da307f63b8..fef7fd52fea05 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/video_intelligence.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/video_intelligence.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudVideoIntelligenceDetectVideoLabelsOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vision.py b/providers/google/src/airflow/providers/google/cloud/operators/vision.py index d13bbb07ed11c..5af37fdf9c943 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vision.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vision.py @@ -41,7 +41,7 @@ from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context MetaData = Sequence[tuple[str, str]] diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py index b7a08e75d40a6..4849cccd6d04a 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py @@ -34,7 +34,7 @@ ) if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigQueryTableExistenceSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py index f1aa0af8713eb..31837b840be0c 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/bigtable.py b/providers/google/src/airflow/providers/google/cloud/sensors/bigtable.py index 45a4e16bc1284..e9335e83f28d4 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/bigtable.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/bigtable.py @@ -33,7 +33,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValidationMixin): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py index b306457de0dc7..149762937468a 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py @@ -38,7 +38,7 @@ from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudComposerDAGRunSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py index c4fe466999066..e8c60522a2ea4 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py @@ -38,7 +38,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/dataflow.py b/providers/google/src/airflow/providers/google/cloud/sensors/dataflow.py index 84f92f602709a..0ccda0ccedd74 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/dataflow.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/dataflow.py @@ -40,7 +40,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DataflowJobStatusSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/dataform.py b/providers/google/src/airflow/providers/google/cloud/sensors/dataform.py index 65a44c1c993dc..f3ca9cf5ee71c 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/dataform.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/dataform.py @@ -27,7 +27,7 @@ from airflow.providers.google.cloud.hooks.dataform import DataformHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DataformWorkflowInvocationStateSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py b/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py index d04806fb40338..3b43d5eb2840d 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py @@ -28,7 +28,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudDataFusionPipelineStateSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/dataplex.py b/providers/google/src/airflow/providers/google/cloud/sensors/dataplex.py index 918883a9d5af0..b6d3c3e8ae278 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/dataplex.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/dataplex.py @@ -25,7 +25,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context from google.api_core.exceptions import GoogleAPICallError from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/dataprep.py b/providers/google/src/airflow/providers/google/cloud/sensors/dataprep.py index 6b6cebe27631b..4d54aa51c4689 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/dataprep.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/dataprep.py @@ -26,7 +26,7 @@ from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook, JobGroupStatuses if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DataprepJobGroupIsFinishedSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/dataproc.py b/providers/google/src/airflow/providers/google/cloud/sensors/dataproc.py index c56a4d01263fa..ecf9a6863ffbb 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/dataproc.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/dataproc.py @@ -32,7 +32,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DataprocJobSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/dataproc_metastore.py b/providers/google/src/airflow/providers/google/cloud/sensors/dataproc_metastore.py index 3cddb5d4dfda1..4fe5e9e85620d 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/dataproc_metastore.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/dataproc_metastore.py @@ -28,7 +28,7 @@ if TYPE_CHECKING: from google.api_core.operation import Operation - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class MetastoreHivePartitionSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/gcs.py b/providers/google/src/airflow/providers/google/cloud/sensors/gcs.py index d8ab4449c0d3e..d6eb0a3363990 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/gcs.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GCSObjectExistenceSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/looker.py b/providers/google/src/airflow/providers/google/cloud/sensors/looker.py index cdc0cd5ed754c..2c8a4a945f9c7 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/looker.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/looker.py @@ -26,7 +26,7 @@ from airflow.providers.google.cloud.hooks.looker import JobStatus, LookerHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LookerCheckPdtBuildSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py b/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py index a0b286ab31890..e1d6f10d4aa4c 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py @@ -33,7 +33,7 @@ from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class PubSubMessageTransformException(AirflowException): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/tasks.py b/providers/google/src/airflow/providers/google/cloud/sensors/tasks.py index 1c2275270b287..5be817cea1c07 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/tasks.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/tasks.py @@ -27,7 +27,7 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class TaskQueueEmptySensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py b/providers/google/src/airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py index 8f8e6804b5b0a..dacc1827a827a 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py @@ -28,7 +28,7 @@ from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class FeatureViewSyncSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/workflows.py b/providers/google/src/airflow/providers/google/cloud/sensors/workflows.py index 89f4309078c45..64e47cf5c0d1e 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/workflows.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/workflows.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class WorkflowExecutionSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/adls_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/adls_to_gcs.py index a445138e1ff0f..d4040c6802bc8 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/adls_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/adls_to_gcs.py @@ -35,7 +35,7 @@ raise AirflowOptionalProviderFeatureException(e) if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class ADLSToGCSOperator(ADLSListOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py index edbfd61b1a2fd..47248e5c4ab97 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py @@ -32,7 +32,7 @@ raise AirflowOptionalProviderFeatureException(e) if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class AzureBlobStorageToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py index 4b890b6c59d59..4db03d1bc8a32 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py @@ -34,7 +34,7 @@ raise AirflowOptionalProviderFeatureException(e) if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class AzureFileShareToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py index d34c71a7dc0fb..201ca4ba9fcf8 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py @@ -28,7 +28,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigQueryToBigQueryOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py index 20f0d1946f96d..0c6eaf396d52f 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py @@ -37,7 +37,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigQueryToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_mssql.py b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_mssql.py index 38fcf2f68d27f..e1706fd68148c 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_mssql.py @@ -30,7 +30,7 @@ from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py index c298d5bfe9b00..76ae2f277bc87 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py @@ -31,7 +31,7 @@ from airflow.providers.postgres.hooks.postgres import PostgresHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_sql.py b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_sql.py index a76f34eae6753..81218f0aa97e2 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_sql.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/bigquery_to_sql.py @@ -29,9 +29,9 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.openlineage.extractors import OperatorLineage - from airflow.utils.context import Context class BigQueryToSqlBaseOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/cassandra_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/cassandra_to_gcs.py index ce56036de6214..ab14e4670f5cc 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/cassandra_to_gcs.py @@ -36,7 +36,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context NotSetType = NewType("NotSetType", object) NOT_SET = NotSetType(object()) diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py index 8814df2c02da3..e170cd626cf2d 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: from facebook_business.adobjects.adsinsights import AdsInsights - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class FlushAction(Enum): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py index f4aed63542d40..dd1fb8e59ba8a 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py @@ -49,7 +49,7 @@ if TYPE_CHECKING: from google.api_core.retry import Retry - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context ALLOWED_FORMATS = [ "CSV", diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py index 296d216b88135..59dff3842f0a8 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_gcs.py @@ -30,7 +30,7 @@ WILDCARD = "*" if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GCSToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_local.py b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_local.py index 0f58a68c943cc..2856107c729b2 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_local.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_local.py @@ -24,7 +24,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context # MAX XCOM Size is 48KB, check discussion: https://github.com/apache/airflow/pull/1618#discussion_r68249677 MAX_XCOM_SIZE = 49344 diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_sftp.py b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_sftp.py index ccb7e4e485f4e..40db84447b592 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_sftp.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_sftp.py @@ -33,7 +33,7 @@ WILDCARD = "*" if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GCSToSFTPOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py index 827f7455910fe..ee0f474c4206b 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py @@ -24,7 +24,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleDriveToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py index 12d903ea52f08..e7fe7d731eb77 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py @@ -23,7 +23,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleDriveToLocalOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/http_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/http_to_gcs.py index b06ac8d22e561..d203b6187c1cb 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/http_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/http_to_gcs.py @@ -31,7 +31,7 @@ from requests.auth import AuthBase - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class HttpToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/local_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/local_to_gcs.py index b6c183e600597..70134563375f8 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/local_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/local_to_gcs.py @@ -28,7 +28,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LocalFilesystemToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py index ba11be3bd5b9c..bd2a0d55f4a00 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/s3_to_gcs.py @@ -57,7 +57,7 @@ from airflow.providers.amazon.aws.operators.s3_list import S3ListOperator # type: ignore[no-redef] if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class S3ToGCSOperator(S3ListOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/salesforce_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/salesforce_to_gcs.py index 5ffe7f59f32cc..c2c42ec68c172 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/salesforce_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/salesforce_to_gcs.py @@ -26,7 +26,7 @@ from airflow.providers.salesforce.hooks.salesforce import SalesforceHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class SalesforceToGcsOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py index 8bd6af7f99105..f76971822d0a8 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/sftp_to_gcs.py @@ -31,7 +31,7 @@ from airflow.providers.sftp.hooks.sftp import SFTPHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context WILDCARD = "*" diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/sheets_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/sheets_to_gcs.py index 6f11b38531f18..25801a54066c4 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/sheets_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/sheets_to_gcs.py @@ -26,7 +26,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleSheetsToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/sql_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/sql_to_gcs.py index e4ec2d730ef81..5fe3242d83839 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/sql_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/sql_to_gcs.py @@ -35,7 +35,7 @@ if TYPE_CHECKING: from airflow.providers.common.compat.openlineage.facet import OutputDataset - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BaseSQLToGCSOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/firebase/operators/firestore.py b/providers/google/src/airflow/providers/google/firebase/operators/firestore.py index 08ad61c7616eb..5a01a1e905a8b 100644 --- a/providers/google/src/airflow/providers/google/firebase/operators/firestore.py +++ b/providers/google/src/airflow/providers/google/firebase/operators/firestore.py @@ -25,7 +25,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class CloudFirestoreExportDatabaseOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/leveldb/operators/leveldb.py b/providers/google/src/airflow/providers/google/leveldb/operators/leveldb.py index 77320d5e6bd7a..f5f6e935f1325 100644 --- a/providers/google/src/airflow/providers/google/leveldb/operators/leveldb.py +++ b/providers/google/src/airflow/providers/google/leveldb/operators/leveldb.py @@ -22,7 +22,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LevelDBOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/marketing_platform/links/analytics_admin.py b/providers/google/src/airflow/providers/google/marketing_platform/links/analytics_admin.py index 6214b8d267ee8..9d2058c74557c 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/links/analytics_admin.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/links/analytics_admin.py @@ -22,8 +22,8 @@ if TYPE_CHECKING: from airflow.models.taskinstancekey import TaskInstanceKey + from airflow.providers.common.compat.sdk import Context from airflow.providers.google.version_compat import BaseOperator - from airflow.utils.context import Context BASE_LINK = "https://analytics.google.com/analytics/web/" diff --git a/providers/google/src/airflow/providers/google/marketing_platform/operators/analytics_admin.py b/providers/google/src/airflow/providers/google/marketing_platform/operators/analytics_admin.py index 4465a6084aa4d..86395d7ed8bcc 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/operators/analytics_admin.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/operators/analytics_admin.py @@ -39,7 +39,7 @@ from google.api_core.retry import Retry from google.protobuf.message import Message - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleAnalyticsAdminListAccountsOperator(GoogleCloudBaseOperator): diff --git a/providers/google/src/airflow/providers/google/marketing_platform/operators/campaign_manager.py b/providers/google/src/airflow/providers/google/marketing_platform/operators/campaign_manager.py index e68a286284872..61ea484a5c9c2 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/operators/campaign_manager.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/operators/campaign_manager.py @@ -33,7 +33,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleCampaignManagerDeleteReportOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py b/providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py index ad7684422270d..f351fd90f6df9 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/operators/display_video.py @@ -30,7 +30,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleDisplayVideo360CreateSDFDownloadTaskOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/marketing_platform/operators/search_ads.py b/providers/google/src/airflow/providers/google/marketing_platform/operators/search_ads.py index 7b22197596c71..415354efb3c5e 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/operators/search_ads.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/operators/search_ads.py @@ -27,7 +27,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class _GoogleSearchAdsBaseOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/marketing_platform/sensors/campaign_manager.py b/providers/google/src/airflow/providers/google/marketing_platform/sensors/campaign_manager.py index cb63c341a5a7c..7325663fd813a 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/sensors/campaign_manager.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/sensors/campaign_manager.py @@ -26,7 +26,7 @@ from airflow.providers.google.marketing_platform.hooks.campaign_manager import GoogleCampaignManagerHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleCampaignManagerReportSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/marketing_platform/sensors/display_video.py b/providers/google/src/airflow/providers/google/marketing_platform/sensors/display_video.py index 8cfebd2054567..628882e88bca7 100644 --- a/providers/google/src/airflow/providers/google/marketing_platform/sensors/display_video.py +++ b/providers/google/src/airflow/providers/google/marketing_platform/sensors/display_video.py @@ -26,7 +26,7 @@ from airflow.providers.google.marketing_platform.hooks.display_video import GoogleDisplayVideo360Hook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleDisplayVideo360GetSDFDownloadOperationSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/suite/sensors/drive.py b/providers/google/src/airflow/providers/google/suite/sensors/drive.py index 0ec7fe721e762..14b4b1e4db06d 100644 --- a/providers/google/src/airflow/providers/google/suite/sensors/drive.py +++ b/providers/google/src/airflow/providers/google/suite/sensors/drive.py @@ -26,7 +26,7 @@ from airflow.providers.google.suite.hooks.drive import GoogleDriveHook if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class GoogleDriveFileExistenceSensor(BaseSensorOperator): diff --git a/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py b/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py index 06854bca89d7b..9cc72bbeb2915 100644 --- a/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py +++ b/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py @@ -29,7 +29,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context WILDCARD = "*" diff --git a/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py b/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py index aa712d773da27..e5b34d727df4b 100644 --- a/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py +++ b/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py @@ -28,7 +28,7 @@ from airflow.providers.google.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LocalFilesystemToGoogleDriveOperator(BaseOperator): diff --git a/providers/google/src/airflow/providers/google/version_compat.py b/providers/google/src/airflow/providers/google/version_compat.py index 3ab92754486f2..23ef9e614c95e 100644 --- a/providers/google/src/airflow/providers/google/version_compat.py +++ b/providers/google/src/airflow/providers/google/version_compat.py @@ -40,17 +40,11 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: # This is needed for DecoratedOperator compatibility if AIRFLOW_V_3_1_PLUS: from airflow.sdk import ( - BaseHook, BaseOperator, ) else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] from airflow.models import BaseOperator + # Explicitly export these imports to protect them from being removed by linters -__all__ = [ - "AIRFLOW_V_3_0_PLUS", - "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseOperator", -] +__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseOperator"] diff --git a/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py b/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py index 999e32a35ecfd..4f8357abc3aa7 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py @@ -41,6 +41,7 @@ from airflow.exceptions import AirflowException from airflow.providers.common.compat.assets import Asset +from airflow.providers.common.compat.sdk import Context from airflow.providers.google.cloud.hooks.bigquery import ( BigQueryAsyncHook, BigQueryHook, @@ -54,11 +55,6 @@ from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import Context -else: - from airflow.utils.context import Context - pytestmark = pytest.mark.filterwarnings("error::airflow.exceptions.AirflowProviderDeprecationWarning") PROJECT_ID = "bq-project" diff --git a/providers/google/tests/unit/google/cloud/links/test_base_link.py b/providers/google/tests/unit/google/cloud/links/test_base_link.py index 6440eb865391d..cd97a304f4f78 100644 --- a/providers/google/tests/unit/google/cloud/links/test_base_link.py +++ b/providers/google/tests/unit/google/cloud/links/test_base_link.py @@ -22,6 +22,7 @@ import pytest +from airflow.providers.common.compat.sdk import XCom from airflow.providers.google.cloud.links.base import BaseGoogleLink from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator @@ -29,9 +30,7 @@ if AIRFLOW_V_3_0_PLUS: from airflow.sdk.execution_time.comms import XComResult - from airflow.sdk.execution_time.xcom import XCom -else: - from airflow.models.xcom import XCom # type: ignore[no-redef] + TEST_LOCATION = "test-location" TEST_CLUSTER_ID = "test-cluster-id" diff --git a/providers/google/tests/unit/google/cloud/operators/test_workflows.py b/providers/google/tests/unit/google/cloud/operators/test_workflows.py index 286fd25fa5783..6afd7ae245b6a 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_workflows.py +++ b/providers/google/tests/unit/google/cloud/operators/test_workflows.py @@ -25,6 +25,7 @@ from google.protobuf.timestamp_pb2 import Timestamp from airflow.models.dagrun import DagRun +from airflow.providers.common.compat.sdk import Context from airflow.providers.google.cloud.operators.workflows import ( WorkflowsCancelExecutionOperator, WorkflowsCreateExecutionOperator, @@ -40,11 +41,6 @@ from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import Context -else: - from airflow.utils.context import Context - BASE_PATH = "airflow.providers.google.cloud.operators.workflows.{}" LOCATION = "europe-west1" WORKFLOW_ID = "workflow_id"