use pickle to serialize data for storing in Redis

Signed-off-by: Jakub Sokołowski <jakub@status.im>
This commit is contained in:
Jakub Sokołowski 2021-03-30 15:58:16 +02:00
parent 0323b8d3bf
commit 85e4bbc6d1
No known key found for this signature in database
GPG Key ID: 4EF064D0E6D63020
1 changed files with 3 additions and 2 deletions

View File

@ -1,4 +1,5 @@
# -*- coding: utf-8 -*-
import pickle
from uuid import uuid4
from typing import Any
from airflow.models.xcom import BaseXCom
@ -14,7 +15,7 @@ class XComRedisBackend(BaseXCom):
hook = RedisHook(redis_conn_id=XComRedisBackend.CONN_ID)
key = str(uuid4())
# We use the default serializer, which pickles or JSONs
hook.get_conn().set(key, BaseXCom.serialize_value(value))
hook.get_conn().set(key, pickle.dumps(value))
# Add prefix to make it clear where the value is stored.
value = XComRedisBackend.PREFIX + key
return BaseXCom.serialize_value(value)
@ -27,5 +28,5 @@ class XComRedisBackend(BaseXCom):
key = result.replace(prefix, "")
hook = RedisHook(redis_conn_id=XComRedisBackend.CONN_ID)
result = hook.get_conn().get(key)
result = BaseXCom.deserialize_value(result)
result = pickle.loads(result)
return result