From 9f70697046599dadac75713fea0f1e586512c6bb Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Wed, 19 Jun 2024 01:33:47 +0300
Subject: [PATCH] chore: use json codec for key value lock (#29285)

---
 superset/utils/lock.py               | 4 ++--
 tests/unit_tests/utils/lock_tests.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/superset/utils/lock.py b/superset/utils/lock.py
index 50bfb6d95..d3c55d459 100644
--- a/superset/utils/lock.py
+++ b/superset/utils/lock.py
@@ -26,7 +26,7 @@ from typing import Any, cast, TypeVar, Union
 
 from superset.exceptions import CreateKeyValueDistributedLockFailedException
 from superset.key_value.exceptions import KeyValueCreateFailedError
-from superset.key_value.types import KeyValueResource, PickleKeyValueCodec
+from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
 from superset.utils import json
 
 LOCK_EXPIRATION = timedelta(seconds=30)
@@ -83,7 +83,7 @@ def KeyValueDistributedLock(  # pylint: disable=invalid-name
         DeleteExpiredKeyValueCommand(resource=KeyValueResource.LOCK).run()
         CreateKeyValueCommand(
             resource=KeyValueResource.LOCK,
-            codec=PickleKeyValueCodec(),
+            codec=JsonKeyValueCodec(),
             key=key,
             value=True,
             expires_on=datetime.now() + LOCK_EXPIRATION,
diff --git a/tests/unit_tests/utils/lock_tests.py b/tests/unit_tests/utils/lock_tests.py
index 270d08236..44eccebc3 100644
--- a/tests/unit_tests/utils/lock_tests.py
+++ b/tests/unit_tests/utils/lock_tests.py
@@ -42,7 +42,7 @@ def test_KeyValueDistributedLock_happy_path(mocker: MockerFixture) -> None:
     DeleteExpiredKeyValueCommand = mocker.patch(
         "superset.commands.key_value.delete_expired.DeleteExpiredKeyValueCommand"
     )
-    PickleKeyValueCodec = mocker.patch("superset.utils.lock.PickleKeyValueCodec")
+    JsonKeyValueCodec = mocker.patch("superset.utils.lock.JsonKeyValueCodec")
 
     with freeze_time("2024-01-01"):
         with KeyValueDistributedLock("ns", a=1, b=2) as key:
@@ -51,7 +51,7 @@ def test_KeyValueDistributedLock_happy_path(mocker: MockerFixture) -> None:
             )
             CreateKeyValueCommand.assert_called_with(
                 resource=KeyValueResource.LOCK,
-                codec=PickleKeyValueCodec(),
+                codec=JsonKeyValueCodec(),
                 key=key,
                 value=True,
                 expires_on=datetime(2024, 1, 1, 0, 0, 30),