Repository: danqing/dqpy

Issue:
Model serialization and caching improvements

Comments:

Patch:
diff --git a/config/local.toml b/config/local.toml
index 0f30dcf..608b2d8 100644
--- a/config/local.toml
+++ b/config/local.toml
@@ -17,6 +17,11 @@ host = "localhost"
 port = 6379
 db = 0
 
+[faulty_redis]
+host = "localhost"
+port = 6380
+db = 0
+
 [sql]
 url = "mysql+pymysql://root@127.0.0.1:3306/dqpy?charset=utf8mb4"
 flavor = "mysql"
diff --git a/docs/conf.py b/docs/conf.py
index d215b10..e70b55e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -59,9 +59,9 @@ author = 'Danqing Liu'
 # built documents.
 #
 # The short X.Y version.
-version = '1.3'
+version = '2.0'
 # The full version, including alpha/beta/rc tags.
-release = '1.3.2'
+release = '2.0.0'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
diff --git a/docs/index.rst b/docs/index.rst
index cbce8ad..af59232 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -12,6 +12,16 @@ Danqing's shared Python library, written with love for `the better Danqing <http
    :maxdepth: 2
    :caption: Table of Contents
 
+Cache
+-----
+
+The cache module provides a simple JSON and Redis based caching decorator.
+
+.. automodule:: dq.cache
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Config
 ------
 
diff --git a/dq/cache.py b/dq/cache.py
new file mode 100644
index 0000000..9e715ec
--- /dev/null
+++ b/dq/cache.py
@@ -0,0 +1,42 @@
+import json
+from functools import wraps
+
+from dq.redis import init_redis, strval
+
+_redis = init_redis('cache')
+
+
+def cache(ttl=600, key_func=None):
+    """Cache decorator.
+
+    This can be applied to any function that returns a raw or JSON-serializable
+    response. To allow caching, the ``cache`` key must be set in the config,
+    namely the redis connection for cache.
+
+    If the function has a keyword argument named ``fresh``, then the decorator
+    gets a fresh copy when it's set to a truthy value.
+
+    If the function has a keyword argument named ``raw``, then the decorator
+    returns the raw (bytes) Redis response as-is, without JSON-deserializing.
+
+    :param number ttl: The TTL in second. Default is 10 minutes.
+    :param func key_func: The key function. This function should take the same
+        arguments as the wrapped function, and return the corresponding cache
+        key as a string.
+    """
+    def memoize(func):
+        @wraps(func)
+        def decorated_func(*args, **kwargs):
+            if not _redis or not key_func:
+                return func(*args, **kwargs)
+            key = key_func(*args, **kwargs)
+            if not kwargs.get('fresh'):
+                resp = _redis.get(key)
+                if resp is not None:
+                    return resp if kwargs.get('raw') else json.loads(resp)
+            resp = func(*args, **kwargs)
+            _redis.setex(key, ttl, strval(resp))
+            return resp
+
+        return decorated_func
+    return memoize
diff --git a/dq/orm.py b/dq/orm.py
index ab897b8..588c18e 100644
--- a/dq/orm.py
+++ b/dq/orm.py
@@ -1,22 +1,13 @@
-import logging
 import json
-import pickle
 from collections import namedtuple
 from uuid import uuid4
 
 import arrow
-import redis
 from sqlalchemy import Column, Integer, String
 from sqlalchemy.ext.declarative import as_declarative
 from sqlalchemy_utils import ArrowType
 
-from dq.config import Config
 from dq.database import query_with_limit_offset, Session
-from dq.logging import error
-
-logger = logging.getLogger(__name__)
-
-SERVICE_NAME = Config.get('meta.name')
 
 # An attribute of a model. The order is the following:
 # - name (string): The name of the attribute (the key)
@@ -162,69 +153,6 @@ class TimeMixin(object):
     deleted_at = Column(ArrowType)
 
 
-class Cache(object):
-    """Redis cache engine. This should not be used externally!
-
-    This is a separate Redis connection from the one in the redis package. The
-    purpose is twofold:
-
-    1. The cache Redis does not allow response decoding, but the user Redis
-       requires it.
-    2. By separating the two, if a project does not have an explicit cache
-       configuration, then cache can be effectively disabled globally.
-    """
-
-    _instance = None
-    _attempted = False
-
-    @classmethod
-    def instance(cls):
-        if cls._instance:
-            return cls._instance
-        if cls._attempted:
-            return None
-        cachecfg = Config.get('cache')
-        if not cachecfg:
-            cls._attempted = True
-            return None
-        try:
-            i = redis.StrictRedis(**cachecfg)
-            # This will attempt to connect to Redis and throw an error if the
-            # connection is invalid.
-            i.info()
-            cls._instance = i
-            return i
-        except Exception:
-            error(logger, 'Unable to connect to cache Redis', None)
-            cls._attempted = True
-            return None
-
-    @classmethod
-    def get(cls, key):
-        """Get the value corresponding to the key in the cache.
-
-        :param string key: The cache key.
-        :returns bytes: The cache value in raw bytes if exists. If the key does
-            not exist or if cache server cannot be reached, returns ``None``.
-        """
-        i = cls.instance()
-        try:
-            return i.get(key) if i else None
-        except Exception as e:
-            error(logger, 'Error querying cache', {'key': key, 'error': e})
-            return None
-
-    @classmethod
-    def set(cls, key, value, ttl):
-        i = cls.instance()
-        if not i:
-            return
-        try:
-            i.setex(key, ttl, pickle.dumps(value))
-        except Exception as e:
-            error(logger, 'Error setting cache', {'key': key, 'error': e})
-
-
 class QueryMixin(object):
     """Query helper functions useful to all models.
 
@@ -232,37 +160,9 @@ class QueryMixin(object):
     class will not need to inherit directly from it.
     """
 
-    @classmethod
-    def cache_key(cls, column, value, contains_deleted=False,
-                  contains_empty=False):
-        key = '{}.cache.{}.{}.{}'.format(
-            SERVICE_NAME, cls.__tablename__, column, value,
-        )
-        if contains_deleted:
-            key += '.del'
-        if contains_empty:
-            key += '.empty'
-        return key
-
-    @classmethod
-    def get_cache(cls, column, value, contains_deleted=False,
-                  contains_empty=False):
-        """Get the object from cache."""
-        key = cls.cache_key(column, value, contains_deleted=contains_deleted,
-                            contains_empty=contains_empty)
-        cached = Cache.get(key)
-        return pickle.loads(cached) if cached else None
-
-    @classmethod
-    def save_to_cache(cls, obj, column, value, ttl=3600,
-                      contains_deleted=False, contains_empty=False):
-        key = cls.cache_key(column, value, contains_deleted=contains_deleted,
-                            contains_empty=contains_empty)
-        Cache.set(key, obj, ttl)
-
     @classmethod
     def get_by(cls, column, value, for_update=False, contains_deleted=False,
-               contains_empty=False, ttl=3600, fresh=False, session=None):
+               contains_empty=False, session=None):
         """Get the object satisfying the query condition.
 
         :param string column: The name of the column to query by.
@@ -273,8 +173,6 @@ class QueryMixin(object):
         :param boolean contains_empty: Whether to contain empty records.
             Default is ``False`` and if value is ``None``, ``None`` will be
             returned.
-        :param int ttl: The TTL of cache. If < 0, no cache will be used.
-        :param boolean fresh: Whether to fetch the response fresh from the DB.
         :param Session session: Optional SQL session.
         :returns QueryMixin: The matching object. This method is designed for
             unique queries and in case of multiple matches, only the first one
@@ -283,14 +181,6 @@ class QueryMixin(object):
         if not contains_empty and value is None:
             return None
 
-        if ttl > 0 and not fresh:
-            cached = cls.get_cache(
-                column, value, contains_deleted=contains_deleted,
-                contains_empty=contains_empty,
-            )
-            if cached:
-                return cached
-
         session = session or Session()
         args = {}
         args[column] = value
@@ -299,13 +189,7 @@ class QueryMixin(object):
         query = session.query(cls).filter_by(**args)
         if for_update:
             query = query.with_for_update()
-        obj = query.first()
-
-        if obj and ttl > 0:
-            cls.save_to_cache(obj, column, value, ttl=ttl,
-                              contains_deleted=contains_deleted,
-                              contains_empty=contains_empty)
-        return obj
+        return query.first()
 
     @classmethod
     def get_multi(cls, column, value, sort_column='updated_at', desc=True,
@@ -378,21 +262,17 @@ class IDBase(QueryMixin, DictMixin, TimeMixin):
     id = Column(Integer, primary_key=True)
 
     @classmethod
-    def get(cls, id, contains_deleted=False, ttl=3600, fresh=False,
-            session=None):
+    def get(cls, id, for_update=False, contains_deleted=False, session=None):
         """Get an object by its ID.
 
         :param int id: The ID of the object.
+        :param boolean for_update: Whether the query is for updating the row.
         :param boolean contains_deleted: Whether to contain deleted objects.
             Default is ``False``.
-        :param int ttl: Cache TTL in seconds. If <= 0 no cache will be used.
-            Default is 3600 (1 hour).
-        :param boolean fresh: Whether to fetch a fresh copy even if cache
-            exists. Default is ``False``.
         :param Session session: Optional SQL session to use.
         """
-        return cls.get_by('id', id, contains_deleted=contains_deleted,
-                          fresh=fresh, session=session)
+        return cls.get_by('id', id, for_update=for_update,
+                          contains_deleted=contains_deleted, session=session)
 
 
 @as_declarative()
@@ -406,18 +286,14 @@ class UUIDBase(QueryMixin, DictMixin, TimeMixin):
     uuid = Column(String, primary_key=True, default=uuid4_string)
 
     @classmethod
-    def get(cls, uuid, contains_deleted=False, ttl=3600, fresh=False,
-            session=None):
+    def get(cls, uuid, for_update=False, contains_deleted=False, session=None):
         """Get an object by its UUID.
 
         :param string uuid: The UUID of the object.
+        :param boolean for_update: Whether the query is for updating the row.
         :param boolean contains_deleted: Whether to contain deleted objects.
             Default is ``False``.
-        :param int ttl: Cache TTL in seconds. If <= 0 no cache will be used.
-            Default is 3600 (1 hour).
-        :param boolean fresh: Whether to fetch a fresh copy even if cache
-            exists. Default is ``False``.
         :param Session session: Optional SQL session to use.
         """
-        return cls.get_by('uuid', uuid, contains_deleted=contains_deleted,
-                          fresh=fresh, session=session)
+        return cls.get_by('uuid', uuid, for_update=for_update,
+                          contains_deleted=contains_deleted, session=session)
diff --git a/dq/redis.py b/dq/redis.py
index e881d98..b9ae3bc 100644
--- a/dq/redis.py
+++ b/dq/redis.py
@@ -1,8 +1,34 @@
 import json
+import logging
 
 import redis
 
 from dq.config import Config
+from dq.logging import error
+
+logger = logging.getLogger(__name__)
+
+
+def init_redis(key):
+    """Initialize a Redis connection.
+
+    :param string key: The config key. The entry should at least contain the
+        host, port and db number of the instance.
+    :returns redis: The redis instance if the config exists and is valid, and
+        None otherwise.
+    """
+    cfg = Config.get(key)
+    if not cfg:
+        return None
+    try:
+        i = redis.StrictRedis(**cfg)
+        # This will attempt to connect to Redis and throw an error if the
+        # connection is invalid.
+        i.info()
+        return i
+    except Exception:
+        error(logger, 'Unable to connect to Redis', None)
+        return None
 
 
 def strval(value):
@@ -34,7 +60,7 @@ def strvals(*values):
 
 class Redis(object):
 
-    _instance = redis.StrictRedis(**Config.get('redis'))
+    _instance = init_redis('redis')
 
     @classmethod
     def exists(cls, key):
@@ -55,6 +81,17 @@ class Redis(object):
         """
         return cls._instance.get(key)
 
+    @classmethod
+    def get_json(cls, key):
+        """Get the value stored at the key as JSON.
+
+        :param string key: The Redis key.
+        :returns object: The value of the key as an unserialized JSON object.
+            If the key does not exist, ``None`` will be returned.
+        """
+        resp = cls.get(key)
+        return json.loads(resp) if resp else None
+
     @classmethod
     def set(cls, key, value):
         """Set the key to the specified value.
diff --git a/setup.py b/setup.py
index cd18a9e..9cd2fe2 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 from setuptools import find_packages, setup
 
-__version__ = '1.3.2'
+__version__ = '2.0.0'
 
 requires = [
     'arrow==0.12.1',


Test patch:
diff --git a/tests/test_cache.py b/tests/test_cache.py
new file mode 100644
index 0000000..0221d26
--- /dev/null
+++ b/tests/test_cache.py
@@ -0,0 +1,45 @@
+import unittest
+
+from dq import cache
+
+
+class TestCache(unittest.TestCase):
+
+    def test_cache_none(self):
+        redis = cache._redis
+        cache._redis = None
+
+        def key_func():
+            return ''
+
+        @cache.cache(key_func=key_func)
+        def some_func():
+            return 'hello'
+
+        assert some_func() == 'hello'
+
+        cache._redis = redis
+
+        @cache.cache()
+        def some_func_2():
+            return 'hello2'
+
+        assert some_func_2() == 'hello2'
+
+    def test_cache_fresh(self):
+        value = 'hello'
+
+        def key_func(fresh=False, raw=False):
+            return 'cache-fresh-key'
+
+        @cache.cache(key_func=key_func)
+        def some_func(fresh=False, raw=False):
+            nonlocal value
+            if value == 'hello':
+                value = 'world'
+                return 'hello'
+            return value
+
+        assert some_func(fresh=True) == 'hello'
+        assert some_func(raw=True) == b'hello'
+        assert some_func(fresh=True) == 'world'
diff --git a/tests/test_orm.py b/tests/test_orm.py
index a437675..036d0e1 100644
--- a/tests/test_orm.py
+++ b/tests/test_orm.py
@@ -1,6 +1,5 @@
 import ast
 import json
-import mock
 import unittest
 from contextlib import suppress
 from uuid import uuid4
@@ -8,7 +7,6 @@ from uuid import uuid4
 import arrow
 
 from dq.database import commit_scope, save_to_database
-from dq.orm import Cache
 from tests.models import Table2, User, UserType
 
 
@@ -18,8 +16,6 @@ class TestORM(unittest.TestCase):
         with suppress(Exception), commit_scope() as session:
             session.query(Table2).delete()
             session.query(User).delete()
-            if Cache._instance:
-                Cache._instance.flushall()
 
     def test_string(self):
         t2 = Table2(id=1, key=1, key2=1, user_type=UserType.admin,
@@ -34,56 +30,6 @@ class TestORM(unittest.TestCase):
             'created_at': 1508544000,
         }
 
-    def test_cache_key(self):
-        assert Table2.cache_key('uuid', '123') == 'dqpy.cache.table2.uuid.123'
-        k2 = Table2.cache_key('id', 123, contains_deleted=True,
-                              contains_empty=True)
-        assert k2 == 'dqpy.cache.table2.id.123.del.empty'
-
-    def test_cached(self):
-        t2 = Table2(id=999, user_uuid=str(uuid4()), key=1, key2=1,
-                    user_type=UserType.regular)
-        save_to_database(t2)
-
-        key = 'dqpy.cache.table2.id.999'
-        assert not Cache.get(key)
-
-        t2 = Table2.get(999)
-        assert t2.key == 1
-        assert Cache.get(key)
-        with commit_scope() as session:
-            session.query(Table2).delete()
-        t2 = Table2.get(999)
-        assert t2.key == 1
-
-    def test_cache_error(self):
-        Cache._instance = 123
-        assert not Cache.get('cornell')
-        Cache.set('cornell', '#1', 123)
-        Cache._instance = None
-
-    @mock.patch('dq.config.Config.get')
-    def test_cache_broken(self, mock_cfg):
-        mock_cfg.return_value = {'port': 1234}
-        Cache._instance = None
-        Cache._attempted = None
-        assert not Cache.instance()
-        assert Cache._attempted
-        assert not Cache.get('cornell')
-        Cache.set('cornell', '#1', 123)
-        Cache._attempted = None
-
-    @mock.patch('dq.config.Config.get')
-    def test_cache_none(self, mock_cfg):
-        mock_cfg.return_value = None
-        Cache._instance = None
-        Cache._attempted = None
-        assert not Cache.instance()
-        assert Cache._attempted
-        assert not Cache.get('cornell')
-        Cache.set('cornell', '#1', 123)
-        Cache._attempted = None
-
     def test_to_dict(self):
         uuid = str(uuid4())
         now = arrow.get()
@@ -146,7 +92,7 @@ class TestORM(unittest.TestCase):
         assert t2.id == 1
         t2.key2 = 10
         save_to_database(t2)
-        t2 = Table2.get_by('user_uuid', uuid, fresh=True)
+        t2 = Table2.get_by('user_uuid', uuid)
         assert t2.key2 == 10
 
     def test_get_by_empty(self):
diff --git a/tests/test_redis.py b/tests/test_redis.py
index 472f1ed..d385f87 100644
--- a/tests/test_redis.py
+++ b/tests/test_redis.py
@@ -3,11 +3,16 @@ import unittest
 from time import sleep
 from uuid import uuid4
 
-from dq.redis import Redis
+from dq.redis import init_redis, Redis
 
 
 class TestRedis(unittest.TestCase):
 
+    def test_init(self):
+        assert init_redis('redis')
+        assert not init_redis('none')
+        assert not init_redis('faulty_redis')
+
     def test_exists(self):
         key = 'dqtest-{}'.format(uuid4())
         assert not Redis.exists(key)
@@ -24,7 +29,7 @@ class TestRedis(unittest.TestCase):
         assert Redis.delete(key)
         assert not Redis.get(key)
 
-    def test_get_json(self):
+    def test_get_json_string(self):
         key = 'dqtest-{}'.format(uuid4())
         assert not Redis.get(key)
         Redis.setex(key, [1, '2', True], 1)
@@ -32,6 +37,14 @@ class TestRedis(unittest.TestCase):
         Redis.setex(key, {'cornell': '#1'}, 1)
         assert Redis.get(key) == '{"cornell": "#1"}'
 
+    def test_get_json(self):
+        key = 'dqtest-{}'.format(uuid4())
+        assert not Redis.get(key)
+        Redis.setex(key, [1, '2', True], 1)
+        assert Redis.get_json(key) == [1, '2', True]
+        Redis.setex(key, {'cornell': '#1'}, 1)
+        assert Redis.get_json(key) == {'cornell': '#1'}
+
     def test_expire(self):
         key = 'dqtest-{}'.format(uuid4())
         assert Redis.set(key, '1')