From 6036ed33b1f6dd851b144ed9bc8f88ac758e5bdd Mon Sep 17 00:00:00 2001
From: Ganesh Mahadevan <ganesh.mahadevan@snowflake.com>
Date: Tue, 27 Aug 2024 15:49:58 -0500
Subject: [PATCH] add JsonNode function for java and scala with test case

---
 .../com/snowflake/snowpark_java/types/Variant.java    |  1 -
 .../scala/com/snowflake/snowpark/types/Variant.scala  | 11 +++++++++++
 .../com/snowflake/snowpark_test/JavaVariantSuite.java |  2 +-
 .../snowflake/snowpark_test/ScalaVariantSuite.scala   |  8 ++++++++
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/snowflake/snowpark_java/types/Variant.java b/src/main/java/com/snowflake/snowpark_java/types/Variant.java
index 9b8630e8..ac14e0eb 100644
--- a/src/main/java/com/snowflake/snowpark_java/types/Variant.java
+++ b/src/main/java/com/snowflake/snowpark_java/types/Variant.java
@@ -383,7 +383,6 @@ public String asJsonString() {
    * @since 1.14.0
    */
   public JsonNode asJsonNode() {
-    verify(VariantTypes.Object);
     return objectToJsonNode(value);
   }
 
diff --git a/src/main/scala/com/snowflake/snowpark/types/Variant.scala b/src/main/scala/com/snowflake/snowpark/types/Variant.scala
index 5ff86f9c..b74a6ebb 100644
--- a/src/main/scala/com/snowflake/snowpark/types/Variant.scala
+++ b/src/main/scala/com/snowflake/snowpark/types/Variant.scala
@@ -44,6 +44,8 @@ private[snowpark] object Variant {
 
     object Object extends VariantType
 
+    object Jsonnode extends VariantType
+
     // internal used when converting from Java
     def getType(name: String): VariantType = name match {
       case "RealNumber" => RealNumber
@@ -381,6 +383,15 @@ class Variant private[snowpark] (
     }
   }
 
+  /**
+   * Converts the variant as valid JsonNode
+   *
+   * @since 1.14.0
+   */
+  def asJsonNode(): JsonNode = {
+    objectToJsonNode(value)
+  }
+
   /**
    * Converts the variant as binary value
    * @since 0.2.0
diff --git a/src/test/java/com/snowflake/snowpark_test/JavaVariantSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaVariantSuite.java
index 8f31a004..482ae8b5 100644
--- a/src/test/java/com/snowflake/snowpark_test/JavaVariantSuite.java
+++ b/src/test/java/com/snowflake/snowpark_test/JavaVariantSuite.java
@@ -308,7 +308,7 @@ public void javaVariantConversion() {
     assert (result.asMap().get("a").get().asString().equals("[1,2]"));
     assert (result.asMap().get("a").get().asJsonString().equals("[1,2]"));
 
-    // Scala variant to Java Variant
+//     Scala variant to Java Variant
     com.snowflake.snowpark.types.Variant sv =
         new com.snowflake.snowpark.types.Variant("{\"a\": [1, 2], \"b\": \"c\"}");
     assert InternalUtils.createVariant(null) == null;
diff --git a/src/test/scala/com/snowflake/snowpark_test/ScalaVariantSuite.scala b/src/test/scala/com/snowflake/snowpark_test/ScalaVariantSuite.scala
index 35e8c572..da9aad43 100644
--- a/src/test/scala/com/snowflake/snowpark_test/ScalaVariantSuite.scala
+++ b/src/test/scala/com/snowflake/snowpark_test/ScalaVariantSuite.scala
@@ -1,6 +1,7 @@
 package com.snowflake.snowpark_test
 
 import com.snowflake.snowpark.types.{Geography, Variant}
+import com.snowflake.snowpark_java.types.InternalUtils
 import org.scalatest.FunSuite
 
 import java.io.UncheckedIOException
@@ -311,4 +312,11 @@ class ScalaVariantSuite extends FunSuite {
 
     assert(v1.toString() == "123")
   }
+
+  test("JsonNode") {
+    val sv = new Variant("{\"a\": [1, 2], \"b\": 3, \"c\": \"xyz\"}")
+    assert(sv.asMap.get("a").get.asJsonNode().isArray)
+    assert(sv.asMap.get("b").get.asJsonNode().asInt().equals(3))
+    assert(sv.asMap.get("c").get.asJsonNode().asText().equals("xyz"))
+  }
 }