models.py 3.35 KB
# -*- coding: utf-8 -*-


import json
from . import Base, db_adapter
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text, TypeDecorator
from sqlalchemy.orm import relation, backref

from datetime import datetime

from server.utils import get_now
from pytz import utc
from dateutil import parser


def relationship(*arg, **kw):
    ret = relation(*arg, **kw)
    db_adapter.commit()
    return ret


def date_serializer(date):
    return long((date - datetime(1970, 1, 1)).total_seconds() * 1000)


def to_dic(inst, cls):
    # add your coversions for things like datetime's
    # and what-not that aren't serializable.
    convert = dict()
    convert[TZDateTime] = date_serializer

    d = dict()
    for c in cls.__table__.columns:
        v = getattr(inst, c.name)
        if c.type.__class__ in convert.keys() and v is not None:
            try:
                func = convert[c.type.__class__]
                d[c.name] = func(v)
            except:
                d[c.name] = "Error:  Failed to covert using ", str(convert[c.type.__class__])
        else:
            d[c.name] = v
    return d


def to_json(inst, cls):
    return json.dumps(to_dic(inst, cls))


class TZDateTime(TypeDecorator):
    """
    usage: remove datetime's tzinfo
    To set all datetime datas are the naive datetime (tzinfo=None)
    in the whole environment
    """
    impl = DateTime

    def process_bind_param(self, value, dialect):
        if value is not None:
            if isinstance(value, basestring) or isinstance(value, str):
                value = parser.parse(value)
            if isinstance(value, datetime):
                if value.tzinfo is not None:
                    value = value.astimezone(utc)
                    value.replace(tzinfo=None)
        return value

    def process_result_value(self, value, dialect):
        if value is not None:
            if isinstance(value, datetime):
                if value.tzinfo is not None:
                    value = value.astimezone(utc)
                    value.replace(tzinfo=None)
        return value


class DBBase(Base):
    """
    DB model base class, providing basic functions
    """
    __abstract__ = True

    def __init__(self, **kwargs):
        super(DBBase, self).__init__(**kwargs)

    def dic(self):
        return to_dic(self, self.__class__)

    def json(self):
        return to_json(self, self.__class__)

    def __repr__(self):
        return '%s: %s' % (self.__class__.__name__, self.json())


class User(DBBase):
    __tablename__ = 'user'

    id = Column(Integer, primary_key=True)
    name = Column(String(64))
    password = Column(String(128))
    role = Column(String(16))
    status = Column(String(32))
    create_time = Column(TZDateTime, default=get_now())
    last_login_time = Column(TZDateTime, default=get_now())

    def __init__(self, **kwargs):
        super(User, self).__init__(**kwargs)


class UserToken(DBBase):
    __tablename__ = 'user_token'

    id = Column(Integer, primary_key=True)
    token = Column(String(64), unique=True, nullable=False)

    user_id = Column(Integer, ForeignKey('user.id', ondelete='CASCADE'))
    user = relationship('User', backref=backref('tokens', lazy='dynamic'))

    create_time = Column(TZDateTime, default=get_now())
    expire_time = Column(TZDateTime, nullable=False)

    def __init__(self, **kwargs):
        super(UserToken, self).__init__(**kwargs)