apache-airflow-providers-microsoft-azure 12.0.0rc2__py3-none-any.whl → 12.1.0__py3-none-any.whl

airflow/providers/microsoft/azure/hooks/container_registry.py

@@ -22,15 +22,14 @@ from __future__ import annotations
 from functools import cached_property
 from typing import Any
 
-from azure.mgmt.containerinstance.models import ImageRegistryCredential
-from azure.mgmt.containerregistry import ContainerRegistryManagementClient
-
 from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
+from azure.mgmt.containerinstance.models import ImageRegistryCredential
+from azure.mgmt.containerregistry import ContainerRegistryManagementClient
 
 
 class AzureContainerRegistryHook(BaseHook):

airflow/providers/microsoft/azure/hooks/container_volume.py

@@ -18,15 +18,14 @@ from __future__ import annotations
 
 from typing import Any
 
-from azure.mgmt.containerinstance.models import AzureFileVolume, Volume
-from azure.mgmt.storage import StorageManagementClient
-
 from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_field,
     get_sync_default_azure_credential,
 )
+from azure.mgmt.containerinstance.models import AzureFileVolume, Volume
+from azure.mgmt.storage import StorageManagementClient
 
 
 class AzureContainerVolumeHook(BaseHook):

airflow/providers/microsoft/azure/hooks/cosmos.py

@@ -30,11 +30,6 @@ import uuid
 from typing import TYPE_CHECKING, Any, Union
 from urllib.parse import urlparse
 
-from azure.cosmos import PartitionKey
-from azure.cosmos.cosmos_client import CosmosClient
-from azure.cosmos.exceptions import CosmosHttpResponseError
-from azure.mgmt.cosmosdb import CosmosDBManagementClient
-
 from airflow.exceptions import AirflowBadRequest, AirflowException
 from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
@@ -42,6 +37,10 @@ from airflow.providers.microsoft.azure.utils import (
     get_field,
     get_sync_default_azure_credential,
 )
+from azure.cosmos import PartitionKey
+from azure.cosmos.cosmos_client import CosmosClient
+from azure.cosmos.exceptions import CosmosHttpResponseError
+from azure.mgmt.cosmosdb import CosmosDBManagementClient
 
 if TYPE_CHECKING:
     PartitionKeyType = Union[str, list[str]]

airflow/providers/microsoft/azure/hooks/data_factory.py

@@ -39,13 +39,6 @@ from functools import wraps
 from typing import IO, TYPE_CHECKING, Any, Callable, TypeVar, Union, cast
 
 from asgiref.sync import sync_to_async
-from azure.identity import ClientSecretCredential, DefaultAzureCredential
-from azure.identity.aio import (
-    ClientSecretCredential as AsyncClientSecretCredential,
-    DefaultAzureCredential as AsyncDefaultAzureCredential,
-)
-from azure.mgmt.datafactory import DataFactoryManagementClient
-from azure.mgmt.datafactory.aio import DataFactoryManagementClient as AsyncDataFactoryManagementClient
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
@@ -54,6 +47,13 @@ from airflow.providers.microsoft.azure.utils import (
     get_async_default_azure_credential,
     get_sync_default_azure_credential,
 )
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
+from azure.identity.aio import (
+    ClientSecretCredential as AsyncClientSecretCredential,
+    DefaultAzureCredential as AsyncDefaultAzureCredential,
+)
+from azure.mgmt.datafactory import DataFactoryManagementClient
+from azure.mgmt.datafactory.aio import DataFactoryManagementClient as AsyncDataFactoryManagementClient
 
 if TYPE_CHECKING:
     from azure.core.polling import LROPoller

airflow/providers/microsoft/azure/hooks/data_lake.py

@@ -20,6 +20,14 @@ from __future__ import annotations
 from functools import cached_property
 from typing import Any, Union
 
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+from airflow.providers.microsoft.azure.utils import (
+    AzureIdentityCredentialAdapter,
+    add_managed_identity_connection_widgets,
+    get_field,
+    get_sync_default_azure_credential,
+)
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
 from azure.datalake.store import core, lib, multithread
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
@@ -32,15 +40,6 @@ from azure.storage.filedatalake import (
     FileSystemProperties,
 )
 
-from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import (
-    AzureIdentityCredentialAdapter,
-    add_managed_identity_connection_widgets,
-    get_field,
-    get_sync_default_azure_credential,
-)
-
 Credentials = Union[ClientSecretCredential, AzureIdentityCredentialAdapter, DefaultAzureCredential]
 
 

airflow/providers/microsoft/azure/hooks/fileshare.py

@@ -19,13 +19,12 @@ from __future__ import annotations
 
 from typing import IO, Any
 
-from azure.storage.fileshare import FileProperties, ShareDirectoryClient, ShareFileClient, ShareServiceClient
-
 from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
     add_managed_identity_connection_widgets,
     get_sync_default_azure_credential,
 )
+from azure.storage.fileshare import FileProperties, ShareDirectoryClient, ShareFileClient, ShareServiceClient
 
 
 class AzureFileShareHook(BaseHook):

airflow/providers/microsoft/azure/hooks/msgraph.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import json
+from ast import literal_eval
 from contextlib import suppress
 from http import HTTPStatus
 from io import BytesIO
@@ -26,8 +27,7 @@ from typing import TYPE_CHECKING, Any
 from urllib.parse import quote, urljoin, urlparse
 
 import httpx
-from azure.identity import ClientSecretCredential
-from httpx import Timeout
+from httpx import AsyncHTTPTransport, Timeout
 from kiota_abstractions.api_error import APIError
 from kiota_abstractions.method import Method
 from kiota_abstractions.request_information import RequestInformation
@@ -43,8 +43,14 @@ from kiota_serialization_text.text_parse_node_factory import TextParseNodeFactor
 from msgraph_core import APIVersion, GraphClientFactory
 from msgraph_core._enums import NationalClouds
 
-from airflow.exceptions import AirflowBadRequest, AirflowException, AirflowNotFoundException
+from airflow.exceptions import (
+    AirflowBadRequest,
+    AirflowConfigException,
+    AirflowException,
+    AirflowNotFoundException,
+)
 from airflow.hooks.base import BaseHook
+from azure.identity import CertificateCredential, ClientSecretCredential
 
 if TYPE_CHECKING:
     from kiota_abstractions.request_adapter import RequestAdapter
@@ -54,6 +60,7 @@ if TYPE_CHECKING:
     from kiota_http.httpx_request_adapter import ResponseType
 
     from airflow.models import Connection
+    from azure.identity._internal.client_credential_base import ClientCredentialBase
 
 
 class DefaultResponseHandler(ResponseHandler):
@@ -107,6 +114,7 @@ class KiotaRequestAdapterHook(BaseHook):
     """
 
     DEFAULT_HEADERS = {"Accept": "application/json;q=1"}
+    DEFAULT_SCOPE = "https://graph.microsoft.com/.default"
     cached_request_adapters: dict[str, tuple[APIVersion, RequestAdapter]] = {}
     conn_type: str = "msgraph"
     conn_name_attr: str = "conn_id"
@@ -119,7 +127,7 @@ class KiotaRequestAdapterHook(BaseHook):
         timeout: float | None = None,
         proxies: dict | None = None,
         host: str = NationalClouds.Global.value,
-        scopes: list[str] | None = None,
+        scopes: str | list[str] | None = None,
         api_version: APIVersion | str | None = None,
     ):
         super().__init__()
@@ -127,7 +135,10 @@ class KiotaRequestAdapterHook(BaseHook):
         self.timeout = timeout
         self.proxies = proxies
         self.host = host
-        self.scopes = scopes or ["https://graph.microsoft.com/.default"]
+        if isinstance(scopes, str):
+            self.scopes = [scopes]
+        else:
+            self.scopes = scopes or [self.DEFAULT_SCOPE]
         self._api_version = self.resolve_api_version_from_value(api_version)
 
     @classmethod
@@ -140,20 +151,21 @@ class KiotaRequestAdapterHook(BaseHook):
         return {
             "tenant_id": StringField(lazy_gettext("Tenant ID"), widget=BS3TextFieldWidget()),
             "api_version": StringField(
-                lazy_gettext("API Version"), widget=BS3TextFieldWidget(), default="v1.0"
+                lazy_gettext("API Version"), widget=BS3TextFieldWidget(), default=APIVersion.v1.value
             ),
             "authority": StringField(lazy_gettext("Authority"), widget=BS3TextFieldWidget()),
+            "certificate_path": StringField(lazy_gettext("Certificate path"), widget=BS3TextFieldWidget()),
+            "certificate_data": StringField(lazy_gettext("Certificate data"), widget=BS3TextFieldWidget()),
             "scopes": StringField(
                 lazy_gettext("Scopes"),
                 widget=BS3TextFieldWidget(),
-                default="https://graph.microsoft.com/.default",
+                default=cls.DEFAULT_SCOPE,
             ),
             "disable_instance_discovery": BooleanField(
                 lazy_gettext("Disable instance discovery"), default=False
             ),
-            "allowed_hosts": StringField(lazy_gettext("Allowed"), widget=BS3TextFieldWidget()),
+            "allowed_hosts": StringField(lazy_gettext("Allowed hosts"), widget=BS3TextFieldWidget()),
             "proxies": StringField(lazy_gettext("Proxies"), widget=BS3TextAreaFieldWidget()),
-            "stream": BooleanField(lazy_gettext("Stream"), default=False),
             "verify": BooleanField(lazy_gettext("Verify"), default=True),
             "trust_env": BooleanField(lazy_gettext("Trust environment"), default=True),
             "base_url": StringField(lazy_gettext("Base URL"), widget=BS3TextFieldWidget()),
@@ -201,24 +213,25 @@ class KiotaRequestAdapterHook(BaseHook):
     @staticmethod
     def format_no_proxy_url(url: str) -> str:
         if "://" not in url:
-            url = f"all://{url}"
+            return f"all://{url}"
         return url
 
     @classmethod
     def to_httpx_proxies(cls, proxies: dict) -> dict:
-        proxies = proxies.copy()
-        if proxies.get("http"):
-            proxies["http://"] = proxies.pop("http")
-        if proxies.get("https"):
-            proxies["https://"] = proxies.pop("https")
-        if proxies.get("no"):
-            for url in proxies.pop("no", "").split(","):
-                proxies[cls.format_no_proxy_url(url.strip())] = None
+        if proxies:
+            proxies = proxies.copy()
+            if proxies.get("http"):
+                proxies["http://"] = AsyncHTTPTransport(proxy=proxies.pop("http"))
+            if proxies.get("https"):
+                proxies["https://"] = AsyncHTTPTransport(proxy=proxies.pop("https"))
+            if proxies.get("no"):
+                for url in proxies.pop("no", "").split(","):
+                    proxies[cls.format_no_proxy_url(url.strip())] = None
         return proxies
 
-    def to_msal_proxies(self, authority: str | None, proxies: dict):
+    def to_msal_proxies(self, authority: str | None, proxies: dict) -> dict | None:
         self.log.debug("authority: %s", authority)
-        if authority:
+        if authority and proxies:
             no_proxies = proxies.get("no")
             self.log.debug("no_proxies: %s", no_proxies)
             if no_proxies:
@@ -241,18 +254,17 @@ class KiotaRequestAdapterHook(BaseHook):
             client_id = connection.login
             client_secret = connection.password
             config = connection.extra_dejson if connection.extra else {}
-            tenant_id = config.get("tenant_id") or config.get("tenantId")
             api_version = self.get_api_version(config)
             host = self.get_host(connection)
             base_url = config.get("base_url", urljoin(host, api_version))
             authority = config.get("authority")
-            proxies = self.proxies or config.get("proxies", {})
-            msal_proxies = self.to_msal_proxies(authority=authority, proxies=proxies)
+            proxies = self.get_proxies(config)
             httpx_proxies = self.to_httpx_proxies(proxies=proxies)
             scopes = config.get("scopes", self.scopes)
+            if isinstance(scopes, str):
+                scopes = scopes.split(",")
             verify = config.get("verify", True)
             trust_env = config.get("trust_env", False)
-            disable_instance_discovery = config.get("disable_instance_discovery", False)
             allowed_hosts = (config.get("allowed_hosts", authority) or "").split(",")
 
             self.log.info(
@@ -262,7 +274,6 @@ class KiotaRequestAdapterHook(BaseHook):
             )
             self.log.info("Host: %s", host)
             self.log.info("Base URL: %s", base_url)
-            self.log.info("Tenant id: %s", tenant_id)
             self.log.info("Client id: %s", client_id)
             self.log.info("Client secret: %s", client_secret)
             self.log.info("API version: %s", api_version)
@@ -271,24 +282,21 @@ class KiotaRequestAdapterHook(BaseHook):
             self.log.info("Timeout: %s", self.timeout)
             self.log.info("Trust env: %s", trust_env)
             self.log.info("Authority: %s", authority)
-            self.log.info("Disable instance discovery: %s", disable_instance_discovery)
             self.log.info("Allowed hosts: %s", allowed_hosts)
             self.log.info("Proxies: %s", proxies)
-            self.log.info("MSAL Proxies: %s", msal_proxies)
             self.log.info("HTTPX Proxies: %s", httpx_proxies)
-            credentials = ClientSecretCredential(
-                tenant_id=tenant_id,  # type: ignore
-                client_id=connection.login,
-                client_secret=connection.password,
+            credentials = self.get_credentials(
+                login=connection.login,
+                password=connection.password,
+                config=config,
                 authority=authority,
-                proxies=msal_proxies,
-                disable_instance_discovery=disable_instance_discovery,
-                connection_verify=verify,
+                verify=verify,
+                proxies=proxies,
             )
             http_client = GraphClientFactory.create_with_default_middleware(
                 api_version=api_version,  # type: ignore
                 client=httpx.AsyncClient(
-                    proxies=httpx_proxies,
+                    mounts=httpx_proxies,
                     timeout=Timeout(timeout=self.timeout),
                     verify=verify,
                     trust_env=trust_env,
@@ -313,6 +321,65 @@ class KiotaRequestAdapterHook(BaseHook):
         self._api_version = api_version
         return request_adapter
 
+    def get_proxies(self, config: dict) -> dict:
+        proxies = self.proxies or config.get("proxies", {})
+        if isinstance(proxies, str):
+            # TODO: Once provider depends on Airflow 2.10 or higher code below won't be needed anymore as
+            #       we could then use the get_extra_dejson method on the connection which deserializes
+            #       nested json. Make sure to use connection.get_extra_dejson(nested=True) instead of
+            #       connection.extra_dejson.
+            with suppress(JSONDecodeError):
+                proxies = json.loads(proxies)
+            with suppress(Exception):
+                proxies = literal_eval(proxies)
+        if not isinstance(proxies, dict):
+            raise AirflowConfigException(
+                f"Proxies must be of type dict, got {type(proxies).__name__} instead!"
+            )
+        return proxies
+
+    def get_credentials(
+        self,
+        login: str | None,
+        password: str | None,
+        config,
+        authority: str | None,
+        verify: bool,
+        proxies: dict,
+    ) -> ClientCredentialBase:
+        tenant_id = config.get("tenant_id") or config.get("tenantId")
+        certificate_path = config.get("certificate_path")
+        certificate_data = config.get("certificate_data")
+        disable_instance_discovery = config.get("disable_instance_discovery", False)
+        msal_proxies = self.to_msal_proxies(authority=authority, proxies=proxies)
+        self.log.info("Tenant id: %s", tenant_id)
+        self.log.info("Certificate path: %s", certificate_path)
+        self.log.info("Certificate data: %s", certificate_data is not None)
+        self.log.info("Authority: %s", authority)
+        self.log.info("Disable instance discovery: %s", disable_instance_discovery)
+        self.log.info("MSAL Proxies: %s", msal_proxies)
+        if certificate_path or certificate_data:
+            return CertificateCredential(
+                tenant_id=tenant_id,  # type: ignore
+                client_id=login,  # type: ignore
+                password=password,
+                certificate_path=certificate_path,
+                certificate_data=certificate_data.encode() if certificate_data else None,
+                authority=authority,
+                proxies=msal_proxies,
+                disable_instance_discovery=disable_instance_discovery,
+                connection_verify=verify,
+            )
+        return ClientSecretCredential(
+            tenant_id=tenant_id,  # type: ignore
+            client_id=login,  # type: ignore
+            client_secret=password,  # type: ignore
+            authority=authority,
+            proxies=msal_proxies,
+            disable_instance_discovery=disable_instance_discovery,
+            connection_verify=verify,
+        )
+
     def test_connection(self):
         """Test HTTP Connection."""
         try:

airflow/providers/microsoft/azure/hooks/synapse.py

@@ -19,11 +19,6 @@ from __future__ import annotations
 import time
 from typing import TYPE_CHECKING, Any, Union
 
-from azure.core.exceptions import ServiceRequestError
-from azure.identity import ClientSecretCredential, DefaultAzureCredential
-from azure.synapse.artifacts import ArtifactsClient
-from azure.synapse.spark import SparkClient
-
 from airflow.exceptions import AirflowException, AirflowTaskTimeout
 from airflow.hooks.base import BaseHook
 from airflow.providers.microsoft.azure.utils import (
@@ -31,6 +26,10 @@ from airflow.providers.microsoft.azure.utils import (
     get_field,
     get_sync_default_azure_credential,
 )
+from azure.core.exceptions import ServiceRequestError
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
+from azure.synapse.artifacts import ArtifactsClient
+from azure.synapse.spark import SparkClient
 
 if TYPE_CHECKING:
     from azure.synapse.artifacts.models import CreateRunResponse, PipelineRun

airflow/providers/microsoft/azure/hooks/wasb.py

@@ -32,6 +32,15 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Any, Union
 
 from asgiref.sync import sync_to_async
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+from airflow.providers.microsoft.azure.utils import (
+    add_managed_identity_connection_widgets,
+    get_async_default_azure_credential,
+    get_sync_default_azure_credential,
+    parse_blob_account_url,
+)
 from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceNotFoundError
 from azure.identity import ClientSecretCredential
 from azure.identity.aio import (
@@ -45,15 +54,6 @@ from azure.storage.blob.aio import (
     ContainerClient as AsyncContainerClient,
 )
 
-from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import (
-    add_managed_identity_connection_widgets,
-    get_async_default_azure_credential,
-    get_sync_default_azure_credential,
-    parse_blob_account_url,
-)
-
 if TYPE_CHECKING:
     from azure.storage.blob._models import BlobProperties
 

airflow/providers/microsoft/azure/log/wasb_task_handler.py

@@ -23,11 +23,10 @@ from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from azure.core.exceptions import HttpResponseError
-
 from airflow.configuration import conf
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
+from azure.core.exceptions import HttpResponseError
 
 if TYPE_CHECKING:
     import logging

airflow/providers/microsoft/azure/operators/adx.py

@@ -28,9 +28,8 @@ from airflow.models import BaseOperator
 from airflow.providers.microsoft.azure.hooks.adx import AzureDataExplorerHook
 
 if TYPE_CHECKING:
-    from azure.kusto.data._models import KustoResultTable
-
     from airflow.utils.context import Context
+    from azure.kusto.data._models import KustoResultTable
 
 
 class AzureDataExplorerQueryOperator(BaseOperator):

airflow/providers/microsoft/azure/operators/asb.py

@@ -19,18 +19,15 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Callable
 
-from azure.core.exceptions import ResourceNotFoundError
-
 from airflow.models import BaseOperator
 from airflow.providers.microsoft.azure.hooks.asb import AdminClientHook, MessageHook
 
 if TYPE_CHECKING:
     import datetime
 
-    from azure.servicebus import ServiceBusMessage
-    from azure.servicebus.management._models import AuthorizationRule
-
     from airflow.utils.context import Context
+    from azure.servicebus import ServiceBusMessage
+    from azure.servicebus.management import AuthorizationRule, CorrelationRuleFilter, SqlRuleFilter
 
     MessageCallback = Callable[[ServiceBusMessage, Context], None]
 
@@ -313,33 +310,23 @@ class AzureServiceBusTopicCreateOperator(BaseOperator):
         # Create the hook
         hook = AdminClientHook(azure_service_bus_conn_id=self.azure_service_bus_conn_id)
 
-        with hook.get_conn() as service_mgmt_conn:
-            try:
-                topic_properties = service_mgmt_conn.get_topic(self.topic_name)
-            except ResourceNotFoundError:
-                topic_properties = None
-            if topic_properties and topic_properties.name == self.topic_name:
-                self.log.info("Topic name already exists")
-                return topic_properties.name
-            topic = service_mgmt_conn.create_topic(
-                topic_name=self.topic_name,
-                default_message_time_to_live=self.default_message_time_to_live,
-                max_size_in_megabytes=self.max_size_in_megabytes,
-                requires_duplicate_detection=self.requires_duplicate_detection,
-                duplicate_detection_history_time_window=self.duplicate_detection_history_time_window,
-                enable_batched_operations=self.enable_batched_operations,
-                size_in_bytes=self.size_in_bytes,
-                filtering_messages_before_publishing=self.filtering_messages_before_publishing,
-                authorization_rules=self.authorization_rules,
-                support_ordering=self.support_ordering,
-                auto_delete_on_idle=self.auto_delete_on_idle,
-                enable_partitioning=self.enable_partitioning,
-                enable_express=self.enable_express,
-                user_metadata=self.user_metadata,
-                max_message_size_in_kilobytes=self.max_message_size_in_kilobytes,
-            )
-            self.log.info("Created Topic %s", topic.name)
-            return topic.name
+        return hook.create_topic(
+            topic_name=self.topic_name,
+            default_message_time_to_live=self.default_message_time_to_live,
+            max_size_in_megabytes=self.max_size_in_megabytes,
+            requires_duplicate_detection=self.requires_duplicate_detection,
+            duplicate_detection_history_time_window=self.duplicate_detection_history_time_window,
+            enable_batched_operations=self.enable_batched_operations,
+            size_in_bytes=self.size_in_bytes,
+            filtering_messages_before_publishing=self.filtering_messages_before_publishing,
+            authorization_rules=self.authorization_rules,
+            support_ordering=self.support_ordering,
+            auto_delete_on_idle=self.auto_delete_on_idle,
+            enable_partitioning=self.enable_partitioning,
+            enable_express=self.enable_express,
+            user_metadata=self.user_metadata,
+            max_message_size_in_kilobytes=self.max_message_size_in_kilobytes,
+        )
 
 
 class AzureServiceBusSubscriptionCreateOperator(BaseOperator):
@@ -378,6 +365,8 @@ class AzureServiceBusSubscriptionCreateOperator(BaseOperator):
     :param auto_delete_on_idle: ISO 8601 time Span idle interval after which the subscription is
         automatically deleted. The minimum duration is 5 minutes. Input value of either
         type ~datetime.timedelta or string in ISO 8601 duration format like "PT300S" is accepted.
+    :param filter_rule: Optional correlation or SQL rule filter to apply on the messages.
+    :param filter_rule_name: Optional rule name to use applying the rule filter to the subscription
     :param azure_service_bus_conn_id: Reference to the
         :ref:`Azure Service Bus connection<howto/connection:azure_service_bus>`.
     """
@@ -402,6 +391,8 @@ class AzureServiceBusSubscriptionCreateOperator(BaseOperator):
         user_metadata: str | None = None,
         forward_dead_lettered_messages_to: str | None = None,
         auto_delete_on_idle: datetime.timedelta | str | None = None,
+        filter_rule: CorrelationRuleFilter | SqlRuleFilter | None = None,
+        filter_rule_name: str | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -419,6 +410,8 @@ class AzureServiceBusSubscriptionCreateOperator(BaseOperator):
         self.forward_dead_lettered_messages_to = forward_dead_lettered_messages_to
         self.auto_delete_on_idle = auto_delete_on_idle
         self.azure_service_bus_conn_id = azure_service_bus_conn_id
+        self.filter_rule = filter_rule
+        self.filter_rule_name = filter_rule_name
 
     def execute(self, context: Context) -> None:
         """Create Subscription in Service Bus namespace, by connecting to Service Bus Admin client."""
@@ -429,24 +422,24 @@ class AzureServiceBusSubscriptionCreateOperator(BaseOperator):
         # Create the hook
         hook = AdminClientHook(azure_service_bus_conn_id=self.azure_service_bus_conn_id)
 
-        with hook.get_conn() as service_mgmt_conn:
-            # create subscription with name
-            subscription = service_mgmt_conn.create_subscription(
-                topic_name=self.topic_name,
-                subscription_name=self.subscription_name,
-                lock_duration=self.lock_duration,
-                requires_session=self.requires_session,
-                default_message_time_to_live=self.default_message_time_to_live,
-                dead_lettering_on_message_expiration=self.dl_on_message_expiration,
-                dead_lettering_on_filter_evaluation_exceptions=self.dl_on_filter_evaluation_exceptions,
-                max_delivery_count=self.max_delivery_count,
-                enable_batched_operations=self.enable_batched_operations,
-                forward_to=self.forward_to,
-                user_metadata=self.user_metadata,
-                forward_dead_lettered_messages_to=self.forward_dead_lettered_messages_to,
-                auto_delete_on_idle=self.auto_delete_on_idle,
-            )
-            self.log.info("Created subscription %s", subscription.name)
+        subscription = hook.create_subscription(
+            topic_name=self.topic_name,
+            subscription_name=self.subscription_name,
+            lock_duration=self.lock_duration,
+            requires_session=self.requires_session,
+            default_message_time_to_live=self.default_message_time_to_live,
+            dead_lettering_on_message_expiration=self.dl_on_message_expiration,
+            dead_lettering_on_filter_evaluation_exceptions=self.dl_on_filter_evaluation_exceptions,
+            max_delivery_count=self.max_delivery_count,
+            enable_batched_operations=self.enable_batched_operations,
+            forward_to=self.forward_to,
+            user_metadata=self.user_metadata,
+            forward_dead_lettered_messages_to=self.forward_dead_lettered_messages_to,
+            auto_delete_on_idle=self.auto_delete_on_idle,
+            filter_rule=self.filter_rule,
+            filter_rule_name=self.filter_rule_name,
+        )
+        self.log.info("Created subscription %s", subscription.name)
 
 
 class AzureServiceBusUpdateSubscriptionOperator(BaseOperator):
@@ -495,18 +488,13 @@ class AzureServiceBusUpdateSubscriptionOperator(BaseOperator):
         """Update Subscription properties, by connecting to Service Bus Admin client."""
         hook = AdminClientHook(azure_service_bus_conn_id=self.azure_service_bus_conn_id)
 
-        with hook.get_conn() as service_mgmt_conn:
-            subscription_prop = service_mgmt_conn.get_subscription(self.topic_name, self.subscription_name)
-            if self.max_delivery_count:
-                subscription_prop.max_delivery_count = self.max_delivery_count
-            if self.dl_on_message_expiration is not None:
-                subscription_prop.dead_lettering_on_message_expiration = self.dl_on_message_expiration
-            if self.enable_batched_operations is not None:
-                subscription_prop.enable_batched_operations = self.enable_batched_operations
-            # update by updating the properties in the model
-            service_mgmt_conn.update_subscription(self.topic_name, subscription_prop)
-            updated_subscription = service_mgmt_conn.get_subscription(self.topic_name, self.subscription_name)
-            self.log.info("Subscription Updated successfully %s", updated_subscription)
+        hook.update_subscription(
+            topic_name=self.topic_name,
+            subscription_name=self.subscription_name,
+            max_delivery_count=self.max_delivery_count,
+            dead_lettering_on_message_expiration=self.dl_on_message_expiration,
+            enable_batched_operations=self.enable_batched_operations,
+        )
 
 
 class ASBReceiveSubscriptionMessageOperator(BaseOperator):

airflow/providers/microsoft/azure/operators/batch.py

@@ -21,11 +21,10 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from azure.batch import models as batch_models
-
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.microsoft.azure.hooks.batch import AzureBatchHook
+from azure.batch import models as batch_models
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context