• Home
• airflow.providers.microsoft.azure
• airflow.providers.microsoft.azure.hooks
• airflow.providers.microsoft.azure.hooks.ai_search

↑↓ Navigate ⏎ Select Esc Close

airflow.providers.microsoft.azure.hooks.ai_search

Classes

AzureAISearchHook

Query and write the documents of one Azure AI Search index, natively async.

Module Contents

class airflow.providers.microsoft.azure.hooks.ai_search.AzureAISearchHook(azure_ai_search_conn_id=default_conn_name, index=None, batch_size=None)

Bases: airflow.providers.common.compat.sdk.BaseHook

Query and write the documents of one Azure AI Search index, natively async.

A thin layer over 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:

• search() is an async generator over every matching document, without the @search.* metadata.

• count() and get_document() (None when the index has no such document).

• upload(), merge(), merge_or_upload() and delete() write in batches of batch_size documents and raise 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: 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.

See also

• 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

Parameters:

• azure_ai_search_conn_id (str) – The Azure AI Search connection id. Default is azure_ai_search_default.

• index (str | None) – The index this hook works on. Overrides the index connection extra.

• batch_size (int | None) – 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'

MAX_BATCH_SIZE = 1000

conn_id = 'azure_ai_search_default'

index = None

batch_size = None

classmethod get_connection_form_widgets()

Return connection widgets to add to connection form.

classmethod get_ui_field_behaviour()

Return custom field behaviour.

async get_async_conn()

Yield the async SearchClient of the index, closed with its credential on exit.

async search(search='*', *, select=None, filter=None, order_by=None)

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.

See also

• 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

Parameters:

• search (str) – Full-text query, * for every document.

• select (collections.abc.Iterable[str] | None) – Fields to return, all retrievable fields when omitted.

• filter (str | None) – OData filter, e.g. "category eq 'report'".

• order_by (collections.abc.Iterable[str] | None) – OData sort expressions, e.g. ["modified desc"].

async count(*, filter=None)

Return the number of documents of the index, or of those matching filter.

Parameters:

filter (str | None) – OData filter, e.g. "category eq 'report'".

async get_document(key, *, select=None)

Return the document with this key, or None when the index has none.

Parameters:

• key (str) – The key of the document.

• select (collections.abc.Iterable[str] | None) – Fields to return, all retrievable fields when omitted.

async upload(documents)

Insert documents, or replace the ones whose key exists (every field is replaced).

Parameters:

documents (collections.abc.Iterable[dict[str, Any]]) – The documents to write.

Returns:

The number of documents written.

Return type:

int

async merge(documents)

Update the given fields of existing documents; a document whose key is missing is rejected.

Parameters:

documents (collections.abc.Iterable[dict[str, Any]]) – The key and the fields to update of each document.

Returns:

The number of documents written.

Return type:

int

async merge_or_upload(documents)

Update the given fields of existing documents, or insert the document when its key is new.

Parameters:

documents (collections.abc.Iterable[dict[str, Any]]) – The key and the fields to write of each document.

Returns:

The number of documents written.

Return type:

int

async delete(documents)

Delete documents.

Parameters:

documents (collections.abc.Iterable[dict[str, Any]]) – The documents to delete; each needs at least its key field.

Returns:

The number of documents deleted.

Return type:

int

Previous Next

Was this entry helpful?