From 9955c08537f1e36b6c62617a1e35c54d8706d5d3 Mon Sep 17 00:00:00 2001
From: Vitya Samoilov <108666580+vsamoilov@users.noreply.github.com>
Date: Wed, 1 Nov 2023 10:23:34 +0100
Subject: [PATCH] [JOBS-11439] Add taskValues support in remoteDbUtils (#406)

## Changes
Support `dbutils.jobs.taskValues` methods in `RemoteDbUtils` and add
additional `checkTaskValue` to allow local testing

## Tests
<!--
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->
Added three test cases in `test_dbutils.py`. Integration tests are not
needed since changes only for remote mode
- [x] `make test` run locally
- [x] `make fmt` applied
- [x] relevant integration tests applied
---
 databricks/sdk/dbutils.py | 26 ++++++++++++++++++++++++++
 tests/test_dbutils.py     | 22 ++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/databricks/sdk/dbutils.py b/databricks/sdk/dbutils.py
index 182b352f3..adfefc6e0 100644
--- a/databricks/sdk/dbutils.py
+++ b/databricks/sdk/dbutils.py
@@ -163,6 +163,31 @@ def listScopes(self) -> typing.List[SecretScope]:
         return [SecretScope(v.name) for v in self._api.list_scopes()]
 
 
+class _JobsUtil:
+    """Remote equivalent of jobs util"""
+
+    class _TaskValuesUtil:
+        """Remote equivalent of task values util"""
+
+        def get(self, taskKey: str, key: str, default: any = None, debugValue: any = None) -> None:
+            """
+            Returns `debugValue` if present, throws an error otherwise as this implementation is always run outside of a job run
+            """
+            if debugValue is None:
+                raise TypeError(
+                    'Must pass debugValue when calling get outside of a job context. debugValue cannot be None.'
+                )
+            return debugValue
+
+        def set(self, key: str, value: any) -> None:
+            """
+            Sets a task value on the current task run
+            """
+
+    def __init__(self) -> None:
+        self.taskValues = self._TaskValuesUtil()
+
+
 class RemoteDbUtils:
 
     def __init__(self, config: 'Config' = None):
@@ -175,6 +200,7 @@ def __init__(self, config: 'Config' = None):
 
         self.fs = _FsUtil(dbfs_ext.DbfsExt(self._client), self.__getattr__)
         self.secrets = _SecretsUtil(workspace.SecretsAPI(self._client))
+        self.jobs = _JobsUtil()
         self._widgets = None
 
     # When we import widget_impl, the init file checks whether user has the
diff --git a/tests/test_dbutils.py b/tests/test_dbutils.py
index 99d363b4a..5a8cb1edf 100644
--- a/tests/test_dbutils.py
+++ b/tests/test_dbutils.py
@@ -204,3 +204,25 @@ def test_secrets_get_and_redacting_logs(dbutils, mocker):
     inner.assert_called_with('GET', '/api/2.0/secrets/get', query={'key': 'bar', 'scope': 'foo'})
 
     assert value == 'hello'
+
+
+def test_jobs_task_values_set(dbutils):
+    dbutils.jobs.taskValues.set('key', 'value')
+
+
+def test_jobs_task_values_get(dbutils):
+    assert dbutils.jobs.taskValues.get('taskKey', 'key', debugValue='debug') == 'debug'
+
+    dbutils.jobs.taskValues.set('key', 'value')
+
+    # Expect `get` to always return the `debugValue`` when calling outside of a job context and not what was previously set using `set`
+    assert dbutils.jobs.taskValues.get('taskKey', 'key', debugValue='debug') == 'debug'
+
+
+def test_jobs_task_values_get_throws(dbutils):
+    try:
+        dbutils.jobs.taskValues.get('taskKey', 'key')
+        assert False
+    except TypeError as e:
+        assert str(
+            e) == 'Must pass debugValue when calling get outside of a job context. debugValue cannot be None.'