apache-airflow-providers-microsoft-azure 10.0.0rc1__py3-none-any.whl → 10.1.0rc1__py3-none-any.whl

airflow/providers/microsoft/azure/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "10.0.0"
+__version__ = "10.1.0"
 
 try:
     from airflow import __version__ as airflow_version
@@ -35,8 +35,8 @@ except ImportError:
     from airflow.version import version as airflow_version
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.6.0"
+    "2.7.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-microsoft-azure:{__version__}` needs Apache Airflow 2.6.0+"
+        f"The package `apache-airflow-providers-microsoft-azure:{__version__}` needs Apache Airflow 2.7.0+"
     )

airflow/providers/microsoft/azure/fs/adls.py

@@ -18,8 +18,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
+from azure.identity import ClientSecretCredential
+
 from airflow.hooks.base import BaseHook
-from airflow.providers.microsoft.azure.utils import get_field
+from airflow.providers.microsoft.azure.utils import get_field, parse_blob_account_url
 
 if TYPE_CHECKING:
     from fsspec import AbstractFileSystem
@@ -35,13 +37,61 @@ def get_fs(conn_id: str | None, storage_options: dict[str, Any] | None = None) -
 
     conn = BaseHook.get_connection(conn_id)
     extras = conn.extra_dejson
+    conn_type = conn.conn_type or "azure_data_lake"
 
-    options = {}
-    fields = ["connection_string", "account_name", "account_key", "sas_token", "tenant"]
-    for field in fields:
-        options[field] = get_field(
-            conn_id=conn_id, conn_type="azure_data_lake", extras=extras, field_name=field
+    # connection string always overrides everything else
+    connection_string = get_field(
+        conn_id=conn_id, conn_type=conn_type, extras=extras, field_name="connection_string"
+    )
+
+    if connection_string:
+        return AzureBlobFileSystem(connection_string=connection_string)
+
+    options: dict[str, Any] = {
+        "account_url": parse_blob_account_url(conn.host, conn.login),
+    }
+
+    # mirror handling of custom field "client_secret_auth_config" from extras. Ignore if missing as AzureBlobFileSystem can handle.
+    tenant_id = get_field(conn_id=conn_id, conn_type=conn_type, extras=extras, field_name="tenant_id")
+    login = conn.login or ""
+    password = conn.password or ""
+    # assumption (from WasbHook) that if tenant_id is set, we want service principal connection
+    if tenant_id:
+        client_secret_auth_config = get_field(
+            conn_id=conn_id, conn_type=conn_type, extras=extras, field_name="client_secret_auth_config"
         )
+        if login:
+            options["client_id"] = login
+        if password:
+            options["client_secret"] = password
+        if client_secret_auth_config and login and password:
+            options["credential"] = ClientSecretCredential(
+                tenant_id=tenant_id, client_id=login, client_secret=password, **client_secret_auth_config
+            )
+
+    # if not service principal, then password is taken to be account admin key
+    if tenant_id is None and password:
+        options["account_key"] = password
+
+    # now take any fields from extras and overlay on these
+    # add empty field to remove defaults
+    fields = [
+        "account_name",
+        "account_key",
+        "sas_token",
+        "tenant_id",
+        "managed_identity_client_id",
+        "workload_identity_client_id",
+        "workload_identity_tenant_id",
+        "anon",
+    ]
+    for field in fields:
+        value = get_field(conn_id=conn_id, conn_type=conn_type, extras=extras, field_name=field)
+        if value is not None:
+            if value == "":
+                options.pop(field, "")
+            else:
+                options[field] = value
 
     options.update(storage_options or {})
 

airflow/providers/microsoft/azure/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Microsoft Azure",
         "description": "`Microsoft Azure <https://azure.microsoft.com/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1709555852,
+        "source-date-epoch": 1714476738,
         "versions": [
+            "10.1.0",
             "10.0.0",
             "9.0.1",
             "9.0.0",
@@ -83,7 +84,7 @@ def get_provider_info():
             "1.0.0",
         ],
         "dependencies": [
-            "apache-airflow>=2.6.0",
+            "apache-airflow>=2.7.0",
             "adlfs>=2023.10.0",
             "azure-batch>=8.0.0",
             "azure-cosmos>=4.6.0",
@@ -105,6 +106,7 @@ def get_provider_info():
             "azure-mgmt-datafactory>=2.0.0",
             "azure-mgmt-containerregistry>=8.0.0",
             "azure-mgmt-containerinstance>=9.0.0",
+            "msgraph-core>=1.0.0",
         ],
         "devel-dependencies": ["pywinrm"],
         "integrations": [
@@ -194,6 +196,13 @@ def get_provider_info():
                 "logo": "/integration-logos/azure/Data Lake Storage.svg",
                 "tags": ["azure"],
             },
+            {
+                "integration-name": "Microsoft Graph API",
+                "external-doc-url": "https://learn.microsoft.com/en-us/graph/use-the-api/",
+                "logo": "/integration-logos/azure/Microsoft-Graph-API.png",
+                "how-to-guide": ["/docs/apache-airflow-providers-microsoft-azure/operators/msgraph.rst"],
+                "tags": ["azure"],
+            },
         ],
         "operators": [
             {
@@ -232,6 +241,10 @@ def get_provider_info():
                 "integration-name": "Microsoft Azure Synapse",
                 "python-modules": ["airflow.providers.microsoft.azure.operators.synapse"],
             },
+            {
+                "integration-name": "Microsoft Graph API",
+                "python-modules": ["airflow.providers.microsoft.azure.operators.msgraph"],
+            },
         ],
         "sensors": [
             {
@@ -246,6 +259,10 @@ def get_provider_info():
                 "integration-name": "Microsoft Azure Data Factory",
                 "python-modules": ["airflow.providers.microsoft.azure.sensors.data_factory"],
             },
+            {
+                "integration-name": "Microsoft Graph API",
+                "python-modules": ["airflow.providers.microsoft.azure.sensors.msgraph"],
+            },
         ],
         "filesystems": ["airflow.providers.microsoft.azure.fs.adls"],
         "hooks": [
@@ -301,6 +318,10 @@ def get_provider_info():
                 "integration-name": "Microsoft Azure Synapse",
                 "python-modules": ["airflow.providers.microsoft.azure.hooks.synapse"],
             },
+            {
+                "integration-name": "Microsoft Graph API",
+                "python-modules": ["airflow.providers.microsoft.azure.hooks.msgraph"],
+            },
         ],
         "triggers": [
             {
@@ -311,6 +332,10 @@ def get_provider_info():
                 "integration-name": "Microsoft Azure Blob Storage",
                 "python-modules": ["airflow.providers.microsoft.azure.triggers.wasb"],
             },
+            {
+                "integration-name": "Microsoft Graph API",
+                "python-modules": ["airflow.providers.microsoft.azure.triggers.msgraph"],
+            },
         ],
         "transfers": [
             {

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

@@ -0,0 +1,358 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# 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 http import HTTPStatus
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Callable
+from urllib.parse import quote, urljoin, urlparse
+
+import httpx
+from azure.identity import ClientSecretCredential
+from httpx import Timeout
+from kiota_abstractions.api_error import APIError
+from kiota_abstractions.method import Method
+from kiota_abstractions.request_information import RequestInformation
+from kiota_abstractions.response_handler import ResponseHandler
+from kiota_authentication_azure.azure_identity_authentication_provider import (
+    AzureIdentityAuthenticationProvider,
+)
+from kiota_http.httpx_request_adapter import HttpxRequestAdapter
+from kiota_http.middleware.options import ResponseHandlerOption
+from msgraph_core import APIVersion, GraphClientFactory
+from msgraph_core._enums import NationalClouds
+
+from airflow.exceptions import AirflowBadRequest, AirflowException, AirflowNotFoundException
+from airflow.hooks.base import BaseHook
+
+if TYPE_CHECKING:
+    from kiota_abstractions.request_adapter import RequestAdapter
+    from kiota_abstractions.request_information import QueryParams
+    from kiota_abstractions.response_handler import NativeResponseType
+    from kiota_abstractions.serialization import ParsableFactory
+    from kiota_http.httpx_request_adapter import ResponseType
+
+    from airflow.models import Connection
+
+
+class CallableResponseHandler(ResponseHandler):
+    """
+    CallableResponseHandler executes the passed callable_function with response as parameter.
+
+    param callable_function: Function that is applied to the response.
+    """
+
+    def __init__(
+        self,
+        callable_function: Callable[[NativeResponseType, dict[str, ParsableFactory | None] | None], Any],
+    ):
+        self.callable_function = callable_function
+
+    async def handle_response_async(
+        self, response: NativeResponseType, error_map: dict[str, ParsableFactory | None] | None = None
+    ) -> Any:
+        """
+        Invoke this callback method when a response is received.
+
+        param response: The type of the native response object.
+        param error_map: The error dict to use in case of a failed request.
+        """
+        value = self.callable_function(response, error_map)
+        if response.status_code not in {200, 201, 202, 204, 302}:
+            message = value or response.reason_phrase
+            status_code = HTTPStatus(response.status_code)
+            if status_code == HTTPStatus.BAD_REQUEST:
+                raise AirflowBadRequest(message)
+            elif status_code == HTTPStatus.NOT_FOUND:
+                raise AirflowNotFoundException(message)
+            raise AirflowException(message)
+        return value
+
+
+class KiotaRequestAdapterHook(BaseHook):
+    """
+    A Microsoft Graph API interaction hook, a Wrapper around KiotaRequestAdapter.
+
+    https://github.com/microsoftgraph/msgraph-sdk-python-core
+
+    :param conn_id: The HTTP Connection ID to run the trigger against.
+    :param timeout: The HTTP timeout being used by the KiotaRequestAdapter (default is None).
+        When no timeout is specified or set to None then no HTTP timeout is applied on each request.
+    :param proxies: A Dict defining the HTTP proxies to be used (default is None).
+    :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
+        You can pass an enum named APIVersion which has 2 possible members v1 and beta,
+        or you can pass a string as "v1.0" or "beta".
+    """
+
+    DEFAULT_HEADERS = {"Accept": "application/json;q=1"}
+    cached_request_adapters: dict[str, tuple[APIVersion, RequestAdapter]] = {}
+    default_conn_name: str = "msgraph_default"
+
+    def __init__(
+        self,
+        conn_id: str = default_conn_name,
+        timeout: float | None = None,
+        proxies: dict | None = None,
+        api_version: APIVersion | str | None = None,
+    ):
+        super().__init__()
+        self.conn_id = conn_id
+        self.timeout = timeout
+        self.proxies = proxies
+        self._api_version = self.resolve_api_version_from_value(api_version)
+
+    @property
+    def api_version(self) -> APIVersion:
+        self.get_conn()  # Make sure config has been loaded through get_conn to have correct api version!
+        return self._api_version
+
+    @staticmethod
+    def resolve_api_version_from_value(
+        api_version: APIVersion | str, default: APIVersion | None = None
+    ) -> APIVersion:
+        if isinstance(api_version, APIVersion):
+            return api_version
+        return next(
+            filter(lambda version: version.value == api_version, APIVersion),
+            default,
+        )
+
+    def get_api_version(self, config: dict) -> APIVersion:
+        if self._api_version is None:
+            return self.resolve_api_version_from_value(
+                api_version=config.get("api_version"), default=APIVersion.v1
+            )
+        return self._api_version
+
+    @staticmethod
+    def get_host(connection: Connection) -> str:
+        if connection.schema and connection.host:
+            return f"{connection.schema}://{connection.host}"
+        return NationalClouds.Global.value
+
+    @staticmethod
+    def format_no_proxy_url(url: str) -> str:
+        if "://" not in url:
+            url = 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
+        return proxies
+
+    def to_msal_proxies(self, authority: str | None, proxies: dict):
+        self.log.info("authority: %s", authority)
+        if authority:
+            no_proxies = proxies.get("no")
+            self.log.info("no_proxies: %s", no_proxies)
+            if no_proxies:
+                for url in no_proxies.split(","):
+                    self.log.info("url: %s", url)
+                    domain_name = urlparse(url).path.replace("*", "")
+                    self.log.info("domain_name: %s", domain_name)
+                    if authority.endswith(domain_name):
+                        return None
+        return proxies
+
+    def get_conn(self) -> RequestAdapter:
+        if not self.conn_id:
+            raise AirflowException("Failed to create the KiotaRequestAdapterHook. No conn_id provided!")
+
+        api_version, request_adapter = self.cached_request_adapters.get(self.conn_id, (None, None))
+
+        if not request_adapter:
+            connection = self.get_connection(conn_id=self.conn_id)
+            client_id = connection.login
+            client_secret = connection.password
+            config = connection.extra_dejson if connection.extra else {}
+            tenant_id = config.get("tenant_id")
+            api_version = self.get_api_version(config)
+            host = self.get_host(connection)
+            base_url = config.get("base_url", urljoin(host, api_version.value))
+            authority = config.get("authority")
+            proxies = self.proxies or config.get("proxies", {})
+            msal_proxies = self.to_msal_proxies(authority=authority, proxies=proxies)
+            httpx_proxies = self.to_httpx_proxies(proxies=proxies)
+            scopes = config.get("scopes", ["https://graph.microsoft.com/.default"])
+            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(
+                "Creating Microsoft Graph SDK client %s for conn_id: %s",
+                api_version.value,
+                self.conn_id,
+            )
+            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.value)
+            self.log.info("Scope: %s", scopes)
+            self.log.info("Verify: %s", verify)
+            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,
+                authority=authority,
+                proxies=msal_proxies,
+                disable_instance_discovery=disable_instance_discovery,
+                connection_verify=verify,
+            )
+            http_client = GraphClientFactory.create_with_default_middleware(
+                api_version=api_version,
+                client=httpx.AsyncClient(
+                    proxies=httpx_proxies,
+                    timeout=Timeout(timeout=self.timeout),
+                    verify=verify,
+                    trust_env=trust_env,
+                ),
+                host=host,
+            )
+            auth_provider = AzureIdentityAuthenticationProvider(
+                credentials=credentials,
+                scopes=scopes,
+                allowed_hosts=allowed_hosts,
+            )
+            request_adapter = HttpxRequestAdapter(
+                authentication_provider=auth_provider,
+                http_client=http_client,
+                base_url=base_url,
+            )
+            self.cached_request_adapters[self.conn_id] = (api_version, request_adapter)
+        self._api_version = api_version
+        return request_adapter
+
+    def test_connection(self):
+        """Test HTTP Connection."""
+        try:
+            self.run()
+            return True, "Connection successfully tested"
+        except Exception as e:
+            return False, str(e)
+
+    async def run(
+        self,
+        url: str = "",
+        response_type: ResponseType | None = None,
+        response_handler: Callable[
+            [NativeResponseType, dict[str, ParsableFactory | None] | None], Any
+        ] = lambda response, error_map: response.json(),
+        path_parameters: dict[str, Any] | None = None,
+        method: str = "GET",
+        query_parameters: dict[str, QueryParams] | None = None,
+        headers: dict[str, str] | None = None,
+        data: dict[str, Any] | str | BytesIO | None = None,
+    ):
+        response = await self.get_conn().send_primitive_async(
+            request_info=self.request_information(
+                url=url,
+                response_type=response_type,
+                response_handler=response_handler,
+                path_parameters=path_parameters,
+                method=method,
+                query_parameters=query_parameters,
+                headers=headers,
+                data=data,
+            ),
+            response_type=response_type,
+            error_map=self.error_mapping(),
+        )
+
+        self.log.debug("response: %s", response)
+
+        return response
+
+    def request_information(
+        self,
+        url: str,
+        response_type: ResponseType | None = None,
+        response_handler: Callable[
+            [NativeResponseType, dict[str, ParsableFactory | None] | None], Any
+        ] = lambda response, error_map: response.json(),
+        path_parameters: dict[str, Any] | None = None,
+        method: str = "GET",
+        query_parameters: dict[str, QueryParams] | None = None,
+        headers: dict[str, str] | None = None,
+        data: dict[str, Any] | str | BytesIO | None = None,
+    ) -> RequestInformation:
+        request_information = RequestInformation()
+        request_information.path_parameters = path_parameters or {}
+        request_information.http_method = Method(method.strip().upper())
+        request_information.query_parameters = self.encoded_query_parameters(query_parameters)
+        if url.startswith("http"):
+            request_information.url = url
+        elif request_information.query_parameters.keys():
+            query = ",".join(request_information.query_parameters.keys())
+            request_information.url_template = f"{{+baseurl}}/{self.normalize_url(url)}{{?{query}}}"
+        else:
+            request_information.url_template = f"{{+baseurl}}/{self.normalize_url(url)}"
+        if not response_type:
+            request_information.request_options[ResponseHandlerOption.get_key()] = ResponseHandlerOption(
+                response_handler=CallableResponseHandler(response_handler)
+            )
+        headers = {**self.DEFAULT_HEADERS, **headers} if headers else self.DEFAULT_HEADERS
+        for header_name, header_value in headers.items():
+            request_information.headers.try_add(header_name=header_name, header_value=header_value)
+        self.log.info("data: %s", data)
+        if isinstance(data, BytesIO) or isinstance(data, bytes) or isinstance(data, str):
+            request_information.content = data
+        elif data:
+            request_information.headers.try_add(
+                header_name=RequestInformation.CONTENT_TYPE_HEADER, header_value="application/json"
+            )
+            request_information.content = json.dumps(data).encode("utf-8")
+        return request_information
+
+    @staticmethod
+    def normalize_url(url: str) -> str | None:
+        if url.startswith("/"):
+            return url.replace("/", "", 1)
+        return url
+
+    @staticmethod
+    def encoded_query_parameters(query_parameters) -> dict:
+        if query_parameters:
+            return {quote(key): value for key, value in query_parameters.items()}
+        return {}
+
+    @staticmethod
+    def error_mapping() -> dict[str, ParsableFactory | None]:
+        return {
+            "4XX": APIError,
+            "5XX": APIError,
+        }

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

@@ -312,6 +312,7 @@ class AzureSynapsePipelineHook(BaseAzureSynapseHook):
             warnings.warn(
                 "The usage of `default_conn_name=azure_synapse_connection` is deprecated and will be removed in future. Please update your code to use the new default connection name: `default_conn_name=azure_synapse_default`. ",
                 AirflowProviderDeprecationWarning,
+                stacklevel=2,
             )
         self._conn: ArtifactsClient | None = None
         self.azure_synapse_workspace_dev_endpoint = azure_synapse_workspace_dev_endpoint

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

@@ -30,7 +30,6 @@ import logging
 import os
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Union
-from urllib.parse import urlparse
 
 from asgiref.sync import sync_to_async
 from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceNotFoundError
@@ -52,6 +51,7 @@ 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:
@@ -167,21 +167,7 @@ class WasbHook(BaseHook):
             # connection_string auth takes priority
             return BlobServiceClient.from_connection_string(connection_string, **extra)
 
-        account_url = conn.host if conn.host else f"https://{conn.login}.blob.core.windows.net/"
-        parsed_url = urlparse(account_url)
-
-        if not parsed_url.netloc:
-            if "." not in parsed_url.path:
-                # if there's no netloc and no dots in the path, then user only
-                # provided the Active Directory ID, not the full URL or DNS name
-                account_url = f"https://{conn.login}.blob.core.windows.net/"
-            else:
-                # if there's no netloc but there are dots in the path, then user
-                # provided the DNS name without the https:// prefix.
-                # Azure storage account name can only be 3 to 24 characters in length
-                # https://learn.microsoft.com/en-us/azure/storage/common/storage-account-overview#storage-account-name
-                acc_name = account_url.split(".")[0][:24]
-                account_url = f"https://{acc_name}." + ".".join(account_url.split(".")[1:])
+        account_url = parse_blob_account_url(conn.host, conn.login)
 
         tenant = self._get_field(extra, "tenant_id")
         if tenant:
@@ -587,21 +573,7 @@ class WasbAsyncHook(WasbHook):
             )
             return self.blob_service_client
 
-        account_url = conn.host if conn.host else f"https://{conn.login}.blob.core.windows.net/"
-        parsed_url = urlparse(account_url)
-
-        if not parsed_url.netloc:
-            if "." not in parsed_url.path:
-                # if there's no netloc and no dots in the path, then user only
-                # provided the Active Directory ID, not the full URL or DNS name
-                account_url = f"https://{conn.login}.blob.core.windows.net/"
-            else:
-                # if there's no netloc but there are dots in the path, then user
-                # provided the DNS name without the https:// prefix.
-                # Azure storage account name can only be 3 to 24 characters in length
-                # https://learn.microsoft.com/en-us/azure/storage/common/storage-account-overview#storage-account-name
-                acc_name = account_url.split(".")[0][:24]
-                account_url = f"https://{acc_name}." + ".".join(account_url.split(".")[1:])
+        account_url = parse_blob_account_url(conn.host, conn.login)
 
         tenant = self._get_field(extra, "tenant_id")
         if tenant:

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

@@ -25,8 +25,10 @@ from typing import TYPE_CHECKING, Any, Sequence
 from azure.mgmt.containerinstance.models import (
     Container,
     ContainerGroup,
+    ContainerGroupDiagnostics,
     ContainerGroupSubnetId,
     ContainerPort,
+    DnsConfiguration,
     EnvironmentVariable,
     IpAddress,
     ResourceRequests,
@@ -90,6 +92,8 @@ class AzureContainerInstancesOperator(BaseOperator):
         Possible values include: 'Always', 'OnFailure', 'Never'
     :param ip_address: The IP address type of the container group.
     :param subnet_ids: The subnet resource IDs for a container group
+    :param dns_config: The DNS configuration for a container group.
+    :param diagnostics: Container group diagnostic information (Log Analytics).
 
     **Example**::
 
@@ -113,6 +117,13 @@ class AzureContainerInstancesOperator(BaseOperator):
             memory_in_gb=14.0,
             cpu=4.0,
             gpu=GpuResource(count=1, sku="K80"),
+            dns_config=["10.0.0.10", "10.0.0.11"],
+            diagnostics={
+                "log_analytics": {
+                    "workspaceId": "workspaceid",
+                    "workspaceKey": "workspaceKey",
+                }
+            },
             command=["/bin/echo", "world"],
             task_id="start_container",
         )
@@ -145,6 +156,8 @@ class AzureContainerInstancesOperator(BaseOperator):
         ip_address: IpAddress | None = None,
         ports: list[ContainerPort] | None = None,
         subnet_ids: list[ContainerGroupSubnetId] | None = None,
+        dns_config: DnsConfiguration | None = None,
+        diagnostics: ContainerGroupDiagnostics | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -183,6 +196,8 @@ class AzureContainerInstancesOperator(BaseOperator):
         self.ip_address = ip_address
         self.ports = ports
         self.subnet_ids = subnet_ids
+        self.dns_config = dns_config
+        self.diagnostics = diagnostics
 
     def execute(self, context: Context) -> int:
         # Check name again in case it was templated.
@@ -256,6 +271,8 @@ class AzureContainerInstancesOperator(BaseOperator):
                 tags=self.tags,
                 ip_address=self.ip_address,
                 subnet_ids=self.subnet_ids,
+                dns_config=self.dns_config,
+                diagnostics=self.diagnostics,
             )
 
             self._ci_hook.create_or_update(self.resource_group, self.name, container_group)