import enum from datetime import datetime from sqlalchemy import (Column, DateTime, Enum, Integer, Unicode, UniqueConstraint, create_engine) from sqlalchemy.ext.declarative import declarative_base from db import SqlDB, get_conn_str Base = declarative_base() engine = create_engine(get_conn_str()) class JobStatus(enum.Enum): UNKNOWN = 0 SCHEDULED = 1 RUNNING = 2 STOPPED = 3 ABORTED = 4 ERROR = 5 CLEARED = 6, DONE = 7 class Job(Base): __tablename__ = 'jobs' __table_args__ = ( None, UniqueConstraint('id'), ) id = Column(Integer, autoincrement=True, primary_key=True) starttime = Column(DateTime, nullable=False) stoptime = Column(DateTime, nullable=True) updated = Column(DateTime, nullable=True) status = Column(Enum(JobStatus), index=True, default=JobStatus.SCHEDULED, nullable=True) comment = Column(Unicode(255), nullable=True) scanner = Column(Unicode(255), nullable=True) target = Column(Unicode(255), nullable=True) @classmethod def starter(cls, func, **kwargs): job_id = kwargs['job_id'] retval = None del kwargs['job_id'] with SqlDB.sql_session() as session: job = session.query(Job).filter(Job.id == job_id).one_or_none() job.status = JobStatus.RUNNING session.commit() try: retval = func(*kwargs) except Exception: job.status = JobStatus.ERROR job.stoptime = datetime.now() session.commit() print("Job raised an exception!") else: job.status = JobStatus.DONE job.stoptime = datetime.now() print("Job finished properly.") return retval def as_dict(self): """Return JSON serializable dict.""" d = {} for col in self.__table__.columns: value = getattr(self, col.name) if issubclass(value.__class__, enum.Enum): value = value.value elif issubclass(value.__class__, Base): continue elif issubclass(value.__class__, datetime): value = str(value) d[col.name] = value return d Base.metadata.create_all(engine)