feat(azure-search): add Azure AI Search instrumentation and semantic conventions - #4392
feat(azure-search): add Azure AI Search instrumentation and semantic conventions#4392sanyamk23 wants to merge 1 commit into
Conversation
|
|
📝 WalkthroughWalkthroughAdded a new Azure AI Search OpenTelemetry instrumentation package. It wraps three Azure Search client types, records request and response attributes, handles span errors, registers semantic conventions, and includes package configuration, documentation, and tests. ChangesAzure AI Search instrumentation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SearchClient
participant AzureSearchInstrumentor
participant Tracer
participant SpanExporter
SearchClient->>AzureSearchInstrumentor: wrapped search operation
AzureSearchInstrumentor->>Tracer: create client span
AzureSearchInstrumentor->>SearchClient: execute original operation
SearchClient-->>AzureSearchInstrumentor: response or exception
AzureSearchInstrumentor->>Tracer: record attributes and status
Tracer->>SpanExporter: export completed span
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment Warning |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py`:
- Around line 176-179: Update the skillset handling branch in the
instrumentation method to stop recording skillset.name as
SpanAttributes.AZURE_SEARCH_INDEXER_NAME; add and use a dedicated
AZURE_SEARCH_SKILLSET_NAME attribute, and update the README entries for
create_skillset, get_skillset, and delete_skillset to document the resulting
behavior.
- Around line 190-205: Update the response-attribute handling for the document
methods in _set_response_attributes to count successful results using each
IndexingResult’s succeeded field instead of accessing r.error. Preserve the
existing total count from results and ensure both AZURE_SEARCH_SUCCEEDED_COUNT
and AZURE_SEARCH_DOCUMENTS_COUNT are set for all listed methods.
In `@packages/opentelemetry-instrumentation-azure-search/README.md`:
- Line 40: The README attribute table is inconsistent with the Azure Search
instrumentation behavior. After updating the input-attribute extraction logic to
capture the index name for SearchIndexClient.analyze_text and the indexer name
for create_skillset, get_skillset, and delete_skillset, update the corresponding
analyze_text and skillset rows in the README table to document the attributes
actually emitted.
In
`@packages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.py`:
- Around line 13-38: Update _make_search_client, _make_index_client, and
_make_indexer_client to construct real patched Azure SDK client instances, or
test doubles that invoke the instrumented SDK methods, instead of
MagicMock(spec=...). Preserve the existing endpoint and index configuration so
calls such as search, get_index, and indexer operations exercise
AzureSearchInstrumentor’s _instrument_method wrapping.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 6e1558ea-d47a-424d-86d5-3053b459f8ac
📒 Files selected for processing (15)
packages/opentelemetry-instrumentation-azure-search/.python-versionpackages/opentelemetry-instrumentation-azure-search/README.mdpackages/opentelemetry-instrumentation-azure-search/opentelemetry/__init__.pypackages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/__init__.pypackages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.pypackages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/config.pypackages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/utils.pypackages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/version.pypackages/opentelemetry-instrumentation-azure-search/poetry.tomlpackages/opentelemetry-instrumentation-azure-search/project.jsonpackages/opentelemetry-instrumentation-azure-search/pyproject.tomlpackages/opentelemetry-instrumentation-azure-search/tests/__init__.pypackages/opentelemetry-instrumentation-azure-search/tests/conftest.pypackages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.pypackages/opentelemetry-semantic-conventions-ai/opentelemetry/semconv_ai/__init__.py
| elif method in ("get_skillset", "delete_skillset", "create_skillset"): | ||
| skillset = kwargs.get("skillset") or (args[0] if args else None) | ||
| if hasattr(skillset, "name"): | ||
| set_span_attribute(span, SpanAttributes.AZURE_SEARCH_INDEXER_NAME, skillset.name) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Skillset name is recorded under the indexer attribute.
For create_skillset, get_skillset, and delete_skillset, the code sets SpanAttributes.AZURE_SEARCH_INDEXER_NAME from skillset.name. A skillset is a distinct Azure Search entity from an indexer, so tagging its name under db.azure_search.indexer_name produces misleading trace data (a create_skillset span will show an "indexer_name" that is actually a skillset name).
This also contradicts the README, which lists "—" (no attributes captured) for create_skillset(), get_skillset(), and delete_skillset() (README.md lines 56-58).
Add a dedicated AZURE_SEARCH_SKILLSET_NAME semantic attribute, or update the README to reflect the current behavior.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py`
around lines 176 - 179, Update the skillset handling branch in the
instrumentation method to stop recording skillset.name as
SpanAttributes.AZURE_SEARCH_INDEXER_NAME; add and use a dedicated
AZURE_SEARCH_SKILLSET_NAME attribute, and update the README entries for
create_skillset, get_skillset, and delete_skillset to document the resulting
behavior.
| elif method in ( | ||
| "index_documents", | ||
| "upload_documents", | ||
| "merge_documents", | ||
| "merge_or_upload_documents", | ||
| "delete_documents", | ||
| ): | ||
| if response is not None and hasattr(response, "results"): | ||
| results = list(response.results) | ||
| succeeded = sum(1 for r in results if not r.error) | ||
| set_span_attribute( | ||
| span, SpanAttributes.AZURE_SEARCH_SUCCEEDED_COUNT, succeeded, | ||
| ) | ||
| set_span_attribute( | ||
| span, SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT, len(results), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does azure-search-documents IndexingResult expose a "succeeded" attribute and not an "error" attribute?
💡 Result:
The IndexingResult class in Azure Search documents does expose a succeeded attribute, but it also includes an error message attribute (typically named errorMessage or error_message) [1][2][3][4]. Specifically, the IndexingResult object contains the following key properties across various language SDKs: succeeded: A boolean value indicating whether the indexing operation for a specific document was successful [1][2][3][4]. errorMessage (or error_message): A string that provides details on why an indexing operation failed; this property is null or empty if the operation succeeded [1][2][3][4]. statusCode: An integer representing the HTTP status code of the operation (e.g., 200, 201, 400, 404, etc.) [1][2][3][4]. key: The key of the document associated with the result [1][2][3][4]. Therefore, while the object does have a succeeded attribute to check the status, it does not lack an error-related attribute; rather, it uses an error message field to convey information when an operation fails [1][2][3][4].
Citations:
- 1: https://learn.microsoft.com/en-us/rest/api/searchservice/documents/?view=rest-searchservice-2026-04-01
- 2: https://azuresdkdocs.z19.web.core.windows.net/javascript/azure-search-documents/12.3.0-beta.1/interfaces/IndexingResult.html
- 3: https://learn.microsoft.com/en-us/python/api/azure-search-documents/azure.search.documents.models.indexingresult?view=azure-python
- 4: https://azuresdkdocs.z19.web.core.windows.net/dotnet/Azure.Search.Documents/11.8.0-beta.1/api/Azure.Search.Documents.Models/Azure.Search.Documents.Models.IndexingResult.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate files =="
git ls-files | rg '(^packages/opentelemetry-instrumentation-azure-search/|^pyproject\.toml|README\.md|.*azure.*search.*)' || true
echo
echo "== instrumentation file excerpt =="
file="packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py"
if [ -f "$file" ]; then
wc -l "$file"
sed -n '130,225p' "$file" | cat -n
fi
echo
echo "== dependency/version references =="
for f in $(git ls-files | rg '(^|/)(pyproject\.toml|uv.lock|poetry\.lock|requirements.*\.txt|setup\.cfg|setup\.py)$'); do
echo "--- $f"
rg -n "azure-search-documents|azure\.search\.documents|opentelemetry-instrumentation-azure-search|azure_search" "$f" || true
done
echo
echo "== tests around these attributes or mocks =="
git ls-files | rg 'test|tests' | xargs rg -n "IndexingResult|r\.error|ERROR_DOCUMENT|succeeded_count|documents_count|AZURE_SEARCH_SUCCEEDED_COUNT|AZURE_SEARCH_DOCUMENTS_COUNT|MagicMock\(error" 2>/dev/null || true
echo
echo "== deterministic source-shape check for exact expression =="
python3 - <<'PY'
from pathlib import Path
p=Path('packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py')
s=p.read_text()
print('__contains r.error:', 'r.error' in s)
print('contains r.succeeded:', 'r.succeeded' in s)
for needle in ['AZURE_SEARCH_SUCCEEDED_COUNT', 'AZURE_SEARCH_DOCUMENTS_COUNT']:
print(f'contains {needle}:', needle in s)
PYRepository: traceloop/openllmetry
Length of output: 15099
Use the SDK success flag instead of IndexingResult.error.
IndexingResult exposes succeeded and error_message; it does not expose error. Reading r.error raises AttributeError on real Azure Search responses, and because _set_response_attributes catches exceptions, AZURE_SEARCH_SUCCEEDED_COUNT and AZURE_SEARCH_DOCUMENTS_COUNT remain unset for index_documents, upload_documents, merge_documents, merge_or_upload_documents, and delete_documents spans.
🐛 Proposed fix using the real `succeeded` field
if response is not None and hasattr(response, "results"):
results = list(response.results)
- succeeded = sum(1 for r in results if not r.error)
+ succeeded = sum(1 for r in results if getattr(r, "succeeded", False))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| elif method in ( | |
| "index_documents", | |
| "upload_documents", | |
| "merge_documents", | |
| "merge_or_upload_documents", | |
| "delete_documents", | |
| ): | |
| if response is not None and hasattr(response, "results"): | |
| results = list(response.results) | |
| succeeded = sum(1 for r in results if not r.error) | |
| set_span_attribute( | |
| span, SpanAttributes.AZURE_SEARCH_SUCCEEDED_COUNT, succeeded, | |
| ) | |
| set_span_attribute( | |
| span, SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT, len(results), | |
| ) | |
| elif method in ( | |
| "index_documents", | |
| "upload_documents", | |
| "merge_documents", | |
| "merge_or_upload_documents", | |
| "delete_documents", | |
| ): | |
| if response is not None and hasattr(response, "results"): | |
| results = list(response.results) | |
| succeeded = sum(1 for r in results if getattr(r, "succeeded", False)) | |
| set_span_attribute( | |
| span, SpanAttributes.AZURE_SEARCH_SUCCEEDED_COUNT, succeeded, | |
| ) | |
| set_span_attribute( | |
| span, SpanAttributes.AZURE_SEARCH_DOCUMENTS_COUNT, len(results), | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py`
around lines 190 - 205, Update the response-attribute handling for the document
methods in _set_response_attributes to count successful results using each
IndexingResult’s succeeded field instead of accessing r.error. Preserve the
existing total count from results and ensure both AZURE_SEARCH_SUCCEEDED_COUNT
and AZURE_SEARCH_DOCUMENTS_COUNT are set for all listed methods.
| | `get_index()` | `azure_search.get_index` | index_name | | ||
| | `list_indexes()` | `azure_search.list_indexes` | — | | ||
| | `get_index_statistics()` | `azure_search.get_index_statistics` | index_doc_count, index_size_bytes | | ||
| | `analyze_text()` | `azure_search.analyze_text` | index_name | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
README attribute tables don't match implementation.
Line 40 documents analyze_text() as capturing index_name, but the input-attribute extraction in __init__.py only sets AZURE_SEARCH_INDEX_NAME for ("get_index", "delete_index", "create_index", "create_or_update_index") — analyze_text is not in that list, and SearchIndexClient instances don't carry the _index_name fallback used for SearchClient.
Lines 56-58 document create_skillset(), get_skillset(), and delete_skillset() as capturing no attributes ("—"), but the implementation does set AZURE_SEARCH_INDEXER_NAME from the skillset name for these methods (see the __init__.py review comment on this).
Update the table once the corresponding implementation is fixed, so the documented attributes match actual span data.
Also applies to: 56-58
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/opentelemetry-instrumentation-azure-search/README.md` at line 40,
The README attribute table is inconsistent with the Azure Search instrumentation
behavior. After updating the input-attribute extraction logic to capture the
index name for SearchIndexClient.analyze_text and the indexer name for
create_skillset, get_skillset, and delete_skillset, update the corresponding
analyze_text and skillset rows in the README table to document the attributes
actually emitted.
| def _make_search_client(index_name="test-index"): | ||
| """Create a mock SearchClient with patched endpoint and index.""" | ||
| from azure.search.documents import SearchClient | ||
|
|
||
| client = MagicMock(spec=SearchClient) | ||
| client._endpoint = "https://test.search.windows.net" | ||
| client._index_name = index_name | ||
| return client | ||
|
|
||
|
|
||
| def _make_index_client(): | ||
| """Create a mock SearchIndexClient.""" | ||
| from azure.search.documents.indexes import SearchIndexClient | ||
|
|
||
| client = MagicMock(spec=SearchIndexClient) | ||
| client._endpoint = "https://test.search.windows.net" | ||
| return client | ||
|
|
||
|
|
||
| def _make_indexer_client(): | ||
| """Create a mock SearchIndexerClient.""" | ||
| from azure.search.documents.indexes import SearchIndexerClient | ||
|
|
||
| client = MagicMock(spec=SearchIndexerClient) | ||
| client._endpoint = "https://test.search.windows.net" | ||
| return client |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🌐 Web query:
Does unittest.mock MagicMock(spec=SomeClass) delegate method calls to the real class implementation or to a class attribute patched by wrapt.wrap_function_wrapper?
💡 Result:
No, unittest.mock.MagicMock(spec=SomeClass) does not delegate method calls to the implementation of a class attribute patched by wrapt.wrap_function_wrapper [1][2]. When you create a MagicMock with a spec (e.g., MagicMock(spec=SomeClass)), the mock object creates a mock interface based on the attributes and methods of the original class at the time the mock is initialized [3][2]. Because MagicMock is designed to be a "test double" that replaces the original object, it is functionally independent of the real class's actual code execution [1][2]. Here is why delegation does not occur: 1. Isolation by Design: The purpose of a MagicMock (especially with a spec) is to intercept calls to methods defined in the spec, not to proxy them to the underlying implementation [1][2]. If you call a method on a MagicMock, the mock simply records the call and returns its own pre-configured return value (or a new child mock) [1][4]. 2. Mocking Interface vs. Behavior: The spec=SomeClass argument only limits the mock's interface—it ensures that if you attempt to access a method that does not exist on SomeClass, the mock raises an AttributeError [3][2]. It does not link the mock's behavior to the logic of the real methods, regardless of whether those methods have been monkey-patched by wrapt [2]. 3. Interaction with wrapt: wrapt.wrap_function_wrapper works by replacing an attribute in the original class (or module) with a wrapper function [5][6]. When you create a MagicMock with spec=SomeClass, the mock inspects the state of SomeClass at that moment. Even if you have successfully patched SomeClass with wrapt, the MagicMock does not execute the patched method; it treats the method name as a target to be mocked [1][2]. If you need the behavior of the real class to persist while only modifying specific methods, you should generally not replace the entire class with a MagicMock. Instead, consider using patch.object to replace only the specific method you wish to change, or use a side_effect function on the mock if you want to conditionally delegate to the original implementation [1][7][2].
Citations:
- 1: https://docs.python.org/3/library/unittest.mock.html
- 2: https://python-testing-debugging.com/advanced-mocking-test-doubles-in-python/deep-dive-into-unittestmock/when-to-use-magicmock-vs-mock-in-python/
- 3: https://docs.python.org/3/library/unittest.mock-examples.html
- 4: https://docs.python.org/release/3.10.19/library/unittest.mock.html
- 5: https://grahamdumpleton.me/posts/2026/05/lazy-monkey-patching-with-wrapt/
- 6: https://grahamdumpleton.me/posts/2015/03/safely-applying-monkey-patches-in-python/
- 7: https://stackoverflow.com/questions/50138219/mock-class-with-original-spec-python
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate relevant files"
fd -a 'test_azure_search_instrumentation.py|__init__.py|conftest.py' . | sed 's#^\./##' | sort
echo
echo "Inspect test helpers and instrumentor setup"
for f in $(fd 'test_azure_search_instrumentation.py|conftest.py|__init__.py' packages/opentelemetry-instrumentation-azure-search); do
echo "===== $f ====="
wc -l "$f"
sed -n '1,220p' "$f"
done
echo
echo "Find AzureSearchInstrumentor and instrumentation helper patterns"
rg -n "AzureSearchInstrumentor|wrap_function_wrapper|instrument\\(|_wrap|search_spans|get_index|create_data_source|indexer|IndexingResult|r\\.error" packages/opentelemetry-instrumentation-azure -S
echo
python3 - <<'PY'
import inspect
from unittest.mock import MagicMock
class Fake:
def method(self):
return "real"
Fake.method = lambda self: "patched class-level"
m = MagicMock(spec=Fake)
print("MagicMock(spec=Fake) call result:", m.method())
print("MagicMock object identity different:", m.method is not Fake.method)
PYRepository: traceloop/openllmetry
Length of output: 32076
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Read remaining azure_search __init__.py"
sed -n '220,327p' packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search/__init__.py
echo
echo "Find AzureSearchInstrumentor instrumentation code"
rg -n "class AzureSearchInstrumentor|def instrument|def _uninstrument|wrap_function_wrapper|_WrappedOperation" packages/opentelemetry-instrumentation-azure-search/opentelemetry/instrumentation/azure_search -S
echo
echo "Search for AzureSearch helpers outside package root"
rg -n "AzureSearchInstrumentor|wrap_function_wrapper|_WrappedOperation" packages/test_*.py packages/opentelemetry-instrumentation-azure-search -S
echo
echo "Static probe: MagicMock(spec=fake_class_with_patch) call bypass"
python3 - <<'PY'
from unittest.mock import MagicMock
import sys
class FakeSdkMethod:
def __init__(self, name):
self.calls = []
self.output = f"real {name}"
def __get__(self, obj, objtype=None):
def call(*args, **kwargs):
self.calls.append((args, kwargs))
return self.output
return call
FakeSdkMethod.search = FakeSdkMethod("search")
FakeSdkMethod.instance_attribute = 1
spec_obj = MagicMock(spec=type("Sdk", (), {"search": lambda self: "spec searched"}))
print("spec_obj.search is spec class attribute?", spec_obj.search is type("Sdk", (), {"search": lambda self: "spec searched"}).search)
print("spec_obj.search callable?", callable(spec_obj.search))
print("spec_obj.search result:", spec_obj.search())
# Mirror the conftest patching model: an existing client-like object gets a patched runtime method later.
class Client:
def search(self):
return "before patch"
Client.search = lambda self: "patched"
spec_client = MagicMock(spec=Client)
spec_client.search().return_value = "config return"
print("spec_client.search is patched runtime method?", spec_client.search is Client.search)
PYRepository: traceloop/openllmetry
Length of output: 6474
Use real patched SDK clients in these tests.
_make_search_client, _make_index_client, and _make_indexer_client build MagicMock(spec=SearchClient) and the index/indexer equivalents, then the tests call client.search(...), client.get_index(...), etc. These mocked methods are independent child mocks, so they do not invoke the AzureSearchInstrumentor wrappage in _instrument_method.
Move to real SearchClient/SearchIndexClient/SearchIndexerClient instances, or another test double whose method call path is exercised by the patched SDK method.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@packages/opentelemetry-instrumentation-azure-search/tests/test_azure_search_instrumentation.py`
around lines 13 - 38, Update _make_search_client, _make_index_client, and
_make_indexer_client to construct real patched Azure SDK client instances, or
test doubles that invoke the instrumented SDK methods, instead of
MagicMock(spec=...). Preserve the existing endpoint and index configuration so
calls such as search, get_index, and indexer operations exercise
AzureSearchInstrumentor’s _instrument_method wrapping.
2 participants
Description
Adds a new
opentelemetry-instrumentation-azure-searchpackage and the corresponding semantic conventions to support tracing for Azure AI Search.Changes
opentelemetry-instrumentation-azure-searchwith instrumentation, config, utils, testssemconv_ai(db.azure_search.*namespace): index name, search text, top, filter, result/document/succeeded counts, indexer name/status, index dimensions/doc count/size, service limit/usage, autocomplete/suggest textSummary by CodeRabbit
New Features
Documentation
Tests