Skip to content

Commit

Permalink
tsp, generate PollingStrategy in emitter (#2757)
Browse files Browse the repository at this point in the history
  • Loading branch information
weidongxu-microsoft authored May 29, 2024
1 parent 83c3851 commit 7da9de5
Show file tree
Hide file tree
Showing 55 changed files with 4,211 additions and 240 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ public class LongRunningMetadata {
private ObjectSchema pollResultType;
private ObjectSchema finalResultType;
private Metadata pollingStrategy;
private String finalResultPropertySerializedName;

/**
* Creates a new instance of the LongRunningMetadata class.
Expand Down Expand Up @@ -70,4 +71,22 @@ public Metadata getPollingStrategy() {
public void setPollingStrategy(Metadata pollingStrategy) {
this.pollingStrategy = pollingStrategy;
}

/**
* Gets the serialized name for the property of final result.
*
* @return the serialized name for the property of final result.
*/
public String getFinalResultPropertySerializedName() {
return finalResultPropertySerializedName;
}

/**
* Sets the serialized name for the property of final result.
*
* @param finalResultPropertySerializedName the serialized name for the property of final result.
*/
public void setFinalResultPropertySerializedName(String finalResultPropertySerializedName) {
this.finalResultPropertySerializedName = finalResultPropertySerializedName;
}
}
4 changes: 2 additions & 2 deletions javagen/src/main/java/com/azure/autorest/Javagen.java
Original file line number Diff line number Diff line change
Expand Up @@ -257,7 +257,7 @@ protected JavaPackage writeToTemplates(CodeModel codeModel, Client client, JavaS

writeClientModels(client, javaPackage, settings);

writeHelperClasses(client, javaPackage, settings);
writeHelperClasses(client, codeModel, javaPackage, settings);

// Unit tests on client model
if (settings.isGenerateTests() && !settings.isDataPlaneClient()) {
Expand Down Expand Up @@ -340,7 +340,7 @@ protected void writeClientModels(Client client, JavaPackage javaPackage, JavaSet
}
}

protected void writeHelperClasses(Client client, JavaPackage javaPackage, JavaSettings settings) {
protected void writeHelperClasses(Client client, CodeModel codeModel, JavaPackage javaPackage, JavaSettings settings) {
}

private void clear() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;
Expand Down Expand Up @@ -1654,8 +1655,8 @@ private static MethodPollingDetails methodPollingDetailsFromMetadata(
ObjectMapper objectMapper = Mappers.getObjectMapper();
IType intermediateType = objectMapper.map(metadata.getPollResultType());
IType finalType = metadata.getFinalResultType() == null
? PrimitiveType.VOID
: objectMapper.map(metadata.getFinalResultType());
? PrimitiveType.VOID
: objectMapper.map(metadata.getFinalResultType());

// PollingDetails would override LongRunningMetadata
if (pollingDetails.getIntermediateType() != null) {
Expand All @@ -1665,19 +1666,33 @@ private static MethodPollingDetails methodPollingDetailsFromMetadata(
finalType = createTypeFromModelName(pollingDetails.getFinalType(), JavaSettings.getInstance());
}

// PollingStrategy
JavaSettings settings = JavaSettings.getInstance();
final String packageName = settings.getPackage(settings.getImplementationSubpackage());
String pollingStrategy = metadata.getPollingStrategy() == null
? pollingDetails.getStrategy()
: String.format(JavaSettings.PollingDetails.DEFAULT_POLLING_STRATEGY_FORMAT, metadata.getPollingStrategy().getLanguage().getJava().getNamespace() + "." + metadata.getPollingStrategy().getLanguage().getJava().getName());
? pollingDetails.getStrategy()
: String.format(JavaSettings.PollingDetails.DEFAULT_POLLING_STRATEGY_FORMAT, packageName + "." + metadata.getPollingStrategy().getLanguage().getJava().getName());
String syncPollingStrategy = metadata.getPollingStrategy() == null
? pollingDetails.getSyncStrategy()
: String.format(JavaSettings.PollingDetails.DEFAULT_POLLING_STRATEGY_FORMAT, metadata.getPollingStrategy().getLanguage().getJava().getNamespace() + ".Sync" + metadata.getPollingStrategy().getLanguage().getJava().getName());
? pollingDetails.getSyncStrategy()
: String.format(JavaSettings.PollingDetails.DEFAULT_POLLING_STRATEGY_FORMAT, packageName + ".Sync" + metadata.getPollingStrategy().getLanguage().getJava().getName());
if (metadata.getPollingStrategy() != null && metadata.getFinalResultPropertySerializedName() != null) {
// add "<property-name>" argument to polling strategy constructor
Function<String, String> addPropertyNameToArguments = (strategy) -> {
strategy = strategy.substring(0, strategy.length() - 1) + ", ";
strategy += ClassType.STRING.defaultValueExpression(metadata.getFinalResultPropertySerializedName());
strategy += ")";
return strategy;
};
pollingStrategy = addPropertyNameToArguments.apply(pollingStrategy);
syncPollingStrategy = addPropertyNameToArguments.apply(syncPollingStrategy);
}

methodPollingDetails = new MethodPollingDetails(
pollingStrategy,
syncPollingStrategy,
intermediateType,
finalType,
pollingDetails.getPollIntervalInSeconds());
pollingStrategy,
syncPollingStrategy,
intermediateType,
finalType,
pollingDetails.getPollIntervalInSeconds());
}
return methodPollingDetails;
}
Expand Down
26 changes: 26 additions & 0 deletions javagen/src/main/java/com/azure/autorest/util/ClientModelUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import com.azure.autorest.extension.base.model.codemodel.CodeModel;
import com.azure.autorest.extension.base.model.codemodel.ConstantSchema;
import com.azure.autorest.extension.base.model.codemodel.KnownMediaType;
import com.azure.autorest.extension.base.model.codemodel.Operation;
import com.azure.autorest.extension.base.model.codemodel.OperationGroup;
import com.azure.autorest.extension.base.model.codemodel.Parameter;
import com.azure.autorest.extension.base.plugin.JavaSettings;
Expand Down Expand Up @@ -58,6 +59,10 @@ public class ClientModelUtil {

public static final String JSON_MERGE_PATCH_HELPER_CLASS_NAME = "JsonMergePatchHelper";

public static final String OPERATION_LOCATION_POLLING_STRATEGY = "OperationLocationPollingStrategy";
public static final String SYNC_OPERATION_LOCATION_POLLING_STRATEGY = "SyncOperationLocationPollingStrategy";
public static final String POLLING_UTILS = "PollingUtils";

private ClientModelUtil() {
}

Expand Down Expand Up @@ -726,6 +731,27 @@ public static Set<String> getExternalPackageNamesUsedInClient(List<ClientModel>
return externalPackageNames;
}

public static boolean requireOperationLocationPollingStrategy(CodeModel codeModel) {
if (!CoreUtils.isNullOrEmpty(codeModel.getClients())) {
for (Client client : codeModel.getClients()) {
if (!CoreUtils.isNullOrEmpty(client.getOperationGroups())) {
for (OperationGroup og : client.getOperationGroups()) {
if (!CoreUtils.isNullOrEmpty(og.getOperations())) {
for (Operation operation : og.getOperations()) {
if (operation.getLroMetadata() != null && operation.getLroMetadata().getPollingStrategy() != null) {
if (OPERATION_LOCATION_POLLING_STRATEGY.equals(operation.getLroMetadata().getPollingStrategy().getLanguage().getJava().getName())) {
return true;
}
}
}
}
}
}
}
}
return false;
}

public static boolean isMultipartModel(ClientModel model) {
return model.getSerializationFormats().contains(KnownMediaType.MULTIPART.value());
}
Expand Down
129 changes: 129 additions & 0 deletions javagen/src/main/resources/OperationLocationPollingStrategy.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
import com.azure.core.exception.AzureException;
import com.azure.core.http.HttpHeader;
import com.azure.core.http.rest.Response;
import com.azure.core.util.BinaryData;
import com.azure.core.util.logging.ClientLogger;
import com.azure.core.util.polling.LongRunningOperationStatus;
import com.azure.core.util.polling.OperationResourcePollingStrategy;
import com.azure.core.util.polling.PollResponse;
import com.azure.core.util.polling.PollingContext;
import com.azure.core.util.polling.PollingStrategyOptions;
import com.azure.core.util.serializer.JsonSerializerProviders;
import com.azure.core.util.serializer.ObjectSerializer;
import com.azure.core.util.serializer.TypeReference;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.time.OffsetDateTime;

// DO NOT modify this helper class

/**
* Implements an operation location polling strategy, from Operation-Location.
*
* @param <T> the type of the response type from a polling call, or BinaryData if raw response body should be kept
* @param <U> the type of the final result object to deserialize into, or BinaryData if raw response body should be
* kept
*/
public final class OperationLocationPollingStrategy<T, U> extends OperationResourcePollingStrategy<T, U> {

private static final ClientLogger LOGGER = new ClientLogger(OperationLocationPollingStrategy.class);

private final ObjectSerializer serializer;
private final String endpoint;
private final String propertyName;

/**
* Creates an instance of the operation resource polling strategy.
*
* @param pollingStrategyOptions options to configure this polling strategy.
* @throws NullPointerException if {@code pollingStrategyOptions} is null.
*/
public OperationLocationPollingStrategy(PollingStrategyOptions pollingStrategyOptions) {
this(pollingStrategyOptions, null);
}

/**
* Creates an instance of the operation resource polling strategy.
*
* @param pollingStrategyOptions options to configure this polling strategy.
* @param propertyName the name of the property to extract final result.
* @throws NullPointerException if {@code pollingStrategyOptions} is null.
*/
public OperationLocationPollingStrategy(PollingStrategyOptions pollingStrategyOptions, String propertyName) {
super(PollingUtils.OPERATION_LOCATION_HEADER, pollingStrategyOptions);
this.propertyName = propertyName;
this.endpoint = pollingStrategyOptions.getEndpoint();
this.serializer = pollingStrategyOptions.getSerializer() != null
? pollingStrategyOptions.getSerializer()
: JsonSerializerProviders.createInstance(true);
}

/**
* {@inheritDoc}
*/
@Override
public Mono<PollResponse<T>> onInitialResponse(Response<?> response, PollingContext<T> pollingContext,
TypeReference<T> pollResponseType) {
// Response<?> is Response<BinaryData>

HttpHeader operationLocationHeader = response.getHeaders().get(PollingUtils.OPERATION_LOCATION_HEADER);
if (operationLocationHeader != null) {
pollingContext.setData(PollingUtils.OPERATION_LOCATION_HEADER.getCaseSensitiveName(),
PollingUtils.getAbsolutePath(operationLocationHeader.getValue(), endpoint, LOGGER));
}
final String httpMethod = response.getRequest().getHttpMethod().name();
pollingContext.setData(PollingUtils.HTTP_METHOD, httpMethod);
pollingContext.setData(PollingUtils.REQUEST_URL, response.getRequest().getUrl().toString());

if (response.getStatusCode() == 200
|| response.getStatusCode() == 201
|| response.getStatusCode() == 202
|| response.getStatusCode() == 204) {
final Duration retryAfter = PollingUtils.getRetryAfterFromHeaders(response.getHeaders(), OffsetDateTime::now);
final Mono<PollResponse<T>> pollResponseMono = PollingUtils.deserializeResponse((BinaryData) response.getValue(), serializer, pollResponseType)
.onErrorResume(exception -> {
LOGGER.info("Failed to parse initial response.");
return Mono.empty();
})
.map(value -> new PollResponse<>(LongRunningOperationStatus.IN_PROGRESS, value, retryAfter));
return pollResponseMono.switchIfEmpty(
Mono.fromSupplier(() -> new PollResponse<>(LongRunningOperationStatus.IN_PROGRESS, null, retryAfter)));
} else {
return Mono.error(new AzureException(String.format(
"Operation failed or cancelled with status code %d," + ", '%s' header: %s, and response body: %s",
response.getStatusCode(), PollingUtils.OPERATION_LOCATION_HEADER, operationLocationHeader,
response.getValue())));
}
}

/**
* {@inheritDoc}
*/
@Override
public Mono<U> getResult(PollingContext<T> pollingContext, TypeReference<U> resultType) {
if (pollingContext.getLatestResponse().getStatus() == LongRunningOperationStatus.FAILED) {
return Mono.error(new AzureException("Long running operation failed."));
} else if (pollingContext.getLatestResponse().getStatus() == LongRunningOperationStatus.USER_CANCELLED) {
return Mono.error(new AzureException("Long running operation cancelled."));
}
if (propertyName != null) {
// take the last poll response body from PollingContext,
// and de-serialize the <propertyName> property as final result
BinaryData latestResponseBody
= BinaryData.fromString(pollingContext.getData(PollingUtils.POLL_RESPONSE_BODY));
return PollingUtils.deserializeResponse(latestResponseBody, serializer, PollingUtils.POST_POLL_RESULT_TYPE_REFERENCE)
.flatMap(value -> {
if (value.get(propertyName) != null) {
return BinaryData.fromObjectAsync(value.get(propertyName))
.flatMap(result -> PollingUtils.deserializeResponse(result, serializer, resultType));
} else {
return Mono.error(new AzureException("Cannot get final result"));
}
})
.switchIfEmpty(Mono.error(new AzureException("Cannot get final result")));
} else {
return super.getResult(pollingContext, resultType);
}
}
}
Loading

0 comments on commit 7da9de5

Please sign in to comment.