From e939ef84983dc15d0bb2917dc5bf9e3cfcef1987 Mon Sep 17 00:00:00 2001 From: Diego Chacon Date: Tue, 20 Feb 2024 16:00:55 -0600 Subject: [PATCH 1/5] Add appName methods to SessionBuilder to register in query_tag. --- .../snowpark_java/SessionBuilder.java | 12 +++++++++++ .../com/snowflake/snowpark/Session.scala | 21 ++++++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java index 210a426e..a3ef0655 100644 --- a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java +++ b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java @@ -80,4 +80,16 @@ 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 + * + * @param appName Name of the app. + * @return A {@code Session} object + * @since 1.11.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..eb6e4ff9 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,19 @@ object Session extends Logging { this } + + /** + * Adds the app name to set in the query_tag after session creation + * + * @param appName Name of the app. + * @return A [[SessionBuilder]] + * @since 1.11.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 +1481,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 } /** From 947ab399e4fe4c474bd24c930191de1692151ce1 Mon Sep 17 00:00:00 2001 From: Diego Chacon Date: Tue, 20 Feb 2024 16:01:13 -0600 Subject: [PATCH 2/5] Add tests --- .../snowflake/snowpark_test/JavaSessionSuite.java | 8 ++++++++ .../com/snowflake/snowpark_test/SessionSuite.scala | 12 ++++++++++++ 2 files changed, 20 insertions(+) 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"))), From 5b55e74f42d4eb4f72138b364083283587ddc84a Mon Sep 17 00:00:00 2001 From: Diego Chacon Date: Wed, 21 Feb 2024 09:12:06 -0600 Subject: [PATCH 3/5] Update appName function documentation --- src/main/java/com/snowflake/snowpark_java/SessionBuilder.java | 2 +- src/main/scala/com/snowflake/snowpark/Session.scala | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java index a3ef0655..64a7775f 100644 --- a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java +++ b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java @@ -85,7 +85,7 @@ public Session getOrCreate() { * Adds the app name to set in the query_tag after session creation * * @param appName Name of the app. - * @return A {@code Session} object + * @return A {@code SessionBuilder} object * @since 1.11.0 */ public SessionBuilder appName(String appName) { diff --git a/src/main/scala/com/snowflake/snowpark/Session.scala b/src/main/scala/com/snowflake/snowpark/Session.scala index eb6e4ff9..11907a32 100644 --- a/src/main/scala/com/snowflake/snowpark/Session.scala +++ b/src/main/scala/com/snowflake/snowpark/Session.scala @@ -1483,7 +1483,7 @@ object Session extends Logging { def create: Session = { val session = createInternal(None) val appName = this.appName - if(appName.isDefined) { + if (appName.isDefined) { session.setQueryTag(s"APPNAME=${appName.get}") } session From 58fa576d5c8b00c9d3aeb69aa359b38dd63489bd Mon Sep 17 00:00:00 2001 From: Diego Chacon Date: Wed, 21 Feb 2024 13:40:02 -0600 Subject: [PATCH 4/5] resolve comments --- src/main/java/com/snowflake/snowpark_java/SessionBuilder.java | 3 ++- src/main/scala/com/snowflake/snowpark/Session.scala | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java index 64a7775f..5929cfff 100644 --- a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java +++ b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java @@ -82,7 +82,8 @@ public Session getOrCreate() { } /** - * Adds the app name to set in the query_tag after session creation + * 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 diff --git a/src/main/scala/com/snowflake/snowpark/Session.scala b/src/main/scala/com/snowflake/snowpark/Session.scala index 11907a32..da7b4e9c 100644 --- a/src/main/scala/com/snowflake/snowpark/Session.scala +++ b/src/main/scala/com/snowflake/snowpark/Session.scala @@ -1410,7 +1410,8 @@ object Session extends Logging { /** - * Adds the app name to set in the query_tag after session creation + * 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]] From 4d61183610a9bd8109642654af875f2e0fbc7401 Mon Sep 17 00:00:00 2001 From: Diego Chacon Date: Mon, 4 Mar 2024 09:14:28 -0600 Subject: [PATCH 5/5] update functions version --- src/main/java/com/snowflake/snowpark_java/SessionBuilder.java | 2 +- src/main/scala/com/snowflake/snowpark/Session.scala | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java index 5929cfff..c590b955 100644 --- a/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java +++ b/src/main/java/com/snowflake/snowpark_java/SessionBuilder.java @@ -87,7 +87,7 @@ public Session getOrCreate() { * * @param appName Name of the app. * @return A {@code SessionBuilder} object - * @since 1.11.0 + * @since 1.12.0 */ public SessionBuilder appName(String appName) { this.builder.appName(appName); diff --git a/src/main/scala/com/snowflake/snowpark/Session.scala b/src/main/scala/com/snowflake/snowpark/Session.scala index da7b4e9c..321ef271 100644 --- a/src/main/scala/com/snowflake/snowpark/Session.scala +++ b/src/main/scala/com/snowflake/snowpark/Session.scala @@ -1415,7 +1415,7 @@ object Session extends Logging { * * @param appName Name of the app. * @return A [[SessionBuilder]] - * @since 1.11.0 + * @since 1.12.0 */ def appName(appName: String): SessionBuilder = { this.appName = Some(appName)