diff --git a/airflow/providers/apache/beam/hooks/beam.py b/airflow/providers/apache/beam/hooks/beam.py index 0644e02b625f0..e1191073c686c 100644 --- a/airflow/providers/apache/beam/hooks/beam.py +++ b/airflow/providers/apache/beam/hooks/beam.py @@ -16,6 +16,8 @@ # specific language governing permissions and limitations # under the License. """This module contains a Apache Beam Hook.""" +from __future__ import annotations + import json import os import select @@ -24,7 +26,7 @@ import subprocess import textwrap from tempfile import TemporaryDirectory -from typing import Callable, List, Optional +from typing import Callable from airflow.exceptions import AirflowConfigException, AirflowException from airflow.hooks.base import BaseHook @@ -50,7 +52,7 @@ class BeamRunnerType: Twister2Runner = "Twister2Runner" -def beam_options_to_args(options: dict) -> List[str]: +def beam_options_to_args(options: dict) -> list[str]: """ Returns a formatted pipeline options from a dictionary of arguments @@ -65,7 +67,7 @@ def beam_options_to_args(options: dict) -> List[str]: if not options: return [] - args: List[str] = [] + args: list[str] = [] for attr, value in options.items(): if value is None or (isinstance(value, bool) and value): args.append(f"--{attr}") @@ -88,14 +90,14 @@ class BeamCommandRunner(LoggingMixin): def __init__( self, - cmd: List[str], - process_line_callback: Optional[Callable[[str], None]] = None, - working_directory: Optional[str] = None, + cmd: list[str], + process_line_callback: Callable[[str], None] | None = None, + working_directory: str | None = None, ) -> None: super().__init__() self.log.info("Running command: %s", " ".join(shlex.quote(c) for c in cmd)) self.process_line_callback = process_line_callback - self.job_id: Optional[str] = None + self.job_id: str | None = None self._proc = subprocess.Popen( cmd, @@ -173,9 +175,9 @@ def __init__( def _start_pipeline( self, variables: dict, - command_prefix: List[str], - process_line_callback: Optional[Callable[[str], None]] = None, - working_directory: Optional[str] = None, + command_prefix: list[str], + process_line_callback: Callable[[str], None] | None = None, + working_directory: str | None = None, ) -> None: cmd = command_prefix + [ f"--runner={self.runner}", @@ -193,11 +195,11 @@ def start_python_pipeline( self, variables: dict, py_file: str, - py_options: List[str], + py_options: list[str], py_interpreter: str = "python3", - py_requirements: Optional[List[str]] = None, + py_requirements: list[str] | None = None, py_system_site_packages: bool = False, - process_line_callback: Optional[Callable[[str], None]] = None, + process_line_callback: Callable[[str], None] | None = None, ): """ Starts Apache Beam python pipeline. @@ -266,8 +268,8 @@ def start_java_pipeline( self, variables: dict, jar: str, - job_class: Optional[str] = None, - process_line_callback: Optional[Callable[[str], None]] = None, + job_class: str | None = None, + process_line_callback: Callable[[str], None] | None = None, ) -> None: """ Starts Apache Beam Java pipeline. @@ -292,7 +294,7 @@ def start_go_pipeline( self, variables: dict, go_file: str, - process_line_callback: Optional[Callable[[str], None]] = None, + process_line_callback: Callable[[str], None] | None = None, should_init_module: bool = False, ) -> None: """ diff --git a/airflow/providers/apache/beam/operators/beam.py b/airflow/providers/apache/beam/operators/beam.py index 2ee6d180de840..5d29805c32811 100644 --- a/airflow/providers/apache/beam/operators/beam.py +++ b/airflow/providers/apache/beam/operators/beam.py @@ -16,11 +16,13 @@ # specific language governing permissions and limitations # under the License. """This module contains Apache Beam operators.""" +from __future__ import annotations + import copy import tempfile from abc import ABC, ABCMeta from contextlib import ExitStack -from typing import TYPE_CHECKING, Callable, List, Optional, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Callable, Sequence from airflow import AirflowException from airflow.models import BaseOperator @@ -47,17 +49,17 @@ class BeamDataflowMixin(metaclass=ABCMeta): :class:`~airflow.providers.apache.beam.operators.beam.BeamRunGoPipelineOperator`. """ - dataflow_hook: Optional[DataflowHook] + dataflow_hook: DataflowHook | None dataflow_config: DataflowConfiguration gcp_conn_id: str - delegate_to: Optional[str] + delegate_to: str | None dataflow_support_impersonation: bool = True def _set_dataflow( self, pipeline_options: dict, - job_name_variable_key: Optional[str] = None, - ) -> Tuple[str, dict, Callable[[str], None]]: + job_name_variable_key: str | None = None, + ) -> tuple[str, dict, Callable[[str], None]]: self.dataflow_hook = self.__set_dataflow_hook() self.dataflow_config.project_id = self.dataflow_config.project_id or self.dataflow_hook.project_id dataflow_job_name = self.__get_dataflow_job_name() @@ -85,7 +87,7 @@ def __get_dataflow_job_name(self) -> str: ) def __get_dataflow_pipeline_options( - self, pipeline_options: dict, job_name: str, job_name_key: Optional[str] = None + self, pipeline_options: dict, job_name: str, job_name_key: str | None = None ) -> dict: pipeline_options = copy.deepcopy(pipeline_options) if job_name_key is not None: @@ -151,11 +153,11 @@ def __init__( self, *, runner: str = "DirectRunner", - default_pipeline_options: Optional[dict] = None, - pipeline_options: Optional[dict] = None, + default_pipeline_options: dict | None = None, + pipeline_options: dict | None = None, gcp_conn_id: str = "google_cloud_default", - delegate_to: Optional[str] = None, - dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None, + delegate_to: str | None = None, + dataflow_config: DataflowConfiguration | dict | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -168,9 +170,9 @@ def __init__( self.dataflow_config = DataflowConfiguration(**dataflow_config) else: self.dataflow_config = dataflow_config or DataflowConfiguration() - self.beam_hook: Optional[BeamHook] = None - self.dataflow_hook: Optional[DataflowHook] = None - self.dataflow_job_id: Optional[str] = None + self.beam_hook: BeamHook | None = None + self.dataflow_hook: DataflowHook | None = None + self.dataflow_job_id: str | None = None if self.dataflow_config and self.runner.lower() != BeamRunnerType.DataflowRunner.lower(): self.log.warning( @@ -180,13 +182,13 @@ def __init__( def _init_pipeline_options( self, format_pipeline_options: bool = False, - job_name_variable_key: Optional[str] = None, - ) -> Tuple[bool, Optional[str], dict, Optional[Callable[[str], None]]]: + job_name_variable_key: str | None = None, + ) -> tuple[bool, str | None, dict, Callable[[str], None] | None]: self.beam_hook = BeamHook(runner=self.runner) pipeline_options = self.default_pipeline_options.copy() - process_line_callback: Optional[Callable[[str], None]] = None + process_line_callback: Callable[[str], None] | None = None is_dataflow = self.runner.lower() == BeamRunnerType.DataflowRunner.lower() - dataflow_job_name: Optional[str] = None + dataflow_job_name: str | None = None if is_dataflow: dataflow_job_name, pipeline_options, process_line_callback = self._set_dataflow( pipeline_options=pipeline_options, @@ -255,15 +257,15 @@ def __init__( *, py_file: str, runner: str = "DirectRunner", - default_pipeline_options: Optional[dict] = None, - pipeline_options: Optional[dict] = None, + default_pipeline_options: dict | None = None, + pipeline_options: dict | None = None, py_interpreter: str = "python3", - py_options: Optional[List[str]] = None, - py_requirements: Optional[List[str]] = None, + py_options: list[str] | None = None, + py_requirements: list[str] | None = None, py_system_site_packages: bool = False, gcp_conn_id: str = "google_cloud_default", - delegate_to: Optional[str] = None, - dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None, + delegate_to: str | None = None, + dataflow_config: DataflowConfiguration | dict | None = None, **kwargs, ) -> None: super().__init__( @@ -285,7 +287,7 @@ def __init__( {"airflow-version": "v" + version.replace(".", "-").replace("+", "-")} ) - def execute(self, context: 'Context'): + def execute(self, context: Context): """Execute the Apache Beam Pipeline.""" ( is_dataflow, @@ -396,12 +398,12 @@ def __init__( *, jar: str, runner: str = "DirectRunner", - job_class: Optional[str] = None, - default_pipeline_options: Optional[dict] = None, - pipeline_options: Optional[dict] = None, + job_class: str | None = None, + default_pipeline_options: dict | None = None, + pipeline_options: dict | None = None, gcp_conn_id: str = "google_cloud_default", - delegate_to: Optional[str] = None, - dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None, + delegate_to: str | None = None, + dataflow_config: DataflowConfiguration | dict | None = None, **kwargs, ) -> None: super().__init__( @@ -416,7 +418,7 @@ def __init__( self.jar = jar self.job_class = job_class - def execute(self, context: 'Context'): + def execute(self, context: Context): """Execute the Apache Beam Pipeline.""" ( is_dataflow, @@ -537,11 +539,11 @@ def __init__( *, go_file: str, runner: str = "DirectRunner", - default_pipeline_options: Optional[dict] = None, - pipeline_options: Optional[dict] = None, + default_pipeline_options: dict | None = None, + pipeline_options: dict | None = None, gcp_conn_id: str = "google_cloud_default", - delegate_to: Optional[str] = None, - dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None, + delegate_to: str | None = None, + dataflow_config: DataflowConfiguration | dict | None = None, **kwargs, ) -> None: super().__init__( @@ -567,7 +569,7 @@ def __init__( {"airflow-version": "v" + version.replace(".", "-").replace("+", "-")} ) - def execute(self, context: 'Context'): + def execute(self, context: Context): """Execute the Apache Beam Pipeline.""" ( is_dataflow, diff --git a/airflow/providers/apache/cassandra/hooks/cassandra.py b/airflow/providers/apache/cassandra/hooks/cassandra.py index 71c360ec472b8..c0d3484f182d2 100644 --- a/airflow/providers/apache/cassandra/hooks/cassandra.py +++ b/airflow/providers/apache/cassandra/hooks/cassandra.py @@ -15,10 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains hook to integrate with Apache Cassandra.""" +from __future__ import annotations -from typing import Any, Dict, Union +from typing import Any, Union from cassandra.auth import PlainTextAuthProvider from cassandra.cluster import Cluster, Session @@ -141,7 +141,7 @@ def shutdown_cluster(self) -> None: self.cluster.shutdown() @staticmethod - def get_lb_policy(policy_name: str, policy_args: Dict[str, Any]) -> Policy: + def get_lb_policy(policy_name: str, policy_args: dict[str, Any]) -> Policy: """ Creates load balancing policy. @@ -188,7 +188,7 @@ def table_exists(self, table: str) -> bool: cluster_metadata = self.get_conn().cluster.metadata return keyspace in cluster_metadata.keyspaces and table in cluster_metadata.keyspaces[keyspace].tables - def record_exists(self, table: str, keys: Dict[str, str]) -> bool: + def record_exists(self, table: str, keys: dict[str, str]) -> bool: """ Checks if a record exists in Cassandra diff --git a/airflow/providers/apache/cassandra/sensors/record.py b/airflow/providers/apache/cassandra/sensors/record.py index f0a407297adfd..f90283d61e8ed 100644 --- a/airflow/providers/apache/cassandra/sensors/record.py +++ b/airflow/providers/apache/cassandra/sensors/record.py @@ -19,8 +19,9 @@ This module contains sensor that check the existence of a record in a Cassandra cluster. """ +from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook from airflow.sensors.base import BaseSensorOperator @@ -58,7 +59,7 @@ class CassandraRecordSensor(BaseSensorOperator): def __init__( self, *, - keys: Dict[str, str], + keys: dict[str, str], table: str, cassandra_conn_id: str = CassandraHook.default_conn_name, **kwargs: Any, @@ -68,7 +69,7 @@ def __init__( self.table = table self.keys = keys - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: self.log.info('Sensor check existence of record: %s', self.keys) hook = CassandraHook(self.cassandra_conn_id) return hook.record_exists(self.table, self.keys) diff --git a/airflow/providers/apache/cassandra/sensors/table.py b/airflow/providers/apache/cassandra/sensors/table.py index 2f5e6681cb4b9..83efd17962eb7 100644 --- a/airflow/providers/apache/cassandra/sensors/table.py +++ b/airflow/providers/apache/cassandra/sensors/table.py @@ -15,11 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """ This module contains sensor that check the existence of a table in a Cassandra cluster. """ +from __future__ import annotations from typing import TYPE_CHECKING, Any, Sequence @@ -65,7 +65,7 @@ def __init__( self.cassandra_conn_id = cassandra_conn_id self.table = table - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: self.log.info('Sensor check existence of table: %s', self.table) hook = CassandraHook(self.cassandra_conn_id) return hook.table_exists(self.table) diff --git a/airflow/providers/apache/drill/hooks/drill.py b/airflow/providers/apache/drill/hooks/drill.py index 218b31bf18ab0..e7400f23e5122 100644 --- a/airflow/providers/apache/drill/hooks/drill.py +++ b/airflow/providers/apache/drill/hooks/drill.py @@ -15,8 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations -from typing import Any, Iterable, Optional, Tuple +from typing import Any, Iterable from sqlalchemy import create_engine from sqlalchemy.engine import Connection @@ -80,8 +81,8 @@ def set_autocommit(self, conn: Connection, autocommit: bool) -> NotImplementedEr def insert_rows( self, table: str, - rows: Iterable[Tuple[str]], - target_fields: Optional[Iterable[str]] = None, + rows: Iterable[tuple[str]], + target_fields: Iterable[str] | None = None, commit_every: int = 1000, replace: bool = False, **kwargs: Any, diff --git a/airflow/providers/apache/drill/operators/drill.py b/airflow/providers/apache/drill/operators/drill.py index 6dad45cc3c323..1aad91967b136 100644 --- a/airflow/providers/apache/drill/operators/drill.py +++ b/airflow/providers/apache/drill/operators/drill.py @@ -15,7 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import TYPE_CHECKING, Iterable, Mapping, Optional, Sequence, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterable, Mapping, Sequence from airflow.models import BaseOperator from airflow.providers.apache.drill.hooks.drill import DrillHook @@ -50,16 +52,16 @@ def __init__( *, sql: str, drill_conn_id: str = 'drill_default', - parameters: Optional[Union[Iterable, Mapping]] = None, + parameters: Iterable | Mapping | None = None, **kwargs, ) -> None: super().__init__(**kwargs) self.sql = sql self.drill_conn_id = drill_conn_id self.parameters = parameters - self.hook: Optional[DrillHook] = None + self.hook: DrillHook | None = None - def execute(self, context: 'Context'): + def execute(self, context: Context): self.log.info('Executing: %s on %s', self.sql, self.drill_conn_id) self.hook = DrillHook(drill_conn_id=self.drill_conn_id) self.hook.run(self.sql, parameters=self.parameters, split_statements=True) diff --git a/airflow/providers/apache/druid/hooks/druid.py b/airflow/providers/apache/druid/hooks/druid.py index 2415bc4bc1262..a5e8ba86540c0 100644 --- a/airflow/providers/apache/druid/hooks/druid.py +++ b/airflow/providers/apache/druid/hooks/druid.py @@ -15,9 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import time -from typing import Any, Dict, Iterable, Optional, Tuple, Union +from typing import Any, Iterable import requests from pydruid.db import connect @@ -46,7 +47,7 @@ def __init__( self, druid_ingest_conn_id: str = 'druid_ingest_default', timeout: int = 1, - max_ingestion_time: Optional[int] = None, + max_ingestion_time: int | None = None, ) -> None: super().__init__() @@ -67,7 +68,7 @@ def get_conn_url(self) -> str: endpoint = conn.extra_dejson.get('endpoint', '') return f"{conn_type}://{host}:{port}/{endpoint}" - def get_auth(self) -> Optional[requests.auth.HTTPBasicAuth]: + def get_auth(self) -> requests.auth.HTTPBasicAuth | None: """ Return username and password from connections tab as requests.auth.HTTPBasicAuth object. @@ -81,7 +82,7 @@ def get_auth(self) -> Optional[requests.auth.HTTPBasicAuth]: else: return None - def submit_indexing_job(self, json_index_spec: Union[Dict[str, Any], str]) -> None: + def submit_indexing_job(self, json_index_spec: dict[str, Any] | str) -> None: """Submit Druid ingestion job""" url = self.get_conn_url() @@ -173,8 +174,8 @@ def set_autocommit(self, conn: connect, autocommit: bool) -> NotImplementedError def insert_rows( self, table: str, - rows: Iterable[Tuple[str]], - target_fields: Optional[Iterable[str]] = None, + rows: Iterable[tuple[str]], + target_fields: Iterable[str] | None = None, commit_every: int = 1000, replace: bool = False, **kwargs: Any, diff --git a/airflow/providers/apache/druid/operators/druid.py b/airflow/providers/apache/druid/operators/druid.py index 9fd8b3595af5c..2424db7a99fa7 100644 --- a/airflow/providers/apache/druid/operators/druid.py +++ b/airflow/providers/apache/druid/operators/druid.py @@ -15,8 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations -from typing import TYPE_CHECKING, Any, Optional, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.models import BaseOperator from airflow.providers.apache.druid.hooks.druid import DruidHook @@ -47,7 +48,7 @@ def __init__( json_index_file: str, druid_ingest_conn_id: str = 'druid_ingest_default', timeout: int = 1, - max_ingestion_time: Optional[int] = None, + max_ingestion_time: int | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -56,7 +57,7 @@ def __init__( self.timeout = timeout self.max_ingestion_time = max_ingestion_time - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: hook = DruidHook( druid_ingest_conn_id=self.conn_id, timeout=self.timeout, diff --git a/airflow/providers/apache/druid/operators/druid_check.py b/airflow/providers/apache/druid/operators/druid_check.py index 75f2b17d0c78f..16c748b24ae44 100644 --- a/airflow/providers/apache/druid/operators/druid_check.py +++ b/airflow/providers/apache/druid/operators/druid_check.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import warnings from airflow.providers.common.sql.operators.sql import SQLCheckOperator diff --git a/airflow/providers/apache/druid/transfers/hive_to_druid.py b/airflow/providers/apache/druid/transfers/hive_to_druid.py index dc5f74109acf1..ea76dbda5141e 100644 --- a/airflow/providers/apache/druid/transfers/hive_to_druid.py +++ b/airflow/providers/apache/druid/transfers/hive_to_druid.py @@ -15,10 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains operator to move data from Hive to Druid.""" +from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.models import BaseOperator from airflow.providers.apache.druid.hooks.druid import DruidHook @@ -74,18 +74,18 @@ def __init__( sql: str, druid_datasource: str, ts_dim: str, - metric_spec: Optional[List[Any]] = None, + metric_spec: list[Any] | None = None, hive_cli_conn_id: str = 'hive_cli_default', druid_ingest_conn_id: str = 'druid_ingest_default', metastore_conn_id: str = 'metastore_default', - hadoop_dependency_coordinates: Optional[List[str]] = None, - intervals: Optional[List[Any]] = None, + hadoop_dependency_coordinates: list[str] | None = None, + intervals: list[Any] | None = None, num_shards: float = -1, target_partition_size: int = -1, query_granularity: str = "NONE", segment_granularity: str = "DAY", - hive_tblproperties: Optional[Dict[Any, Any]] = None, - job_properties: Optional[Dict[Any, Any]] = None, + hive_tblproperties: dict[Any, Any] | None = None, + job_properties: dict[Any, Any] | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -105,7 +105,7 @@ def __init__( self.hive_tblproperties = hive_tblproperties or {} self.job_properties = job_properties - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: hive = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id) self.log.info("Extracting data from Hive") hive_table = 'druid.' + context['task_instance_key_str'].replace('.', '_') @@ -152,7 +152,7 @@ def execute(self, context: "Context") -> None: hql = f"DROP TABLE IF EXISTS {hive_table}" hive.run_cli(hql) - def construct_ingest_query(self, static_path: str, columns: List[str]) -> Dict[str, Any]: + def construct_ingest_query(self, static_path: str, columns: list[str]) -> dict[str, Any]: """ Builds an ingest query for an HDFS TSV load. @@ -176,7 +176,7 @@ def construct_ingest_query(self, static_path: str, columns: List[str]) -> Dict[s # or a metric, as the dimension columns dimensions = [c for c in columns if c not in metric_names and c != self.ts_dim] - ingest_query_dict: Dict[str, Any] = { + ingest_query_dict: dict[str, Any] = { "type": "index_hadoop", "spec": { "dataSchema": { diff --git a/airflow/providers/apache/hdfs/hooks/hdfs.py b/airflow/providers/apache/hdfs/hooks/hdfs.py index 6d1ce3d010125..1041ef2423ea9 100644 --- a/airflow/providers/apache/hdfs/hooks/hdfs.py +++ b/airflow/providers/apache/hdfs/hooks/hdfs.py @@ -16,7 +16,9 @@ # specific language governing permissions and limitations # under the License. """Hook for HDFS operations""" -from typing import Any, Optional +from __future__ import annotations + +from typing import Any from airflow.configuration import conf from airflow.exceptions import AirflowException @@ -49,7 +51,7 @@ class HDFSHook(BaseHook): hook_name = 'HDFS' def __init__( - self, hdfs_conn_id: str = 'hdfs_default', proxy_user: Optional[str] = None, autoconfig: bool = False + self, hdfs_conn_id: str = 'hdfs_default', proxy_user: str | None = None, autoconfig: bool = False ): super().__init__() if not snakebite_loaded: diff --git a/airflow/providers/apache/hdfs/hooks/webhdfs.py b/airflow/providers/apache/hdfs/hooks/webhdfs.py index dc27a80237d58..c791d6c5acd94 100644 --- a/airflow/providers/apache/hdfs/hooks/webhdfs.py +++ b/airflow/providers/apache/hdfs/hooks/webhdfs.py @@ -16,9 +16,11 @@ # specific language governing permissions and limitations # under the License. """Hook for Web HDFS""" +from __future__ import annotations + import logging import socket -from typing import Any, Optional +from typing import Any import requests from hdfs import HdfsError, InsecureClient @@ -50,7 +52,7 @@ class WebHDFSHook(BaseHook): :param proxy_user: The user used to authenticate. """ - def __init__(self, webhdfs_conn_id: str = 'webhdfs_default', proxy_user: Optional[str] = None): + def __init__(self, webhdfs_conn_id: str = 'webhdfs_default', proxy_user: str | None = None): super().__init__() self.webhdfs_conn_id = webhdfs_conn_id self.proxy_user = proxy_user @@ -95,7 +97,7 @@ def _find_valid_server(self) -> Any: return None def _get_client( - self, namenode: str, port: int, login: str, password: Optional[str], schema: str, extra_dejson: dict + self, namenode: str, port: int, login: str, password: str | None, schema: str, extra_dejson: dict ) -> Any: connection_str = f'http://{namenode}' session = requests.Session() diff --git a/airflow/providers/apache/hdfs/sensors/hdfs.py b/airflow/providers/apache/hdfs/sensors/hdfs.py index a445bb688e722..02bbe1089f92c 100644 --- a/airflow/providers/apache/hdfs/sensors/hdfs.py +++ b/airflow/providers/apache/hdfs/sensors/hdfs.py @@ -15,10 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import logging import re import sys -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Pattern, Sequence, Type +from typing import TYPE_CHECKING, Any, Pattern, Sequence from airflow import settings from airflow.providers.apache.hdfs.hooks.hdfs import HDFSHook @@ -53,10 +55,10 @@ def __init__( *, filepath: str, hdfs_conn_id: str = 'hdfs_default', - ignored_ext: Optional[List[str]] = None, + ignored_ext: list[str] | None = None, ignore_copying: bool = True, - file_size: Optional[int] = None, - hook: Type[HDFSHook] = HDFSHook, + file_size: int | None = None, + hook: type[HDFSHook] = HDFSHook, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -70,7 +72,7 @@ def __init__( self.hook = hook @staticmethod - def filter_for_filesize(result: List[Dict[Any, Any]], size: Optional[int] = None) -> List[Dict[Any, Any]]: + def filter_for_filesize(result: list[dict[Any, Any]], size: int | None = None) -> list[dict[Any, Any]]: """ Will test the filepath result and test if its size is at least self.filesize @@ -87,8 +89,8 @@ def filter_for_filesize(result: List[Dict[Any, Any]], size: Optional[int] = None @staticmethod def filter_for_ignored_ext( - result: List[Dict[Any, Any]], ignored_ext: List[str], ignore_copying: bool - ) -> List[Dict[Any, Any]]: + result: list[dict[Any, Any]], ignored_ext: list[str], ignore_copying: bool + ) -> list[dict[Any, Any]]: """ Will filter if instructed to do so the result to remove matching criteria @@ -110,7 +112,7 @@ def filter_for_ignored_ext( log.debug('HdfsSensor.poke: after ext filter result is %s', result) return result - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: """Get a snakebite client connection and check for file.""" sb_client = self.hook(self.hdfs_conn_id).get_conn() self.log.info('Poking for file %s', self.filepath) @@ -144,7 +146,7 @@ def __init__(self, regex: Pattern[str], *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.regex = regex - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: """ Poke matching files in a directory with self.regex @@ -177,7 +179,7 @@ def __init__(self, be_empty: bool = False, *args: Any, **kwargs: Any): super().__init__(*args, **kwargs) self.be_empty = be_empty - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: """ Poke for a non empty directory diff --git a/airflow/providers/apache/hdfs/sensors/web_hdfs.py b/airflow/providers/apache/hdfs/sensors/web_hdfs.py index adacdefecad07..cf2c516c10ad3 100644 --- a/airflow/providers/apache/hdfs/sensors/web_hdfs.py +++ b/airflow/providers/apache/hdfs/sensors/web_hdfs.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + from typing import TYPE_CHECKING, Any, Sequence from airflow.sensors.base import BaseSensorOperator @@ -33,7 +35,7 @@ def __init__(self, *, filepath: str, webhdfs_conn_id: str = 'webhdfs_default', * self.filepath = filepath self.webhdfs_conn_id = webhdfs_conn_id - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: from airflow.providers.apache.hdfs.hooks.webhdfs import WebHDFSHook hook = WebHDFSHook(self.webhdfs_conn_id) diff --git a/airflow/providers/apache/hive/hooks/hive.py b/airflow/providers/apache/hive/hooks/hive.py index f5e8c24d65705..d700dbaf9d4fd 100644 --- a/airflow/providers/apache/hive/hooks/hive.py +++ b/airflow/providers/apache/hive/hooks/hive.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import contextlib import os import re @@ -24,7 +26,7 @@ import warnings from collections import OrderedDict from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import Any, Dict, Iterable, List, Mapping, Optional, Union +from typing import Any, Iterable, Mapping import pandas import unicodecsv as csv @@ -40,7 +42,7 @@ HIVE_QUEUE_PRIORITIES = ['VERY_HIGH', 'HIGH', 'NORMAL', 'LOW', 'VERY_LOW'] -def get_context_from_env_var() -> Dict[Any, Any]: +def get_context_from_env_var() -> dict[Any, Any]: """ Extract context from env variable, e.g. dag_id, task_id and execution_date, so that they can be used inside BashOperator and PythonOperator. @@ -86,10 +88,10 @@ class HiveCliHook(BaseHook): def __init__( self, hive_cli_conn_id: str = default_conn_name, - run_as: Optional[str] = None, - mapred_queue: Optional[str] = None, - mapred_queue_priority: Optional[str] = None, - mapred_job_name: Optional[str] = None, + run_as: str | None = None, + mapred_queue: str | None = None, + mapred_queue_priority: str | None = None, + mapred_job_name: str | None = None, ) -> None: super().__init__() conn = self.get_connection(hive_cli_conn_id) @@ -124,7 +126,7 @@ def _get_proxy_user(self) -> str: return f"hive.server2.proxy.user={proxy_user_value}" return proxy_user_value # The default proxy user (undefined) - def _prepare_cli_cmd(self) -> List[Any]: + def _prepare_cli_cmd(self) -> list[Any]: """This function creates the command list from available information""" conn = self.conn hive_bin = 'hive' @@ -157,7 +159,7 @@ def _prepare_cli_cmd(self) -> List[Any]: return [hive_bin] + cmd_extra + hive_params_list @staticmethod - def _prepare_hiveconf(d: Dict[Any, Any]) -> List[Any]: + def _prepare_hiveconf(d: dict[Any, Any]) -> list[Any]: """ This function prepares a list of hiveconf params from a dictionary of key value pairs. @@ -178,9 +180,9 @@ def _prepare_hiveconf(d: Dict[Any, Any]) -> List[Any]: def run_cli( self, hql: str, - schema: Optional[str] = None, + schema: str | None = None, verbose: bool = True, - hive_conf: Optional[Dict[Any, Any]] = None, + hive_conf: dict[Any, Any] | None = None, ) -> Any: """ Run an hql statement using the hive cli. If hive_conf is specified @@ -302,7 +304,7 @@ def load_df( self, df: pandas.DataFrame, table: str, - field_dict: Optional[Dict[Any, Any]] = None, + field_dict: dict[Any, Any] | None = None, delimiter: str = ',', encoding: str = 'utf8', pandas_kwargs: Any = None, @@ -325,7 +327,7 @@ def load_df( :param kwargs: passed to self.load_file """ - def _infer_field_types_from_df(df: pandas.DataFrame) -> Dict[Any, Any]: + def _infer_field_types_from_df(df: pandas.DataFrame) -> dict[Any, Any]: dtype_kind_hive_type = { 'b': 'BOOLEAN', # boolean 'i': 'BIGINT', # signed integer @@ -372,12 +374,12 @@ def load_file( filepath: str, table: str, delimiter: str = ",", - field_dict: Optional[Dict[Any, Any]] = None, + field_dict: dict[Any, Any] | None = None, create: bool = True, overwrite: bool = True, - partition: Optional[Dict[str, Any]] = None, + partition: dict[str, Any] | None = None, recreate: bool = False, - tblproperties: Optional[Dict[str, Any]] = None, + tblproperties: dict[str, Any] | None = None, ) -> None: """ Loads a local file into Hive @@ -470,14 +472,14 @@ def __init__(self, metastore_conn_id: str = default_conn_name) -> None: self.conn = self.get_connection(metastore_conn_id) self.metastore = self.get_metastore_client() - def __getstate__(self) -> Dict[str, Any]: + def __getstate__(self) -> dict[str, Any]: # This is for pickling to work despite the thrift hive client not # being picklable state = dict(self.__dict__) del state['metastore'] return state - def __setstate__(self, d: Dict[str, Any]) -> None: + def __setstate__(self, d: dict[str, Any]) -> None: self.__dict__.update(d) self.__dict__['metastore'] = self.get_metastore_client() @@ -617,9 +619,7 @@ def get_databases(self, pattern: str = '*') -> Any: with self.metastore as client: return client.get_databases(pattern) - def get_partitions( - self, schema: str, table_name: str, partition_filter: Optional[str] = None - ) -> List[Any]: + def get_partitions(self, schema: str, table_name: str, partition_filter: str | None = None) -> list[Any]: """ Returns a list of all partitions in a table. Works only for tables with less than 32767 (java short max val). @@ -655,7 +655,7 @@ def get_partitions( @staticmethod def _get_max_partition_from_part_specs( - part_specs: List[Any], partition_key: Optional[str], filter_map: Optional[Dict[str, Any]] + part_specs: list[Any], partition_key: str | None, filter_map: dict[str, Any] | None ) -> Any: """ Helper method to get max partition of partitions with partition_key @@ -701,8 +701,8 @@ def max_partition( self, schema: str, table_name: str, - field: Optional[str] = None, - filter_map: Optional[Dict[Any, Any]] = None, + field: str | None = None, + filter_map: dict[Any, Any] | None = None, ) -> Any: """ Returns the maximum value for all partitions with given field in a table. @@ -806,10 +806,10 @@ class HiveServer2Hook(DbApiHook): hook_name = 'Hive Server 2 Thrift' supports_autocommit = False - def get_conn(self, schema: Optional[str] = None) -> Any: + def get_conn(self, schema: str | None = None) -> Any: """Returns a Hive connection object.""" - username: Optional[str] = None - password: Optional[str] = None + username: str | None = None + password: str | None = None db = self.get_connection(self.hiveserver2_conn_id) # type: ignore @@ -857,10 +857,10 @@ def get_conn(self, schema: Optional[str] = None) -> Any: def _get_results( self, - sql: Union[str, List[str]], + sql: str | list[str], schema: str = 'default', - fetch_size: Optional[int] = None, - hive_conf: Optional[Union[Iterable, Mapping]] = None, + fetch_size: int | None = None, + hive_conf: Iterable | Mapping | None = None, ) -> Any: from pyhive.exc import ProgrammingError @@ -911,11 +911,11 @@ def _get_results( def get_results( self, - sql: Union[str, List[str]], + sql: str | list[str], schema: str = 'default', - fetch_size: Optional[int] = None, - hive_conf: Optional[Union[Iterable, Mapping]] = None, - ) -> Dict[str, Any]: + fetch_size: int | None = None, + hive_conf: Iterable | Mapping | None = None, + ) -> dict[str, Any]: """ Get results of the provided hql in target schema. @@ -940,7 +940,7 @@ def to_csv( lineterminator: str = '\r\n', output_header: bool = True, fetch_size: int = 1000, - hive_conf: Optional[Dict[Any, Any]] = None, + hive_conf: dict[Any, Any] | None = None, ) -> None: """ Execute hql in target schema and write results to a csv file. @@ -982,7 +982,7 @@ def to_csv( self.log.info("Done. Loaded a total of %s rows.", i) def get_records( - self, sql: Union[str, List[str]], parameters: Optional[Union[Iterable, Mapping]] = None, **kwargs + self, sql: str | list[str], parameters: Iterable | Mapping | None = None, **kwargs ) -> Any: """ Get a set of records from a Hive query. You can optionally pass 'schema' kwarg @@ -1005,7 +1005,7 @@ def get_pandas_df( # type: ignore self, sql: str, schema: str = 'default', - hive_conf: Optional[Dict[Any, Any]] = None, + hive_conf: dict[Any, Any] | None = None, **kwargs, ) -> pandas.DataFrame: """ diff --git a/airflow/providers/apache/hive/operators/hive.py b/airflow/providers/apache/hive/operators/hive.py index 45cae0fa4e31f..6ab3bdd296796 100644 --- a/airflow/providers/apache/hive/operators/hive.py +++ b/airflow/providers/apache/hive/operators/hive.py @@ -15,9 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import os import re -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.configuration import conf from airflow.models import BaseOperator @@ -78,13 +80,13 @@ def __init__( hql: str, hive_cli_conn_id: str = 'hive_cli_default', schema: str = 'default', - hiveconfs: Optional[Dict[Any, Any]] = None, + hiveconfs: dict[Any, Any] | None = None, hiveconf_jinja_translate: bool = False, - script_begin_tag: Optional[str] = None, + script_begin_tag: str | None = None, run_as_owner: bool = False, - mapred_queue: Optional[str] = None, - mapred_queue_priority: Optional[str] = None, - mapred_job_name: Optional[str] = None, + mapred_queue: str | None = None, + mapred_queue_priority: str | None = None, + mapred_job_name: str | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -114,7 +116,7 @@ def __init__( # `None` initial value, later it will be populated by the execute method. # This also makes `on_kill` implementation consistent since it assumes `self.hook` # is defined. - self.hook: Optional[HiveCliHook] = None + self.hook: HiveCliHook | None = None def get_hook(self) -> HiveCliHook: """Get Hive cli hook""" @@ -132,7 +134,7 @@ def prepare_template(self) -> None: if self.script_begin_tag and self.script_begin_tag in self.hql: self.hql = "\n".join(self.hql.split(self.script_begin_tag)[1:]) - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: self.log.info('Executing: %s', self.hql) self.hook = self.get_hook() diff --git a/airflow/providers/apache/hive/operators/hive_stats.py b/airflow/providers/apache/hive/operators/hive_stats.py index a82040e8466e7..2659554f2785c 100644 --- a/airflow/providers/apache/hive/operators/hive_stats.py +++ b/airflow/providers/apache/hive/operators/hive_stats.py @@ -15,10 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import json import warnings from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence +from typing import TYPE_CHECKING, Any, Callable, List, Sequence from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -66,9 +68,9 @@ def __init__( *, table: str, partition: Any, - extra_exprs: Optional[Dict[str, Any]] = None, - excluded_columns: Optional[List[str]] = None, - assignment_func: Optional[Callable[[str, str], Optional[Dict[Any, Any]]]] = None, + extra_exprs: dict[str, Any] | None = None, + excluded_columns: list[str] | None = None, + assignment_func: Callable[[str, str], dict[Any, Any] | None] | None = None, metastore_conn_id: str = 'metastore_default', presto_conn_id: str = 'presto_default', mysql_conn_id: str = 'airflow_db', @@ -95,7 +97,7 @@ def __init__( self.ds = '{{ ds }}' self.dttm = '{{ execution_date.isoformat() }}' - def get_default_exprs(self, col: str, col_type: str) -> Dict[Any, Any]: + def get_default_exprs(self, col: str, col_type: str) -> dict[Any, Any]: """Get default expressions""" if col in self.excluded_columns: return {} @@ -114,7 +116,7 @@ def get_default_exprs(self, col: str, col_type: str) -> Dict[Any, Any]: return exp - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: metastore = HiveMetastoreHook(metastore_conn_id=self.metastore_conn_id) table = metastore.get_table(table_name=self.table) field_types = {col.name: col.type for col in table.sd.cols} diff --git a/airflow/providers/apache/hive/sensors/hive_partition.py b/airflow/providers/apache/hive/sensors/hive_partition.py index f03dcb18f52e1..8d589f3d68329 100644 --- a/airflow/providers/apache/hive/sensors/hive_partition.py +++ b/airflow/providers/apache/hive/sensors/hive_partition.py @@ -15,7 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import TYPE_CHECKING, Any, Optional, Sequence +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence from airflow.providers.apache.hive.hooks.hive import HiveMetastoreHook from airflow.sensors.base import BaseSensorOperator @@ -53,7 +55,7 @@ def __init__( self, *, table: str, - partition: Optional[str] = "ds='{{ ds }}'", + partition: str | None = "ds='{{ ds }}'", metastore_conn_id: str = 'metastore_default', schema: str = 'default', poke_interval: int = 60 * 3, @@ -67,7 +69,7 @@ def __init__( self.partition = partition self.schema = schema - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: if '.' in self.table: self.schema, self.table = self.table.split('.') self.log.info('Poking for table %s.%s, partition %s', self.schema, self.table, self.partition) diff --git a/airflow/providers/apache/hive/sensors/metastore_partition.py b/airflow/providers/apache/hive/sensors/metastore_partition.py index ef1d166f024fe..640e2a672f3cd 100644 --- a/airflow/providers/apache/hive/sensors/metastore_partition.py +++ b/airflow/providers/apache/hive/sensors/metastore_partition.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + from typing import TYPE_CHECKING, Any, Sequence from airflow.providers.common.sql.sensors.sql import SqlSensor @@ -65,7 +67,7 @@ def __init__( # constructor below and apply_defaults will no longer throw an exception. super().__init__(**kwargs) - def poke(self, context: "Context") -> Any: + def poke(self, context: Context) -> Any: if self.first_poke: self.first_poke = False if '.' in self.table: diff --git a/airflow/providers/apache/hive/sensors/named_hive_partition.py b/airflow/providers/apache/hive/sensors/named_hive_partition.py index da93e92f41898..4043fd90dc8af 100644 --- a/airflow/providers/apache/hive/sensors/named_hive_partition.py +++ b/airflow/providers/apache/hive/sensors/named_hive_partition.py @@ -15,7 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import TYPE_CHECKING, Any, List, Sequence, Tuple +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence from airflow.sensors.base import BaseSensorOperator @@ -44,7 +46,7 @@ class NamedHivePartitionSensor(BaseSensorOperator): def __init__( self, *, - partition_names: List[str], + partition_names: list[str], metastore_conn_id: str = 'metastore_default', poke_interval: int = 60 * 3, hook: Any = None, @@ -65,7 +67,7 @@ def __init__( ) @staticmethod - def parse_partition_name(partition: str) -> Tuple[Any, ...]: + def parse_partition_name(partition: str) -> tuple[Any, ...]: """Get schema, table, and partition info.""" first_split = partition.split('.', 1) if len(first_split) == 1: @@ -92,7 +94,7 @@ def poke_partition(self, partition: str) -> Any: self.log.info('Poking for %s.%s/%s', schema, table, partition) return self.hook.check_for_named_partition(schema, table, partition) - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: number_of_partitions = len(self.partition_names) poke_index_start = self.next_index_to_poke diff --git a/airflow/providers/apache/hive/transfers/hive_to_mysql.py b/airflow/providers/apache/hive/transfers/hive_to_mysql.py index 9c01b3162b655..d1fe670d35825 100644 --- a/airflow/providers/apache/hive/transfers/hive_to_mysql.py +++ b/airflow/providers/apache/hive/transfers/hive_to_mysql.py @@ -15,10 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains an operator to move data from Hive to MySQL.""" +from __future__ import annotations + from tempfile import NamedTemporaryFile -from typing import TYPE_CHECKING, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Sequence from airflow.models import BaseOperator from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook @@ -75,10 +76,10 @@ def __init__( mysql_table: str, hiveserver2_conn_id: str = 'hiveserver2_default', mysql_conn_id: str = 'mysql_default', - mysql_preoperator: Optional[str] = None, - mysql_postoperator: Optional[str] = None, + mysql_preoperator: str | None = None, + mysql_postoperator: str | None = None, bulk_load: bool = False, - hive_conf: Optional[Dict] = None, + hive_conf: dict | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -91,7 +92,7 @@ def __init__( self.bulk_load = bulk_load self.hive_conf = hive_conf - def execute(self, context: 'Context'): + def execute(self, context: Context): hive = HiveServer2Hook(hiveserver2_conn_id=self.hiveserver2_conn_id) self.log.info("Extracting data from Hive: %s", self.sql) diff --git a/airflow/providers/apache/hive/transfers/hive_to_samba.py b/airflow/providers/apache/hive/transfers/hive_to_samba.py index 63a811c8ce366..143384ac8f928 100644 --- a/airflow/providers/apache/hive/transfers/hive_to_samba.py +++ b/airflow/providers/apache/hive/transfers/hive_to_samba.py @@ -15,8 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains an operator to move data from Hive to Samba.""" +from __future__ import annotations from tempfile import NamedTemporaryFile from typing import TYPE_CHECKING, Sequence @@ -64,7 +64,7 @@ def __init__( self.destination_filepath = destination_filepath self.hql = hql.strip().rstrip(';') - def execute(self, context: 'Context'): + def execute(self, context: Context): with NamedTemporaryFile() as tmp_file: self.log.info("Fetching file from Hive") hive = HiveServer2Hook(hiveserver2_conn_id=self.hiveserver2_conn_id) diff --git a/airflow/providers/apache/hive/transfers/mssql_to_hive.py b/airflow/providers/apache/hive/transfers/mssql_to_hive.py index b1b6e0b3c4907..67adb954b0467 100644 --- a/airflow/providers/apache/hive/transfers/mssql_to_hive.py +++ b/airflow/providers/apache/hive/transfers/mssql_to_hive.py @@ -15,12 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains an operator to move data from MSSQL to Hive.""" +from __future__ import annotations from collections import OrderedDict from tempfile import NamedTemporaryFile -from typing import TYPE_CHECKING, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Sequence import pymssql import unicodecsv as csv @@ -77,11 +77,11 @@ def __init__( hive_table: str, create: bool = True, recreate: bool = False, - partition: Optional[Dict] = None, + partition: dict | None = None, delimiter: str = chr(1), mssql_conn_id: str = 'mssql_default', hive_cli_conn_id: str = 'hive_cli_default', - tblproperties: Optional[Dict] = None, + tblproperties: dict | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -106,7 +106,7 @@ def type_map(cls, mssql_type: int) -> str: } return map_dict.get(mssql_type, 'STRING') - def execute(self, context: "Context"): + def execute(self, context: Context): mssql = MsSqlHook(mssql_conn_id=self.mssql_conn_id) self.log.info("Dumping Microsoft SQL Server query results to local file") with mssql.get_conn() as conn: diff --git a/airflow/providers/apache/hive/transfers/mysql_to_hive.py b/airflow/providers/apache/hive/transfers/mysql_to_hive.py index ccdaf7b682fc1..2bc0e7537cfd8 100644 --- a/airflow/providers/apache/hive/transfers/mysql_to_hive.py +++ b/airflow/providers/apache/hive/transfers/mysql_to_hive.py @@ -15,12 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains an operator to move data from MySQL to Hive.""" +from __future__ import annotations from collections import OrderedDict from tempfile import NamedTemporaryFile -from typing import TYPE_CHECKING, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Sequence import MySQLdb import unicodecsv as csv @@ -80,14 +80,14 @@ def __init__( hive_table: str, create: bool = True, recreate: bool = False, - partition: Optional[Dict] = None, + partition: dict | None = None, delimiter: str = chr(1), - quoting: Optional[str] = None, + quoting: str | None = None, quotechar: str = '"', - escapechar: Optional[str] = None, + escapechar: str | None = None, mysql_conn_id: str = 'mysql_default', hive_cli_conn_id: str = 'hive_cli_default', - tblproperties: Optional[Dict] = None, + tblproperties: dict | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -125,7 +125,7 @@ def type_map(cls, mysql_type: int) -> str: } return type_map.get(mysql_type, 'STRING') - def execute(self, context: "Context"): + def execute(self, context: Context): hive = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id) mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id) diff --git a/airflow/providers/apache/hive/transfers/s3_to_hive.py b/airflow/providers/apache/hive/transfers/s3_to_hive.py index 9f3827c56f2a5..dc7c577502464 100644 --- a/airflow/providers/apache/hive/transfers/s3_to_hive.py +++ b/airflow/providers/apache/hive/transfers/s3_to_hive.py @@ -15,15 +15,15 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains an operator to move data from an S3 bucket to Hive.""" +from __future__ import annotations import bz2 import gzip import os import tempfile from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, Sequence from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -94,21 +94,21 @@ def __init__( self, *, s3_key: str, - field_dict: Dict, + field_dict: dict, hive_table: str, delimiter: str = ',', create: bool = True, recreate: bool = False, - partition: Optional[Dict] = None, + partition: dict | None = None, headers: bool = False, check_headers: bool = False, wildcard_match: bool = False, aws_conn_id: str = 'aws_default', - verify: Optional[Union[bool, str]] = None, + verify: bool | str | None = None, hive_cli_conn_id: str = 'hive_cli_default', input_compressed: bool = False, - tblproperties: Optional[Dict] = None, - select_expression: Optional[str] = None, + tblproperties: dict | None = None, + select_expression: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -132,7 +132,7 @@ def __init__( if self.check_headers and not (self.field_dict is not None and self.headers): raise AirflowException("To check_headers provide field_dict and headers") - def execute(self, context: 'Context'): + def execute(self, context: Context): # Downloading file from S3 s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify) hive_hook = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id) @@ -162,7 +162,7 @@ def execute(self, context: 'Context'): if self.delimiter: option['FieldDelimiter'] = self.delimiter - input_serialization: Dict[str, Any] = {'CSV': option} + input_serialization: dict[str, Any] = {'CSV': option} if self.input_compressed: input_serialization['CompressionType'] = 'GZIP' diff --git a/airflow/providers/apache/hive/transfers/vertica_to_hive.py b/airflow/providers/apache/hive/transfers/vertica_to_hive.py index e3f432d9ad70d..66e615c7a6938 100644 --- a/airflow/providers/apache/hive/transfers/vertica_to_hive.py +++ b/airflow/providers/apache/hive/transfers/vertica_to_hive.py @@ -15,12 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains an operator to move data from Vertica to Hive.""" +from __future__ import annotations from collections import OrderedDict from tempfile import NamedTemporaryFile -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Any, Sequence import unicodecsv as csv @@ -72,7 +72,7 @@ def __init__( hive_table: str, create: bool = True, recreate: bool = False, - partition: Optional[Dict] = None, + partition: dict | None = None, delimiter: str = chr(1), vertica_conn_id: str = 'vertica_default', hive_cli_conn_id: str = 'hive_cli_default', @@ -106,7 +106,7 @@ def type_map(cls, vertica_type): } return type_map.get(vertica_type, 'STRING') - def execute(self, context: 'Context'): + def execute(self, context: Context): hive = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id) vertica = VerticaHook(vertica_conn_id=self.vertica_conn_id) diff --git a/airflow/providers/apache/kylin/hooks/kylin.py b/airflow/providers/apache/kylin/hooks/kylin.py index 53a8bf5d909c8..5f41a459e54c6 100644 --- a/airflow/providers/apache/kylin/hooks/kylin.py +++ b/airflow/providers/apache/kylin/hooks/kylin.py @@ -15,8 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - -from typing import Optional +from __future__ import annotations from kylinpy import exceptions, kylinpy @@ -36,8 +35,8 @@ class KylinHook(BaseHook): def __init__( self, kylin_conn_id: str = 'kylin_default', - project: Optional[str] = None, - dsn: Optional[str] = None, + project: str | None = None, + dsn: str | None = None, ): super().__init__() self.kylin_conn_id = kylin_conn_id diff --git a/airflow/providers/apache/kylin/operators/kylin_cube.py b/airflow/providers/apache/kylin/operators/kylin_cube.py index 18bd0966e0c13..8ffd98129ac70 100644 --- a/airflow/providers/apache/kylin/operators/kylin_cube.py +++ b/airflow/providers/apache/kylin/operators/kylin_cube.py @@ -15,10 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import time from datetime import datetime -from typing import TYPE_CHECKING, Optional, Sequence +from typing import TYPE_CHECKING, Sequence from kylinpy import kylinpy @@ -101,15 +102,15 @@ def __init__( self, *, kylin_conn_id: str = 'kylin_default', - project: Optional[str] = None, - cube: Optional[str] = None, - dsn: Optional[str] = None, - command: Optional[str] = None, - start_time: Optional[str] = None, - end_time: Optional[str] = None, - offset_start: Optional[str] = None, - offset_end: Optional[str] = None, - segment_name: Optional[str] = None, + project: str | None = None, + cube: str | None = None, + dsn: str | None = None, + command: str | None = None, + start_time: str | None = None, + end_time: str | None = None, + offset_start: str | None = None, + offset_end: str | None = None, + segment_name: str | None = None, is_track_job: bool = False, interval: int = 60, timeout: int = 60 * 60 * 24, @@ -133,7 +134,7 @@ def __init__( self.eager_error_status = eager_error_status self.jobs_error_status = [stat.upper() for stat in eager_error_status] - def execute(self, context: 'Context'): + def execute(self, context: Context): _hook = KylinHook(kylin_conn_id=self.kylin_conn_id, project=self.project, dsn=self.dsn) diff --git a/airflow/providers/apache/livy/hooks/livy.py b/airflow/providers/apache/livy/hooks/livy.py index 84b2fd12ade6d..05ef32bd58d22 100644 --- a/airflow/providers/apache/livy/hooks/livy.py +++ b/airflow/providers/apache/livy/hooks/livy.py @@ -14,12 +14,13 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains the Apache Livy hook.""" +from __future__ import annotations + import json import re from enum import Enum -from typing import Any, Dict, List, Optional, Sequence, Union +from typing import Any, Sequence import requests @@ -74,16 +75,16 @@ class LivyHook(HttpHook, LoggingMixin): def __init__( self, livy_conn_id: str = default_conn_name, - extra_options: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - auth_type: Optional[Any] = None, + extra_options: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + auth_type: Any | None = None, ) -> None: super().__init__(http_conn_id=livy_conn_id) self.extra_headers = extra_headers or {} self.extra_options = extra_options or {} self.auth_type = auth_type or self.auth_type - def get_conn(self, headers: Optional[Dict[str, Any]] = None) -> Any: + def get_conn(self, headers: dict[str, Any] | None = None) -> Any: """ Returns http session for use with requests @@ -100,9 +101,9 @@ def run_method( self, endpoint: str, method: str = 'GET', - data: Optional[Any] = None, - headers: Optional[Dict[str, Any]] = None, - retry_args: Optional[Dict[str, Any]] = None, + data: Any | None = None, + headers: dict[str, Any] | None = None, + retry_args: dict[str, Any] | None = None, ) -> Any: """ Wrapper for HttpHook, allows to change method on the same HttpHook @@ -173,7 +174,7 @@ def post_batch(self, *args: Any, **kwargs: Any) -> Any: return batch_id - def get_batch(self, session_id: Union[int, str]) -> Any: + def get_batch(self, session_id: int | str) -> Any: """ Fetch info about the specified batch @@ -196,9 +197,7 @@ def get_batch(self, session_id: Union[int, str]) -> Any: return response.json() - def get_batch_state( - self, session_id: Union[int, str], retry_args: Optional[Dict[str, Any]] = None - ) -> BatchState: + def get_batch_state(self, session_id: int | str, retry_args: dict[str, Any] | None = None) -> BatchState: """ Fetch the state of the specified batch @@ -228,7 +227,7 @@ def get_batch_state( raise AirflowException(f"Unable to get state for batch with id: {session_id}") return BatchState(jresp['state']) - def delete_batch(self, session_id: Union[int, str]) -> Any: + def delete_batch(self, session_id: int | str) -> Any: """ Delete the specified batch @@ -253,7 +252,7 @@ def delete_batch(self, session_id: Union[int, str]) -> Any: return response.json() - def get_batch_logs(self, session_id: Union[int, str], log_start_position, log_batch_size) -> Any: + def get_batch_logs(self, session_id: int | str, log_start_position, log_batch_size) -> Any: """ Gets the session logs for a specified batch. :param session_id: identifier of the batch sessions @@ -278,7 +277,7 @@ def get_batch_logs(self, session_id: Union[int, str], log_start_position, log_ba ) return response.json() - def dump_batch_logs(self, session_id: Union[int, str]) -> Any: + def dump_batch_logs(self, session_id: int | str) -> Any: """ Dumps the session logs for a specified batch @@ -301,7 +300,7 @@ def dump_batch_logs(self, session_id: Union[int, str]) -> Any: self.log.info(log_line) @staticmethod - def _validate_session_id(session_id: Union[int, str]) -> None: + def _validate_session_id(session_id: int | str) -> None: """ Validate session id is a int @@ -313,7 +312,7 @@ def _validate_session_id(session_id: Union[int, str]) -> None: raise TypeError("'session_id' must be an integer") @staticmethod - def _parse_post_response(response: Dict[Any, Any]) -> Any: + def _parse_post_response(response: dict[Any, Any]) -> Any: """ Parse batch response for batch id @@ -324,7 +323,7 @@ def _parse_post_response(response: Dict[Any, Any]) -> Any: return response.get('id') @staticmethod - def _parse_request_response(response: Dict[Any, Any], parameter) -> Any: + def _parse_request_response(response: dict[Any, Any], parameter) -> Any: """ Parse batch response for batch id @@ -337,21 +336,21 @@ def _parse_request_response(response: Dict[Any, Any], parameter) -> Any: @staticmethod def build_post_batch_body( file: str, - args: Optional[Sequence[Union[str, int, float]]] = None, - class_name: Optional[str] = None, - jars: Optional[List[str]] = None, - py_files: Optional[List[str]] = None, - files: Optional[List[str]] = None, - archives: Optional[List[str]] = None, - name: Optional[str] = None, - driver_memory: Optional[str] = None, - driver_cores: Optional[Union[int, str]] = None, - executor_memory: Optional[str] = None, - executor_cores: Optional[int] = None, - num_executors: Optional[Union[int, str]] = None, - queue: Optional[str] = None, - proxy_user: Optional[str] = None, - conf: Optional[Dict[Any, Any]] = None, + args: Sequence[str | int | float] | None = None, + class_name: str | None = None, + jars: list[str] | None = None, + py_files: list[str] | None = None, + files: list[str] | None = None, + archives: list[str] | None = None, + name: str | None = None, + driver_memory: str | None = None, + driver_cores: int | str | None = None, + executor_memory: str | None = None, + executor_cores: int | None = None, + num_executors: int | str | None = None, + queue: str | None = None, + proxy_user: str | None = None, + conf: dict[Any, Any] | None = None, ) -> Any: """ Build the post batch request body. @@ -376,7 +375,7 @@ def build_post_batch_body( :return: request body :rtype: dict """ - body: Dict[str, Any] = {'file': file} + body: dict[str, Any] = {'file': file} if proxy_user: body['proxyUser'] = proxy_user @@ -425,7 +424,7 @@ def _validate_size_format(size: str) -> bool: return True @staticmethod - def _validate_list_of_stringables(vals: Sequence[Union[str, int, float]]) -> bool: + def _validate_list_of_stringables(vals: Sequence[str | int | float]) -> bool: """ Check the values in the provided list can be converted to strings. @@ -442,7 +441,7 @@ def _validate_list_of_stringables(vals: Sequence[Union[str, int, float]]) -> boo return True @staticmethod - def _validate_extra_conf(conf: Dict[Any, Any]) -> bool: + def _validate_extra_conf(conf: dict[Any, Any]) -> bool: """ Check configuration values are either strings or ints. diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py index ce2910449d609..b6c2f67bfba7c 100644 --- a/airflow/providers/apache/livy/operators/livy.py +++ b/airflow/providers/apache/livy/operators/livy.py @@ -14,10 +14,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains the Apache Livy operator.""" +from __future__ import annotations + from time import sleep -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, Sequence from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -64,27 +65,27 @@ def __init__( self, *, file: str, - class_name: Optional[str] = None, - args: Optional[Sequence[Union[str, int, float]]] = None, - conf: Optional[Dict[Any, Any]] = None, - jars: Optional[Sequence[str]] = None, - py_files: Optional[Sequence[str]] = None, - files: Optional[Sequence[str]] = None, - driver_memory: Optional[str] = None, - driver_cores: Optional[Union[int, str]] = None, - executor_memory: Optional[str] = None, - executor_cores: Optional[Union[int, str]] = None, - num_executors: Optional[Union[int, str]] = None, - archives: Optional[Sequence[str]] = None, - queue: Optional[str] = None, - name: Optional[str] = None, - proxy_user: Optional[str] = None, + class_name: str | None = None, + args: Sequence[str | int | float] | None = None, + conf: dict[Any, Any] | None = None, + jars: Sequence[str] | None = None, + py_files: Sequence[str] | None = None, + files: Sequence[str] | None = None, + driver_memory: str | None = None, + driver_cores: int | str | None = None, + executor_memory: str | None = None, + executor_cores: int | str | None = None, + num_executors: int | str | None = None, + archives: Sequence[str] | None = None, + queue: str | None = None, + name: str | None = None, + proxy_user: str | None = None, livy_conn_id: str = 'livy_default', - livy_conn_auth_type: Optional[Any] = None, + livy_conn_auth_type: Any | None = None, polling_interval: int = 0, - extra_options: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - retry_args: Optional[Dict[str, Any]] = None, + extra_options: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + retry_args: dict[str, Any] | None = None, **kwargs: Any, ) -> None: @@ -115,8 +116,8 @@ def __init__( self._extra_options = extra_options or {} self._extra_headers = extra_headers or {} - self._livy_hook: Optional[LivyHook] = None - self._batch_id: Union[int, str] + self._livy_hook: LivyHook | None = None + self._batch_id: int | str self.retry_args = retry_args def get_hook(self) -> LivyHook: @@ -135,7 +136,7 @@ def get_hook(self) -> LivyHook: ) return self._livy_hook - def execute(self, context: "Context") -> Any: + def execute(self, context: Context) -> Any: self._batch_id = self.get_hook().post_batch(**self.spark_params) if self._polling_interval > 0: @@ -143,7 +144,7 @@ def execute(self, context: "Context") -> Any: return self._batch_id - def poll_for_termination(self, batch_id: Union[int, str]) -> None: + def poll_for_termination(self, batch_id: int | str) -> None: """ Pool Livy for batch termination. diff --git a/airflow/providers/apache/livy/sensors/livy.py b/airflow/providers/apache/livy/sensors/livy.py index fb3e60c82865b..319a60945fcb0 100644 --- a/airflow/providers/apache/livy/sensors/livy.py +++ b/airflow/providers/apache/livy/sensors/livy.py @@ -14,9 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - """This module contains the Apache Livy sensor.""" -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence from airflow.providers.apache.livy.hooks.livy import LivyHook from airflow.sensors.base import BaseSensorOperator @@ -39,17 +40,17 @@ class LivySensor(BaseSensorOperator): def __init__( self, *, - batch_id: Union[int, str], + batch_id: int | str, livy_conn_id: str = 'livy_default', - livy_conn_auth_type: Optional[Any] = None, - extra_options: Optional[Dict[str, Any]] = None, + livy_conn_auth_type: Any | None = None, + extra_options: dict[str, Any] | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) self.batch_id = batch_id self._livy_conn_id = livy_conn_id self._livy_conn_auth_type = livy_conn_auth_type - self._livy_hook: Optional[LivyHook] = None + self._livy_hook: LivyHook | None = None self._extra_options = extra_options or {} def get_hook(self) -> LivyHook: @@ -67,7 +68,7 @@ def get_hook(self) -> LivyHook: ) return self._livy_hook - def poke(self, context: "Context") -> bool: + def poke(self, context: Context) -> bool: batch_id = self.batch_id status = self.get_hook().get_batch_state(batch_id) diff --git a/airflow/providers/apache/pig/hooks/pig.py b/airflow/providers/apache/pig/hooks/pig.py index ae9db33c3db53..b63d8fba63caf 100644 --- a/airflow/providers/apache/pig/hooks/pig.py +++ b/airflow/providers/apache/pig/hooks/pig.py @@ -15,9 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import subprocess from tempfile import NamedTemporaryFile, TemporaryDirectory -from typing import Any, List, Optional +from typing import Any from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook @@ -45,7 +47,7 @@ def __init__(self, pig_cli_conn_id: str = default_conn_name) -> None: self.conn = conn self.sub_process = None - def run_cli(self, pig: str, pig_opts: Optional[str] = None, verbose: bool = True) -> Any: + def run_cli(self, pig: str, pig_opts: str | None = None, verbose: bool = True) -> Any: """ Run an pig script using the pig cli @@ -60,7 +62,7 @@ def run_cli(self, pig: str, pig_opts: Optional[str] = None, verbose: bool = True f.flush() fname = f.name pig_bin = 'pig' - cmd_extra: List[str] = [] + cmd_extra: list[str] = [] pig_cmd = [pig_bin] diff --git a/airflow/providers/apache/pig/operators/pig.py b/airflow/providers/apache/pig/operators/pig.py index 5a285530dc637..fb12ae0169b91 100644 --- a/airflow/providers/apache/pig/operators/pig.py +++ b/airflow/providers/apache/pig/operators/pig.py @@ -15,8 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import re -from typing import TYPE_CHECKING, Any, Optional, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.models import BaseOperator from airflow.providers.apache.pig.hooks.pig import PigCliHook @@ -52,7 +54,7 @@ def __init__( pig: str, pig_cli_conn_id: str = 'pig_cli_default', pigparams_jinja_translate: bool = False, - pig_opts: Optional[str] = None, + pig_opts: str | None = None, **kwargs: Any, ) -> None: @@ -61,13 +63,13 @@ def __init__( self.pig = pig self.pig_cli_conn_id = pig_cli_conn_id self.pig_opts = pig_opts - self.hook: Optional[PigCliHook] = None + self.hook: PigCliHook | None = None def prepare_template(self): if self.pigparams_jinja_translate: self.pig = re.sub(r"(\$([a-zA-Z_][a-zA-Z0-9_]*))", r"{{ \g<2> }}", self.pig) - def execute(self, context: 'Context'): + def execute(self, context: Context): self.log.info('Executing: %s', self.pig) self.hook = PigCliHook(pig_cli_conn_id=self.pig_cli_conn_id) self.hook.run_cli(pig=self.pig, pig_opts=self.pig_opts) diff --git a/airflow/providers/apache/pinot/hooks/pinot.py b/airflow/providers/apache/pinot/hooks/pinot.py index 90b3804d98856..69e608e945c05 100644 --- a/airflow/providers/apache/pinot/hooks/pinot.py +++ b/airflow/providers/apache/pinot/hooks/pinot.py @@ -15,10 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import os import subprocess -from typing import Any, Iterable, List, Mapping, Optional, Union +from typing import Any, Iterable, Mapping from pinotdb import connect @@ -102,24 +103,24 @@ def add_table(self, file_path: str, with_exec: bool = True) -> Any: def create_segment( self, - generator_config_file: Optional[str] = None, - data_dir: Optional[str] = None, - segment_format: Optional[str] = None, - out_dir: Optional[str] = None, - overwrite: Optional[str] = None, - table_name: Optional[str] = None, - segment_name: Optional[str] = None, - time_column_name: Optional[str] = None, - schema_file: Optional[str] = None, - reader_config_file: Optional[str] = None, - enable_star_tree_index: Optional[str] = None, - star_tree_index_spec_file: Optional[str] = None, - hll_size: Optional[str] = None, - hll_columns: Optional[str] = None, - hll_suffix: Optional[str] = None, - num_threads: Optional[str] = None, - post_creation_verification: Optional[str] = None, - retry: Optional[str] = None, + generator_config_file: str | None = None, + data_dir: str | None = None, + segment_format: str | None = None, + out_dir: str | None = None, + overwrite: str | None = None, + table_name: str | None = None, + segment_name: str | None = None, + time_column_name: str | None = None, + schema_file: str | None = None, + reader_config_file: str | None = None, + enable_star_tree_index: str | None = None, + star_tree_index_spec_file: str | None = None, + hll_size: str | None = None, + hll_columns: str | None = None, + hll_suffix: str | None = None, + num_threads: str | None = None, + post_creation_verification: str | None = None, + retry: str | None = None, ) -> Any: """Create Pinot segment by run CreateSegment command""" cmd = ["CreateSegment"] @@ -180,7 +181,7 @@ def create_segment( self.run_cli(cmd) - def upload_segment(self, segment_dir: str, table_name: Optional[str] = None) -> Any: + def upload_segment(self, segment_dir: str, table_name: str | None = None) -> Any: """ Upload Segment with run UploadSegment command @@ -196,7 +197,7 @@ def upload_segment(self, segment_dir: str, table_name: Optional[str] = None) -> cmd += ["-tableName", table_name] self.run_cli(cmd) - def run_cli(self, cmd: List[str], verbose: bool = True) -> str: + def run_cli(self, cmd: list[str], verbose: bool = True) -> str: """ Run command with pinot-admin.sh @@ -276,7 +277,7 @@ def get_uri(self) -> str: return f'{conn_type}://{host}/{endpoint}' def get_records( - self, sql: Union[str, List[str]], parameters: Optional[Union[Iterable, Mapping]] = None, **kwargs + self, sql: str | list[str], parameters: Iterable | Mapping | None = None, **kwargs ) -> Any: """ Executes the sql and returns a set of records. @@ -289,9 +290,7 @@ def get_records( cur.execute(sql) return cur.fetchall() - def get_first( - self, sql: Union[str, List[str]], parameters: Optional[Union[Iterable, Mapping]] = None - ) -> Any: + def get_first(self, sql: str | list[str], parameters: Iterable | Mapping | None = None) -> Any: """ Executes the sql and returns the first resulting row. @@ -310,7 +309,7 @@ def insert_rows( self, table: str, rows: str, - target_fields: Optional[str] = None, + target_fields: str | None = None, commit_every: int = 1000, replace: bool = False, **kwargs: Any, diff --git a/airflow/providers/apache/spark/hooks/spark_jdbc.py b/airflow/providers/apache/spark/hooks/spark_jdbc.py index e95fae3dbbabc..13abd9053ab82 100644 --- a/airflow/providers/apache/spark/hooks/spark_jdbc.py +++ b/airflow/providers/apache/spark/hooks/spark_jdbc.py @@ -15,9 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import os -from typing import Any, Dict, Optional +from typing import Any from airflow.exceptions import AirflowException from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook @@ -94,32 +95,32 @@ def __init__( self, spark_app_name: str = 'airflow-spark-jdbc', spark_conn_id: str = default_conn_name, - spark_conf: Optional[Dict[str, Any]] = None, - spark_py_files: Optional[str] = None, - spark_files: Optional[str] = None, - spark_jars: Optional[str] = None, - num_executors: Optional[int] = None, - executor_cores: Optional[int] = None, - executor_memory: Optional[str] = None, - driver_memory: Optional[str] = None, + spark_conf: dict[str, Any] | None = None, + spark_py_files: str | None = None, + spark_files: str | None = None, + spark_jars: str | None = None, + num_executors: int | None = None, + executor_cores: int | None = None, + executor_memory: str | None = None, + driver_memory: str | None = None, verbose: bool = False, - principal: Optional[str] = None, - keytab: Optional[str] = None, + principal: str | None = None, + keytab: str | None = None, cmd_type: str = 'spark_to_jdbc', - jdbc_table: Optional[str] = None, + jdbc_table: str | None = None, jdbc_conn_id: str = 'jdbc-default', - jdbc_driver: Optional[str] = None, - metastore_table: Optional[str] = None, + jdbc_driver: str | None = None, + metastore_table: str | None = None, jdbc_truncate: bool = False, - save_mode: Optional[str] = None, - save_format: Optional[str] = None, - batch_size: Optional[int] = None, - fetch_size: Optional[int] = None, - num_partitions: Optional[int] = None, - partition_column: Optional[str] = None, - lower_bound: Optional[str] = None, - upper_bound: Optional[str] = None, - create_table_column_types: Optional[str] = None, + save_mode: str | None = None, + save_format: str | None = None, + batch_size: int | None = None, + fetch_size: int | None = None, + num_partitions: int | None = None, + partition_column: str | None = None, + lower_bound: str | None = None, + upper_bound: str | None = None, + create_table_column_types: str | None = None, *args: Any, **kwargs: Any, ): @@ -154,7 +155,7 @@ def __init__( self._create_table_column_types = create_table_column_types self._jdbc_connection = self._resolve_jdbc_connection() - def _resolve_jdbc_connection(self) -> Dict[str, Any]: + def _resolve_jdbc_connection(self) -> dict[str, Any]: conn_data = {'url': '', 'schema': '', 'conn_prefix': '', 'user': '', 'password': ''} try: conn = self.get_connection(self._jdbc_conn_id) @@ -173,7 +174,7 @@ def _resolve_jdbc_connection(self) -> Dict[str, Any]: ) return conn_data - def _build_jdbc_application_arguments(self, jdbc_conn: Dict[str, Any]) -> Any: + def _build_jdbc_application_arguments(self, jdbc_conn: dict[str, Any]) -> Any: arguments = [] arguments += ["-cmdType", self._cmd_type] if self._jdbc_connection['url']: diff --git a/airflow/providers/apache/spark/hooks/spark_jdbc_script.py b/airflow/providers/apache/spark/hooks/spark_jdbc_script.py index c354de6ab0ca4..b72053f6af986 100644 --- a/airflow/providers/apache/spark/hooks/spark_jdbc_script.py +++ b/airflow/providers/apache/spark/hooks/spark_jdbc_script.py @@ -15,9 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import argparse -from typing import Any, List, Optional +from typing import Any from pyspark.sql import SparkSession @@ -121,7 +122,7 @@ def spark_read_from_jdbc( reader.load().write.saveAsTable(metastore_table, format=save_format, mode=save_mode) -def _parse_arguments(args: Optional[List[str]] = None) -> Any: +def _parse_arguments(args: list[str] | None = None) -> Any: parser = argparse.ArgumentParser(description='Spark-JDBC') parser.add_argument('-cmdType', dest='cmd_type', action='store') parser.add_argument('-url', dest='url', action='store') diff --git a/airflow/providers/apache/spark/hooks/spark_sql.py b/airflow/providers/apache/spark/hooks/spark_sql.py index 2621bf9a4099b..1a6ae39e8a9eb 100644 --- a/airflow/providers/apache/spark/hooks/spark_sql.py +++ b/airflow/providers/apache/spark/hooks/spark_sql.py @@ -15,9 +15,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import subprocess -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any from airflow.exceptions import AirflowException, AirflowNotFoundException from airflow.hooks.base import BaseHook @@ -57,22 +58,22 @@ class SparkSqlHook(BaseHook): def __init__( self, sql: str, - conf: Optional[str] = None, + conf: str | None = None, conn_id: str = default_conn_name, - total_executor_cores: Optional[int] = None, - executor_cores: Optional[int] = None, - executor_memory: Optional[str] = None, - keytab: Optional[str] = None, - principal: Optional[str] = None, - master: Optional[str] = None, + total_executor_cores: int | None = None, + executor_cores: int | None = None, + executor_memory: str | None = None, + keytab: str | None = None, + principal: str | None = None, + master: str | None = None, name: str = 'default-name', - num_executors: Optional[int] = None, + num_executors: int | None = None, verbose: bool = True, - yarn_queue: Optional[str] = None, + yarn_queue: str | None = None, ) -> None: super().__init__() - options: Dict = {} - conn: Optional[Connection] = None + options: dict = {} + conn: Connection | None = None try: conn = self.get_connection(conn_id) @@ -109,7 +110,7 @@ def __init__( def get_conn(self) -> Any: pass - def _prepare_command(self, cmd: Union[str, List[str]]) -> List[str]: + def _prepare_command(self, cmd: str | list[str]) -> list[str]: """ Construct the spark-sql command to execute. Verbose output is enabled as default. diff --git a/airflow/providers/apache/spark/hooks/spark_submit.py b/airflow/providers/apache/spark/hooks/spark_submit.py index fa09c0b6c9d7d..b0be7298f40a3 100644 --- a/airflow/providers/apache/spark/hooks/spark_submit.py +++ b/airflow/providers/apache/spark/hooks/spark_submit.py @@ -15,13 +15,14 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import contextlib import os import re import subprocess import time -from typing import Any, Dict, Iterator, List, Optional, Union +from typing import Any, Iterator from airflow.configuration import conf as airflow_conf from airflow.exceptions import AirflowException @@ -85,7 +86,7 @@ class SparkSubmitHook(BaseHook, LoggingMixin): hook_name = 'Spark' @staticmethod - def get_ui_field_behaviour() -> Dict[str, Any]: + def get_ui_field_behaviour() -> dict[str, Any]: """Returns custom field behaviour""" return { "hidden_fields": ['schema', 'login', 'password'], @@ -94,31 +95,31 @@ def get_ui_field_behaviour() -> Dict[str, Any]: def __init__( self, - conf: Optional[Dict[str, Any]] = None, + conf: dict[str, Any] | None = None, conn_id: str = 'spark_default', - files: Optional[str] = None, - py_files: Optional[str] = None, - archives: Optional[str] = None, - driver_class_path: Optional[str] = None, - jars: Optional[str] = None, - java_class: Optional[str] = None, - packages: Optional[str] = None, - exclude_packages: Optional[str] = None, - repositories: Optional[str] = None, - total_executor_cores: Optional[int] = None, - executor_cores: Optional[int] = None, - executor_memory: Optional[str] = None, - driver_memory: Optional[str] = None, - keytab: Optional[str] = None, - principal: Optional[str] = None, - proxy_user: Optional[str] = None, + files: str | None = None, + py_files: str | None = None, + archives: str | None = None, + driver_class_path: str | None = None, + jars: str | None = None, + java_class: str | None = None, + packages: str | None = None, + exclude_packages: str | None = None, + repositories: str | None = None, + total_executor_cores: int | None = None, + executor_cores: int | None = None, + executor_memory: str | None = None, + driver_memory: str | None = None, + keytab: str | None = None, + principal: str | None = None, + proxy_user: str | None = None, name: str = 'default-name', - num_executors: Optional[int] = None, + num_executors: int | None = None, status_poll_interval: int = 1, - application_args: Optional[List[Any]] = None, - env_vars: Optional[Dict[str, Any]] = None, + application_args: list[Any] | None = None, + env_vars: dict[str, Any] | None = None, verbose: bool = False, - spark_binary: Optional[str] = None, + spark_binary: str | None = None, ) -> None: super().__init__() self._conf = conf or {} @@ -145,9 +146,9 @@ def __init__( self._application_args = application_args self._env_vars = env_vars self._verbose = verbose - self._submit_sp: Optional[Any] = None - self._yarn_application_id: Optional[str] = None - self._kubernetes_driver_pod: Optional[str] = None + self._submit_sp: Any | None = None + self._yarn_application_id: str | None = None + self._kubernetes_driver_pod: str | None = None self._spark_binary = spark_binary self._connection = self._resolve_connection() @@ -159,10 +160,10 @@ def __init__( ) self._should_track_driver_status = self._resolve_should_track_driver_status() - self._driver_id: Optional[str] = None - self._driver_status: Optional[str] = None - self._spark_exit_code: Optional[int] = None - self._env: Optional[Dict[str, Any]] = None + self._driver_id: str | None = None + self._driver_status: str | None = None + self._spark_exit_code: int | None = None + self._env: dict[str, Any] | None = None def _resolve_should_track_driver_status(self) -> bool: """ @@ -172,7 +173,7 @@ def _resolve_should_track_driver_status(self) -> bool: """ return 'spark://' in self._connection['master'] and self._connection['deploy_mode'] == 'cluster' - def _resolve_connection(self) -> Dict[str, Any]: + def _resolve_connection(self) -> dict[str, Any]: # Build from connection master or default to yarn if not available conn_data = { 'master': 'yarn', @@ -212,7 +213,7 @@ def _resolve_connection(self) -> Dict[str, Any]: def get_conn(self) -> Any: pass - def _get_spark_binary_path(self) -> List[str]: + def _get_spark_binary_path(self) -> list[str]: # If the spark_home is passed then build the spark-submit executable path using # the spark_home; otherwise assume that spark-submit is present in the path to # the executing user @@ -225,7 +226,7 @@ def _get_spark_binary_path(self) -> List[str]: return connection_cmd - def _mask_cmd(self, connection_cmd: Union[str, List[str]]) -> str: + def _mask_cmd(self, connection_cmd: str | list[str]) -> str: # Mask any password related fields in application args with key value pair # where key contains password (case insensitive), e.g. HivePassword='abc' connection_cmd_masked = re.sub( @@ -249,7 +250,7 @@ def _mask_cmd(self, connection_cmd: Union[str, List[str]]) -> str: return connection_cmd_masked - def _build_spark_submit_command(self, application: str) -> List[str]: + def _build_spark_submit_command(self, application: str) -> list[str]: """ Construct the spark-submit command to execute. @@ -335,7 +336,7 @@ def _build_spark_submit_command(self, application: str) -> List[str]: return connection_cmd - def _build_track_driver_status_command(self) -> List[str]: + def _build_track_driver_status_command(self) -> list[str]: """ Construct the command to poll the driver status. @@ -574,7 +575,7 @@ def _start_driver_status_tracking(self) -> None: f"returncode = {returncode}" ) - def _build_spark_driver_kill_command(self) -> List[str]: + def _build_spark_driver_kill_command(self) -> list[str]: """ Construct the spark-submit command to kill a driver. :return: full command to kill a driver diff --git a/airflow/providers/apache/spark/operators/spark_jdbc.py b/airflow/providers/apache/spark/operators/spark_jdbc.py index 87f244be50899..ab12c0fe57916 100644 --- a/airflow/providers/apache/spark/operators/spark_jdbc.py +++ b/airflow/providers/apache/spark/operators/spark_jdbc.py @@ -15,8 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# -from typing import TYPE_CHECKING, Any, Dict, Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, Any from airflow.providers.apache.spark.hooks.spark_jdbc import SparkJDBCHook from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator @@ -98,32 +99,32 @@ def __init__( *, spark_app_name: str = 'airflow-spark-jdbc', spark_conn_id: str = 'spark-default', - spark_conf: Optional[Dict[str, Any]] = None, - spark_py_files: Optional[str] = None, - spark_files: Optional[str] = None, - spark_jars: Optional[str] = None, - num_executors: Optional[int] = None, - executor_cores: Optional[int] = None, - executor_memory: Optional[str] = None, - driver_memory: Optional[str] = None, + spark_conf: dict[str, Any] | None = None, + spark_py_files: str | None = None, + spark_files: str | None = None, + spark_jars: str | None = None, + num_executors: int | None = None, + executor_cores: int | None = None, + executor_memory: str | None = None, + driver_memory: str | None = None, verbose: bool = False, - principal: Optional[str] = None, - keytab: Optional[str] = None, + principal: str | None = None, + keytab: str | None = None, cmd_type: str = 'spark_to_jdbc', - jdbc_table: Optional[str] = None, + jdbc_table: str | None = None, jdbc_conn_id: str = 'jdbc-default', - jdbc_driver: Optional[str] = None, - metastore_table: Optional[str] = None, + jdbc_driver: str | None = None, + metastore_table: str | None = None, jdbc_truncate: bool = False, - save_mode: Optional[str] = None, - save_format: Optional[str] = None, - batch_size: Optional[int] = None, - fetch_size: Optional[int] = None, - num_partitions: Optional[int] = None, - partition_column: Optional[str] = None, - lower_bound: Optional[str] = None, - upper_bound: Optional[str] = None, - create_table_column_types: Optional[str] = None, + save_mode: str | None = None, + save_format: str | None = None, + batch_size: int | None = None, + fetch_size: int | None = None, + num_partitions: int | None = None, + partition_column: str | None = None, + lower_bound: str | None = None, + upper_bound: str | None = None, + create_table_column_types: str | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -155,9 +156,9 @@ def __init__( self._lower_bound = lower_bound self._upper_bound = upper_bound self._create_table_column_types = create_table_column_types - self._hook: Optional[SparkJDBCHook] = None + self._hook: SparkJDBCHook | None = None - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: """Call the SparkSubmitHook to run the provided spark job""" if self._hook is None: self._hook = self._get_hook() diff --git a/airflow/providers/apache/spark/operators/spark_sql.py b/airflow/providers/apache/spark/operators/spark_sql.py index 33f19e9c43287..a71be5fd1ae10 100644 --- a/airflow/providers/apache/spark/operators/spark_sql.py +++ b/airflow/providers/apache/spark/operators/spark_sql.py @@ -15,8 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# -from typing import TYPE_CHECKING, Any, Optional, Sequence +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence from airflow.models import BaseOperator from airflow.providers.apache.spark.hooks.spark_sql import SparkSqlHook @@ -59,18 +60,18 @@ def __init__( self, *, sql: str, - conf: Optional[str] = None, + conf: str | None = None, conn_id: str = 'spark_sql_default', - total_executor_cores: Optional[int] = None, - executor_cores: Optional[int] = None, - executor_memory: Optional[str] = None, - keytab: Optional[str] = None, - principal: Optional[str] = None, - master: Optional[str] = None, + total_executor_cores: int | None = None, + executor_cores: int | None = None, + executor_memory: str | None = None, + keytab: str | None = None, + principal: str | None = None, + master: str | None = None, name: str = 'default-name', - num_executors: Optional[int] = None, + num_executors: int | None = None, verbose: bool = True, - yarn_queue: Optional[str] = None, + yarn_queue: str | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -87,9 +88,9 @@ def __init__( self._num_executors = num_executors self._verbose = verbose self._yarn_queue = yarn_queue - self._hook: Optional[SparkSqlHook] = None + self._hook: SparkSqlHook | None = None - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: """Call the SparkSqlHook to run the provided sql query""" if self._hook is None: self._hook = self._get_hook() diff --git a/airflow/providers/apache/spark/operators/spark_submit.py b/airflow/providers/apache/spark/operators/spark_submit.py index db1114cf201dc..c79508f0a05f5 100644 --- a/airflow/providers/apache/spark/operators/spark_submit.py +++ b/airflow/providers/apache/spark/operators/spark_submit.py @@ -15,8 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence from airflow.models import BaseOperator from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook @@ -94,31 +95,31 @@ def __init__( self, *, application: str = '', - conf: Optional[Dict[str, Any]] = None, + conf: dict[str, Any] | None = None, conn_id: str = 'spark_default', - files: Optional[str] = None, - py_files: Optional[str] = None, - archives: Optional[str] = None, - driver_class_path: Optional[str] = None, - jars: Optional[str] = None, - java_class: Optional[str] = None, - packages: Optional[str] = None, - exclude_packages: Optional[str] = None, - repositories: Optional[str] = None, - total_executor_cores: Optional[int] = None, - executor_cores: Optional[int] = None, - executor_memory: Optional[str] = None, - driver_memory: Optional[str] = None, - keytab: Optional[str] = None, - principal: Optional[str] = None, - proxy_user: Optional[str] = None, + files: str | None = None, + py_files: str | None = None, + archives: str | None = None, + driver_class_path: str | None = None, + jars: str | None = None, + java_class: str | None = None, + packages: str | None = None, + exclude_packages: str | None = None, + repositories: str | None = None, + total_executor_cores: int | None = None, + executor_cores: int | None = None, + executor_memory: str | None = None, + driver_memory: str | None = None, + keytab: str | None = None, + principal: str | None = None, + proxy_user: str | None = None, name: str = 'arrow-spark', - num_executors: Optional[int] = None, + num_executors: int | None = None, status_poll_interval: int = 1, - application_args: Optional[List[Any]] = None, - env_vars: Optional[Dict[str, Any]] = None, + application_args: list[Any] | None = None, + env_vars: dict[str, Any] | None = None, verbose: bool = False, - spark_binary: Optional[str] = None, + spark_binary: str | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -147,10 +148,10 @@ def __init__( self._env_vars = env_vars self._verbose = verbose self._spark_binary = spark_binary - self._hook: Optional[SparkSubmitHook] = None + self._hook: SparkSubmitHook | None = None self._conn_id = conn_id - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: """Call the SparkSubmitHook to run the provided spark job""" if self._hook is None: self._hook = self._get_hook() diff --git a/airflow/providers/apache/sqoop/hooks/sqoop.py b/airflow/providers/apache/sqoop/hooks/sqoop.py index 65ed7500cb5ae..27b736ffcec68 100644 --- a/airflow/providers/apache/sqoop/hooks/sqoop.py +++ b/airflow/providers/apache/sqoop/hooks/sqoop.py @@ -15,12 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# - """This module contains a sqoop 1.x hook""" +from __future__ import annotations + import subprocess from copy import deepcopy -from typing import Any, Dict, List, Optional +from typing import Any from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook @@ -57,10 +57,10 @@ def __init__( self, conn_id: str = default_conn_name, verbose: bool = False, - num_mappers: Optional[int] = None, - hcatalog_database: Optional[str] = None, - hcatalog_table: Optional[str] = None, - properties: Optional[Dict[str, Any]] = None, + num_mappers: int | None = None, + hcatalog_database: str | None = None, + hcatalog_table: str | None = None, + properties: dict[str, Any] | None = None, ) -> None: # No mutable types in the default parameters super().__init__() @@ -83,7 +83,7 @@ def __init__( def get_conn(self) -> Any: return self.conn - def cmd_mask_password(self, cmd_orig: List[str]) -> List[str]: + def cmd_mask_password(self, cmd_orig: list[str]) -> list[str]: """Mask command password for safety""" cmd = deepcopy(cmd_orig) try: @@ -93,7 +93,7 @@ def cmd_mask_password(self, cmd_orig: List[str]) -> List[str]: self.log.debug("No password in sqoop cmd") return cmd - def popen(self, cmd: List[str], **kwargs: Any) -> None: + def popen(self, cmd: list[str], **kwargs: Any) -> None: """ Remote Popen @@ -112,7 +112,7 @@ def popen(self, cmd: List[str], **kwargs: Any) -> None: if sub_process.returncode: raise AirflowException(f"Sqoop command failed: {masked_cmd}") - def _prepare_command(self, export: bool = False) -> List[str]: + def _prepare_command(self, export: bool = False) -> list[str]: sqoop_cmd_type = "export" if export else "import" connection_cmd = ["sqoop", sqoop_cmd_type] @@ -158,7 +158,7 @@ def _prepare_command(self, export: bool = False) -> List[str]: return connection_cmd @staticmethod - def _get_export_format_argument(file_type: str = 'text') -> List[str]: + def _get_export_format_argument(file_type: str = 'text') -> list[str]: if file_type == "avro": return ["--as-avrodatafile"] elif file_type == "sequence": @@ -172,14 +172,14 @@ def _get_export_format_argument(file_type: str = 'text') -> List[str]: def _import_cmd( self, - target_dir: Optional[str], + target_dir: str | None, append: bool, file_type: str, - split_by: Optional[str], - direct: Optional[bool], + split_by: str | None, + direct: bool | None, driver: Any, extra_import_options: Any, - ) -> List[str]: + ) -> list[str]: cmd = self._prepare_command(export=False) @@ -211,16 +211,16 @@ def _import_cmd( def import_table( self, table: str, - target_dir: Optional[str] = None, + target_dir: str | None = None, append: bool = False, file_type: str = "text", - columns: Optional[str] = None, - split_by: Optional[str] = None, - where: Optional[str] = None, + columns: str | None = None, + split_by: str | None = None, + where: str | None = None, direct: bool = False, driver: Any = None, - extra_import_options: Optional[Dict[str, Any]] = None, - schema: Optional[str] = None, + extra_import_options: dict[str, Any] | None = None, + schema: str | None = None, ) -> Any: """ Imports table from remote location to target dir. Arguments are @@ -257,13 +257,13 @@ def import_table( def import_query( self, query: str, - target_dir: Optional[str] = None, + target_dir: str | None = None, append: bool = False, file_type: str = "text", - split_by: Optional[str] = None, - direct: Optional[bool] = None, - driver: Optional[Any] = None, - extra_import_options: Optional[Dict[str, Any]] = None, + split_by: str | None = None, + direct: bool | None = None, + driver: Any | None = None, + extra_import_options: dict[str, Any] | None = None, ) -> Any: """ Imports a specific query from the rdbms to hdfs @@ -288,21 +288,21 @@ def import_query( def _export_cmd( self, table: str, - export_dir: Optional[str] = None, - input_null_string: Optional[str] = None, - input_null_non_string: Optional[str] = None, - staging_table: Optional[str] = None, + export_dir: str | None = None, + input_null_string: str | None = None, + input_null_non_string: str | None = None, + staging_table: str | None = None, clear_staging_table: bool = False, - enclosed_by: Optional[str] = None, - escaped_by: Optional[str] = None, - input_fields_terminated_by: Optional[str] = None, - input_lines_terminated_by: Optional[str] = None, - input_optionally_enclosed_by: Optional[str] = None, + enclosed_by: str | None = None, + escaped_by: str | None = None, + input_fields_terminated_by: str | None = None, + input_lines_terminated_by: str | None = None, + input_optionally_enclosed_by: str | None = None, batch: bool = False, relaxed_isolation: bool = False, - extra_export_options: Optional[Dict[str, Any]] = None, - schema: Optional[str] = None, - ) -> List[str]: + extra_export_options: dict[str, Any] | None = None, + schema: str | None = None, + ) -> list[str]: cmd = self._prepare_command(export=True) @@ -359,20 +359,20 @@ def _export_cmd( def export_table( self, table: str, - export_dir: Optional[str] = None, - input_null_string: Optional[str] = None, - input_null_non_string: Optional[str] = None, - staging_table: Optional[str] = None, + export_dir: str | None = None, + input_null_string: str | None = None, + input_null_non_string: str | None = None, + staging_table: str | None = None, clear_staging_table: bool = False, - enclosed_by: Optional[str] = None, - escaped_by: Optional[str] = None, - input_fields_terminated_by: Optional[str] = None, - input_lines_terminated_by: Optional[str] = None, - input_optionally_enclosed_by: Optional[str] = None, + enclosed_by: str | None = None, + escaped_by: str | None = None, + input_fields_terminated_by: str | None = None, + input_lines_terminated_by: str | None = None, + input_optionally_enclosed_by: str | None = None, batch: bool = False, relaxed_isolation: bool = False, - extra_export_options: Optional[Dict[str, Any]] = None, - schema: Optional[str] = None, + extra_export_options: dict[str, Any] | None = None, + schema: str | None = None, ) -> None: """ Exports Hive table to remote location. Arguments are copies of direct diff --git a/airflow/providers/apache/sqoop/operators/sqoop.py b/airflow/providers/apache/sqoop/operators/sqoop.py index 3ad7c5fe1a464..0cdebb2c44ba8 100644 --- a/airflow/providers/apache/sqoop/operators/sqoop.py +++ b/airflow/providers/apache/sqoop/operators/sqoop.py @@ -15,11 +15,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# """This module contains a sqoop 1 operator""" +from __future__ import annotations + import os import signal -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -118,37 +119,37 @@ def __init__( *, conn_id: str = 'sqoop_default', cmd_type: str = 'import', - table: Optional[str] = None, - query: Optional[str] = None, - target_dir: Optional[str] = None, + table: str | None = None, + query: str | None = None, + target_dir: str | None = None, append: bool = False, file_type: str = 'text', - columns: Optional[str] = None, - num_mappers: Optional[int] = None, - split_by: Optional[str] = None, - where: Optional[str] = None, - export_dir: Optional[str] = None, - input_null_string: Optional[str] = None, - input_null_non_string: Optional[str] = None, - staging_table: Optional[str] = None, + columns: str | None = None, + num_mappers: int | None = None, + split_by: str | None = None, + where: str | None = None, + export_dir: str | None = None, + input_null_string: str | None = None, + input_null_non_string: str | None = None, + staging_table: str | None = None, clear_staging_table: bool = False, - enclosed_by: Optional[str] = None, - escaped_by: Optional[str] = None, - input_fields_terminated_by: Optional[str] = None, - input_lines_terminated_by: Optional[str] = None, - input_optionally_enclosed_by: Optional[str] = None, + enclosed_by: str | None = None, + escaped_by: str | None = None, + input_fields_terminated_by: str | None = None, + input_lines_terminated_by: str | None = None, + input_optionally_enclosed_by: str | None = None, batch: bool = False, direct: bool = False, - driver: Optional[Any] = None, + driver: Any | None = None, verbose: bool = False, relaxed_isolation: bool = False, - properties: Optional[Dict[str, Any]] = None, - hcatalog_database: Optional[str] = None, - hcatalog_table: Optional[str] = None, + properties: dict[str, Any] | None = None, + hcatalog_database: str | None = None, + hcatalog_table: str | None = None, create_hcatalog_table: bool = False, - extra_import_options: Optional[Dict[str, Any]] = None, - extra_export_options: Optional[Dict[str, Any]] = None, - schema: Optional[str] = None, + extra_import_options: dict[str, Any] | None = None, + extra_export_options: dict[str, Any] | None = None, + schema: str | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) @@ -184,10 +185,10 @@ def __init__( self.properties = properties self.extra_import_options = extra_import_options or {} self.extra_export_options = extra_export_options or {} - self.hook: Optional[SqoopHook] = None + self.hook: SqoopHook | None = None self.schema = schema - def execute(self, context: "Context") -> None: + def execute(self, context: Context) -> None: """Execute sqoop job""" if self.hook is None: self.hook = self._get_hook() diff --git a/tests/providers/apache/beam/hooks/test_beam.py b/tests/providers/apache/beam/hooks/test_beam.py index a69560fe84926..431bc24021f47 100644 --- a/tests/providers/apache/beam/hooks/test_beam.py +++ b/tests/providers/apache/beam/hooks/test_beam.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import copy import os diff --git a/tests/providers/apache/beam/operators/test_beam.py b/tests/providers/apache/beam/operators/test_beam.py index 23700dc71bc85..f8654ed5bbdd1 100644 --- a/tests/providers/apache/beam/operators/test_beam.py +++ b/tests/providers/apache/beam/operators/test_beam.py @@ -14,7 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import unittest from unittest import mock from unittest.mock import MagicMock diff --git a/tests/providers/apache/cassandra/hooks/test_cassandra.py b/tests/providers/apache/cassandra/hooks/test_cassandra.py index 8b588e737b948..a8f88d3c420c5 100644 --- a/tests/providers/apache/cassandra/hooks/test_cassandra.py +++ b/tests/providers/apache/cassandra/hooks/test_cassandra.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from unittest import mock diff --git a/tests/providers/apache/cassandra/sensors/test_record.py b/tests/providers/apache/cassandra/sensors/test_record.py index d3252366d4d32..825042daa7656 100644 --- a/tests/providers/apache/cassandra/sensors/test_record.py +++ b/tests/providers/apache/cassandra/sensors/test_record.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from unittest.mock import patch diff --git a/tests/providers/apache/cassandra/sensors/test_table.py b/tests/providers/apache/cassandra/sensors/test_table.py index af3e2128fc07f..ad86d5838e4e1 100644 --- a/tests/providers/apache/cassandra/sensors/test_table.py +++ b/tests/providers/apache/cassandra/sensors/test_table.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from unittest.mock import patch diff --git a/tests/providers/apache/drill/hooks/test_drill.py b/tests/providers/apache/drill/hooks/test_drill.py index 3e4e5c9af0b99..8087268c9836d 100644 --- a/tests/providers/apache/drill/hooks/test_drill.py +++ b/tests/providers/apache/drill/hooks/test_drill.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import unittest from unittest.mock import MagicMock diff --git a/tests/providers/apache/drill/operators/test_drill.py b/tests/providers/apache/drill/operators/test_drill.py index 32cdddf1e33cc..4832287bbbbd7 100644 --- a/tests/providers/apache/drill/operators/test_drill.py +++ b/tests/providers/apache/drill/operators/test_drill.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest diff --git a/tests/providers/apache/druid/hooks/test_druid.py b/tests/providers/apache/druid/hooks/test_druid.py index a30766adccb72..1a3a5c18400b8 100644 --- a/tests/providers/apache/druid/hooks/test_druid.py +++ b/tests/providers/apache/druid/hooks/test_druid.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import unittest from unittest.mock import MagicMock, patch diff --git a/tests/providers/apache/druid/operators/test_druid.py b/tests/providers/apache/druid/operators/test_druid.py index e9000cf5ee624..de02024385139 100644 --- a/tests/providers/apache/druid/operators/test_druid.py +++ b/tests/providers/apache/druid/operators/test_druid.py @@ -15,7 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import json from airflow.providers.apache.druid.operators.druid import DruidOperator diff --git a/tests/providers/apache/druid/transfers/test_hive_to_druid.py b/tests/providers/apache/druid/transfers/test_hive_to_druid.py index 944aa157f1a5a..2265d187b12d8 100644 --- a/tests/providers/apache/druid/transfers/test_hive_to_druid.py +++ b/tests/providers/apache/druid/transfers/test_hive_to_druid.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import unittest diff --git a/tests/providers/apache/hdfs/hooks/test_hdfs.py b/tests/providers/apache/hdfs/hooks/test_hdfs.py index 58bab173f6c00..8db618aefb673 100644 --- a/tests/providers/apache/hdfs/hooks/test_hdfs.py +++ b/tests/providers/apache/hdfs/hooks/test_hdfs.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import json import unittest diff --git a/tests/providers/apache/hdfs/hooks/test_webhdfs.py b/tests/providers/apache/hdfs/hooks/test_webhdfs.py index 473ef025956b8..c9e28c5d909d4 100644 --- a/tests/providers/apache/hdfs/hooks/test_webhdfs.py +++ b/tests/providers/apache/hdfs/hooks/test_webhdfs.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from unittest.mock import call, patch diff --git a/tests/providers/apache/hdfs/sensors/test_hdfs.py b/tests/providers/apache/hdfs/sensors/test_hdfs.py index c6630cfead2ad..86b99e14c8206 100644 --- a/tests/providers/apache/hdfs/sensors/test_hdfs.py +++ b/tests/providers/apache/hdfs/sensors/test_hdfs.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import logging import re import unittest diff --git a/tests/providers/apache/hdfs/sensors/test_web_hdfs.py b/tests/providers/apache/hdfs/sensors/test_web_hdfs.py index 9af9442c334f3..991d53fcd2713 100644 --- a/tests/providers/apache/hdfs/sensors/test_web_hdfs.py +++ b/tests/providers/apache/hdfs/sensors/test_web_hdfs.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations from unittest import mock diff --git a/tests/providers/apache/hive/__init__.py b/tests/providers/apache/hive/__init__.py index 25212cc10caaa..03c4ebcc6e885 100644 --- a/tests/providers/apache/hive/__init__.py +++ b/tests/providers/apache/hive/__init__.py @@ -15,9 +15,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations from datetime import datetime -from typing import Optional from unittest import TestCase from unittest.mock import MagicMock @@ -160,7 +160,7 @@ def __init__(self, *args, **kwargs): class MockSubProcess: PIPE = -1 STDOUT = -2 - returncode: Optional[int] = None + returncode: int | None = None def __init__(self, *args, **kwargs): self.stdout = MockStdOut(*args, **kwargs) diff --git a/tests/providers/apache/hive/hooks/test_hive.py b/tests/providers/apache/hive/hooks/test_hive.py index c1fe5bd2235eb..b06cf7f15a9a9 100644 --- a/tests/providers/apache/hive/hooks/test_hive.py +++ b/tests/providers/apache/hive/hooks/test_hive.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import datetime import itertools diff --git a/tests/providers/apache/hive/operators/test_hive.py b/tests/providers/apache/hive/operators/test_hive.py index 97027a1ea60dd..a26159269e26a 100644 --- a/tests/providers/apache/hive/operators/test_hive.py +++ b/tests/providers/apache/hive/operators/test_hive.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import os import unittest diff --git a/tests/providers/apache/hive/operators/test_hive_stats.py b/tests/providers/apache/hive/operators/test_hive_stats.py index 3c2260b799d98..b45cb90b8bbcd 100644 --- a/tests/providers/apache/hive/operators/test_hive_stats.py +++ b/tests/providers/apache/hive/operators/test_hive_stats.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import os import re diff --git a/tests/providers/apache/hive/sensors/test_hdfs.py b/tests/providers/apache/hive/sensors/test_hdfs.py index aa4fca0138ce2..429849bcd1877 100644 --- a/tests/providers/apache/hive/sensors/test_hdfs.py +++ b/tests/providers/apache/hive/sensors/test_hdfs.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import os import unittest diff --git a/tests/providers/apache/hive/sensors/test_hive_partition.py b/tests/providers/apache/hive/sensors/test_hive_partition.py index f28c0e090a1fa..4ddc3be1dc9e8 100644 --- a/tests/providers/apache/hive/sensors/test_hive_partition.py +++ b/tests/providers/apache/hive/sensors/test_hive_partition.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import os import unittest diff --git a/tests/providers/apache/hive/sensors/test_metastore_partition.py b/tests/providers/apache/hive/sensors/test_metastore_partition.py index 3e6a0449ab252..e9372e3eeb195 100644 --- a/tests/providers/apache/hive/sensors/test_metastore_partition.py +++ b/tests/providers/apache/hive/sensors/test_metastore_partition.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import os import unittest diff --git a/tests/providers/apache/hive/sensors/test_named_hive_partition.py b/tests/providers/apache/hive/sensors/test_named_hive_partition.py index bd579ab173f6b..7ddd82df7a3a1 100644 --- a/tests/providers/apache/hive/sensors/test_named_hive_partition.py +++ b/tests/providers/apache/hive/sensors/test_named_hive_partition.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import os import unittest from datetime import timedelta diff --git a/tests/providers/apache/hive/transfers/test_hive_to_mysql.py b/tests/providers/apache/hive/transfers/test_hive_to_mysql.py index 7e056a17bab5d..c4fa7e7f68719 100644 --- a/tests/providers/apache/hive/transfers/test_hive_to_mysql.py +++ b/tests/providers/apache/hive/transfers/test_hive_to_mysql.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import os import re import unittest diff --git a/tests/providers/apache/hive/transfers/test_hive_to_samba.py b/tests/providers/apache/hive/transfers/test_hive_to_samba.py index 6620943758b70..ccd80389c9743 100644 --- a/tests/providers/apache/hive/transfers/test_hive_to_samba.py +++ b/tests/providers/apache/hive/transfers/test_hive_to_samba.py @@ -15,6 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + import os import unittest from unittest.mock import MagicMock, Mock, PropertyMock, patch diff --git a/tests/providers/apache/hive/transfers/test_mssql_to_hive.py b/tests/providers/apache/hive/transfers/test_mssql_to_hive.py index 8fae259a35c47..8398341b54fab 100644 --- a/tests/providers/apache/hive/transfers/test_mssql_to_hive.py +++ b/tests/providers/apache/hive/transfers/test_mssql_to_hive.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - +from __future__ import annotations import unittest from collections import OrderedDict diff --git a/tests/providers/apache/hive/transfers/test_mysql_to_hive.py b/tests/providers/apache/hive/transfers/test_mysql_to_hive.py index fd5bd05c92f5d..54abfdc55b62e 100644 --- a/tests/providers/apache/hive/transfers/test_mysql_to_hive.py +++ b/tests/providers/apache/hive/transfers/test_mysql_to_hive.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import textwrap from collections import OrderedDict diff --git a/tests/providers/apache/hive/transfers/test_s3_to_hive.py b/tests/providers/apache/hive/transfers/test_s3_to_hive.py index 26d7430fdb49e..a808f17eb89f0 100644 --- a/tests/providers/apache/hive/transfers/test_s3_to_hive.py +++ b/tests/providers/apache/hive/transfers/test_s3_to_hive.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import bz2 import errno diff --git a/tests/providers/apache/hive/transfers/test_vertica_to_hive.py b/tests/providers/apache/hive/transfers/test_vertica_to_hive.py index d71860aad76dd..e6da9c24595dc 100644 --- a/tests/providers/apache/hive/transfers/test_vertica_to_hive.py +++ b/tests/providers/apache/hive/transfers/test_vertica_to_hive.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import datetime import unittest diff --git a/tests/providers/apache/kylin/hooks/test_kylin.py b/tests/providers/apache/kylin/hooks/test_kylin.py index c802197d8c357..f7009bba0475c 100644 --- a/tests/providers/apache/kylin/hooks/test_kylin.py +++ b/tests/providers/apache/kylin/hooks/test_kylin.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import unittest from unittest.mock import MagicMock, patch diff --git a/tests/providers/apache/kylin/operators/test_kylin_cube.py b/tests/providers/apache/kylin/operators/test_kylin_cube.py index bb17a40567671..a72958ae2449a 100644 --- a/tests/providers/apache/kylin/operators/test_kylin_cube.py +++ b/tests/providers/apache/kylin/operators/test_kylin_cube.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from datetime import datetime diff --git a/tests/providers/apache/livy/hooks/test_livy.py b/tests/providers/apache/livy/hooks/test_livy.py index 7d663e0fd0b6d..3ee78e369a735 100644 --- a/tests/providers/apache/livy/hooks/test_livy.py +++ b/tests/providers/apache/livy/hooks/test_livy.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import json import unittest diff --git a/tests/providers/apache/livy/operators/test_livy.py b/tests/providers/apache/livy/operators/test_livy.py index f1c43bf5f1759..374e45f252a8c 100644 --- a/tests/providers/apache/livy/operators/test_livy.py +++ b/tests/providers/apache/livy/operators/test_livy.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import logging import unittest diff --git a/tests/providers/apache/livy/sensors/test_livy.py b/tests/providers/apache/livy/sensors/test_livy.py index 654f120ff8c8b..50512341053ea 100644 --- a/tests/providers/apache/livy/sensors/test_livy.py +++ b/tests/providers/apache/livy/sensors/test_livy.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. - +from __future__ import annotations import unittest from unittest.mock import patch diff --git a/tests/providers/apache/pig/hooks/test_pig.py b/tests/providers/apache/pig/hooks/test_pig.py index 12fbc50885293..db66a683698e2 100644 --- a/tests/providers/apache/pig/hooks/test_pig.py +++ b/tests/providers/apache/pig/hooks/test_pig.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from unittest import mock diff --git a/tests/providers/apache/pig/operators/test_pig.py b/tests/providers/apache/pig/operators/test_pig.py index e391beb88b28f..fd891808fe056 100644 --- a/tests/providers/apache/pig/operators/test_pig.py +++ b/tests/providers/apache/pig/operators/test_pig.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import unittest from unittest import mock diff --git a/tests/providers/apache/pinot/hooks/test_pinot.py b/tests/providers/apache/pinot/hooks/test_pinot.py index 9c363f3d52ee2..476dfce56c447 100644 --- a/tests/providers/apache/pinot/hooks/test_pinot.py +++ b/tests/providers/apache/pinot/hooks/test_pinot.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import io import os diff --git a/tests/providers/apache/spark/hooks/test_spark_jdbc.py b/tests/providers/apache/spark/hooks/test_spark_jdbc.py index 75271c4b93074..25d136c083066 100644 --- a/tests/providers/apache/spark/hooks/test_spark_jdbc.py +++ b/tests/providers/apache/spark/hooks/test_spark_jdbc.py @@ -15,7 +15,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations + import unittest from airflow.models import Connection diff --git a/tests/providers/apache/spark/hooks/test_spark_jdbc_script.py b/tests/providers/apache/spark/hooks/test_spark_jdbc_script.py index c0ac0130ccab6..c3d45af48d8ca 100644 --- a/tests/providers/apache/spark/hooks/test_spark_jdbc_script.py +++ b/tests/providers/apache/spark/hooks/test_spark_jdbc_script.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + from unittest import mock import pytest diff --git a/tests/providers/apache/spark/hooks/test_spark_sql.py b/tests/providers/apache/spark/hooks/test_spark_sql.py index 4cdc64c1f3510..7a7a5ebf35c22 100644 --- a/tests/providers/apache/spark/hooks/test_spark_sql.py +++ b/tests/providers/apache/spark/hooks/test_spark_sql.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import io import unittest diff --git a/tests/providers/apache/spark/hooks/test_spark_submit.py b/tests/providers/apache/spark/hooks/test_spark_submit.py index 041d03812e15c..190625e8bba15 100644 --- a/tests/providers/apache/spark/hooks/test_spark_submit.py +++ b/tests/providers/apache/spark/hooks/test_spark_submit.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations import io import os diff --git a/tests/providers/apache/spark/operators/test_spark_jdbc.py b/tests/providers/apache/spark/operators/test_spark_jdbc.py index f8c1a564aca22..0c7b8c149f84a 100644 --- a/tests/providers/apache/spark/operators/test_spark_jdbc.py +++ b/tests/providers/apache/spark/operators/test_spark_jdbc.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import unittest diff --git a/tests/providers/apache/spark/operators/test_spark_sql.py b/tests/providers/apache/spark/operators/test_spark_sql.py index 6d1e664eb1bb2..6cd99fcb976ff 100644 --- a/tests/providers/apache/spark/operators/test_spark_sql.py +++ b/tests/providers/apache/spark/operators/test_spark_sql.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import datetime import unittest diff --git a/tests/providers/apache/spark/operators/test_spark_submit.py b/tests/providers/apache/spark/operators/test_spark_submit.py index 6cbdd0cf512d2..151713429b46c 100644 --- a/tests/providers/apache/spark/operators/test_spark_submit.py +++ b/tests/providers/apache/spark/operators/test_spark_submit.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import unittest from datetime import timedelta diff --git a/tests/providers/apache/sqoop/hooks/test_sqoop.py b/tests/providers/apache/sqoop/hooks/test_sqoop.py index c6a40ed85b6d3..048466595f78f 100644 --- a/tests/providers/apache/sqoop/hooks/test_sqoop.py +++ b/tests/providers/apache/sqoop/hooks/test_sqoop.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import collections import json diff --git a/tests/providers/apache/sqoop/operators/test_sqoop.py b/tests/providers/apache/sqoop/operators/test_sqoop.py index 9b8b205700364..35cec3b1be37a 100644 --- a/tests/providers/apache/sqoop/operators/test_sqoop.py +++ b/tests/providers/apache/sqoop/operators/test_sqoop.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# +from __future__ import annotations import datetime import unittest