This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 935bac04f06 Add AzureAISearchHook to the Microsoft Azure provider
(#74290)
935bac04f06 is described below
commit 935bac04f06a125aca6cc432442af2a8de5b7cc7
Author: David Blain <[email protected]>
AuthorDate: Sun Oct 18 23:28:18 2026 +0200
Add AzureAISearchHook to the Microsoft Azure provider (#74290)
* Add AzureAISearchHook to the Microsoft Azure provider
The Azure provider had no integration with Azure AI Search, the vector and
full-text store behind most retrieval-augmented generation on Azure. A Dag
that keeps such an index up to date had to build the SDK client from the
connection itself, and to repeat for every write what the SDK leaves to the
caller: batching by the service limit of 1000 documents, and checking the
per-document results, because a partially failed write (HTTP 207) does not
raise.
The hook is natively async, so that an async task can read and write an
index next to its other awaits. It reads the connection extra through
get_async_extra_dejson, which makes no blocking call to the Task SDK on the
event loop.
* Add OData to the spelling wordlist
The docs spell check rejected the word in the docstrings of
AzureAISearchHook,
which name the OData filter and sort expressions of Azure AI Search.
* Add an example Dag and a guide for AzureAISearchHook
A hook without an operator has no page in the provider documentation that
shows how it is used from a task. The system test creates its own index,
writes, searches, updates, reads and deletes documents through the hook,
and the guide includes those tasks, with the Airflow 3.2 requirement for
async callables.
* Register the Azure AI Search guide in provider.yaml
The provider.yaml check expects every page under docs/operators to be the
how-to guide of an integration.
* Import DAG, task and TriggerRule from common.compat.sdk in the example
The compat module already covers the Airflow 2 and 3 import paths, as the
other provider examples use it.
* Update
providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_search.py
Co-authored-by: Aaron Chen <[email protected]>
* Exclude azure-search-documents 12.0.0 from the Azure provider
On a 413 the async SearchClient of 12.0.0 splits the batch, writes the
first half and retries the second half with a keyword argument its
_index_documents_actions does not accept, so an oversized write raises
TypeError after a partial write. 12.1.0b2 fixes it; 11.6.0 is not
affected. The lock moves to 11.6.0 until 12.1.0 is released.
* Document the 100,000-result limit of AzureAISearchHook.search and test
get_document's select
The service pages search results with $skip, which cannot go above
100,000, so an unbounded search over a larger index fails partway
through. The docstring and the guide now say so and show how to read
such an index by ranges of a sortable, filterable key.
The fake client's get_document now honours selected_fields, and a new
test checks that the select argument reaches the client.
* Link the Azure AI Search hook and guide to the Microsoft documentation
Every limit and behaviour of the service the hook and its guide state
(1000 documents and 16 MB per write request, the per-document results
of a write, the skip limit of 100,000 and the key-range paging, the
data-plane roles) now links to the Microsoft page that documents it.
---------
Co-authored-by: Aaron Chen <[email protected]>
---
docs/spelling_wordlist.txt | 1 +
providers/microsoft/azure/README.rst | 7 +-
.../azure/docs/connections/azure_ai_search.rst | 92 +++++
providers/microsoft/azure/docs/index.rst | 7 +-
.../microsoft/azure/docs/operators/ai_search.rst | 121 +++++++
providers/microsoft/azure/provider.yaml | 48 +++
providers/microsoft/azure/pyproject.toml | 3 +
.../providers/microsoft/azure/exceptions.py | 4 +
.../providers/microsoft/azure/get_provider_info.py | 48 +++
.../providers/microsoft/azure/hooks/ai_search.py | 355 +++++++++++++++++++
.../microsoft/azure/example_azure_ai_search.py | 175 +++++++++
.../unit/microsoft/azure/hooks/test_ai_search.py | 394 +++++++++++++++++++++
uv.lock | 17 +
13 files changed, 1266 insertions(+), 6 deletions(-)
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index 660c5b82560..029b8e21ba4 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -1214,6 +1214,7 @@ objectORfile
objectstorage
observability
od
+OData
odbc
odps
OIDC
diff --git a/providers/microsoft/azure/README.rst
b/providers/microsoft/azure/README.rst
index 78e314bf46d..2907a0d0f11 100644
--- a/providers/microsoft/azure/README.rst
+++ b/providers/microsoft/azure/README.rst
@@ -50,9 +50,9 @@ The package supports the following python versions:
3.11,3.12,3.13,3.14
Requirements
------------
-========================================== ===================
+========================================== =====================
PIP package Version required
-========================================== ===================
+========================================== =====================
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.21.0``
``adlfs`` ``>=2026.5.0``
@@ -66,6 +66,7 @@ PIP package Version required
``azure-keyvault-secrets`` ``>=4.10.0``
``azure-mgmt-datalake-store`` ``>=0.5.0``
``azure-mgmt-resource`` ``>=25.0.0``
+``azure-search-documents`` ``>=11.6.0,!=12.0.0``
``azure-storage-blob`` ``>=12.28.0``
``azure-mgmt-storage`` ``>=24.0.1``
``azure-storage-file-share`` ``>=12.24.0``
@@ -86,7 +87,7 @@ PIP package Version required
``microsoft-kiota-abstractions`` ``>=1.9.4,<2.0.0``
``microsoft-kiota-authentication-azure`` ``>=1.9.4,<2.0.0``
``msal-extensions`` ``>=1.3.0``
-========================================== ===================
+========================================== =====================
Optional cross provider package dependencies
--------------------------------------------
diff --git a/providers/microsoft/azure/docs/connections/azure_ai_search.rst
b/providers/microsoft/azure/docs/connections/azure_ai_search.rst
new file mode 100644
index 00000000000..de9e295ac8a
--- /dev/null
+++ b/providers/microsoft/azure/docs/connections/azure_ai_search.rst
@@ -0,0 +1,92 @@
+.. 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.
+
+
+.. _howto/connection:azure_ai_search:
+
+Microsoft Azure AI Search
+=========================
+
+The Microsoft Azure AI Search connection type enables the Azure AI Search
integrations.
+
+Authenticating to Azure AI Search
+---------------------------------
+
+There are three ways to connect to Azure AI Search using Airflow.
+
+1. Use `token credentials
+
<https://learn.microsoft.com/en-us/azure/developer/python/sdk/authentication/overview>`_
+ i.e. add the credentials of a service principal (client_id, secret, tenant)
to the Airflow connection.
+2. Use an `API key`_ by setting it as the password, without a client_id.
+ A query key is enough to read documents, writing them needs an admin key.
+3. Fallback on DefaultAzureCredential_ when the connection has neither.
+ This includes a mechanism to try different options to authenticate: Managed
System Identity, environment variables, authentication through Azure CLI...
+ Set ``managed_identity_client_id`` and ``workload_identity_tenant_id`` to
use a specific managed identity.
+
+With a service principal or ``DefaultAzureCredential`` the identity needs the
+`Search Index Data Reader`_ role on the search service, or ``Search Index Data
Contributor`` to write documents.
+
+Default Connection IDs
+----------------------
+
+All hooks related to Microsoft Azure AI Search use ``azure_ai_search_default``
by default.
+
+Configuring the Connection
+--------------------------
+
+Search Service Endpoint
+ Specify the endpoint of the search service, e.g.
``https://<service>.search.windows.net``.
+ ``https://`` is added to a bare host name.
+
+Azure Client ID (optional)
+ Specify the ``client_id`` of the service principal.
+ Leave it empty to authenticate with an API key or with
DefaultAzureCredential_.
+
+Azure Secret or API Key (optional)
+ Specify the ``secret`` of the service principal, or an `API key`_ when no
client_id is set.
+ It can be left out to fall back on DefaultAzureCredential_.
+
+Azure Tenant ID (optional)
+ Specify the ``tenantId`` of the service principal.
+
+Default Index (optional)
+ Specify the ``index`` the hook works on when it is given none.
+
+Batch Size (optional)
+ Specify the ``batch_size``: the number of documents the hook sends per
write request, from 1 to 1000 (the default).
+
+Managed Identity Client ID (optional)
+ The client ID of a user-assigned managed identity. If provided with
``workload_identity_tenant_id``, they'll pass to DefaultAzureCredential_.
+
+Workload Identity Tenant ID (optional)
+ ID of the application's Microsoft Entra tenant. Also called its
"directory" ID. If provided with ``managed_identity_client_id``, they'll pass
to DefaultAzureCredential_.
+
+When specifying the connection in environment variable you should specify
+it using URI syntax.
+
+Note that all components of the URI should be URL-encoded.
+
+For example:
+
+.. code-block:: bash
+
+ export
AIRFLOW_CONN_AZURE_AI_SEARCH_DEFAULT='azure-ai-search://:api%[email protected]?index=my-index'
+
+
+.. _API key:
https://learn.microsoft.com/en-us/azure/search/search-security-api-keys
+.. _Search Index Data Reader:
https://learn.microsoft.com/en-us/azure/search/search-security-rbac
+.. _DefaultAzureCredential:
https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential
diff --git a/providers/microsoft/azure/docs/index.rst
b/providers/microsoft/azure/docs/index.rst
index 9d6277fbcef..f074d98c6b2 100644
--- a/providers/microsoft/azure/docs/index.rst
+++ b/providers/microsoft/azure/docs/index.rst
@@ -107,9 +107,9 @@ Requirements
The minimum Apache Airflow version supported by this provider distribution is
``2.11.0``.
-========================================== ===================
+========================================== =====================
PIP package Version required
-========================================== ===================
+========================================== =====================
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.21.0``
``adlfs`` ``>=2026.5.0``
@@ -123,6 +123,7 @@ PIP package Version required
``azure-keyvault-secrets`` ``>=4.10.0``
``azure-mgmt-datalake-store`` ``>=0.5.0``
``azure-mgmt-resource`` ``>=25.0.0``
+``azure-search-documents`` ``>=11.6.0,!=12.0.0``
``azure-storage-blob`` ``>=12.28.0``
``azure-mgmt-storage`` ``>=24.0.1``
``azure-storage-file-share`` ``>=12.24.0``
@@ -143,7 +144,7 @@ PIP package Version required
``microsoft-kiota-abstractions`` ``>=1.9.4,<2.0.0``
``microsoft-kiota-authentication-azure`` ``>=1.9.4,<2.0.0``
``msal-extensions`` ``>=1.3.0``
-========================================== ===================
+========================================== =====================
Optional cross provider package dependencies
--------------------------------------------
diff --git a/providers/microsoft/azure/docs/operators/ai_search.rst
b/providers/microsoft/azure/docs/operators/ai_search.rst
new file mode 100644
index 00000000000..3f460c94949
--- /dev/null
+++ b/providers/microsoft/azure/docs/operators/ai_search.rst
@@ -0,0 +1,121 @@
+.. 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.
+
+
+
+.. _howto/hook:AzureAISearchHook:
+
+Azure AI Search
+===============
+
+`Azure AI Search <https://learn.microsoft.com/en-us/azure/search/>`__ is the
full-text and vector
+store behind retrieval-augmented generation on Azure. The
+:class:`~airflow.providers.microsoft.azure.hooks.ai_search.AzureAISearchHook`
reads and writes the
+documents of one index, from an async task: every method is a coroutine, and
+`async Python callables
<https://airflow.apache.org/docs/apache-airflow-providers-standard/stable/operators/python.html>`__
+are supported since Airflow 3.2. Creating and changing indexes is left to the
+``azure-search-documents`` SDK.
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include:: /operators/_partials/prerequisite_tasks.rst
+
+The hook uses the ``azure_ai_search_default`` connection by default, see
+:ref:`howto/connection:azure_ai_search`. The index comes from the hook
(``index=``) or from the
+``index`` extra of the connection.
+
+Writing documents
+-----------------
+
+``upload`` inserts documents and replaces the ones whose key exists. The
documents are sent in batches
+of ``batch_size`` (1000 by default, the
+`limit of the service
<https://learn.microsoft.com/en-us/azure/search/search-limits-quotas-capacity#api-request-limits>`__),
+and the hook raises
+:class:`~airflow.providers.microsoft.azure.exceptions.AzureAISearchIndexingError`
with the keys and
+reasons when the service rejects some of them, from the
+`per-document results
<https://learn.microsoft.com/en-us/rest/api/searchservice/documents/index>`__
+of the write.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_search.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_hook_azure_ai_search_upload]
+ :end-before: [END howto_hook_azure_ai_search_upload]
+
+Searching documents
+-------------------
+
+``search`` is an async generator over every matching document, following the
pages of the service.
+Select the fields you need, and filter and sort with
+`OData expressions
<https://learn.microsoft.com/en-us/azure/search/search-query-odata-filter>`__.
+
+The service pages the results with
+`skip
<https://learn.microsoft.com/en-us/rest/api/searchservice/documents/search-post>`__,
which stops
+at 100,000 documents: past that many matches the search fails. To read a
larger index, sort on a
+sortable and filterable key (``order_by=["id"]``) and filter on ranges of it
+(``filter="id gt '<last key read>'"``), one range per call, as described in
+`Paging through a large number of results
<https://learn.microsoft.com/en-us/azure/search/search-pagination-page-layout#paging-through-a-large-number-of-results>`__.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_search.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_hook_azure_ai_search_search]
+ :end-before: [END howto_hook_azure_ai_search_search]
+
+Updating and counting documents
+-------------------------------
+
+``merge`` changes the given fields of existing documents and
``merge_or_upload`` also inserts the
+documents whose key is new. ``count`` returns the number of documents of the
index, or of those
+matching a filter.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_search.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_hook_azure_ai_search_merge]
+ :end-before: [END howto_hook_azure_ai_search_merge]
+
+Reading one document
+--------------------
+
+``get_document`` returns the document with a key, or ``None`` when the index
has none.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_search.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_hook_azure_ai_search_get_document]
+ :end-before: [END howto_hook_azure_ai_search_get_document]
+
+Deleting documents
+------------------
+
+``delete`` removes documents by key.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_search.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_hook_azure_ai_search_delete]
+ :end-before: [END howto_hook_azure_ai_search_delete]
+
+Reference
+---------
+
+For further information, look at:
+
+* `Azure AI Search documentation
<https://learn.microsoft.com/en-us/azure/search/>`__
+* `azure-search-documents SDK
<https://learn.microsoft.com/en-us/python/api/overview/azure/search-documents-readme>`__
diff --git a/providers/microsoft/azure/provider.yaml
b/providers/microsoft/azure/provider.yaml
index c2130f80bb7..7c43d87a67e 100644
--- a/providers/microsoft/azure/provider.yaml
+++ b/providers/microsoft/azure/provider.yaml
@@ -195,6 +195,12 @@ integrations:
- /docs/apache-airflow-providers-microsoft-azure/operators/ai_agents.rst
logo: /docs/integration-logos/Microsoft-Azure.png
tags: [azure]
+ - integration-name: Microsoft Azure AI Search
+ external-doc-url: https://learn.microsoft.com/en-us/azure/search/
+ how-to-guide:
+ - /docs/apache-airflow-providers-microsoft-azure/operators/ai_search.rst
+ logo: /docs/integration-logos/Microsoft-Azure.png
+ tags: [azure]
- integration-name: Microsoft Azure Service Bus
external-doc-url: https://azure.microsoft.com/en-us/services/service-bus/
logo: /docs/integration-logos/Service-Bus.svg
@@ -318,6 +324,9 @@ hooks:
- integration-name: Microsoft Azure AI Foundry Hosted Agents
python-modules:
- airflow.providers.microsoft.azure.hooks.ai_agents
+ - integration-name: Microsoft Azure AI Search
+ python-modules:
+ - airflow.providers.microsoft.azure.hooks.ai_search
- integration-name: Microsoft Azure Data Lake Storage
python-modules:
- airflow.providers.microsoft.azure.hooks.data_lake
@@ -507,6 +516,45 @@ connection-types:
label: Workload Identity Tenant ID
schema:
type: ["string", "null"]
+ - hook-class-name:
airflow.providers.microsoft.azure.hooks.ai_search.AzureAISearchHook
+ hook-name: "Azure AI Search"
+ connection-type: azure_ai_search
+ ui-field-behaviour:
+ hidden-fields: ["schema", "port"]
+ relabeling:
+ host: Search Service Endpoint
+ login: Azure Client ID
+ password: Azure Secret or API Key
+ placeholders:
+ host: https://<service>.search.windows.net
+ login: client_id (service principal only)
+ password: secret (service principal), or an API key when no client_id
is set
+ tenantId: tenantId (service principal only)
+ index: index used when the hook is given none
+ batch_size: documents per write request, 1 to 1000 (default 1000)
+ conn-fields:
+ tenantId:
+ label: Azure Tenant ID
+ schema:
+ type: ["string", "null"]
+ index:
+ label: Default Index
+ schema:
+ type: ["string", "null"]
+ batch_size:
+ label: Batch Size
+ schema:
+ type: ["integer", "null"]
+ minimum: 1
+ maximum: 1000
+ managed_identity_client_id:
+ label: Managed Identity Client ID
+ schema:
+ type: ["string", "null"]
+ workload_identity_tenant_id:
+ label: Workload Identity Tenant ID
+ schema:
+ type: ["string", "null"]
- hook-class-name:
airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook
hook-name: "Azure Compute"
connection-type: azure_compute
diff --git a/providers/microsoft/azure/pyproject.toml
b/providers/microsoft/azure/pyproject.toml
index 12d89dc4af0..dc08d81376b 100644
--- a/providers/microsoft/azure/pyproject.toml
+++ b/providers/microsoft/azure/pyproject.toml
@@ -72,6 +72,9 @@ dependencies = [
"azure-keyvault-secrets>=4.10.0",
"azure-mgmt-datalake-store>=0.5.0",
"azure-mgmt-resource>=25.0.0",
+ # 12.0.0 is excluded: on a 413 its async client writes the first half of
the batch, then retries
+ # the second half with a keyword it does not accept (TypeError). Fixed
since 12.1.0b2.
+ "azure-search-documents>=11.6.0,!=12.0.0",
"azure-storage-blob>=12.28.0",
"azure-mgmt-storage>=24.0.1",
"azure-storage-file-share>=12.24.0",
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
index 7235826189a..461463e100c 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/exceptions.py
@@ -58,3 +58,7 @@ class
AzureBatchJobPreparationTaskMissingError(AirflowException):
class AzureBatchRequiredParametersMissingError(AirflowException):
"""Raised when one or more required pool, job or task parameters are
missing."""
+
+
+class AzureAISearchIndexingError(AirflowException):
+ """Raised when Azure AI Search rejects some documents of a write."""
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
index eac58f76d7b..dd4a470f34a 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
@@ -119,6 +119,13 @@ def get_provider_info():
"logo": "/docs/integration-logos/Microsoft-Azure.png",
"tags": ["azure"],
},
+ {
+ "integration-name": "Microsoft Azure AI Search",
+ "external-doc-url":
"https://learn.microsoft.com/en-us/azure/search/";,
+ "how-to-guide":
["/docs/apache-airflow-providers-microsoft-azure/operators/ai_search.rst"],
+ "logo": "/docs/integration-logos/Microsoft-Azure.png",
+ "tags": ["azure"],
+ },
{
"integration-name": "Microsoft Azure Service Bus",
"external-doc-url":
"https://azure.microsoft.com/en-us/services/service-bus/";,
@@ -282,6 +289,10 @@ def get_provider_info():
"integration-name": "Microsoft Azure AI Foundry Hosted Agents",
"python-modules":
["airflow.providers.microsoft.azure.hooks.ai_agents"],
},
+ {
+ "integration-name": "Microsoft Azure AI Search",
+ "python-modules":
["airflow.providers.microsoft.azure.hooks.ai_search"],
+ },
{
"integration-name": "Microsoft Azure Data Lake Storage",
"python-modules":
["airflow.providers.microsoft.azure.hooks.data_lake"],
@@ -499,6 +510,43 @@ def get_provider_info():
},
},
},
+ {
+ "hook-class-name":
"airflow.providers.microsoft.azure.hooks.ai_search.AzureAISearchHook",
+ "hook-name": "Azure AI Search",
+ "connection-type": "azure_ai_search",
+ "ui-field-behaviour": {
+ "hidden-fields": ["schema", "port"],
+ "relabeling": {
+ "host": "Search Service Endpoint",
+ "login": "Azure Client ID",
+ "password": "Azure Secret or API Key",
+ },
+ "placeholders": {
+ "host": "https://<service>.search.windows.net",
+ "login": "client_id (service principal only)",
+ "password": "secret (service principal), or an API key
when no client_id is set",
+ "tenantId": "tenantId (service principal only)",
+ "index": "index used when the hook is given none",
+ "batch_size": "documents per write request, 1 to 1000
(default 1000)",
+ },
+ },
+ "conn-fields": {
+ "tenantId": {"label": "Azure Tenant ID", "schema":
{"type": ["string", "null"]}},
+ "index": {"label": "Default Index", "schema": {"type":
["string", "null"]}},
+ "batch_size": {
+ "label": "Batch Size",
+ "schema": {"type": ["integer", "null"], "minimum": 1,
"maximum": 1000},
+ },
+ "managed_identity_client_id": {
+ "label": "Managed Identity Client ID",
+ "schema": {"type": ["string", "null"]},
+ },
+ "workload_identity_tenant_id": {
+ "label": "Workload Identity Tenant ID",
+ "schema": {"type": ["string", "null"]},
+ },
+ },
+ },
{
"hook-class-name":
"airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook",
"hook-name": "Azure Compute",
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_search.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_search.py
new file mode 100644
index 00000000000..c67890acd05
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_search.py
@@ -0,0 +1,355 @@
+#
+# 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
+
+from collections.abc import AsyncGenerator, AsyncIterator, Iterable
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+from azure.core.credentials import AzureKeyCredential
+from azure.core.exceptions import ResourceNotFoundError
+from azure.identity.aio import ClientSecretCredential
+from azure.search.documents.aio import SearchClient
+
+from airflow.providers.common.compat.connection import get_async_connection,
get_async_extra_dejson
+from airflow.providers.common.compat.sdk import BaseHook
+from airflow.providers.microsoft.azure.exceptions import
AzureAISearchIndexingError
+from airflow.providers.microsoft.azure.utils import (
+ add_managed_identity_connection_widgets,
+ get_async_default_azure_credential,
+ get_field,
+)
+
+if TYPE_CHECKING:
+ from azure.core.credentials_async import AsyncTokenCredential
+
+ from airflow.sdk import Connection
+
+
+class AzureAISearchHook(BaseHook):
+ """
+ Query and write the documents of one Azure AI Search index, natively async.
+
+ A thin layer over :class:`azure.search.documents.aio.SearchClient`. The
SDK pages the search
+ results and retries throttled requests (429 and 503) with their
``Retry-After``; the hook binds
+ the client to one index and adds what every Dag would otherwise repeat:
+
+ * :meth:`search` is an async generator over every matching document,
without the ``@search.*``
+ metadata.
+ * :meth:`count` and :meth:`get_document` (``None`` when the index has no
such document).
+ * :meth:`upload`, :meth:`merge`, :meth:`merge_or_upload` and
:meth:`delete` write in batches of
+ ``batch_size`` documents and raise
+
:class:`~airflow.providers.microsoft.azure.exceptions.AzureAISearchIndexingError`
when the
+ service rejects any document, naming the keys and reasons. The SDK
itself returns the rejected
+ documents of a partially failed write without raising.
+
+ Other errors are the SDK's:
:class:`azure.core.exceptions.HttpResponseError` with its
+ ``status_code``, and ``ServiceRequestError`` or ``ServiceResponseError``
when the connection fails.
+
+ Every method is a coroutine, so the hook is meant for async tasks and
triggers. The credential is
+ chosen in this order:
+
+ 1. Login, password and ``tenantId``: a service principal.
+ 2. Password alone: an API key (a query key is enough to read, writing
needs an admin key).
+ 3. Neither: ``DefaultAzureCredential``, i.e. a managed or workload
identity, or the environment.
+
+ With Microsoft Entra ID the identity needs the *Search Index Data Reader*
role, or
+ *Search Index Data Contributor* to write.
+
+ .. seealso::
+ - Azure AI Search: https://learn.microsoft.com/en-us/azure/search/
+ - Roles for Microsoft Entra ID:
https://learn.microsoft.com/en-us/azure/search/search-security-rbac
+ - Limits of a write request (1000 documents, 16 MB):
+
https://learn.microsoft.com/en-us/azure/search/search-limits-quotas-capacity#api-request-limits
+ - Per-document results of a write (HTTP 207):
+
https://learn.microsoft.com/en-us/rest/api/searchservice/documents/index
+
+ :param azure_ai_search_conn_id: The :ref:`Azure AI Search connection
id<howto/connection:azure_ai_search>`.
+ Default is ``azure_ai_search_default``.
+ :param index: The index this hook works on. Overrides the ``index``
connection extra.
+ :param batch_size: Documents per write request, at most 1000 (the limit of
the service). Overrides
+ the ``batch_size`` connection extra. Lower it for large documents: the
service also refuses a
+ request over about 16 MB.
+ """
+
+ conn_name_attr = "azure_ai_search_conn_id"
+ default_conn_name = "azure_ai_search_default"
+ conn_type = "azure_ai_search"
+ hook_name = "Azure AI Search"
+
+ # The service accepts at most 1000 documents per write request, see
+ #
https://learn.microsoft.com/en-us/azure/search/search-limits-quotas-capacity#api-request-limits
+ MAX_BATCH_SIZE = 1000
+
+ def __init__(
+ self,
+ azure_ai_search_conn_id: str = default_conn_name,
+ index: str | None = None,
+ batch_size: int | None = None,
+ ) -> None:
+ super().__init__()
+ self.conn_id = azure_ai_search_conn_id
+ self.index = index
+ # Checked here, so that a wrong value fails when the Dag builds the
hook.
+ self.batch_size = None if batch_size is None else
self._valid_batch_size(batch_size, "batch_size")
+ # Resolved from the hook or the connection each time a client is built.
+ self._index: str | None = index
+ self._batch_size = self.batch_size or self.MAX_BATCH_SIZE
+
+ @classmethod
+ @add_managed_identity_connection_widgets
+ def get_connection_form_widgets(cls) -> dict[str, Any]:
+ """Return connection widgets to add to connection form."""
+ from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
+ from flask_babel import lazy_gettext
+ from wtforms import IntegerField, StringField
+ from wtforms.validators import NumberRange, Optional
+
+ return {
+ "tenantId": StringField(lazy_gettext("Azure Tenant ID"),
widget=BS3TextFieldWidget()),
+ "index": StringField(lazy_gettext("Default Index"),
widget=BS3TextFieldWidget()),
+ "batch_size": IntegerField(
+ lazy_gettext("Batch Size"),
+ widget=BS3TextFieldWidget(),
+ validators=[Optional(), NumberRange(min=1,
max=cls.MAX_BATCH_SIZE)],
+ ),
+ }
+
+ @classmethod
+ def get_ui_field_behaviour(cls) -> dict[str, Any]:
+ """Return custom field behaviour."""
+ return {
+ "hidden_fields": ["schema", "port"],
+ "relabeling": {
+ "host": "Search Service Endpoint",
+ "login": "Azure Client ID",
+ "password": "Azure Secret or API Key",
+ },
+ "placeholders": {
+ "host": "https://<service>.search.windows.net",
+ "login": "client_id (service principal only)",
+ "password": "secret (service principal), or an API key when no
client_id is set",
+ "tenantId": "tenantId (service principal only)",
+ "index": "index used when the hook is given none",
+ "batch_size": "documents per write request, 1 to 1000 (default
1000)",
+ },
+ }
+
+ def _get_field(self, extras: dict[str, Any], field_name: str) -> Any:
+ return get_field(
+ conn_id=self.conn_id,
+ conn_type=self.conn_type,
+ extras=extras,
+ field_name=field_name,
+ )
+
+ def _get_credential(
+ self, conn: Connection, extras: dict[str, Any]
+ ) -> AzureKeyCredential | AsyncTokenCredential:
+ tenant = self._get_field(extras, "tenantId")
+ if conn.login:
+ if not (conn.password and tenant):
+ raise ValueError(
+ "Azure Client ID, Azure Secret, and Azure Tenant ID must
all be provided when "
+ "authenticating with a service principal."
+ )
+ self.log.info("Getting connection using specific credentials.")
+ return ClientSecretCredential(tenant_id=tenant,
client_id=conn.login, client_secret=conn.password)
+ if conn.password:
+ self.log.info("Getting connection using an API key.")
+ return AzureKeyCredential(conn.password)
+ self.log.info("Using DefaultAzureCredential as credential.")
+ return get_async_default_azure_credential(
+ managed_identity_client_id=self._get_field(extras,
"managed_identity_client_id"),
+ workload_identity_tenant_id=self._get_field(extras,
"workload_identity_tenant_id"),
+ )
+
+ @classmethod
+ def _valid_batch_size(cls, value: Any, source: str) -> int:
+ """Return ``value`` as a batch size, or raise naming where it came
from."""
+ # bool is a subclass of int: refuse it, True would silently mean 1.
+ if isinstance(value, bool):
+ raise ValueError(f"{source} must be an integer, got {value!r}")
+ if isinstance(value, str) and value.strip().isdigit():
+ value = int(value)
+ if not isinstance(value, int):
+ raise ValueError(f"{source} must be an integer, got {value!r}")
+ if not 1 <= value <= cls.MAX_BATCH_SIZE:
+ raise ValueError(f"{source} must be between 1 and
{cls.MAX_BATCH_SIZE}, got {value}")
+ return value
+
+ @asynccontextmanager
+ async def get_async_conn(self) -> AsyncGenerator[SearchClient, None]:
+ """Yield the async ``SearchClient`` of the index, closed with its
credential on exit."""
+ conn = await get_async_connection(self.conn_id)
+ # Masks the extra's secrets without a synchronous call to the Task SDK
on the event loop.
+ extras = await get_async_extra_dejson(conn)
+
+ self._index = self.index or self._get_field(extras, "index")
+ if not self._index:
+ raise ValueError(
+ f"No index for connection {self.conn_id!r}: pass index to the
hook or set it in the connection."
+ )
+ host = (conn.host or "").rstrip("/")
+ if not host:
+ raise ValueError(f"Connection {self.conn_id!r} has no search
service endpoint.")
+ endpoint = host if "://" in host else f"https://{host}";
+
+ if self.batch_size is not None:
+ self._batch_size = self.batch_size
+ elif (batch_size := self._get_field(extras, "batch_size")) not in
(None, ""):
+ self._batch_size = self._valid_batch_size(
+ batch_size, f"extra['batch_size'] of connection
{self.conn_id!r}"
+ )
+ else:
+ self._batch_size = self.MAX_BATCH_SIZE
+
+ credential = self._get_credential(conn, extras)
+ try:
+ async with SearchClient(endpoint, self._index, credential) as
client:
+ yield client
+ finally:
+ # The token credentials hold their own HTTP session, the key
credential has none.
+ if not isinstance(credential, AzureKeyCredential):
+ await credential.close()
+
+ async def search(
+ self,
+ search: str = "*",
+ *,
+ select: Iterable[str] | None = None,
+ filter: str | None = None,
+ order_by: Iterable[str] | None = None,
+ ) -> AsyncIterator[dict[str, Any]]:
+ """
+ Yield every document of the index matching ``search`` and ``filter``.
+
+ The service pages the results with ``$skip``, which stops at 100,000:
past that many matches
+ the search fails with an ``HttpResponseError``. To read a larger
index, order by a sortable and
+ filterable key and filter on ranges of it, e.g. ``order_by=["id"]`` and
+ ``filter="id gt '<last key read>'"``, one range per call.
+
+ .. seealso::
+ - Paging through a large number of results:
+
https://learn.microsoft.com/en-us/azure/search/search-pagination-page-layout#paging-through-a-large-number-of-results
+ - The ``skip`` limit of the search request:
+
https://learn.microsoft.com/en-us/rest/api/searchservice/documents/search-post
+
+ :param search: Full-text query, ``*`` for every document.
+ :param select: Fields to return, all retrievable fields when omitted.
+ :param filter: OData filter, e.g. ``"category eq 'report'"``.
+ :param order_by: OData sort expressions, e.g. ``["modified desc"]``.
+ """
+ async with self.get_async_conn() as client:
+ results = await client.search(
+ search_text=search,
+ select=list(select) if select else None,
+ filter=filter,
+ order_by=list(order_by) if order_by else None,
+ )
+ async for document in results:
+ yield {key: value for key, value in document.items() if not
key.startswith("@search.")}
+
+ async def count(self, *, filter: str | None = None) -> int:
+ """
+ Return the number of documents of the index, or of those matching
``filter``.
+
+ :param filter: OData filter, e.g. ``"category eq 'report'"``.
+ """
+ async with self.get_async_conn() as client:
+ if not filter:
+ return await client.get_document_count()
+ results = await client.search(search_text="*", filter=filter,
include_total_count=True, top=0)
+ return await results.get_count()
+
+ async def get_document(self, key: str, *, select: Iterable[str] | None =
None) -> dict[str, Any] | None:
+ """
+ Return the document with this key, or ``None`` when the index has none.
+
+ :param key: The key of the document.
+ :param select: Fields to return, all retrievable fields when omitted.
+ """
+ async with self.get_async_conn() as client:
+ try:
+ document = await client.get_document(key,
selected_fields=list(select) if select else None)
+ except ResourceNotFoundError:
+ return None
+ return dict(document)
+
+ async def upload(self, documents: Iterable[dict[str, Any]]) -> int:
+ """
+ Insert documents, or replace the ones whose key exists (every field is
replaced).
+
+ :param documents: The documents to write.
+ :return: The number of documents written.
+ """
+ return await self._write("upload_documents", documents)
+
+ async def merge(self, documents: Iterable[dict[str, Any]]) -> int:
+ """
+ Update the given fields of existing documents; a document whose key is
missing is rejected.
+
+ :param documents: The key and the fields to update of each document.
+ :return: The number of documents written.
+ """
+ return await self._write("merge_documents", documents)
+
+ async def merge_or_upload(self, documents: Iterable[dict[str, Any]]) ->
int:
+ """
+ Update the given fields of existing documents, or insert the document
when its key is new.
+
+ :param documents: The key and the fields to write of each document.
+ :return: The number of documents written.
+ """
+ return await self._write("merge_or_upload_documents", documents)
+
+ async def delete(self, documents: Iterable[dict[str, Any]]) -> int:
+ """
+ Delete documents.
+
+ :param documents: The documents to delete; each needs at least its key
field.
+ :return: The number of documents deleted.
+ """
+ return await self._write("delete_documents", documents)
+
+ async def _write(self, method: str, documents: Iterable[dict[str, Any]])
-> int:
+ written = 0
+ batch: list[dict[str, Any]] = []
+ async with self.get_async_conn() as client:
+ for document in documents:
+ batch.append(document)
+ if len(batch) == self._batch_size:
+ written += await self._write_batch(client, method, batch)
+ batch = []
+ if batch:
+ written += await self._write_batch(client, method, batch)
+ return written
+
+ async def _write_batch(self, client: SearchClient, method: str, batch:
list[dict[str, Any]]) -> int:
+ # The SDK returns one result per document and does not raise when the
service rejected
+ # some of them (HTTP 207).
+ results = await getattr(client, method)(documents=batch)
+ failed = [result for result in results if not result.succeeded]
+ if failed:
+ details = "; ".join(
+ f"{result.key}: {result.status_code} {result.error_message}"
for result in failed[:10]
+ )
+ raise AzureAISearchIndexingError(
+ f"Index {self._index!r} rejected {len(failed)} of {len(batch)}
documents ({details})"
+ )
+ return len(batch)
diff --git
a/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_search.py
b/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_search.py
new file mode 100644
index 00000000000..e8262a0056d
--- /dev/null
+++
b/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_search.py
@@ -0,0 +1,175 @@
+#
+# 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.
+"""
+System test for AzureAISearchHook.
+
+Requires a real Azure AI Search service, reachable through the
``azure_ai_search_default``
+connection (``AIRFLOW_CONN_AZURE_AI_SEARCH_DEFAULT``) with an admin API key as
password: the
+test creates its own index, writes and reads documents through the hook and
deletes the index.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+from airflow.providers.common.compat.sdk import DAG, TriggerRule, task
+from airflow.providers.microsoft.azure.hooks.ai_search import AzureAISearchHook
+
+DAG_ID = "example_azure_ai_search"
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") or "default"
+INDEX = f"airflow-system-test-{ENV_ID}"
+DOCUMENTS = [
+ {"id": "1", "title": "Annual report 2025", "category": "report"},
+ {"id": "2", "title": "Quarterly report Q1", "category": "report"},
+ {"id": "3", "title": "Meeting notes", "category": "memo"},
+]
+
+
+@task
+def create_index():
+ """Create the index of this test; the hook works on documents, not on
indexes."""
+ from azure.core.credentials import AzureKeyCredential
+ from azure.search.documents.indexes import SearchIndexClient
+ from azure.search.documents.indexes.models import (
+ SearchableField,
+ SearchFieldDataType,
+ SearchIndex,
+ SimpleField,
+ )
+
+ conn =
AzureAISearchHook.get_connection(AzureAISearchHook.default_conn_name)
+ client = SearchIndexClient(
+ f"https://{conn.host.removeprefix('https://')}",
AzureKeyCredential(conn.password)
+ )
+ client.create_or_update_index(
+ SearchIndex(
+ name=INDEX,
+ fields=[
+ SimpleField(name="id", type=SearchFieldDataType.String,
key=True, sortable=True),
+ SearchableField(name="title", type=SearchFieldDataType.String),
+ SimpleField(name="category", type=SearchFieldDataType.String,
filterable=True),
+ ],
+ )
+ )
+
+
+# [START howto_hook_azure_ai_search_upload]
+@task
+async def upload_documents() -> int:
+ """Insert the documents, replacing the ones whose key already exists."""
+ hook = AzureAISearchHook(index=INDEX)
+ return await hook.upload(DOCUMENTS)
+
+
+# [END howto_hook_azure_ai_search_upload]
+
+
+# [START howto_hook_azure_ai_search_search]
+@task(retries=3)
+async def find_reports() -> list[str]:
+ """Return the keys of the reports; retries while the service still indexes
the write."""
+ hook = AzureAISearchHook(index=INDEX)
+ reports = [
+ document["id"]
+ async for document in hook.search(select=["id"], filter="category eq
'report'", order_by=["id"])
+ ]
+ if len(reports) < 2:
+ raise ValueError(f"Expected 2 reports, found {reports}")
+ return reports
+
+
+# [END howto_hook_azure_ai_search_search]
+
+
+# [START howto_hook_azure_ai_search_merge]
+@task
+async def archive_reports(keys: list[str]) -> int:
+ """Change one field of the reports, then count the documents that have
it."""
+ hook = AzureAISearchHook(index=INDEX)
+ await hook.merge([{"id": key, "category": "archive"} for key in keys])
+ return await hook.count(filter="category eq 'archive'")
+
+
+# [END howto_hook_azure_ai_search_merge]
+
+
+# [START howto_hook_azure_ai_search_get_document]
+@task
+async def read_first_document(keys: list[str]) -> dict:
+ hook = AzureAISearchHook(index=INDEX)
+ document = await hook.get_document(keys[0], select=["id", "title"])
+ if document is None:
+ raise ValueError(f"Document {keys[0]!r} is missing")
+ return document
+
+
+# [END howto_hook_azure_ai_search_get_document]
+
+
+# [START howto_hook_azure_ai_search_delete]
+@task
+async def delete_documents() -> int:
+ """Delete the documents by key; only the key field is needed."""
+ hook = AzureAISearchHook(index=INDEX)
+ return await hook.delete({"id": document["id"]} for document in DOCUMENTS)
+
+
+# [END howto_hook_azure_ai_search_delete]
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_index():
+ from azure.core.credentials import AzureKeyCredential
+ from azure.search.documents.indexes import SearchIndexClient
+
+ conn =
AzureAISearchHook.get_connection(AzureAISearchHook.default_conn_name)
+ client = SearchIndexClient(
+ f"https://{conn.host.removeprefix('https://')}",
AzureKeyCredential(conn.password)
+ )
+ client.delete_index(INDEX)
+
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+ tags=["example"],
+) as dag:
+ reports = find_reports()
+ (
+ create_index()
+ >> upload_documents()
+ >> reports
+ >> archive_reports(reports)
+ >> read_first_document(reports)
+ >> delete_documents()
+ >> delete_index()
+ )
+
+ from tests_common.test_utils.watcher import watcher
+
+ # This test needs watcher in order to properly mark success/failure
+ # when "tearDown" task with trigger rule is part of the DAG
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+# Needed to run the example DAG with pytest (see:
contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_search.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_search.py
new file mode 100644
index 00000000000..a96164b4aaa
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_search.py
@@ -0,0 +1,394 @@
+#
+# 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
+
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+from azure.core.credentials import AzureKeyCredential
+from azure.core.exceptions import ResourceNotFoundError
+
+from airflow.models import Connection
+from airflow.providers.microsoft.azure.exceptions import
AzureAISearchIndexingError
+from airflow.providers.microsoft.azure.hooks.ai_search import AzureAISearchHook
+
+MODULE = "airflow.providers.microsoft.azure.hooks.ai_search"
+CONN_ID = "azure_ai_search_test"
+HOST = "test.search.windows.net"
+ENDPOINT = f"https://{HOST}";
+
+
+class FakeResults:
+ """Stands for the SDK's ``AsyncSearchItemPaged``: async iteration and
``get_count()``."""
+
+ def __init__(self, documents: list[dict]) -> None:
+ self.documents = documents
+
+ def __aiter__(self):
+ async def generate():
+ for document in self.documents:
+ yield {"@search.score": 1.0, **document}
+
+ return generate()
+
+ async def get_count(self) -> int:
+ return len(self.documents)
+
+
+class FakeSearchClient:
+ """Stands for the SDK's async ``SearchClient`` and records how the hook
builds and uses it."""
+
+ def __init__(self, endpoint, index_name, credential, documents=(),
reject=()) -> None:
+ self.endpoint = endpoint
+ self.index_name = index_name
+ self.credential = credential
+ self.documents = list(documents)
+ self.reject = set(reject)
+ self.searches: list[dict] = []
+ self.writes: list[tuple[str, list[dict]]] = []
+ self.closed = False
+
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, *exc):
+ self.closed = True
+
+ async def search(self, search_text=None, **kwargs):
+ self.searches.append({"search_text": search_text, **kwargs})
+ documents = self.documents
+ if kwargs.get("filter") == "category eq 'report'":
+ documents = [d for d in documents if d.get("category") == "report"]
+ if kwargs.get("select"):
+ documents = [{field: d.get(field) for field in kwargs["select"]}
for d in documents]
+ return FakeResults(documents)
+
+ async def get_document_count(self) -> int:
+ return len(self.documents)
+
+ async def get_document(self, key, selected_fields=None):
+ for document in self.documents:
+ if document["id"] == key:
+ if selected_fields:
+ return {field: document.get(field) for field in
selected_fields}
+ return document
+ raise ResourceNotFoundError("not found")
+
+ def __getattr__(self, name):
+ if not name.endswith("_documents"):
+ raise AttributeError(name)
+
+ async def write(documents):
+ self.writes.append((name, list(documents)))
+ return [
+ SimpleNamespace(
+ key=d["id"],
+ succeeded=d["id"] not in self.reject,
+ status_code=400 if d["id"] in self.reject else 200,
+ error_message="bad field" if d["id"] in self.reject else
None,
+ )
+ for d in documents
+ ]
+
+ return write
+
+
+def docs(count: int) -> list[dict]:
+ return [{"id": str(i), "title": f"title {i}", "category": None} for i in
range(count)]
+
+
[email protected]
+def search_service():
+ """Patch the connection and the SDK client; returns a namespace to set
them up and inspect them."""
+ service = SimpleNamespace(documents=[], reject=set(), clients=[])
+
+ def set_connection(**kwargs):
+ kwargs.setdefault("host", HOST)
+ kwargs.setdefault("extra", {"index": "idx"})
+ mock_get_connection.return_value = Connection(conn_id=CONN_ID,
conn_type="azure_ai_search", **kwargs)
+
+ def client(*args):
+ fake = FakeSearchClient(*args, documents=service.documents,
reject=service.reject)
+ service.clients.append(fake)
+ return fake
+
+ service.set_connection = set_connection
+ with (
+ mock.patch(f"{MODULE}.get_async_connection", autospec=True) as
mock_get_connection,
+ mock.patch(f"{MODULE}.SearchClient", side_effect=client),
+ ):
+ set_connection(password="api-key")
+ yield service
+
+
+class TestAzureAISearchHook:
+ def test_connection_form_widgets(self):
+ pytest.importorskip("flask_appbuilder")
+ widgets = AzureAISearchHook.get_connection_form_widgets()
+
+ assert set(widgets) == {
+ "tenantId",
+ "index",
+ "batch_size",
+ "managed_identity_client_id",
+ "workload_identity_tenant_id",
+ }
+
+ def test_ui_field_behaviour(self):
+ behaviour = AzureAISearchHook.get_ui_field_behaviour()
+
+ assert behaviour["hidden_fields"] == ["schema", "port"]
+ assert behaviour["relabeling"] == {
+ "host": "Search Service Endpoint",
+ "login": "Azure Client ID",
+ "password": "Azure Secret or API Key",
+ }
+
+ @pytest.mark.parametrize("batch_size", [0, 1001, -5, True, 2.5, "abc"])
+ def test_batch_size_is_refused_when_the_hook_is_built(self, batch_size):
+ with pytest.raises(ValueError, match="batch_size must be"):
+ AzureAISearchHook(CONN_ID, index="idx", batch_size=batch_size)
+
+
[email protected]
+class TestAzureAISearchHookAsync:
+ async def test_client_is_bound_to_the_index_with_an_api_key(self,
search_service):
+ hook = AzureAISearchHook(CONN_ID, index="my-index")
+
+ async with hook.get_async_conn() as client:
+ assert not client.closed
+
+ assert client.endpoint == ENDPOINT
+ assert client.index_name == "my-index"
+ assert isinstance(client.credential, AzureKeyCredential)
+ assert client.credential.key == "api-key"
+ assert client.closed
+
+ @pytest.mark.parametrize("host", [ENDPOINT, f"{ENDPOINT}/", HOST])
+ async def test_endpoint_from_the_host(self, search_service, host):
+ search_service.set_connection(host=host, password="api-key")
+
+ await AzureAISearchHook(CONN_ID).count()
+
+ assert search_service.clients[0].endpoint == ENDPOINT
+
+ async def test_index_from_the_extra(self, search_service):
+ await AzureAISearchHook(CONN_ID).count()
+
+ assert search_service.clients[0].index_name == "idx"
+
+ async def test_index_from_the_prefixed_extra(self, search_service):
+ search_service.set_connection(password="api-key",
extra={"extra__azure_ai_search__index": "legacy"})
+
+ await AzureAISearchHook(CONN_ID).count()
+
+ assert search_service.clients[0].index_name == "legacy"
+
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ async def test_service_principal_with_login_password_and_tenant(self,
mock_credential, search_service):
+ search_service.set_connection(
+ login="client-id", password="secret", extra={"index": "idx",
"tenantId": "tenant"}
+ )
+
+ await AzureAISearchHook(CONN_ID).count()
+
+ mock_credential.assert_called_once_with(
+ tenant_id="tenant", client_id="client-id", client_secret="secret"
+ )
+ assert search_service.clients[0].credential is
mock_credential.return_value
+ mock_credential.return_value.close.assert_awaited_once()
+
+ @pytest.mark.parametrize(
+ ("password", "extra"),
+ [
+ pytest.param("secret", {"index": "idx"}, id="no-tenant"),
+ pytest.param(None, {"index": "idx", "tenantId": "tenant"},
id="no-secret"),
+ ],
+ )
+ async def test_incomplete_service_principal(self, search_service,
password, extra):
+ search_service.set_connection(login="client-id", password=password,
extra=extra)
+
+ with pytest.raises(ValueError, match="must all be provided"):
+ await AzureAISearchHook(CONN_ID).count()
+
+ async def test_password_without_client_id_is_an_api_key(self,
search_service):
+ search_service.set_connection(password="api-key", extra={"index":
"idx", "tenantId": "tenant"})
+
+ await AzureAISearchHook(CONN_ID).count()
+
+ assert isinstance(search_service.clients[0].credential,
AzureKeyCredential)
+
+ @mock.patch(f"{MODULE}.get_async_default_azure_credential", autospec=True)
+ async def test_default_azure_credential_without_secret(self, mock_default,
search_service):
+ mock_default.return_value.close = mock.AsyncMock()
+ search_service.set_connection(
+ extra={
+ "index": "idx",
+ "managed_identity_client_id": "client-id",
+ "workload_identity_tenant_id": "tenant",
+ }
+ )
+
+ await AzureAISearchHook(CONN_ID).count()
+
+ mock_default.assert_called_once_with(
+ managed_identity_client_id="client-id",
workload_identity_tenant_id="tenant"
+ )
+ assert search_service.clients[0].credential is
mock_default.return_value
+ mock_default.return_value.close.assert_awaited_once()
+
+ async def test_missing_index(self, search_service):
+ search_service.set_connection(password="api-key", extra={})
+
+ with pytest.raises(ValueError, match="No index"):
+ await AzureAISearchHook(CONN_ID).count()
+
+ async def test_missing_endpoint(self, search_service):
+ search_service.set_connection(host="", password="api-key")
+
+ with pytest.raises(ValueError, match="no search service endpoint"):
+ await AzureAISearchHook(CONN_ID).count()
+
+ @mock.patch(f"{MODULE}.get_async_extra_dejson", autospec=True)
+ async def test_extra_is_read_through_the_async_helper(self, mock_extra,
search_service):
+ # Connection.extra_dejson masks through a synchronous call to the Task
SDK, which is not
+ # safe on the event loop of an async task.
+ mock_extra.return_value = {"index": "from-helper"}
+
+ await AzureAISearchHook(CONN_ID).count()
+
+ mock_extra.assert_awaited_once()
+ assert search_service.clients[0].index_name == "from-helper"
+
+ async def test_search_yields_every_document_without_search_metadata(self,
search_service):
+ search_service.documents = docs(2500)
+
+ found = [document async for document in
AzureAISearchHook(CONN_ID).search(select=["id", "title"])]
+
+ assert [document["id"] for document in found] == [str(i) for i in
range(2500)]
+ assert found[0] == {"id": "0", "title": "title 0"}
+ assert search_service.clients[0].searches == [
+ {"search_text": "*", "select": ["id", "title"], "filter": None,
"order_by": None}
+ ]
+ assert search_service.clients[0].closed
+
+ async def test_search_with_a_filter_and_an_order(self, search_service):
+ search_service.documents = [{"id": "1", "category": "report"}, {"id":
"2", "category": None}]
+
+ found = [
+ document
+ async for document in AzureAISearchHook(CONN_ID).search(
+ "annual", filter="category eq 'report'", order_by=("id desc",)
+ )
+ ]
+
+ assert [document["id"] for document in found] == ["1"]
+ assert search_service.clients[0].searches == [
+ {
+ "search_text": "annual",
+ "select": None,
+ "filter": "category eq 'report'",
+ "order_by": ["id desc"],
+ }
+ ]
+
+ async def test_count(self, search_service):
+ search_service.documents = docs(42)
+
+ assert await AzureAISearchHook(CONN_ID).count() == 42
+ assert search_service.clients[0].searches == []
+
+ async def test_count_with_a_filter(self, search_service):
+ search_service.documents = [{"id": "1", "category": "report"}, {"id":
"2", "category": None}]
+
+ assert await AzureAISearchHook(CONN_ID).count(filter="category eq
'report'") == 1
+ assert search_service.clients[0].searches == [
+ {"search_text": "*", "filter": "category eq 'report'",
"include_total_count": True, "top": 0}
+ ]
+
+ async def test_get_document(self, search_service):
+ search_service.documents = docs(3)
+ hook = AzureAISearchHook(CONN_ID)
+
+ assert await hook.get_document("2") == {"id": "2", "title": "title 2",
"category": None}
+ assert await hook.get_document("9") is None
+
+ async def test_get_document_with_select(self, search_service):
+ search_service.documents = docs(3)
+
+ document = await AzureAISearchHook(CONN_ID).get_document("2",
select=["id", "title"])
+
+ assert document == {"id": "2", "title": "title 2"}
+
+ @pytest.mark.parametrize(
+ ("method", "sdk_method"),
+ [
+ ("upload", "upload_documents"),
+ ("merge", "merge_documents"),
+ ("merge_or_upload", "merge_or_upload_documents"),
+ ("delete", "delete_documents"),
+ ],
+ )
+ async def test_writes_in_batches_of_1000(self, search_service, method,
sdk_method):
+ written = await getattr(AzureAISearchHook(CONN_ID), method)(docs(2001))
+
+ assert written == 2001
+ assert len(search_service.clients) == 1
+ assert [(name, len(batch)) for name, batch in
search_service.clients[0].writes] == [
+ (sdk_method, 1000),
+ (sdk_method, 1000),
+ (sdk_method, 1),
+ ]
+
+ async def test_writes_a_generator(self, search_service):
+ written = await AzureAISearchHook(CONN_ID,
batch_size=2).upload(document for document in docs(3))
+
+ assert written == 3
+ assert [len(batch) for _, batch in search_service.clients[0].writes]
== [2, 1]
+
+ async def test_rejected_documents_raise(self, search_service):
+ search_service.reject = {"1"}
+
+ with pytest.raises(
+ AzureAISearchIndexingError,
+ match=r"Index 'idx' rejected 1 of 3 documents \(1: 400 bad
field\)",
+ ):
+ await AzureAISearchHook(CONN_ID).merge(docs(3))
+
+ async def test_batch_size_from_the_hook(self, search_service):
+ search_service.set_connection(password="api-key", extra={"index":
"idx", "batch_size": 3})
+
+ await AzureAISearchHook(CONN_ID, batch_size=2).upload(docs(5))
+
+ assert [len(batch) for _, batch in search_service.clients[0].writes]
== [2, 2, 1]
+
+ @pytest.mark.parametrize("batch_size", [3, "3"])
+ async def test_batch_size_from_the_extra(self, search_service, batch_size):
+ search_service.set_connection(password="api-key", extra={"index":
"idx", "batch_size": batch_size})
+
+ await AzureAISearchHook(CONN_ID).upload(docs(5))
+
+ assert [len(batch) for _, batch in search_service.clients[0].writes]
== [3, 2]
+
+ @pytest.mark.parametrize("batch_size", ["abc", True, 2.5, 0, 5000])
+ async def test_batch_size_in_the_extra_is_validated_with_its_source(self,
search_service, batch_size):
+ search_service.set_connection(password="api-key", extra={"index":
"idx", "batch_size": batch_size})
+
+ with pytest.raises(ValueError, match=rf"extra\['batch_size'\] of
connection '{CONN_ID}' must be"):
+ await AzureAISearchHook(CONN_ID).upload(docs(1))
diff --git a/uv.lock b/uv.lock
index f05397da96e..d8cd255e072 100644
--- a/uv.lock
+++ b/uv.lock
@@ -6645,6 +6645,7 @@ dependencies = [
{ name = "azure-mgmt-datalake-store" },
{ name = "azure-mgmt-resource" },
{ name = "azure-mgmt-storage" },
+ { name = "azure-search-documents" },
{ name = "azure-servicebus" },
{ name = "azure-storage-blob" },
{ name = "azure-storage-file-datalake" },
@@ -6727,6 +6728,7 @@ requires-dist = [
{ name = "azure-mgmt-datalake-store", specifier = ">=0.5.0" },
{ name = "azure-mgmt-resource", specifier = ">=25.0.0" },
{ name = "azure-mgmt-storage", specifier = ">=24.0.1" },
+ { name = "azure-search-documents", specifier = ">=11.6.0,!=12.0.0" },
{ name = "azure-servicebus", specifier = ">=7.14.3" },
{ name = "azure-storage-blob", specifier = ">=12.28.0" },
{ name = "azure-storage-file-datalake", specifier = ">=12.23.0" },
@@ -10288,6 +10290,21 @@ wheels = [
{ url =
"https://files.pythonhosted.org/packages/c4/0c/c562be95a9a2ed52454f598571cf300b1114d0db2aa27f5b8ed3bb9cd0c0/azure_nspkg-3.0.2-py3-none-any.whl";,
hash =
"sha256:31a060caca00ed1ebd369fc7fe01a56768c927e404ebc92268f4d9d636435e28", size
= 1530, upload-time = "2018-10-02T22:48:44.806Z" },
]
+[[package]]
+name = "azure-search-documents"
+version = "11.6.0"
+source = { registry = "https://pypi.org/simple"; }
+dependencies = [
+ { name = "azure-common" },
+ { name = "azure-core" },
+ { name = "isodate" },
+ { name = "typing-extensions" },
+]
+sdist = { url =
"https://files.pythonhosted.org/packages/cf/68/9d59a0bed5fd9581b45444e8abc3ecda97e0466ae0f03affc7cddfb9fa74/azure_search_documents-11.6.0.tar.gz";,
hash =
"sha256:fcc807076ff82024be576ffccb0d0f3261e5c2a112a6666b86ec70bbdb2e1d64", size
= 311194, upload-time = "2025-10-09T22:04:03.655Z" }
+wheels = [
+ { url =
"https://files.pythonhosted.org/packages/c5/4c/d74e5c3ccc0b9ead0e400a2d70ded67554b56a5d799aaa8bf5baaacf4aea/azure_search_documents-11.6.0-py3-none-any.whl";,
hash =
"sha256:c3eb2deaf7926844e99a881830861225ef68e8b3bc067a76019e87fc7f5586dc", size
= 307935, upload-time = "2025-10-09T22:04:05.008Z" },
+]
+
[[package]]
name = "azure-servicebus"
version = "7.14.3"
