From 5df24aff925d61a60ae92843630d6962fa7eff0a Mon Sep 17 00:00:00 2001
From: Chuhaa <chuhaashanan@yahoo.com>
Date: Wed, 15 Jun 2022 08:44:28 +0530
Subject: [PATCH] Fix record type in stream response

---
 .../ballerinax/cosmosdb/RecordIteratorUtils.java   | 14 +++++++++-----
 cosmosdb/iterator.bal                              |  6 ++++--
 cosmosdb/tests/test.bal                            |  2 +-
 3 files changed, 14 insertions(+), 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
index 2467942b..3e350348 100644
--- 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
index d93a483f..f5e94ee4 100644
--- 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
index c2b4fb05..60def3a7 100644
--- 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 {