shahar1 commented on code in PR #74290: URL: https://github.com/apache/airflow/pull/74290#discussion_r4306410526
########## providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_search.py: ########## @@ -0,0 +1,385 @@ +# +# 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: + 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 Review Comment: **nit:** `FakeSearchClient.get_document` ignores `selected_fields`, and no test passes `select`. If `get_document(key, select=[...])` stopped forwarding the fields, these tests would still pass. Could you add one assertion that `selected_fields == ["id", "title"]` reaches the client? ########## providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_search.py: ########## @@ -0,0 +1,338 @@ +# +# 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:: + https://learn.microsoft.com/en-us/azure/search/ + + :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. + 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``. Review Comment: **nit:** "every document" stops being true above 100,000 matches. The SDK follows the service's `nextPageParameters`, which page with `$skip`. The SDK's own model documents that limit (`SearchRequest.skip`, 11.6.0): > The number of search results to skip. This value cannot be greater than 100,000. If you need to scan documents in sequence, but cannot use $skip due to this limitation, consider using $orderby on a totally-ordered key and $filter with a range query instead. An unbounded scan of a larger index (like the `indexed_fingerprints` task in the PR description) fails partway through with an `HttpResponseError`. A note in this docstring and in `docs/operators/ai_search.rst` would cover it, for example: "The service follows at most 100,000 results; to scan a larger index, order by the key and filter on key ranges." -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
