Repository: PrefectHQ/prefect

Issue:
Redis Tasks
Using https://github.com/andymccurdy/redis-py let's add some tasks for communicating with Redis. These tasks should first connect to a Redis instance and then execute their respective commands. As a first pass (and to keep this minimal) let's add tasks for both `set` and `get` then branch out from there if there are more obvious basic commands we want to support. Other than those two we should also add a task for `execute_command` in which a task would take in a custom Redis command to execute against a db.

- [ ] set
- [ ] get
- [ ] execute_command



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -19,6 +19,7 @@
     "kubernetes": ["kubernetes >= 8.0.1, < 9.0", "dask-kubernetes == 0.7.0"],
     "rss": ["feedparser >= 5.0.1, < 6.0"],
     "postgres": ["psycopg2-binary >= 2.8.2"],
+    "redis": ["redis >= 3.2.1"],
     "spacy": ["spacy >= 2.0.0, < 3.0.0"],
     "templates": ["jinja2 >= 2.0, < 3.0"],
     "viz": ["graphviz >= 0.8.3"],
diff --git a/src/prefect/tasks/redis/__init__.py b/src/prefect/tasks/redis/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/redis/__init__.py
@@ -0,0 +1,11 @@
+"""
+This module contains a collection of tasks for interacting with Redis via
+the redis-py library.
+"""
+
+try:
+    from prefect.tasks.redis.redis_tasks import RedisSet, RedisGet, RedisExecute
+except ImportError:
+    raise ImportError(
+        'Using `prefect.tasks.redis` requires Prefect to be installed with the "redis" extra.'
+    )
diff --git a/src/prefect/tasks/redis/redis_tasks.py b/src/prefect/tasks/redis/redis_tasks.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/redis/redis_tasks.py
@@ -0,0 +1,244 @@
+from prefect import Task
+from prefect.client import Secret
+from prefect.utilities.tasks import defaults_from_attrs
+
+import redis
+
+
+class RedisSet(Task):
+    """
+    Task for setting a Redis key-value pair.
+
+    Args:
+        - host (str, optional): name of Redis host, defaults to 'localhost'
+        - port (int, optional): Redis port, defaults to 6379
+        - db (int, optional): redis database index, defaults to 0
+        - password_secret (str, optional): the name of the Prefect Secret
+            which stores your Redis credentials
+        - redis_key (str, optional): Redis key to be set, can be provided at initialization or runtime
+        - redis_val (Redis native type, optional): Redis val to be set, can be provided at initialization or runtime,
+            Redis native types include strings, dictionaries, lists, sets, and sorted sets
+        - redis_connection_params (dict, optional): key-value pairs passed to the redis.Redis connection
+            initializer
+        - ex (int, optional): if provided, sets an expire flag, in seconds, on 'redis_key' set
+        - px (int, optional): if provided, sets an expire flag, in milliseconds, on 'redis_key' set
+        - nx (int, optional): if set to True, set the value at 'redis_key' to 'redis_val' only
+            if it does not exist, defaults to False
+        - xx (int, optional): if set to True, set the value at 'redis_key' to 'redis_val' only
+            if it already exists, defaults to False
+        - **kwargs (dict, optional): additional keyword arguments to pass to the
+            Task constructor
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password_secret: str = "REDIS_PASSWORD",
+        redis_key: str = None,
+        redis_val=None,
+        redis_connection_params: dict = None,
+        ex: int = None,
+        px: int = None,
+        nx: bool = False,
+        xx: bool = False,
+        **kwargs
+    ):
+        self.host = host
+        self.port = port
+        self.db = db
+        self.password_secret = password_secret
+        self.redis_connection_params = redis_connection_params or {}
+
+        self.redis_key = redis_key
+        self.redis_val = redis_val
+        self.ex = ex
+        self.px = px
+        self.nx = nx
+        self.xx = xx
+        super().__init__(**kwargs)
+
+    @defaults_from_attrs("redis_key", "redis_val", "ex", "px", "nx", "xx")
+    def run(
+        self,
+        redis_key: str = None,
+        redis_val=None,
+        ex: int = None,
+        px: int = None,
+        nx: bool = False,
+        xx: bool = False,
+    ):
+        """
+        Task run method. Sets Redis key-value pair.
+
+        Args:
+            - redis_key (str, optional): Redis key to be set, can be provided at initialization or runtime
+            - redis_val (Redis native type, optional): Redis val to be set, can be provided at initialization or runtime,
+                Redis native types include strings, dictionaries, lists, sets, and sorted sets
+            - ex (int, optional): if provided, sets an expire flag, in seconds, on 'redis_key' set
+            - px (int, optional): if provided, sets an expire flag, in milliseconds, on 'redis_key' set
+            - nx (int, optional): if set to True, set the value at 'redis_key' to 'redis_val' only
+                if it does not exist, defaults to False
+            - xx (int, optional): if set to True, set the value at 'redis_key' to 'redis_val' only
+                if it already exists, defaults to False
+
+        Returns:
+            - bool: status of set operation
+
+        Raises:
+            - ValueError: if redis_key or redis_val is not provided
+        """
+        if None in (redis_key, redis_val):
+            raise ValueError("redis_key and redis_val must be provided")
+
+        ## connect to redis
+        password = Secret(self.password_secret).get()
+        connection = redis.Redis(
+            host=self.host,
+            port=self.port,
+            db=self.db,
+            password=password,
+            **self.redis_connection_params
+        )
+
+        result = connection.set(
+            name=redis_key, value=redis_val, ex=ex, px=px, nx=nx, xx=xx
+        )
+
+        return result
+
+
+class RedisGet(Task):
+    """
+    Task for getting a value based on key from a Redis connection.
+
+    Args:
+        - host (str, optional): name of Redis host, defaults to 'localhost'
+        - port (int, optional): Redis port, defaults to 6379
+        - db (int, optional): redis database index, defaults to 0
+        - password_secret (str, optional): the name of the Prefect Secret
+            which stores your Redis password
+        - redis_connection_params (dict, optional): key-value pairs passed to the redis.Redis connection
+            initializer
+        - redis_key (str, optional): Redis key to get value, can be provided at initialization or runtime
+        - **kwargs (dict, optional): additional keyword arguments to pass to the
+            Task constructor
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password_secret: str = "REDIS_PASSWORD",
+        redis_connection_params: dict = None,
+        redis_key: str = None,
+        **kwargs
+    ):
+        self.host = host
+        self.port = port
+        self.db = db
+        self.password_secret = password_secret
+        self.redis_connection_params = redis_connection_params or {}
+
+        self.redis_key = redis_key
+        super().__init__(**kwargs)
+
+    @defaults_from_attrs("redis_key")
+    def run(self, redis_key: str = None):
+        """
+        Task run method.
+
+        Args:
+            - redis_key (str, optional): Redis key to get value, can be provided at initialization or runtime
+
+        Returns:
+            - value: value associated with redis_key
+
+        Raises:
+            - ValueError: if redis_key is not provided
+        """
+        if not redis_key:
+            raise ValueError("redis_key must be provided")
+
+        ## connect to redis
+        password = Secret(self.password_secret).get()
+        connection = redis.Redis(
+            host=self.host,
+            port=self.port,
+            db=self.db,
+            password=password,
+            **self.redis_connection_params
+        )
+
+        result = connection.get(name=redis_key)
+
+        return result
+
+
+class RedisExecute(Task):
+    """
+    Task for executing a command against a Redis connection
+
+    Args:
+        - host (str, optional): name of Redis host, defaults to 'localhost'
+        - port (int, optional): Redis port, defaults to 6379
+        - db (int, optional): redis database index, defaults to 0
+        - password_secret (str, optional): the name of the Prefect Secret
+            which stores your Redis credentials
+        - redis_connection_params (dict, optional): key-value pairs passed to the redis.Redis connection
+            initializer
+        - redis_cmd (str, optional): Redis command to execute, must be provided at initialization or
+            runtime
+        - **kwargs (dict, optional): additional keyword arguments to pass to the
+            Task constructor
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password_secret: str = "REDIS_PASSWORD",
+        redis_connection_params: dict = None,
+        redis_cmd: str = None,
+        **kwargs
+    ):
+        self.host = host
+        self.port = port
+        self.db = db
+        self.password_secret = password_secret
+        self.redis_connection_params = redis_connection_params or {}
+
+        self.redis_cmd = redis_cmd
+        super().__init__(**kwargs)
+
+    @defaults_from_attrs("redis_cmd")
+    def run(self, redis_cmd: str = None):
+        """
+        Task run method. Executes a command against a Redis connection.
+
+        Args:
+            - redis_cmd (str, optional): Redis command to execute, must be provided at initialization or
+                runtime
+
+        Returns:
+            - result: result of executed Redis command
+        """
+        if not redis_cmd:
+            raise ValueError("A redis command must be specified")
+
+        ## connect to redis
+        password = Secret(self.password_secret).get()
+        connection = redis.Redis(
+            host=self.host,
+            port=self.port,
+            db=self.db,
+            password=password,
+            **self.redis_connection_params
+        )
+
+        result = connection.execute_command(redis_cmd)
+
+        return result


Test patch:
diff --git a/tests/tasks/redis/__init__.py b/tests/tasks/redis/__init__.py
new file mode 100644
--- /dev/null
+++ b/tests/tasks/redis/__init__.py
@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("redis")
diff --git a/tests/tasks/redis/test_redis_tasks.py b/tests/tasks/redis/test_redis_tasks.py
new file mode 100644
--- /dev/null
+++ b/tests/tasks/redis/test_redis_tasks.py
@@ -0,0 +1,110 @@
+from unittest.mock import MagicMock
+
+import prefect
+from prefect.tasks.redis import RedisSet, RedisGet, RedisExecute
+from prefect.utilities.configuration import set_temporary_config
+
+import pytest
+
+
+class TestRedisSet:
+    def test_construction(self):
+        task = RedisSet()
+        assert task.host == "localhost"
+
+    def test_raises_key_val_not_provided(self):
+        task = RedisSet()
+
+        ## raises if neither provided
+        with pytest.raises(ValueError) as exc:
+            task.run()
+        assert "redis_key and redis_val must be provided" == str(exc.value)
+
+        ## raises if only one arg is missing
+        with pytest.raises(ValueError) as exc:
+            task.run(redis_key="foo")
+        assert "redis_key and redis_val must be provided" == str(exc.value)
+        with pytest.raises(ValueError) as exc:
+            task.run(redis_val="bar")
+        assert "redis_key and redis_val must be provided" == str(exc.value)
+
+    def test_creds_are_pulled_from_secret(self, monkeypatch):
+        task = RedisSet()
+        redis = MagicMock()
+        monkeypatch.setattr("prefect.tasks.redis.redis_tasks.redis.Redis", redis)
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with prefect.context(secrets=dict(REDIS_PASSWORD="42")):
+                task.run(redis_key="foo", redis_val="bar")
+        assert redis.call_args[1]["password"] == 42
+
+    def test_redis_params_passed_to_connection(self, monkeypatch):
+        redis_params = {"custom_parameter": "value"}
+        task = RedisSet(redis_connection_params=redis_params)
+        redis = MagicMock()
+        monkeypatch.setattr("prefect.tasks.redis.redis_tasks.redis.Redis", redis)
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with prefect.context(secrets=dict(REDIS_PASSWORD="42")):
+                task.run(redis_key="foo", redis_val="bar")
+        assert redis.call_args[1]["custom_parameter"] == "value"
+
+
+class TestRedisGet:
+    def test_construction(self):
+        task = RedisGet()
+        assert task.host == "localhost"
+
+    def test_raises_key_val_not_provided(self):
+        task = RedisGet()
+        with pytest.raises(ValueError) as exc:
+            task.run()
+        assert "redis_key must be provided" == str(exc.value)
+
+    def test_creds_are_pulled_from_secret(self, monkeypatch):
+        task = RedisGet()
+        redis = MagicMock()
+        monkeypatch.setattr("prefect.tasks.redis.redis_tasks.redis.Redis", redis)
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with prefect.context(secrets=dict(REDIS_PASSWORD="42")):
+                task.run(redis_key="foo")
+        assert redis.call_args[1]["password"] == 42
+
+    def test_redis_params_passed_to_connection(self, monkeypatch):
+        redis_params = {"custom_parameter": "value"}
+        task = RedisGet(redis_connection_params=redis_params)
+        redis = MagicMock()
+        monkeypatch.setattr("prefect.tasks.redis.redis_tasks.redis.Redis", redis)
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with prefect.context(secrets=dict(REDIS_PASSWORD="42")):
+                task.run(redis_key="foo")
+        assert redis.call_args[1]["custom_parameter"] == "value"
+
+
+class TestRedisExecute:
+    def test_construction(self):
+        task = RedisExecute()
+        assert task.host == "localhost"
+
+    def test_raises_if_command_not_provided(self):
+        task = RedisExecute()
+        with pytest.raises(ValueError) as exc:
+            task.run()
+        assert "A redis command must be specified" == str(exc.value)
+
+    def test_creds_are_pulled_from_secret(self, monkeypatch):
+        task = RedisExecute()
+        redis = MagicMock()
+        monkeypatch.setattr("prefect.tasks.redis.redis_tasks.redis.Redis", redis)
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with prefect.context(secrets=dict(REDIS_PASSWORD="42")):
+                task.run(redis_cmd="GET foo")
+        assert redis.call_args[1]["password"] == 42
+
+    def test_redis_params_passed_to_connection(self, monkeypatch):
+        redis_params = {"custom_parameter": "value"}
+        task = RedisExecute(redis_connection_params=redis_params)
+        redis = MagicMock()
+        monkeypatch.setattr("prefect.tasks.redis.redis_tasks.redis.Redis", redis)
+        with set_temporary_config({"cloud.use_local_secrets": True}):
+            with prefect.context(secrets=dict(REDIS_PASSWORD="42")):
+                task.run(redis_cmd="GET foo")
+        assert redis.call_args[1]["custom_parameter"] == "value"