diff --git a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java index 210a426e..c590b955 100644 --- 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 index 950e90b0..321ef271 100644 --- a/src/main/scala/com/snowflake/snowpark/Session.scala +++ b/src/main/scala/com/snowflake/snowpark/Session.scala @@ -1394,6 +1394,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 = { @@ -1407,6 +1408,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. * @@ -1467,7 +1482,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/test/java/com/snowflake/snowpark_test/JavaSessionSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaSessionSuite.java index ee00443c..6726dd44 100644 --- 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 index 6ab7d5db..9b82ba7b 100644 --- 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"))),