Merge branch 'main' into snow-1183003

snowflakedb · Mar 8, 2024 · 8d3626d · 8d3626d
2 parents 9934af0 + e968927
commit 8d3626d
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 1 deletion.
diff --git a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java
@@ -80,4 +80,17 @@ public Session create() {
   public Session getOrCreate() {
     return new Session(this.builder.getOrCreate());
   }
+
+  /**
+   * Adds the app name to set in the query_tag after session creation.
+   * The query tag will be set with this format 'APPNAME=${appName}'.
+   *
+   * @param appName Name of the app.
+   * @return A {@code SessionBuilder} object
+   * @since 1.12.0
+   */
+  public SessionBuilder appName(String appName) {
+    this.builder.appName(appName);
+    return this;
+  }
 }
diff --git a/src/main/scala/com/snowflake/snowpark/Session.scala b/src/main/scala/com/snowflake/snowpark/Session.scala
@@ -1396,6 +1396,7 @@ object Session extends Logging {
     private var options: Map[String, String] = Map()
 
     private var isScalaAPI: Boolean = true
+    private var appName: Option[String] = None
 
     // used by Java API only
     private[snowpark] def setJavaAPI(): SessionBuilder = {
@@ -1409,6 +1410,20 @@ object Session extends Logging {
       this
     }
 
+
+    /**
+     * Adds the app name to set in the query_tag after session creation.
+     * The query tag will be set with this format 'APPNAME=${appName}'.
+     *
+     * @param appName Name of the app.
+     * @return A [[SessionBuilder]]
+     * @since 1.12.0
+     */
+    def appName(appName: String): SessionBuilder = {
+      this.appName = Some(appName)
+      this
+    }
+
     /**
      * Adds the configuration properties in the specified file to the SessionBuilder configuration.
      *
@@ -1469,7 +1484,12 @@ object Session extends Logging {
      * @since 0.1.0
      */
     def create: Session = {
-      createInternal(None)
+      val session = createInternal(None)
+      val appName = this.appName
+      if (appName.isDefined) {
+        session.setQueryTag(s"APPNAME=${appName.get}")
+      }
+      session
     }
 
     /**

diff --git a/src/main/scala/com/snowflake/snowpark/functions.scala b/src/main/scala/com/snowflake/snowpark/functions.scala
@@ -1017,6 +1017,46 @@ object functions {
 
   /**
    * Splits a given string with a given separator and returns the result in an array of strings.
+   * To specify a string separator, use the lit() function.
+   *
+   * Example 1:
+   * {{{
+   *   val df = session.createDataFrame(
+   *                  Seq(("many-many-words", "-"), ("hello--hello", "--"))).toDF("V", "D")
+   *   df.select(split(col("V"), col("D"))).show()
+   * }}}
+   *   -------------------------
+   *   |"SPLIT(""V"", ""D"")"  |
+   *   -------------------------
+   *   |[                      |
+   *   |  "many",              |
+   *   |  "many",              |
+   *   |  "words"              |
+   *   |]                      |
+   *   |[                      |
+   *   |  "hello",             |
+   *   |  "hello"              |
+   *   |]                      |
+   *   -------------------------
+   *
+   * Example 2:
+   * {{{
+   *   val df = session.createDataFrame(Seq("many-many-words", "hello-hi-hello")).toDF("V")
+   *   df.select(split(col("V"), lit("-"))).show()
+   * }}}
+   *   -------------------------
+   *   |"SPLIT(""V"", ""D"")"  |
+   *   -------------------------
+   *   |[                      |
+   *   |  "many",              |
+   *   |  "many",              |
+   *   |  "words"              |
+   *   |]                      |
+   *   |[                      |
+   *   |  "hello",             |
+   *   |  "hello"              |
+   *   |]                      |
+   *   -------------------------
    *
    * @group str_func
    * @since 0.1.0

diff --git a/src/test/java/com/snowflake/snowpark_test/JavaSessionSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaSessionSuite.java
@@ -88,6 +88,14 @@ public void getOrCreate() {
     assert (actualSessionInfo.equals(expectedSessionInfo));
   }
 
+  @Test
+  public void appName() {
+    String appName = "my-app";
+    String expectedAppName = String.format("APPNAME=%s", appName);
+    Session session = Session.builder().configFile(defaultProfile).appName(appName).create();
+    assert (expectedAppName.equals(session.getQueryTag().get()));
+  }
+
   @Test
   public void getSessionInfo() {
     String result = getSession().getSessionInfo();

diff --git a/src/test/scala/com/snowflake/snowpark_test/SessionSuite.scala b/src/test/scala/com/snowflake/snowpark_test/SessionSuite.scala
@@ -251,6 +251,18 @@ class SessionSuite extends SNTestBase {
 
   }
 
+  test("Set an app name in the query tag") {
+    val appName = "my_app"
+    val expectedAppName = s"APPNAME=$appName"
+    val newSession = Session.builder.appName(appName).configFile(defaultProfile).create
+    assert(getParameterValue("query_tag", newSession) == expectedAppName)
+  }
+
+  test("The app name is not defined") {
+    val newSession = Session.builder.configFile(defaultProfile).create
+    assert(getParameterValue("query_tag", newSession) == "")
+  }
+
   test("generator") {
     checkAnswer(
       session.generator(3, Seq(lit(1).as("a"), lit(2).as("b"))),