tsp, add polling option, allow it to override LroMetadata (#2404)

Azure · Nov 14, 2023 · 0d6f86e · 0d6f86e
1 parent 12bbfc0
commit 0d6f86e
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 5 deletions.
diff --git a/javagen/src/main/java/com/azure/autorest/mapper/ClientMethodMapper.java b/javagen/src/main/java/com/azure/autorest/mapper/ClientMethodMapper.java
@@ -415,7 +415,7 @@ private List<ClientMethod> createClientMethods(Operation operation, boolean isPr
 
                     if (pollingDetails != null) {
                         // try lroMetadata
-                        methodPollingDetails = methodPollingDetailsFromMetadata(operation, pollingDetails, settings);
+                        methodPollingDetails = methodPollingDetailsFromMetadata(operation, pollingDetails);
 
                         // fallback to JavaSettings.PollingDetails
                         if (methodPollingDetails == null) {
@@ -1654,16 +1654,15 @@ private static String formatReturnTypeDescription(String description) {
 
     private static MethodPollingDetails methodPollingDetailsFromMetadata(
         Operation operation,
-        JavaSettings.PollingDetails pollingDetails,
-        JavaSettings settings) {
+        JavaSettings.PollingDetails pollingDetails) {
 
         if (pollingDetails == null || operation.getConvenienceApi() == null) {
             return null;
         }
 
         MethodPollingDetails methodPollingDetails = null;
         if (operation.getLroMetadata() != null) {
-            // Only Typespec would have longRunningMetadata
+            // Only TypeSpec would have longRunningMetadata
 
             LongRunningMetadata metadata = operation.getLroMetadata();
             ObjectMapper objectMapper = Mappers.getObjectMapper();
@@ -1672,6 +1671,32 @@ private static MethodPollingDetails methodPollingDetailsFromMetadata(
                     ? PrimitiveType.Void
                     : objectMapper.map(metadata.getFinalResultType());
 
+            // merge with PollingDetails on intermediateType and finalType
+            if (pollingDetails.getIntermediateType() != null) {
+                String intermediateTypeName;
+                String intermediateTypePackage;
+                if (pollingDetails.getIntermediateType().contains(".")) {
+                    intermediateTypeName = ANYTHING_THEN_PERIOD.matcher(pollingDetails.getIntermediateType()).replaceAll("");
+                    intermediateTypePackage = pollingDetails.getIntermediateType().replace("." + intermediateTypeName, "");
+                } else {
+                    intermediateTypeName = pollingDetails.getIntermediateType();
+                    intermediateTypePackage = JavaSettings.getInstance().getPackage();
+                }
+                intermediateType = new ClassType.Builder().packageName(intermediateTypePackage).name(intermediateTypeName).build();
+            }
+            if (pollingDetails.getFinalType() != null) {
+                String finalTypeName;
+                String finalTypePackage;
+                if (pollingDetails.getFinalType().contains(".")) {
+                    finalTypeName = ANYTHING_THEN_PERIOD.matcher(pollingDetails.getFinalType()).replaceAll("");
+                    finalTypePackage = pollingDetails.getFinalType().replace("." + finalTypeName, "");
+                } else {
+                    finalTypeName = pollingDetails.getFinalType();
+                    finalTypePackage = JavaSettings.getInstance().getPackage();
+                }
+                finalType = new ClassType.Builder().packageName(finalTypePackage).name(finalTypeName).build();
+            }
+
             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());

diff --git a/typespec-extension/src/emitter.ts b/typespec-extension/src/emitter.ts
@@ -37,6 +37,7 @@ export interface EmitterOptions {
   "custom-types"?: string;
   "custom-types-subpackage"?: string;
   "customization-class"?: string;
+  "polling"?: any;
 
   "arm"?: boolean;
 
@@ -82,6 +83,7 @@ const EmitterOptionsSchema: JSONSchemaType<EmitterOptions> = {
     "custom-types": { type: "string", nullable: true },
     "custom-types-subpackage": { type: "string", nullable: true },
     "customization-class": { type: "string", nullable: true },
+    "polling": { type: "object", additionalProperties: true, nullable: true },
 
     "arm": { type: "boolean", nullable: true },
 

diff --git a/typespec-extension/src/main/java/com/azure/autorest/TypeSpecPlugin.java b/typespec-extension/src/main/java/com/azure/autorest/TypeSpecPlugin.java
@@ -205,10 +205,12 @@ public TypeSpecPlugin(EmitterOptions options, boolean sdkIntegration) {
         if (options.getCustomTypeSubpackage() != null) {
             SETTINGS_MAP.put("custom-types-subpackage", options.getCustomTypeSubpackage());
         }
-
         if (options.getCustomizationClass() != null) {
             SETTINGS_MAP.put("customization-class", options.getCustomizationClass());
         }
+        if (emitterOptions.getPolling() != null) {
+            SETTINGS_MAP.put("polling", options.getPolling());
+        }
 
         if (options.getBranded() == Boolean.FALSE) {
             SETTINGS_MAP.put("branded", options.getBranded());

diff --git a/typespec-extension/src/main/java/com/azure/typespec/model/EmitterOptions.java b/typespec-extension/src/main/java/com/azure/typespec/model/EmitterOptions.java
@@ -3,14 +3,17 @@
 
 package com.azure.typespec.model;
 
+import com.azure.autorest.extension.base.plugin.JavaSettings;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.databind.DeserializationContext;
 import com.fasterxml.jackson.databind.JsonDeserializer;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 public class EmitterOptions {
     @JsonDeserialize(using = EmptyStringToNullDeserializer.class)
@@ -58,6 +61,9 @@ public class EmitterOptions {
     @JsonProperty(value="customization-class")
     private String customizationClass;
 
+    @JsonProperty(value = "polling")
+    private Map<String, JavaSettings.PollingDetails> polling = new HashMap<>();
+
     @JsonProperty(value = "arm")
     private Boolean arm = false;
 
@@ -123,6 +129,14 @@ public String getCustomizationClass() {
         return customizationClass;
     }
 
+    public Map<String, JavaSettings.PollingDetails> getPolling() {
+        return polling;
+    }
+
+    public void setPolling(Map<String, JavaSettings.PollingDetails> polling) {
+        this.polling = polling;
+    }
+
     public Boolean getBranded() {
         return branded;
     }