Fix record type in stream response

ballerina-platform · Jun 15, 2022 · 5df24af · 5df24af
1 parent e6285ee
commit 5df24af
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/cosmosdb-native/src/main/java/io/ballerinax/cosmosdb/RecordIteratorUtils.java b/cosmosdb-native/src/main/java/io/ballerinax/cosmosdb/RecordIteratorUtils.java
@@ -19,15 +19,15 @@
 package io.ballerinax.cosmosdb;
 
 import com.azure.cosmos.models.CosmosStoredProcedureProperties;
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.ballerina.runtime.api.PredefinedTypes;
 import io.ballerina.runtime.api.creators.TypeCreator;
 import io.ballerina.runtime.api.creators.ValueCreator;
-import io.ballerina.runtime.api.types.MapType;
 import io.ballerina.runtime.api.types.RecordType;
+import io.ballerina.runtime.api.types.UnionType;
 import io.ballerina.runtime.api.values.BObject;
-import io.ballerina.runtime.internal.JsonParser;
+import io.ballerina.runtime.api.values.BTypedesc;
+import org.ballerinalang.langlib.value.FromJsonStringWithType;
 
 import java.util.HashMap;
 import java.util.Iterator;
@@ -57,10 +57,14 @@ public static Object nextResult(BObject recordIterator) {
             return null;
         } else {
             Iterator<Object> results = (Iterator<Object>) recordIterator.getNativeData(Constants.OBJECT_ITERATOR);
-            ObjectMapper mapper = new ObjectMapper();
             if (results.hasNext()) {
                 try {
-                    return JsonParser.parse(mapper.writeValueAsString(results.next()));
+                    String result = new ObjectMapper().writeValueAsString(results.next());
+                    RecordType recordType = (RecordType) recordIterator.getNativeData(Constants.RECORD_TYPE);
+                    UnionType responseType = TypeCreator.createUnionType(recordType, PredefinedTypes.TYPE_ERROR,
+                            PredefinedTypes.TYPE_NULL);
+                    BTypedesc responseTypedescValue = ValueCreator.createTypedescValue(responseType);
+                    return FromJsonStringWithType.fromJsonStringWithType(fromString(result), responseTypedescValue);
                 } catch (Exception e) {
                     return BallerinaErrorGenerator.createBallerinaDatabaseError(e);
                 }

diff --git a/cosmosdb/iterator.bal b/cosmosdb/iterator.bal
@@ -28,14 +28,16 @@ public class ResultIterator {
     public isolated function next() returns record {|record {} value;|}|Error? {
         record {}|Error? result;
         result = nextResult(self);
-        if (result is record {}) {
+        if result is record {} {
             record {|
                 record {} value;
             |} streamRecord = {value: result};
             return streamRecord;
-        } else {
+        } else if result is Error {
             self.err = result;
             return self.err;
+        } else {
+            return result;
         }
     }
 }

diff --git a/cosmosdb/tests/test.bal b/cosmosdb/tests/test.bal
@@ -484,9 +484,9 @@ function testGetDocumentList() returns error? {
     stream<Person, error?> result = check azureCosmosClient->getDocumentList(databaseId, containerId, 
     valueOfPartitionKey);
     check result.forEach(isolated function(Person queryResult) {
+        test:assertTrue(queryResult is Person);
         test:assertEquals(1234, queryResult.AccountNumber);
     });
-
 }
 
 @test:Config {