Consume AirflowOptionalProviderFeatureException from compat sdk in providers (apache#60335)

* Consume AirflowOptionalProviderFeatureException from compat sdk in providers

* fixing static checks

* fixing static checks

Author: amoghrajesh

providers/amazon/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -27,15 +27,14 @@
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
 from airflow.cli.cli_config import CLICommand
-from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities
 from airflow.providers.amazon.aws.auth_manager.avp.facade import (
     AwsAuthManagerAmazonVerifiedPermissionsFacade,
     IsAuthorizedRequest,
 )
 from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.providers.common.compat.sdk import conf
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, conf
 
 if TYPE_CHECKING:
     from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod

providers/amazon/src/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py

@@ -22,9 +22,8 @@
 
 import boto3
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.auth_manager.constants import CONF_REGION_NAME_KEY, CONF_SECTION_NAME
-from airflow.providers.common.compat.sdk import conf
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, conf
 from airflow.utils import cli as cli_utils
 
 try:

providers/amazon/src/airflow/providers/amazon/aws/auth_manager/user.py

@@ -16,7 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
 try:
     from airflow.api_fastapi.auth.managers.models.base_user import BaseUser

providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py

@@ -156,7 +156,7 @@ def _get_conn_params(self) -> dict[str, str | None]:
 
     def get_uri(self) -> str:
         """Overridden to use the Athena dialect as driver name."""
-        from airflow.exceptions import AirflowOptionalProviderFeatureException
+        from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
         if URL is None:
             raise AirflowOptionalProviderFeatureException(

providers/amazon/src/airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -29,9 +29,8 @@
 except ImportError:
     URL = create_engine = None  # type: ignore[assignment,misc]
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 
 if TYPE_CHECKING:

providers/amazon/src/airflow/providers/amazon/aws/queues/sqs.py

@@ -19,8 +19,8 @@
 import re
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
 try:
     from airflow.providers.common.messaging.providers.base_provider import BaseMessageQueueProvider

providers/amazon/src/airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py

@@ -28,7 +28,7 @@
 try:
     from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
 except ModuleNotFoundError as e:
-    from airflow.exceptions import AirflowOptionalProviderFeatureException
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
     raise AirflowOptionalProviderFeatureException(e)
 

providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -183,7 +183,7 @@ def _fix_dtypes(df: pd.DataFrame, file_format: FILE_FORMAT) -> None:
             import numpy as np
             import pandas as pd
         except ImportError as e:
-            from airflow.exceptions import AirflowOptionalProviderFeatureException
+            from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
             raise AirflowOptionalProviderFeatureException(e)
 

providers/amazon/tests/unit/amazon/aws/hooks/test_hooks_signature.py

@@ -23,8 +23,8 @@
 
 import pytest
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
 BASE_AWS_HOOKS = ["AwsGenericHook", "AwsBaseHook"]
 ALLOWED_THICK_HOOKS_PARAMETERS: dict[str, set[str]] = {

providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py

@@ -378,7 +378,7 @@ def start_go_pipeline(
         try:
             from airflow.providers.google.go_module_utils import init_module, install_dependencies
         except ImportError:
-            from airflow.exceptions import AirflowOptionalProviderFeatureException
+            from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
             raise AirflowOptionalProviderFeatureException(
                 "Failed to import apache-airflow-google-provider. To start a go pipeline, please install the"