Change to continuation token handling

Author: alzimmermsft

sdk/digitaltwins/azure-digitaltwins-core/src/main/java/com/azure/digitaltwins/core/DigitalTwinsAsyncClient.java

@@ -46,12 +46,17 @@
 import com.azure.digitaltwins.core.models.UpdateComponentOptions;
 import com.azure.digitaltwins.core.models.UpdateDigitalTwinOptions;
 import com.azure.digitaltwins.core.models.UpdateRelationshipOptions;
+import com.azure.json.JsonProviders;
+import com.azure.json.JsonReader;
+import com.azure.json.JsonWriter;
 import reactor.core.publisher.Mono;
 
 import java.io.IOException;
+import java.io.StringWriter;
 import java.io.UncheckedIOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
 
@@ -1541,7 +1546,23 @@ <T> Mono<PagedResponse<T>> queryNextPage(String nextLink, Class<T> clazz, QueryO
             context = Context.NONE;
         }
 
-        QuerySpecification querySpecification = new QuerySpecification().setContinuationToken(nextLink);
+        // The nextLink continuation token is a stringified JSON object. Attempt to deserialize the nextLink into any
+        // Object and inspect whether is what a stringified JSON value (List = array, Map = object, String = string).
+        // If it was a stringified value, use nextLink as-is. If not, serialize the nextLink into a JSON string.
+        String nextLinkToUse;
+        try (JsonReader jsonReader = JsonProviders.createReader(nextLink)) {
+            Object untyped = jsonReader.readUntyped();
+
+            if (untyped instanceof List || untyped instanceof Map || untyped instanceof String) {
+                nextLinkToUse = nextLink;
+            } else {
+                nextLinkToUse = serializeNextLink(nextLink);
+            }
+        } catch (IOException ex) {
+            nextLinkToUse = serializeNextLink(nextLink);
+        }
+
+        QuerySpecification querySpecification = new QuerySpecification().setContinuationToken(nextLinkToUse);
 
         return protocolLayer.getQueries()
             .queryTwinsWithResponseAsync(querySpecification, OptionsConverter.toProtocolLayerOptions(options), context)
@@ -1557,6 +1578,15 @@ <T> Mono<PagedResponse<T>> queryNextPage(String nextLink, Class<T> clazz, QueryO
                 objectPagedResponse.getDeserializedHeaders()));
     }
 
+    private static String serializeNextLink(String nextLink) {
+        try (StringWriter writer = new StringWriter(); JsonWriter jsonWriter = JsonProviders.createWriter(writer)) {
+            jsonWriter.writeString(nextLink);
+            return writer.toString();
+        } catch (IOException ex) {
+            throw LOGGER.logExceptionAsError(new UncheckedIOException(ex));
+        }
+    }
+
     //endregion Query APIs
 
     //region Event Route APIs

sdk/digitaltwins/azure-digitaltwins-core/src/main/java/com/azure/digitaltwins/core/implementation/models/QuerySpecification.java

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 // Code generated by Microsoft (R) AutoRest Code Generator.
-
 package com.azure.digitaltwins.core.implementation.models;
 
 import com.azure.core.annotation.Fluent;
@@ -16,6 +15,7 @@
  */
 @Fluent
 public final class QuerySpecification implements JsonSerializable<QuerySpecification> {
+
     /*
      * The query to execute. This value is ignored if a continuation token is provided.
      */
@@ -34,7 +34,7 @@ public QuerySpecification() {
 
     /**
      * Get the query property: The query to execute. This value is ignored if a continuation token is provided.
-     * 
+     *
      * @return the query value.
      */
     public String getQuery() {
@@ -43,7 +43,7 @@ public String getQuery() {
 
     /**
      * Set the query property: The query to execute. This value is ignored if a continuation token is provided.
-     * 
+     *
      * @param query the query value to set.
      * @return the QuerySpecification object itself.
      */
@@ -55,7 +55,7 @@ public QuerySpecification setQuery(String query) {
     /**
      * Get the continuationToken property: A token which is used to retrieve the next set of results from a previous
      * query.
-     * 
+     *
      * @return the continuationToken value.
      */
     public String getContinuationToken() {
@@ -65,7 +65,7 @@ public String getContinuationToken() {
     /**
      * Set the continuationToken property: A token which is used to retrieve the next set of results from a previous
      * query.
-     * 
+     *
      * @param continuationToken the continuationToken value to set.
      * @return the QuerySpecification object itself.
      */
@@ -81,13 +81,15 @@ public QuerySpecification setContinuationToken(String continuationToken) {
     public JsonWriter toJson(JsonWriter jsonWriter) throws IOException {
         jsonWriter.writeStartObject();
         jsonWriter.writeStringField("query", this.query);
-        jsonWriter.writeStringField("continuationToken", this.continuationToken);
+        if (this.continuationToken != null) {
+            jsonWriter.writeRawField("continuationToken", this.continuationToken);
+        }
         return jsonWriter.writeEndObject();
     }
 
     /**
      * Reads an instance of QuerySpecification from the JsonReader.
-     * 
+     *
      * @param jsonReader The JsonReader being read.
      * @return An instance of QuerySpecification if the JsonReader was pointing to an instance of it, or null if it was
      * pointing to JSON null.
@@ -99,7 +101,6 @@ public static QuerySpecification fromJson(JsonReader jsonReader) throws IOExcept
             while (reader.nextToken() != JsonToken.END_OBJECT) {
                 String fieldName = reader.getFieldName();
                 reader.nextToken();
-
                 if ("query".equals(fieldName)) {
                     deserializedQuerySpecification.query = reader.getString();
                 } else if ("continuationToken".equals(fieldName)) {
@@ -108,7 +109,6 @@ public static QuerySpecification fromJson(JsonReader jsonReader) throws IOExcept
                     reader.skipChildren();
                 }
             }
-
             return deserializedQuerySpecification;
         });
     }

sdk/digitaltwins/azure-digitaltwins-core/swagger/Update-Codegeneration.ps1

@@ -0,0 +1 @@
+& (Join-Path $PSScriptRoot ".." ".." ".." ".." eng scripts Invoke-Codegeneration.ps1) -Directory $PSScriptRoot

sdk/digitaltwins/azure-digitaltwins-core/swagger/autorest.md

@@ -47,6 +47,7 @@ implementation-subpackage: implementation
 models-subpackage: implementation.models
 custom-types-subpackage: models
 required-fields-as-ctor-args: true
+customization-class: src/main/java/DigitalTwinsCustomization.java
 ```
 
 ## This directive removes the specified enum values from the swagger so the code generator will expose IfNonMatch header as an option instead of always attaching it to requests with its only default value.

sdk/digitaltwins/azure-digitaltwins-core/swagger/pom.xml

@@ -0,0 +1,21 @@
+<?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 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+
+  <parent>
+    <groupId>com.azure</groupId>
+    <artifactId>azure-code-customization-parent</artifactId>
+    <version>1.0.0-beta.1</version> <!-- {x-version-update;com.azure:azure-code-customization-parent;current} -->
+    <relativePath>../../../parents/azure-code-customization-parent</relativePath>
+  </parent>
+
+  <name>Microsoft Azure Digital Twins code customization</name>
+  <description>This package contains code customization for Azure Digital Twins</description>
+
+  <groupId>com.azure.tools</groupId>
+  <artifactId>azure-digitaltwins-core-autorest-customization</artifactId>
+  <version>1.0.0-beta.1</version>
+  <packaging>jar</packaging>
+</project>

sdk/digitaltwins/azure-digitaltwins-core/swagger/src/main/java/DigitalTwinsCustomization.java

@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import com.azure.autorest.customization.ClassCustomization;
+import com.azure.autorest.customization.Customization;
+import com.azure.autorest.customization.LibraryCustomization;
+import com.azure.autorest.customization.PackageCustomization;
+import com.github.javaparser.StaticJavaParser;
+import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
+import com.github.javaparser.ast.body.MethodDeclaration;
+import org.slf4j.Logger;
+
+/**
+ * This class contains the customization code to customize the AutoRest generated code for Digital Twins.
+ */
+public class DigitalTwinsCustomization extends Customization {
+
+    @Override
+    public void customize(LibraryCustomization customization, Logger logger) {
+        PackageCustomization implModels = customization.getPackage("com.azure.digitaltwins.core.implementation.models");
+
+        customizeQuerySpecification(implModels.getClass("QuerySpecification"));
+    }
+
+    /**
+     * Customization for Digital Twins {@code QuerySpecification}.
+     * <p>
+     * Digital Twins designed continuation tokens to use a JSON stringified value, rather than a full object type. So,
+     * QuerySpecification's handling of its {@code continuationToken} field in JSON serialization needs to write the
+     * value as raw JSON rather than a JSON string. If written as a JSON string it would strigify the continuation token
+     * a second time, resulting in incorrect JSON being sent to the service, resulting in the operation failing.
+     */
+    private static void customizeQuerySpecification(ClassCustomization customization) {
+        customization.customizeAst(ast -> ast.getClassByName(customization.getClassName()).ifPresent(clazz -> {
+            MethodDeclaration toJsonMethod = clazz.getMethodsByName("toJson").get(0);
+            String toJsonMethodBodyString = toJsonMethod.getBody().get().toString();
+
+            // Replace 'jsonWriter.writeStringField("continuationToken", this.continuationToken)' with
+            // if (this.continuationToken != null) { jsonWriter.writeRawField("continuationToken", this.continuationToken); }
+            // to have the continuationToken written as raw JSON as needed by the service to prevent double
+            // stringifying.
+            toJsonMethodBodyString = toJsonMethodBodyString.replace(
+                "jsonWriter.writeStringField(\"continuationToken\", this.continuationToken);",
+                "if (this.continuationToken != null) {jsonWriter.writeRawField(\"continuationToken\", this.continuationToken);}");
+            toJsonMethod.setBody(StaticJavaParser.parseBlock(toJsonMethodBodyString));
+        }));
+    }
+}