From ae1331403b06b842844e3a6b5e56bb26c0b40e37 Mon Sep 17 00:00:00 2001
From: Dominik Przybysz <dominik.przybysz@snowflake.com>
Date: Wed, 17 Jul 2024 11:29:31 +0200
Subject: [PATCH] SNOW-1531307: Add variant bind parameters sample

---
 .../client/jdbc/PreparedStatement1IT.java     | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/src/test/java/net/snowflake/client/jdbc/PreparedStatement1IT.java b/src/test/java/net/snowflake/client/jdbc/PreparedStatement1IT.java
index 56bef419f..09452649a 100644
--- a/src/test/java/net/snowflake/client/jdbc/PreparedStatement1IT.java
+++ b/src/test/java/net/snowflake/client/jdbc/PreparedStatement1IT.java
@@ -25,6 +25,7 @@
 import java.sql.Types;
 import java.util.Map;
 import java.util.Properties;
+import java.util.UUID;
 import net.snowflake.client.ConditionalIgnoreRule;
 import net.snowflake.client.RunningOnGithubAction;
 import net.snowflake.client.category.TestCategoryStatement;
@@ -790,4 +791,74 @@ public void manualTestForPreparedStatementLogging() throws SQLException {
       }
     }
   }
+
+  /**
+   * This test shows how to use batch requests with variant columns. Syntax <code>
+   * insert into ... Select column1, ... from values (?,...)</code> is crucial there.
+   */
+  @Test
+  public void shouldUseVariantBindParametersSNOW1531307() throws SQLException {
+    try (Connection con = getConnection();
+        Statement stmt = con.createStatement()) {
+      String tableName = "SNOW1531307";
+      createTableForVariantBindParametersSNOW1531307(stmt, tableName);
+      executePreparedStatementWithBatchesSNOW1531307(con, tableName, 20);
+    }
+  }
+
+  /**
+   * This test shows how to use batch requests with variant columns. Syntax <code>
+   * insert into ... Select column1, ... from values (?,...)</code> is crucial there.
+   * CLIENT_STAGE_ARRAY_BINDING_THRESHOLD parameter is used on session to set the minimal number of
+   * bind parameters to create stage file with provided parameters - it may be set only by
+   * administrator.
+   */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void shouldUseVariantBindParametersWithStageSNOW1531307() throws SQLException {
+    try (Connection con = getConnection();
+        Statement stmt = con.createStatement()) {
+      String tableName = "SNOW1531307";
+      createTableForVariantBindParametersSNOW1531307(stmt, tableName);
+
+      try {
+        stmt.execute("ALTER SESSION SET CLIENT_STAGE_ARRAY_BINDING_THRESHOLD = 10");
+        executePreparedStatementWithBatchesSNOW1531307(con, tableName, 20);
+      } finally {
+        // we are planning to reuse connection so let's cleanup
+        stmt.execute("ALTER SESSION UNSET CLIENT_STAGE_ARRAY_BINDING_THRESHOLD");
+      }
+    }
+  }
+
+  private static void createTableForVariantBindParametersSNOW1531307(
+      Statement stmt, String tableName) throws SQLException {
+    stmt.execute(
+        "CREATE OR REPLACE TABLE "
+            + tableName
+            + "(id text, array variant, json1 variant, json2 variant, str variant, num variant)");
+  }
+
+  private static void executePreparedStatementWithBatchesSNOW1531307(
+      Connection con, String tableName, int numberOfBatches) throws SQLException {
+    try (PreparedStatement ps =
+        con.prepareStatement(
+            "insert into "
+                + tableName
+                + "(id, array, json1, json2, str, num) "
+                + "Select column1, parse_json(column2), parse_json(column3), parse_json(column4), to_variant(column5), to_variant(column6) "
+                + "from values(?, ?, ?, ?, ?, ?)")) {
+      for (int i = 0; i < numberOfBatches; i++) {
+        int parameterIndex = 1;
+        ps.setString(parameterIndex++, UUID.randomUUID().toString());
+        ps.setObject(parameterIndex++, "[1,2]");
+        ps.setObject(parameterIndex++, null);
+        ps.setObject(parameterIndex++, "{\"a\": 5}");
+        ps.setObject(parameterIndex++, "bla");
+        ps.setObject(parameterIndex++, 5);
+        ps.addBatch();
+      }
+      assertEquals(numberOfBatches, ps.executeBatch().length);
+    }
+  }
 }