handler.py 3.35 KB
import json
import time
import datetime


class DateTimeJSONEncoder(json.JSONEncoder):
    """
    JSONEncoder subclass that knows how to encode date/time, decimal types, and
    UUIDs.
    """
    def default(self, o):
        # See "Date Time String Format" in the ECMA-262 specification.
        if isinstance(o, datetime.datetime):
            r = o.isoformat()
            if o.microsecond:
                r = r[:23] + r[26:]
            if r.endswith('+00:00'):
                r = r[:-6] + 'Z'
            return r
        else:
            return super().default(o)


def dict_str_to_int(res):
    for k, v in res.items():
        res[k] = int(v)


class RedisHandler:

    def __init__(self, redis):
        self.redis = redis
        self.time_expires = datetime.timedelta(hours=24)
        self.time_format = '%a %b %d %H:%M:%S %Y'
        self.prefix = 'bwm_ocr'
        self.common_queue_key = '{0}:common_queue'.format(self.prefix)
        self.priority_queue_key = '{0}:priority_queue'.format(self.prefix)
        self.zip_queue_key = '{0}:zip_queue'.format(self.prefix)
        self.session_id_key = '{0}:session_id'.format(self.prefix)
        self.cms_token_key = '{0}:cms_token'.format(self.prefix)
        self.ecm_token_key = '{0}:ecm_token'.format(self.prefix)
        self.login_limit_key = '{0}:login_limit'.format(self.prefix)

    def enqueue(self, tasks, is_priority=False, is_zip=False):
        # 1
        if is_zip:
            key = self.zip_queue_key
        elif is_priority:
            key = self.priority_queue_key
        else:
            key = self.common_queue_key
        return self.redis.lpush(key, tasks)

    def dequeue_zip(self):
        # task or None
        task = self.redis.rpop(self.zip_queue_key)
        return task

    def dequeue(self):
        # task or None
        task = self.redis.rpop(self.priority_queue_key)
        is_priority = True
        if task is None:
            task = self.redis.rpop(self.common_queue_key)
            is_priority = False
        return task, is_priority

    def set_session_id(self, session_id):
        return self.redis.set(self.session_id_key, session_id)

    def get_session_id(self):
        return self.redis.get(self.session_id_key)

    def get_cms_token(self):
        return self.redis.get(self.cms_token_key)

    def set_cms_token(self, token, expires=None):
        return self.redis.set(self.cms_token_key, token, expires)

    def get_ecm_token(self):
        return self.redis.get(self.ecm_token_key)

    def set_ecm_token(self, token, expires=None):
        return self.redis.set(self.ecm_token_key, token, expires)

    def get_login_times(self, user_name):
        if user_name == '':
            return None
        return self.redis.get('{0}:{1}'.format(self.login_limit_key, user_name))

    def set_login_times(self, user_name, expires=None):
        if user_name == '':
            return
        key = '{0}:{1}'.format(self.login_limit_key, user_name)
        self.redis.incr(key)
        if isinstance(expires, int):
            self.redis.expire(key, expires)

    def get_token_key(self, token_str):
        return '{0}:token:{1}'.format(self.prefix, token_str)

    def set_token(self, token_str, username, expires=1800):
        return self.redis.set(self.get_token_key(token_str), username, expires)

    def get_token(self, token_str):
        return self.redis.get(self.get_token_key(token_str))