align adapter-azureaisearch to API changes

This commit is contained in:
Jonathan Shook 2024-10-24 14:20:30 -05:00
parent a7e5580203
commit 109ca07b2d
17 changed files with 145 additions and 126 deletions

View File

@ -15,45 +15,50 @@
-->
<!--<?xml version="1.0" encoding="UTF-8"?>-->
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<artifactId>adapter-azure-aisearch</artifactId>
<packaging>jar</packaging>
<artifactId>adapter-azure-aisearch</artifactId>
<packaging>jar</packaging>
<parent>
<artifactId>mvn-defaults</artifactId>
<groupId>io.nosqlbench</groupId>
<version>${revision}</version>
<relativePath>../../mvn-defaults</relativePath>
</parent>
<parent>
<artifactId>mvn-defaults</artifactId>
<groupId>io.nosqlbench</groupId>
<version>${revision}</version>
<relativePath>../../mvn-defaults</relativePath>
</parent>
<name>${project.artifactId}</name>
<description>
An nosqlbench adapter driver module for the Azure AI Search database.
</description>
<name>${project.artifactId}</name>
<description>
An nosqlbench adapter driver module for the Azure AI Search database.
</description>
<dependencies>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>nb-annotations</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapters-api</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-search-documents</artifactId>
</dependency>
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-identity</artifactId>
</dependency>
</dependencies>
<dependencies>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>nb-annotations</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapters-api</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-search-documents</artifactId>
<version>11.7.0</version>
</dependency>
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-identity</artifactId>
<version>1.13.2</version>
</dependency>
</dependencies>
</project>

View File

@ -17,6 +17,7 @@ package io.nosqlbench.adapter.azureaisearch;
import java.util.function.Function;
import java.util.function.LongFunction;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp;
import io.nosqlbench.adapters.api.activityimpl.OpMapper;
@ -31,19 +32,19 @@ import io.nosqlbench.nb.api.labels.NBLabels;
import static io.nosqlbench.adapter.azureaisearch.AzureAISearchAdapterUtils.AZURE_AI_SEARCH;
@Service(value = DriverAdapter.class, selector = AZURE_AI_SEARCH)
public class AzureAISearchDriverAdapter extends BaseDriverAdapter<AzureAISearchBaseOp<?>, AzureAISearchSpace> {
public class AzureAISearchDriverAdapter extends BaseDriverAdapter<AzureAISearchBaseOp, AzureAISearchSpace> {
public AzureAISearchDriverAdapter(NBComponent parentComponent, NBLabels labels) {
super(parentComponent, labels);
}
@Override
public OpMapper<AzureAISearchBaseOp<?>> getOpMapper() {
public OpMapper<AzureAISearchBaseOp,AzureAISearchSpace> getOpMapper() {
return new AzureAISearchOpMapper(this);
}
@Override
public Function<String, ? extends AzureAISearchSpace> getSpaceInitializer(NBConfiguration cfg) {
public LongFunction<AzureAISearchSpace> getSpaceInitializer(NBConfiguration cfg) {
return (s) -> new AzureAISearchSpace(s, cfg);
}

View File

@ -32,44 +32,53 @@ import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.nosqlbench.engine.api.templating.TypeAndTarget;
public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp<?>> {
private static final Logger logger = LogManager.getLogger(AzureAISearchOpMapper.class);
private final AzureAISearchDriverAdapter adapter;
import java.util.function.LongFunction;
/**
* Create a new {@code AzureAISearchOpMapper} implementing the {@link OpMapper}.
* interface.
*
* @param adapter The associated {@link AzureAISearchDriverAdapter}
*/
public AzureAISearchOpMapper(AzureAISearchDriverAdapter adapter) {
this.adapter = adapter;
}
public class AzureAISearchOpMapper implements OpMapper<AzureAISearchBaseOp, AzureAISearchSpace> {
private static final Logger logger = LogManager.getLogger(AzureAISearchOpMapper.class);
private final AzureAISearchDriverAdapter adapter;
/**
* Given an instance of a {@link ParsedOp} returns the appropriate
* {@link AzureAISearchBaseOpDispenser} subclass.
*
* @param op The {@link ParsedOp} to be evaluated
* @return The correct {@link AzureAISearchBaseOpDispenser} subclass based on
* the op type
*/
@Override
public OpDispenser<? extends AzureAISearchBaseOp<?>> apply(ParsedOp op) {
TypeAndTarget<AzureAISearchOpType, String> typeAndTarget = op.getTypeAndTarget(AzureAISearchOpType.class,
String.class, "type", "target");
logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'");
/**
* Create a new {@code AzureAISearchOpMapper} implementing the {@link OpMapper}.
* interface.
*
* @param adapter
* The associated {@link AzureAISearchDriverAdapter}
*/
public AzureAISearchOpMapper(AzureAISearchDriverAdapter adapter) {
this.adapter = adapter;
}
return switch (typeAndTarget.enumId) {
case delete_index -> new AzureAISearchDeleteIndexOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_or_update_index ->
new AzureAISearchCreateOrUpdateIndexOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_indexes -> new AzureAISearchListIndexesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case upload_documents -> new AzureAISearchUploadDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
case search_documents -> new AzureAISearchSearchDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
/**
* Given an instance of a {@link ParsedOp} returns the appropriate
* {@link AzureAISearchBaseOpDispenser} subclass.
*
* @param op
* The {@link ParsedOp} to be evaluated
* @param spaceInitF
* @return The correct {@link AzureAISearchBaseOpDispenser} subclass based on
* the op type
*/
@Override
public OpDispenser<AzureAISearchBaseOp> apply(ParsedOp op, LongFunction<AzureAISearchSpace> spaceInitF) {
TypeAndTarget<AzureAISearchOpType, String> typeAndTarget = op.getTypeAndTarget(AzureAISearchOpType.class,
String.class, "type", "target");
logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'");
return switch (typeAndTarget.enumId) {
case delete_index -> new AzureAISearchDeleteIndexOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_or_update_index ->
new AzureAISearchCreateOrUpdateIndexOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_indexes -> new AzureAISearchListIndexesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case upload_documents ->
new AzureAISearchUploadDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
case search_documents ->
new AzureAISearchSearchDocumentsOpDispenser(adapter, op, typeAndTarget.targetFunction);
// default -> throw new RuntimeException(
// "Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " + "mapping parsed op " + op);
};
}
};
}
}

View File

@ -20,6 +20,7 @@ import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import io.nosqlbench.adapters.api.activityimpl.uniform.BaseSpace;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@ -50,9 +51,8 @@ import io.nosqlbench.nb.api.config.standard.Param;
* "https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/search/azure-search-documents/">Azure
* AI Search Java searchIndexClient</a>
*/
public class AzureAISearchSpace implements AutoCloseable {
public class AzureAISearchSpace extends BaseSpace {
private final static Logger logger = LogManager.getLogger(AzureAISearchSpace.class);
private final String name;
private final NBConfiguration cfg;
protected SearchIndexClient searchIndexClient;
@ -62,11 +62,11 @@ public class AzureAISearchSpace implements AutoCloseable {
* contextual information needed to interact with the <b>Azure AI Search</b>
* database instance.
*
* @param name The name of this space
* @param idx The name of this space
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/
public AzureAISearchSpace(String name, NBConfiguration cfg) {
this.name = name;
public AzureAISearchSpace(long idx, NBConfiguration cfg) {
super(idx);
this.cfg = cfg;
}

View File

@ -26,33 +26,32 @@ import com.azure.search.documents.indexes.SearchIndexClient;
import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.CycleOp;
public abstract class AzureAISearchBaseOp<T> implements CycleOp<Object> {
public abstract class AzureAISearchBaseOp<REQUEST, RESULT> implements CycleOp<RESULT> {
protected final static Logger logger = LogManager.getLogger(AzureAISearchBaseOp.class);
protected final SearchIndexClient searchIndexClient;
protected final T request;
protected final LongFunction<Object> apiCall;
protected final REQUEST request;
protected final LongFunction<RESULT> apiCall;
public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, T requestParam) {
public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, REQUEST requestParam) {
this.searchIndexClient = searchIndexClient;
this.request = requestParam;
this.apiCall = this::applyOp;
}
public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, T requestParam, LongFunction<Object> call) {
public AzureAISearchBaseOp(SearchIndexClient searchIndexClient, REQUEST requestParam, LongFunction<RESULT> call) {
this.searchIndexClient = searchIndexClient;
this.request = requestParam;
this.apiCall = call;
}
@Override
public final Object apply(long value) {
public final RESULT apply(long value) {
logger.trace(() -> "applying op: " + this);
try {
Object result = applyOp(value);
RESULT result = applyOp(value);
return result;
} catch (Exception e) {
if (e instanceof RuntimeException rte) {
@ -63,7 +62,7 @@ public abstract class AzureAISearchBaseOp<T> implements CycleOp<Object> {
}
};
public abstract Object applyOp(long value);
public abstract RESULT applyOp(long value);
@Override
public String toString() {

View File

@ -17,11 +17,13 @@ package io.nosqlbench.adapter.azureaisearch.ops;
import com.azure.search.documents.SearchClient;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.util.AutocompletePagedIterable;
import com.azure.search.documents.util.SearchPagedIterable;
public abstract class AzureAISearchClientBaseOp<T> extends AzureAISearchBaseOp<T> {
public abstract class AzureAISearchClientBaseOp<REQUEST, RESULT> extends AzureAISearchBaseOp<REQUEST, RESULT> {
protected final SearchClient searchClient;
public AzureAISearchClientBaseOp(SearchIndexClient searchIdxClient, SearchClient searchClnt, T requestParam) {
public AzureAISearchClientBaseOp(SearchIndexClient searchIdxClient, SearchClient searchClnt, REQUEST requestParam) {
super(searchIdxClient, requestParam);
this.searchClient = searchClnt;
}

View File

@ -18,14 +18,14 @@ package io.nosqlbench.adapter.azureaisearch.ops;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.models.SearchIndex;
public class AzureAISearchCreateOrUpdateIndexOp extends AzureAISearchBaseOp<SearchIndex> {
public class AzureAISearchCreateOrUpdateIndexOp extends AzureAISearchBaseOp<SearchIndex,SearchIndex> {
public AzureAISearchCreateOrUpdateIndexOp(SearchIndexClient client, SearchIndex request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
public SearchIndex applyOp(long value) {
SearchIndex createResponse = null;
try {
createResponse = searchIndexClient.createOrUpdateIndex(request);

View File

@ -17,14 +17,14 @@ package io.nosqlbench.adapter.azureaisearch.ops;
import com.azure.search.documents.indexes.SearchIndexClient;
public class AzureAISearchDeleteIndexOp extends AzureAISearchBaseOp<String> {
public class AzureAISearchDeleteIndexOp extends AzureAISearchBaseOp<String,String> {
public AzureAISearchDeleteIndexOp(SearchIndexClient client, String request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
public String applyOp(long value) {
try {
searchIndexClient.deleteIndex(request);
logger.debug("Successfully deleted the index: {}", request);

View File

@ -19,14 +19,14 @@ import com.azure.core.http.rest.PagedIterable;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.models.SearchIndex;
public class AzureAISearchListIndexesOp extends AzureAISearchBaseOp<String> {
public class AzureAISearchListIndexesOp extends AzureAISearchBaseOp<String,String> {
public AzureAISearchListIndexesOp(SearchIndexClient client, String request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
public String applyOp(long value) {
try {
PagedIterable<SearchIndex> response = searchIndexClient.listIndexes();
response.forEach((index) -> {

View File

@ -19,9 +19,10 @@ import com.azure.core.util.Context;
import com.azure.search.documents.SearchClient;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.models.SearchOptions;
import com.azure.search.documents.util.AutocompletePagedIterable;
import com.azure.search.documents.util.SearchPagedIterable;
public class AzureAISearchSearchDocumentsOp extends AzureAISearchClientBaseOp<SearchOptions> {
public class AzureAISearchSearchDocumentsOp extends AzureAISearchClientBaseOp<SearchOptions,SearchPagedIterable> {
public AzureAISearchSearchDocumentsOp(SearchIndexClient searchIndexClient, SearchClient searchClient,
SearchOptions request) {
@ -29,7 +30,7 @@ public class AzureAISearchSearchDocumentsOp extends AzureAISearchClientBaseOp<Se
}
@Override
public Object applyOp(long value) {
public SearchPagedIterable applyOp(long value) {
SearchPagedIterable searchDocsResponse = null;
try {
searchDocsResponse = searchClient.search(null, // we've not implemented other complex searches yet here.

View File

@ -22,7 +22,7 @@ import com.azure.search.documents.SearchDocument;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.models.IndexDocumentsResult;
public class AzureAISearchUploadDocumentsOp extends AzureAISearchClientBaseOp<SearchDocument> {
public class AzureAISearchUploadDocumentsOp extends AzureAISearchClientBaseOp<SearchDocument,IndexDocumentsResult> {
public AzureAISearchUploadDocumentsOp(SearchIndexClient searchIndexClient, SearchClient searchClient,
SearchDocument request) {
@ -30,7 +30,7 @@ public class AzureAISearchUploadDocumentsOp extends AzureAISearchClientBaseOp<Se
}
@Override
public Object applyOp(long value) {
public IndexDocumentsResult applyOp(long value) {
IndexDocumentsResult uploadDocsResponse = null;
try {
uploadDocsResponse = searchClient.uploadDocuments(List.of(request));

View File

@ -28,13 +28,13 @@ import io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapters.api.templating.ParsedOp;
public abstract class AzureAISearchBaseOpDispenser<T>
extends BaseOpDispenser<AzureAISearchBaseOp<T>, AzureAISearchSpace> {
public abstract class AzureAISearchBaseOpDispenser<REQUEST,RESULT>
extends BaseOpDispenser<AzureAISearchBaseOp, AzureAISearchSpace> {
protected final LongFunction<AzureAISearchSpace> azureAISearchSpaceFunction;
protected final LongFunction<SearchIndexClient> clientFunction;
private final LongFunction<? extends AzureAISearchBaseOp<T>> opF;
private final LongFunction<T> paramF;
private final LongFunction<? extends AzureAISearchBaseOp<REQUEST,RESULT>> opF;
private final LongFunction<REQUEST> paramF;
@SuppressWarnings("rawtypes")
protected AzureAISearchBaseOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
@ -57,14 +57,14 @@ public abstract class AzureAISearchBaseOpDispenser<T>
return (AzureAISearchDriverAdapter) adapter;
}
public abstract LongFunction<T> getParamFunc(LongFunction<SearchIndexClient> clientF, ParsedOp op,
public abstract LongFunction<REQUEST> getParamFunc(LongFunction<SearchIndexClient> clientF, ParsedOp op,
LongFunction<String> targetF);
public abstract LongFunction<AzureAISearchBaseOp<T>> createOpFunc(LongFunction<T> paramF,
public abstract LongFunction<AzureAISearchBaseOp<REQUEST,RESULT>> createOpFunc(LongFunction<REQUEST> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF);
@Override
public AzureAISearchBaseOp<T> getOp(long value) {
public AzureAISearchBaseOp<REQUEST,RESULT> getOp(long value) {
return opF.apply(value);
}

View File

@ -53,7 +53,7 @@ import io.nosqlbench.nb.api.errors.OpConfigError;
* "https://learn.microsoft.com/en-us/java/api/com.azure.search.documents?view=azure-java-stable">Index
* docs</a>
*/
public class AzureAISearchCreateOrUpdateIndexOpDispenser extends AzureAISearchBaseOpDispenser<SearchIndex> {
public class AzureAISearchCreateOrUpdateIndexOpDispenser extends AzureAISearchBaseOpDispenser<SearchIndex,SearchIndex> {
private SearchField searchField;
private VectorSearchProfile vsProfile;
@ -290,7 +290,7 @@ public class AzureAISearchCreateOrUpdateIndexOpDispenser extends AzureAISearchBa
}
@Override
public LongFunction<AzureAISearchBaseOp<SearchIndex>> createOpFunc(LongFunction<SearchIndex> paramF,
public LongFunction<AzureAISearchBaseOp<SearchIndex,SearchIndex>> createOpFunc(LongFunction<SearchIndex> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchCreateOrUpdateIndexOp(clientF.apply(l), paramF.apply(l));
}

View File

@ -34,7 +34,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
* "https://learn.microsoft.com/en-us/rest/api/searchservice/">REST
* API</a>.
*/
public class AzureAISearchDeleteIndexOpDispenser extends AzureAISearchBaseOpDispenser<String> {
public class AzureAISearchDeleteIndexOpDispenser extends AzureAISearchBaseOpDispenser<String,String> {
public AzureAISearchDeleteIndexOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) {
@ -48,7 +48,7 @@ public class AzureAISearchDeleteIndexOpDispenser extends AzureAISearchBaseOpDisp
}
@Override
public LongFunction<AzureAISearchBaseOp<String>> createOpFunc(LongFunction<String> paramF,
public LongFunction<AzureAISearchBaseOp<String,String>> createOpFunc(LongFunction<String> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchDeleteIndexOp(clientF.apply(l), paramF.apply(l));
}

View File

@ -27,7 +27,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
/**
* Lists the indexes available.
*/
public class AzureAISearchListIndexesOpDispenser extends AzureAISearchBaseOpDispenser<String> {
public class AzureAISearchListIndexesOpDispenser extends AzureAISearchBaseOpDispenser<String,String> {
public AzureAISearchListIndexesOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) {
super(adapter, op, targetF);
@ -40,7 +40,7 @@ public class AzureAISearchListIndexesOpDispenser extends AzureAISearchBaseOpDisp
}
@Override
public LongFunction<AzureAISearchBaseOp<String>> createOpFunc(LongFunction<String> paramF,
public LongFunction<AzureAISearchBaseOp<String,String>> createOpFunc(LongFunction<String> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchListIndexesOp(clientF.apply(l), paramF.apply(l));
}

View File

@ -26,6 +26,7 @@ import com.azure.search.documents.models.VectorQuery;
import com.azure.search.documents.models.VectorSearchOptions;
import com.azure.search.documents.models.VectorizedQuery;
import com.azure.search.documents.util.SearchPagedIterable;
import io.nosqlbench.adapter.azureaisearch.AzureAISearchDriverAdapter;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchSearchDocumentsOp;
@ -40,7 +41,7 @@ import io.nosqlbench.nb.api.errors.OpConfigError;
* "https://learn.microsoft.com/en-us/azure/search/vector-search-how-to-query?tabs=query-2024-07-01%2Cfilter-2024-07-01%2Cbuiltin-portal#vector-query-request">How
* to query/vector search</a>
*/
public class AzureAISearchSearchDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchOptions> {
public class AzureAISearchSearchDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchOptions, SearchPagedIterable> {
public AzureAISearchSearchDocumentsOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) {
super(adapter, op, targetF);
@ -70,7 +71,7 @@ public class AzureAISearchSearchDocumentsOpDispenser extends AzureAISearchBaseOp
}
@Override
public LongFunction<AzureAISearchBaseOp<SearchOptions>> createOpFunc(LongFunction<SearchOptions> paramF,
public LongFunction<AzureAISearchBaseOp<SearchOptions,SearchPagedIterable>> createOpFunc(LongFunction<SearchOptions> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchSearchDocumentsOp(clientF.apply(l),
clientF.apply(l).getSearchClient(targetF.apply(l)), paramF.apply(l));

View File

@ -21,6 +21,7 @@ import java.util.function.LongFunction;
import com.azure.search.documents.SearchDocument;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.models.IndexDocumentsResult;
import io.nosqlbench.adapter.azureaisearch.AzureAISearchDriverAdapter;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchBaseOp;
import io.nosqlbench.adapter.azureaisearch.ops.AzureAISearchUploadDocumentsOp;
@ -34,7 +35,7 @@ import io.nosqlbench.adapters.api.templating.ParsedOp;
* "https://learn.microsoft.com/en-us/java/api/com.azure.search.documents?view=azure-java-stable">Index
* docs</a>
*/
public class AzureAISearchUploadDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchDocument> {
public class AzureAISearchUploadDocumentsOpDispenser extends AzureAISearchBaseOpDispenser<SearchDocument, IndexDocumentsResult> {
public AzureAISearchUploadDocumentsOpDispenser(AzureAISearchDriverAdapter adapter, ParsedOp op,
LongFunction<String> targetF) {
super(adapter, op, targetF);
@ -55,7 +56,7 @@ public class AzureAISearchUploadDocumentsOpDispenser extends AzureAISearchBaseOp
}
@Override
public LongFunction<AzureAISearchBaseOp<SearchDocument>> createOpFunc(LongFunction<SearchDocument> paramF,
public LongFunction<AzureAISearchBaseOp<SearchDocument,IndexDocumentsResult>> createOpFunc(LongFunction<SearchDocument> paramF,
LongFunction<SearchIndexClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> new AzureAISearchUploadDocumentsOp(clientF.apply(l),
clientF.apply(l).getSearchClient(targetF.apply(l)), paramF.apply(l));