From 94b621f42156b11f368b99046c098dabd414d936 Mon Sep 17 00:00:00 2001 From: Hemit Shah Date: Tue, 9 Sep 2025 00:47:13 +0000 Subject: [PATCH] SNOW-1062684: Finalize test changes after server-side support --- scripts/utils.sh | 4 ++-- .../snowpark/internal/JavaUtils.scala | 5 +++++ .../JavaStoredProcedureSuite.java | 5 ++++- .../snowflake/snowpark_test/UDFTestBase.java | 7 ++++++- .../com/snowflake/snowpark/TestUtils.scala | 12 +++--------- .../snowpark/UDFClasspathSuite.scala | 19 ++++++++++++++----- .../snowflake/snowpark/UDFInternalSuite.scala | 14 ++++++++++---- .../snowpark/UDFRegistrationSuite.scala | 14 ++++++++++++-- .../snowpark/UDTFInternalSuite.scala | 12 +++++++++++- .../snowpark_test/PermanentUDFSuite.scala | 7 +++++++ .../snowpark_test/PermanentUDTFSuite.scala | 6 ++++++ .../snowpark_test/StoredProcedureSuite.scala | 13 ++++++++++--- .../snowflake/snowpark_test/UDFSuite.scala | 7 +++++++ .../snowflake/snowpark_test/UDTFSuite.scala | 7 +++++++ .../snowpark_test/UdxOpenTelemetrySuite.scala | 10 ++++++++++ 15 files changed, 114 insertions(+), 28 deletions(-) diff --git a/scripts/utils.sh b/scripts/utils.sh index 2476e360..8b05e132 100644 --- a/scripts/utils.sh +++ b/scripts/utils.sh @@ -77,8 +77,8 @@ run_test_suites() { '++ 2.12.20 OtherTests:testOnly * -- -l SampleDataTest' \ '++ 2.13.16 OtherTests:testOnly * -- -l SampleDataTest' \ '++ 2.12.20 UDFTests:testOnly * -- -l SampleDataTest' \ - '++ 2.13.16 UDFTests:testOnly * -- -l SampleDataTest -l com.snowflake.snowpark.UDFPackageTest' \ + '++ 2.13.16 UDFTests:testOnly * -- -l SampleDataTest' \ '++ 2.12.20 UDTFTests:testOnly * -- -l SampleDataTest' \ - '++ 2.13.16 UDTFTests:testOnly * -- -l SampleDataTest -l com.snowflake.snowpark.UDFPackageTest' \ + '++ 2.13.16 UDTFTests:testOnly * -- -l SampleDataTest' \ +SprocTests:test } diff --git a/src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala b/src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala index 6406f15c..9b9c25ec 100644 --- a/src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala +++ b/src/main/scala/com/snowflake/snowpark/internal/JavaUtils.scala @@ -25,6 +25,7 @@ import com.snowflake.snowpark.{ } import java.io._ +import com.snowflake.snowpark.internal.Utils.{ScalaCompatVersion, SnowparkPackageName} import com.snowflake.snowpark.types.{Geography, Geometry, Variant} import com.snowflake.snowpark_java.types.InternalUtils import com.snowflake.snowpark_java.udtf._ @@ -33,6 +34,10 @@ import scala.collection.{JavaConverters, mutable} import scala.collection.JavaConverters._ object JavaUtils { + def snowparkScalaCompatVersion(): String = ScalaCompatVersion + + def snowparkPackageName(): String = SnowparkPackageName + def session_setJavaAPI(builder: SessionBuilder): SessionBuilder = builder.setJavaAPI() diff --git a/src/test/java/com/snowflake/snowpark_test/JavaStoredProcedureSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaStoredProcedureSuite.java index 50d9a1cc..9037b666 100644 --- a/src/test/java/com/snowflake/snowpark_test/JavaStoredProcedureSuite.java +++ b/src/test/java/com/snowflake/snowpark_test/JavaStoredProcedureSuite.java @@ -1,6 +1,7 @@ package com.snowflake.snowpark_test; import com.snowflake.snowpark.SnowparkClientException; +import com.snowflake.snowpark.internal.JavaUtils; import com.snowflake.snowpark_java.Row; import com.snowflake.snowpark_java.Session; import com.snowflake.snowpark_java.StoredProcedure; @@ -37,7 +38,9 @@ public void call() { + "returns STRING\n" + "language scala\n" + "runtime_version=2.12\n" - + "packages=('com.snowflake:snowpark:latest')\n" + + "packages=('" + + JavaUtils.snowparkPackageName() + + ":latest')\n" + "handler='Test.run'\n" + "as\n" + "$$\n" diff --git a/src/test/java/com/snowflake/snowpark_test/UDFTestBase.java b/src/test/java/com/snowflake/snowpark_test/UDFTestBase.java index 60da787c..68524d2d 100644 --- a/src/test/java/com/snowflake/snowpark_test/UDFTestBase.java +++ b/src/test/java/com/snowflake/snowpark_test/UDFTestBase.java @@ -1,6 +1,7 @@ package com.snowflake.snowpark_test; import com.snowflake.snowpark.TestUtils; +import com.snowflake.snowpark.internal.JavaUtils; import com.snowflake.snowpark_java.Session; public abstract class UDFTestBase extends TestFunctions { @@ -18,7 +19,11 @@ public Session getSession() { } protected Session createSession() { - return Session.builder().configFile(defaultProfile).create(); + Session newSession = Session.builder().configFile(defaultProfile).create(); + if (JavaUtils.snowparkScalaCompatVersion().equals("2.13")) { + newSession.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect(); + } + return newSession; } protected void runQuery(String sql) { diff --git a/src/test/scala/com/snowflake/snowpark/TestUtils.scala b/src/test/scala/com/snowflake/snowpark/TestUtils.scala index e74b1702..0b82a48f 100644 --- a/src/test/scala/com/snowflake/snowpark/TestUtils.scala +++ b/src/test/scala/com/snowflake/snowpark/TestUtils.scala @@ -125,7 +125,7 @@ object TestUtils extends Logging { sess.removeDependency(snClassDir) if (usePackages) { sess.removePackage(Utils.clientPackageName) - sess.addPackage("com.snowflake:snowpark:latest") + sess.addPackage(s"${Utils.SnowparkPackageName}:latest") } else { // Replace scoverage classes in classpath because they are instrumented. sess.addDependency(snClassDir.replace("scoverage-", "")) @@ -147,10 +147,7 @@ object TestUtils extends Logging { classOf[org.scalatest.compatible.Assertion], classOf[org.scalactic.TripleEquals], // scalactic jar classOf[io.opentelemetry.exporters.inmemory.InMemorySpanExporter], - classOf[io.opentelemetry.sdk.trace.export.SpanExporter], - classOf[scala.Product], - classOf[scala.reflect.api.Universe], - classOf[scala.tools.nsc.Global]) + classOf[io.opentelemetry.sdk.trace.export.SpanExporter]) .flatMap(UDFClassPath.getPathForClass(_)) .foreach(path => { val file = new File(path) @@ -179,10 +176,7 @@ object TestUtils extends Logging { classOf[BeforeAndAfterAll], // scala test jar classOf[org.scalactic.TripleEquals], // scalactic jar classOf[io.opentelemetry.exporters.inmemory.InMemorySpanExporter], - classOf[io.opentelemetry.sdk.trace.export.SpanExporter], - classOf[scala.Product], - classOf[scala.reflect.api.Universe], - classOf[scala.tools.nsc.Global]) + classOf[io.opentelemetry.sdk.trace.export.SpanExporter]) .flatMap(UDFClassPath.getPathForClass(_)) .foreach(path => { val file = new File(path) diff --git a/src/test/scala/com/snowflake/snowpark/UDFClasspathSuite.scala b/src/test/scala/com/snowflake/snowpark/UDFClasspathSuite.scala index 3c3f388d..d475d08f 100644 --- a/src/test/scala/com/snowflake/snowpark/UDFClasspathSuite.scala +++ b/src/test/scala/com/snowflake/snowpark/UDFClasspathSuite.scala @@ -3,7 +3,6 @@ package com.snowflake.snowpark import java.io.File import java.net.URLClassLoader import com.snowflake.snowpark.internal.ScalaFunctions._toUdf -import com.snowflake.snowpark.internal.Utils.clientPackageName import com.snowflake.snowpark.internal.{UDFClassPath, UDXRegistrationHandler, Utils} import org.mockito.ArgumentMatchers.any import org.mockito.Mockito.{atLeastOnce, never, reset, spy, times, verify} @@ -15,9 +14,19 @@ import scala.util.Random @UDFTest class UDFClasspathSuite extends SNTestBase { - override def beforeAll(): Unit = { - super.beforeAll() + override def beforeAll: Unit = { + super.beforeAll TestUtils.addDepsToClassPath(session) + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true") + } + } + + override def afterAll: Unit = { + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false") + } + super.afterAll } test("Test that jars uploaded to different stages") { @@ -97,8 +106,8 @@ class UDFClasspathSuite extends SNTestBase { val fixedPath = path.replace("scoverage-", "") // Remove snowpark jar from classpath, The code will catch error and add the path and retry newSession.removeDependency(fixedPath) - newSession.removePackage("com.snowflake:snowpark:latest") - newSession.removePackage(clientPackageName) + newSession.removePackage(s"${Utils.SnowparkPackageName}:latest") + newSession.removePackage(Utils.clientPackageName) val func = "func_" + Random.nextInt().abs ignoreClassNotFoundForScoverageClasses { diff --git a/src/test/scala/com/snowflake/snowpark/UDFInternalSuite.scala b/src/test/scala/com/snowflake/snowpark/UDFInternalSuite.scala index a8ed7861..f3884eeb 100644 --- a/src/test/scala/com/snowflake/snowpark/UDFInternalSuite.scala +++ b/src/test/scala/com/snowflake/snowpark/UDFInternalSuite.scala @@ -21,10 +21,16 @@ class UDFInternalSuite extends TestData { TestUtils.addDepsToClassPath(session, Some(stageName)) TestUtils.addDepsToClassPath(newSession, Some(stageName)) } + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true") + } } override def afterAll: Unit = { dropStage(stageName) + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false") + } super.afterAll } @@ -46,7 +52,7 @@ class UDFInternalSuite extends TestData { verify(mockSession, times(1)).addDependency(path) } verify(mockSession, times(1)).removeDependency(path) - verify(mockSession, times(1)).addPackage("com.snowflake:snowpark:latest") + verify(mockSession, times(1)).addPackage(s"${Utils.SnowparkPackageName}:latest") } test("Test permanent udf not failing back to upload jar", JavaStoredProcExclude) { @@ -80,7 +86,7 @@ class UDFInternalSuite extends TestData { verify(mockSession, times(1)).addDependency(path) } verify(mockSession, times(1)).removeDependency(path) - verify(mockSession, times(1)).addPackage("com.snowflake:snowpark:latest") + verify(mockSession, times(1)).addPackage(s"${Utils.SnowparkPackageName}:latest") } test("Test add version logic", JavaStoredProcExclude) { @@ -189,7 +195,7 @@ class PackageUDFSuite extends UDFSuite { val snClassDir = UDFClassPath.getPathForClass(classOf[Session]).get session.removeDependency(snClassDir.replace("scoverage-", "")) session.removePackage(Utils.clientPackageName) - session.addPackage("com.snowflake:snowpark:latest") + session.addPackage(s"${Utils.SnowparkPackageName}:latest") } override def afterAll: Unit = { @@ -207,7 +213,7 @@ class PackageUDTFSuite extends UDTFSuite { val snClassDir = UDFClassPath.getPathForClass(classOf[Session]).get session.removeDependency(snClassDir.replace("scoverage-", "")) session.removePackage(Utils.clientPackageName) - session.addPackage("com.snowflake:snowpark:latest") + session.addPackage(s"${Utils.SnowparkPackageName}:latest") } override def afterAll: Unit = { diff --git a/src/test/scala/com/snowflake/snowpark/UDFRegistrationSuite.scala b/src/test/scala/com/snowflake/snowpark/UDFRegistrationSuite.scala index 43d40c02..fe133496 100644 --- a/src/test/scala/com/snowflake/snowpark/UDFRegistrationSuite.scala +++ b/src/test/scala/com/snowflake/snowpark/UDFRegistrationSuite.scala @@ -4,7 +4,7 @@ import com.snowflake.snowpark.internal.Utils.clientPackageName import java.io.{BufferedOutputStream, File, FileOutputStream} import java.nio.file.{Files, NoSuchFileException} -import com.snowflake.snowpark.internal.{JavaUtils, UDFClassPath} +import com.snowflake.snowpark.internal.{JavaUtils, UDFClassPath, Utils} import scala.reflect.internal.util.BatchSourceFile import scala.reflect.io.{AbstractFile, VirtualDirectory} @@ -17,12 +17,22 @@ class UDFRegistrationSuite extends SNTestBase with FileUtils { private val tempStage = this.getClass.getSimpleName + Random.alphanumeric.take(5).mkString("") private val stagePrefix = "udfJar" - override def beforeAll(): Unit = { + override def beforeAll: Unit = { super.beforeAll() session.runQuery(s"create or replace temporary stage $tempStage") if (!isStoredProc(session)) { TestUtils.addDepsToClassPath(session) } + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true") + } + } + + override def afterAll: Unit = { + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false") + } + super.afterAll } test("Test that jar files are uploaded to stage correctly") { diff --git a/src/test/scala/com/snowflake/snowpark/UDTFInternalSuite.scala b/src/test/scala/com/snowflake/snowpark/UDTFInternalSuite.scala index 4abb4d30..d1d7ef0b 100644 --- a/src/test/scala/com/snowflake/snowpark/UDTFInternalSuite.scala +++ b/src/test/scala/com/snowflake/snowpark/UDTFInternalSuite.scala @@ -1,6 +1,6 @@ package com.snowflake.snowpark -import com.snowflake.snowpark.internal.ScalaFunctions +import com.snowflake.snowpark.internal.{ScalaFunctions, Utils} import com.snowflake.snowpark.types._ import com.snowflake.snowpark.udtf._ @@ -13,6 +13,16 @@ class UDTFInternalSuite extends SNTestBase { if (!isStoredProc(session)) { TestUtils.addDepsToClassPath(session) } + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true") + } + } + + override def afterAll: Unit = { + if (Utils.ScalaCompatVersion == "2.13") { + session.runQuery("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false") + } + super.afterAll } test("Unit test for UDTF0") { diff --git a/src/test/scala/com/snowflake/snowpark_test/PermanentUDFSuite.scala b/src/test/scala/com/snowflake/snowpark_test/PermanentUDFSuite.scala index 08fdd84f..1b6f4bde 100644 --- a/src/test/scala/com/snowflake/snowpark_test/PermanentUDFSuite.scala +++ b/src/test/scala/com/snowflake/snowpark_test/PermanentUDFSuite.scala @@ -1,6 +1,7 @@ package com.snowflake.snowpark_test import com.snowflake.snowpark.TestUtils.{removeFile, writeFile} +import com.snowflake.snowpark.internal.Utils.ScalaCompatVersion import com.snowflake.snowpark.functions._ import com.snowflake.snowpark._ import net.snowflake.client.jdbc.SnowflakeSQLException @@ -30,12 +31,18 @@ class PermanentUDFSuite extends TestData { // In stored procs mode, there is only one session TestUtils.addDepsToClassPath(newSession, Some(stageName)) } + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect() + } } override def afterAll: Unit = { dropStage(stageName) removeFile(tempDirectory1.getAbsolutePath, session) removeFile(tempDirectory2.getAbsolutePath, session) + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect() + } super.afterAll } diff --git a/src/test/scala/com/snowflake/snowpark_test/PermanentUDTFSuite.scala b/src/test/scala/com/snowflake/snowpark_test/PermanentUDTFSuite.scala index c293d830..4cbfdbb5 100644 --- a/src/test/scala/com/snowflake/snowpark_test/PermanentUDTFSuite.scala +++ b/src/test/scala/com/snowflake/snowpark_test/PermanentUDTFSuite.scala @@ -27,10 +27,16 @@ class PermanentUDTFSuite extends TestData { TestUtils.addDepsToClassPath(session, Some(stageName)) TestUtils.addDepsToClassPath(newSession, Some(stageName)) } + if (Utils.ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect() + } } override def afterAll: Unit = { dropStage(stageName) + if (Utils.ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect() + } super.afterAll } diff --git a/src/test/scala/com/snowflake/snowpark_test/StoredProcedureSuite.scala b/src/test/scala/com/snowflake/snowpark_test/StoredProcedureSuite.scala index 12a2ab37..f411127f 100644 --- a/src/test/scala/com/snowflake/snowpark_test/StoredProcedureSuite.scala +++ b/src/test/scala/com/snowflake/snowpark_test/StoredProcedureSuite.scala @@ -1,6 +1,7 @@ package com.snowflake.snowpark_test import com.snowflake.snowpark._ +import com.snowflake.snowpark.internal.Utils.{ScalaCompatVersion, SnowparkPackageName} import net.snowflake.client.jdbc.SnowflakeSQLException import java.sql.{Date, Timestamp} @@ -15,10 +16,16 @@ class StoredProcedureSuite extends SNTestBase { if (!isStoredProc(session)) { TestUtils.addDepsToClassPath(session, Some(testStage)) } + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect() + } } override def afterAll: Unit = { dropStage(testStage) + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect() + } super.afterAll } @@ -29,7 +36,7 @@ class StoredProcedureSuite extends SNTestBase { |returns STRING |language scala |runtime_version=2.12 - |packages=('com.snowflake:snowpark:latest') + |packages=('${SnowparkPackageName}:latest') |handler='Test.run' |as |$$$$ @@ -56,7 +63,7 @@ class StoredProcedureSuite extends SNTestBase { |returns STRING |language scala |runtime_version=2.12 - |packages=('com.snowflake:snowpark:latest') + |packages=('${SnowparkPackageName}:latest') |handler='Test.run' |as |$$$$ @@ -81,7 +88,7 @@ class StoredProcedureSuite extends SNTestBase { |returns STRING |language scala |runtime_version=2.12 - |packages=('com.snowflake:snowpark:latest') + |packages=('${SnowparkPackageName}:latest') |handler='Test.run' |as |$$$$ diff --git a/src/test/scala/com/snowflake/snowpark_test/UDFSuite.scala b/src/test/scala/com/snowflake/snowpark_test/UDFSuite.scala index 5152716c..1d125315 100644 --- a/src/test/scala/com/snowflake/snowpark_test/UDFSuite.scala +++ b/src/test/scala/com/snowflake/snowpark_test/UDFSuite.scala @@ -1,6 +1,7 @@ package com.snowflake.snowpark_test import com.snowflake.snowpark._ +import com.snowflake.snowpark.internal.Utils.ScalaCompatVersion import com.snowflake.snowpark.functions.{col, _} import com.snowflake.snowpark.types.{Geography, Geometry, Variant} @@ -53,6 +54,9 @@ trait UDFSuite extends TestData { if (!isStoredProc(session)) { TestUtils.addDepsToClassPath(session) } + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect() + } } override def afterAll: Unit = { @@ -63,6 +67,9 @@ trait UDFSuite extends TestData { dropTable(tableName) dropTable(semiStructuredTable) runQuery(s"DROP STAGE IF EXISTS $tmpStageName", session) + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect() + } super.afterAll() } diff --git a/src/test/scala/com/snowflake/snowpark_test/UDTFSuite.scala b/src/test/scala/com/snowflake/snowpark_test/UDTFSuite.scala index ad4fab89..5d0ceb48 100644 --- a/src/test/scala/com/snowflake/snowpark_test/UDTFSuite.scala +++ b/src/test/scala/com/snowflake/snowpark_test/UDTFSuite.scala @@ -31,11 +31,18 @@ class UDTFSuite extends TestData { .toDF("c1", "c2") .write .saveAsTable(wordCountTableName) + + if (Utils.ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect() + } } override def afterAll: Unit = { dropTable(wordCountTableName) dropTable(tableName) + if (Utils.ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect() + } super.afterAll } diff --git a/src/test/scala/com/snowflake/snowpark_test/UdxOpenTelemetrySuite.scala b/src/test/scala/com/snowflake/snowpark_test/UdxOpenTelemetrySuite.scala index 8f21b2fd..056f43d8 100644 --- a/src/test/scala/com/snowflake/snowpark_test/UdxOpenTelemetrySuite.scala +++ b/src/test/scala/com/snowflake/snowpark_test/UdxOpenTelemetrySuite.scala @@ -1,5 +1,6 @@ package com.snowflake.snowpark_test +import com.snowflake.snowpark.internal.Utils.ScalaCompatVersion import com.snowflake.snowpark.types.{IntegerType, StructField, StructType} import com.snowflake.snowpark.udtf.UDTF0 import com.snowflake.snowpark.{OpenTelemetryEnabled, Row, Session, TestUtils, functions} @@ -10,6 +11,15 @@ class UdxOpenTelemetrySuite extends OpenTelemetryEnabled { if (!isStoredProc(session)) { TestUtils.addDepsToClassPath(session) } + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=true").collect() + } + } + + override def afterAll(): Unit = { + if (ScalaCompatVersion == "2.13") { + session.sql("alter session set ENABLE_SCALA_UDF_RUNTIME_2_13=false").collect() + } } test("udf") {