From 6bf6f0f3257f43947e6107f4abc21cdddc826bc6 Mon Sep 17 00:00:00 2001
From: Vitya Samoilov <victor.samoilov@databricks.com>
Date: Wed, 1 Nov 2023 11:31:29 +0100
Subject: [PATCH] Get/SetContext throw error

---
 .../sdk/scala/dbutils/SdkDBUtilsImpl.scala    |  6 ++---
 .../sdk/scala/dbutils/SdkTaskValuesTest.scala | 23 +++++++++----------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/databricks-dbutils-scala/src/main/scala/com/databricks/sdk/scala/dbutils/SdkDBUtilsImpl.scala b/databricks-dbutils-scala/src/main/scala/com/databricks/sdk/scala/dbutils/SdkDBUtilsImpl.scala
index 8ed5fcd..2cbe755 100644
--- a/databricks-dbutils-scala/src/main/scala/com/databricks/sdk/scala/dbutils/SdkDBUtilsImpl.scala
+++ b/databricks-dbutils-scala/src/main/scala/com/databricks/sdk/scala/dbutils/SdkDBUtilsImpl.scala
@@ -160,8 +160,6 @@ private class SdkJobsUtils extends JobsUtils with NoHelp {
 }
 
 private class SdkTaskValues extends TaskValuesUtils with NoHelp {
-  private var commandContext =
-    CommandContext(rootRunId = None, currentRunId = None, jobGroup = None, tags = Map.empty, extraContext = Map.empty)
 
   /**
    * Sets a task value on the current task run. This method is a no-op if used outside of the job context.
@@ -195,10 +193,10 @@ private class SdkTaskValues extends TaskValuesUtils with NoHelp {
   override def getJson(taskKey: String, key: String): Seq[String] = Seq.empty
 
   override def getContext(): CommandContext = {
-    commandContext
+    throw new NotImplementedError("getContext is not supported outside of DBR.")
   }
 
   override def setContext(context: CommandContext): Unit = {
-    commandContext = context
+    throw new NotImplementedError("setContext is not supported outside of DBR.")
   }
 }
diff --git a/databricks-dbutils-scala/src/test/scala/com/databricks/sdk/scala/dbutils/SdkTaskValuesTest.scala b/databricks-dbutils-scala/src/test/scala/com/databricks/sdk/scala/dbutils/SdkTaskValuesTest.scala
index a4849f6..a10bf3b 100644
--- a/databricks-dbutils-scala/src/test/scala/com/databricks/sdk/scala/dbutils/SdkTaskValuesTest.scala
+++ b/databricks-dbutils-scala/src/test/scala/com/databricks/sdk/scala/dbutils/SdkTaskValuesTest.scala
@@ -5,17 +5,8 @@ import org.scalatest.matchers.should.Matchers.{be, convertToAnyShouldWrapper}
 
 class SdkTaskValuesTest extends AnyFlatSpec {
   private val taskValues = new SdkTaskValues
-  "When outside of DBR, TaskValuesUtils.getContext" should "return default context initially" in {
-    taskValues.getContext should be(
-      CommandContext(
-        rootRunId = None,
-        currentRunId = None,
-        jobGroup = None,
-        tags = Map.empty,
-        extraContext = Map.empty))
-  }
 
-  "When outside of DBR, TaskValuesUtils.getContext" should "return context, which was set" in {
+  "When outside of DBR, TaskValuesUtils.getContext and TaskValuesUtils.setContext" should "throw error" in {
     val context = CommandContext(
       rootRunId = None,
       currentRunId = None,
@@ -23,9 +14,17 @@ class SdkTaskValuesTest extends AnyFlatSpec {
       tags = Map.empty,
       extraContext = Map("test" -> "test"))
 
-    taskValues.setContext(context)
+    val setE = intercept[NotImplementedError] {
+      taskValues.setContext(context)
+    }
+
+    setE.getMessage should be("setContext is not supported outside of DBR.")
+
+    val getE = intercept[NotImplementedError] {
+      taskValues.getContext()
+    }
 
-    taskValues.getContext should be(context)
+    getE.getMessage should be("getContext is not supported outside of DBR.")
   }
 
   "When outside of DBR, TaskValuesUtils.get and TaskValues.set" should "throw error" in {