stacktach-winchester/winchester/models.py

from datetime import datetime
from decimal import Decimal
import calendar
from enum import IntEnum

import timex

from sqlalchemy import event
from sqlalchemy import and_, or_
from sqlalchemy import literal_column
from sqlalchemy import Column, Table, ForeignKey, Index, UniqueConstraint
from sqlalchemy import Float, Boolean, Text, DateTime, Integer, String
from sqlalchemy import cast, null, case
from sqlalchemy.orm.interfaces import PropComparator
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.dialects.mysql import DECIMAL
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.orm import composite
from sqlalchemy.orm import backref
from sqlalchemy.orm import relationship
from sqlalchemy.orm.collections import attribute_mapped_collection
from sqlalchemy.types import TypeDecorator, DATETIME


class Datatype(IntEnum):
    none = 0
    string = 1
    int = 2
    float = 3
    datetime = 4
    timerange = 5


class StreamState(IntEnum):
    active = 1
    firing = 2
    expiring = 3
    error = 4
    expire_error = 5
    completed = 6


class DBException(Exception):
    pass


class InvalidTraitType(DBException):
    pass


def dt_to_decimal(dt):
    t_sec = calendar.timegm(dt.utctimetuple()) + (dt.microsecond/1e6)
    return Decimal("%.6f" % t_sec)


def decimal_to_dt(decimal_timestamp):
    return datetime.utcfromtimestamp(float(decimal_timestamp))


class PreciseTimestamp(TypeDecorator):
    """Represents a timestamp precise to the microsecond."""

    impl = DATETIME

    def load_dialect_impl(self, dialect):
        if dialect.name == 'mysql':
            return dialect.type_descriptor(DECIMAL(precision=20,
                                                   scale=6,
                                                   asdecimal=True))
        return dialect.type_descriptor(DATETIME())

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        elif dialect.name == 'mysql':
            return dt_to_decimal(value)
        return value

    def process_result_value(self, value, dialect):
        if value is None:
            return value
        elif dialect.name == 'mysql':
            return decimal_to_dt(value)
        return value


class DBTimeRange(object):
    def __init__(self, begin, end):
        self.begin = begin
        self.end = end

    def __composite_values__(self):
        return self.begin, self.end

    def __repr__(self):
        return "DBTimeRange(begin=%r, end=%r)" % (self.begin, self.end)

    def __eq__(self, other):
        return isinstance(other, DBTimeRange) and \
            other.begin == self.begin and \
            other.end == self.end

    def __ne__(self, other):
        return not self.__eq__(other)


class ProxiedDictMixin(object):
    """Adds obj[name] access to a mapped class.

    This class basically proxies dictionary access to an attribute
    called ``_proxied``.  The class which inherits this class
    should have an attribute called ``_proxied`` which points to a dictionary.

    """

    def __len__(self):
        return len(self._proxied)

    def __iter__(self):
        return iter(self._proxied)

    def __getitem__(self, name):
        return self._proxied[name]

    def __contains__(self, name):
        return name in self._proxied

    def __setitem__(self, name, value):
        self._proxied[name] = value

    def __delitem__(self, name):
        del self._proxied[name]


class PolymorphicVerticalProperty(object):
    """A name/value pair with polymorphic value storage."""

    ATTRIBUTE_MAP = {Datatype.none: None}
    PY_TYPE_MAP = {unicode: Datatype.string,
                   int: Datatype.int,
                   float: Datatype.float,
                   datetime: Datatype.datetime,
                   DBTimeRange: Datatype.timerange}

    def __init__(self, name, value=None):
        self.name = name
        self.value = value

    @classmethod
    def get_type_value(cls, value):
        if value is None:
            return Datatype.none, None
        if isinstance(value, str):
            value = value.decode('utf8', 'ignore')
        if isinstance(value, timex.Timestamp):
            value = value.timestamp
        if isinstance(value, timex.TimeRange):
            value = DBTimeRange(value.begin, value.end)
        if type(value) in cls.PY_TYPE_MAP:
            return cls.PY_TYPE_MAP[type(value)], value
        return None, value

    @hybrid_property
    def value(self):
        if self.type not in self.ATTRIBUTE_MAP:
            raise InvalidTraitType("Invalid trait type in db for %s: %s" % (self.name, self.type))
        attribute = self.ATTRIBUTE_MAP[self.type]
        if attribute is None:
            return None
        if self.type == Datatype.timerange:
            val = getattr(self, attribute)
            return timex.TimeRange(val.begin, val.end)
        else:
            return getattr(self, attribute)

    @value.setter
    def value(self, value):
        datatype, value = self.get_type_value(value)
        if datatype not in self.ATTRIBUTE_MAP:
            raise InvalidTraitType("Invalid trait type for %s: %s" % (self.name, datatype))
        attribute = self.ATTRIBUTE_MAP[datatype]
        self.type = int(datatype)
        if attribute is not None:
            setattr(self, attribute, value)

    @value.deleter
    def value(self):
        self._set_value(None)

    @value.comparator
    class value(PropComparator):
        """A comparator for .value, builds a polymorphic comparison.

        """
        def __init__(self, cls):
            self.cls = cls

        def __eq__(self, other):
            dtype, value = self.cls.get_type_value(other)
            if dtype is None:
                dtype = Datatype.string
            if dtype == Datatype.none:
                return self.cls.type == int(Datatype.none)
            attr = getattr(self.cls, self.cls.ATTRIBUTE_MAP[dtype])
            return and_(attr == value, self.cls.type == int(dtype))

        def __ne__(self, other):
            dtype, value = self.cls.get_type_value(other)
            if dtype is None:
                dtype = Datatype.string
            if dtype == Datatype.none:
                return self.cls.type != int(Datatype.none)
            attr = getattr(self.cls, self.cls.ATTRIBUTE_MAP[dtype])
            return and_(attr != value, self.cls.type == int(dtype))

    def __repr__(self):
        return '<%s %r=%r>' % (self.__class__.__name__, self.name, self.value)


Base = declarative_base()


class Trait(PolymorphicVerticalProperty, Base):
    __tablename__ = 'trait'
    __table_args__ = (
        Index('ix_trait_t_int', 't_int'),
        Index('ix_trait_t_string', 't_string'),
        Index('ix_trait_t_datetime', 't_datetime'),
        Index('ix_trait_t_float', 't_float'),
    )
    event_id = Column(Integer, ForeignKey('event.id'), primary_key=True)
    name = Column(String(100), primary_key=True)
    type = Column(Integer)

    ATTRIBUTE_MAP = {Datatype.none: None,
                     Datatype.string: 't_string',
                     Datatype.int: 't_int',
                     Datatype.float: 't_float',
                     Datatype.datetime: 't_datetime',}

    t_string = Column(String(255), nullable=True, default=None)
    t_float = Column(Float, nullable=True, default=None)
    t_int = Column(Integer, nullable=True, default=None)
    t_datetime = Column(PreciseTimestamp(),
                      nullable=True, default=None)

    def __repr__(self):
        return "<Trait(%s) %s=%s/%s/%s/%s on %s>" % (self.name,
                                                     self.type,
                                                     self.t_string,
                                                     self.t_float,
                                                     self.t_int,
                                                     self.t_datetime,
                                                     self.event_id)


class EventType(Base):
    """Types of event records."""
    __tablename__ = 'event_type'

    id = Column(Integer, primary_key=True)
    desc = Column(String(255), unique=True)

    def __init__(self, event_type):
        self.desc = event_type

    def __repr__(self):
        return "<EventType: %s>" % self.desc


class Event(ProxiedDictMixin, Base):
    __tablename__ = 'event'
    __table_args__ = (
        Index('ix_event_message_id', 'message_id'),
        Index('ix_event_type_id', 'event_type_id'),
        Index('ix_event_generated', 'generated')
    )
    id = Column(Integer, primary_key=True)
    message_id = Column(String(50), unique=True)
    generated = Column(PreciseTimestamp())

    event_type_id = Column(Integer, ForeignKey('event_type.id'))
    event_type = relationship("EventType", backref=backref('event_type'))

    traits = relationship("Trait",
                    collection_class=attribute_mapped_collection('name'))
    _proxied = association_proxy("traits", "value",
                            creator=lambda name, value: Trait(name=name, value=value))

    @property
    def event_type_string(self):
        return self.event_type.desc

    @property
    def as_dict(self):
        d = dict(self._proxied)
        d['message_id'] = self.message_id
        d['event_type'] = self.event_type_string
        d['timestamp'] = self.generated
        return d

    def __init__(self, message_id, event_type, generated):

         self.message_id = message_id
         self.event_type = event_type
         self.generated = generated

    def __repr__(self):
        return "<Event %s ('Event : %s %s, Generated: %s')>" % (self.id,
                                                              self.message_id,
                                                              self.event_type,
                                                              self.generated)


stream_event_table = Table('streamevent', Base.metadata,
        Column('stream_id', Integer, ForeignKey('stream.id'), primary_key=True),
        Column('event_id', Integer,
                ForeignKey('event.id'),
                primary_key=True)
        )


class Stream(ProxiedDictMixin, Base):
    __tablename__ = 'stream'

    __table_args__ = (
        Index('ix_stream_name', 'name'),
        Index('ix_stream_state', 'state'),
        Index('ix_stream_expire_timestamp', 'expire_timestamp'),
        Index('ix_stream_fire_timestamp', 'fire_timestamp')
    )
    id = Column(Integer, primary_key=True)
    first_event = Column(PreciseTimestamp(), nullable=False)
    last_event = Column(PreciseTimestamp(), nullable=False)
    expire_timestamp = Column(PreciseTimestamp())
    fire_timestamp = Column(PreciseTimestamp())
    name = Column(String(255), nullable=False)
    state = Column(Integer, default=StreamState.active, nullable=False)
    state_serial_no = Column(Integer, default=0, nullable=False)

    distinguished_by = relationship("DistinguishingTrait",
                    collection_class=attribute_mapped_collection('name'))
    _proxied = association_proxy("distinguished_by", "value",
                    creator=lambda name, value: DistinguishingTrait(name=name, value=value))

    events = relationship(Event, secondary=stream_event_table,
                          order_by=Event.generated)

    @property
    def distinguished_by_dict(self):
        return dict(self._proxied)

    def __init__(self, name, first_event, last_event=None, expire_timestamp=None,
                fire_timestamp=None, state=None, state_serial_no=None):
        self.name = name
        self.first_event = first_event
        if last_event is None:
            last_event = first_event
        self.last_event = last_event
        self.expire_timestamp = expire_timestamp
        self.fire_timestamp = fire_timestamp
        if state is None:
            state = StreamState.active
        self.state = state
        if state_serial_no is None:
            state_serial_no = 0
        self.state_serial_no = state_serial_no


class DistinguishingTrait(PolymorphicVerticalProperty, Base):
    __tablename__ = 'dist_trait'
    __table_args__ = (
        Index('ix_dist_trait_dt_int', 'dt_int'),
        Index('ix_dist_trait_dt_float', 'dt_float'),
        Index('ix_dist_trait_dt_string', 'dt_string'),
        Index('ix_dist_trait_dt_datetime', 'dt_datetime'),
        Index('ix_dist_trait_dt_timerange_begin', 'dt_timerange_begin'),
        Index('ix_dist_trait_dt_timerange_end', 'dt_timerange_end'),
    )
    stream_id = Column(Integer, ForeignKey('stream.id'), primary_key=True)
    name = Column(String(100), primary_key=True)
    type = Column(Integer)


    ATTRIBUTE_MAP = {Datatype.none:     None,
                     Datatype.string:   'dt_string',
                     Datatype.int:      'dt_int',
                     Datatype.float:    'dt_float',
                     Datatype.datetime: 'dt_datetime',
                     Datatype.timerange:'dt_timerange',
                    }

    dt_string = Column(String(255), nullable=True, default=None)
    dt_float = Column(Float, nullable=True, default=None)
    dt_int = Column(Integer, nullable=True, default=None)
    dt_datetime = Column(PreciseTimestamp(),
                      nullable=True, default=None)
    dt_timerange_begin = Column(PreciseTimestamp(), nullable=True, default=None)
    dt_timerange_end = Column(PreciseTimestamp(), nullable=True, default=None)

    dt_timerange = composite(DBTimeRange, dt_timerange_begin, dt_timerange_end)

    @property
    def as_dict(self):
        return {self.name: self.value}

    def __repr__(self):
        return "<DistinguishingTrait(%s) %s=%s/%s/%s/%s/(%s to %s) on %s>" % (self.name,
                                                             self.type,
                                                             self.dt_string,
                                                             self.dt_float,
                                                             self.dt_int,
                                                             self.dt_datetime,
                                                             self.dt_timerange_begin,
                                                             self.dt_timerange_end,
                                                             self.stream_id)