task.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. from datetime import datetime
  2. import json
  3. from app import db
  4. class Task(db.Model):
  5. """Task model for scan tasks"""
  6. __tablename__ = 'tasks'
  7. id = db.Column(db.Integer, primary_key=True)
  8. name = db.Column(db.String(200), nullable=False)
  9. status = db.Column(
  10. db.Enum('pending', 'running', 'completed', 'failed', name='task_status'),
  11. default='pending',
  12. index=True
  13. )
  14. progress = db.Column(db.Integer, default=0)
  15. created_by = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
  16. created_at = db.Column(db.DateTime, default=datetime.utcnow, index=True)
  17. started_at = db.Column(db.DateTime)
  18. completed_at = db.Column(db.DateTime)
  19. celery_task_id = db.Column(db.String(100), index=True)
  20. # Task configuration (JSON)
  21. _credential_ids = db.Column('credential_ids', db.Text)
  22. _regions = db.Column('regions', db.Text)
  23. _project_metadata = db.Column('project_metadata', db.Text)
  24. # Relationships
  25. created_by_user = db.relationship('User', back_populates='tasks')
  26. logs = db.relationship('TaskLog', back_populates='task', cascade='all, delete-orphan')
  27. report = db.relationship('Report', back_populates='task', uselist=False, cascade='all, delete-orphan')
  28. @property
  29. def credential_ids(self) -> list:
  30. """Get credential IDs as list"""
  31. if self._credential_ids:
  32. return json.loads(self._credential_ids)
  33. return []
  34. @credential_ids.setter
  35. def credential_ids(self, value: list):
  36. """Set credential IDs from list"""
  37. self._credential_ids = json.dumps(value)
  38. @property
  39. def regions(self) -> list:
  40. """Get regions as list"""
  41. if self._regions:
  42. return json.loads(self._regions)
  43. return []
  44. @regions.setter
  45. def regions(self, value: list):
  46. """Set regions from list"""
  47. self._regions = json.dumps(value)
  48. @property
  49. def project_metadata(self) -> dict:
  50. """Get project metadata as dict"""
  51. if self._project_metadata:
  52. return json.loads(self._project_metadata)
  53. return {}
  54. @project_metadata.setter
  55. def project_metadata(self, value: dict):
  56. """Set project metadata from dict"""
  57. self._project_metadata = json.dumps(value)
  58. def to_dict(self) -> dict:
  59. """Convert task to dictionary"""
  60. return {
  61. 'id': self.id,
  62. 'name': self.name,
  63. 'status': self.status,
  64. 'progress': self.progress,
  65. 'created_by': self.created_by,
  66. 'created_at': self.created_at.isoformat() if self.created_at else None,
  67. 'started_at': self.started_at.isoformat() if self.started_at else None,
  68. 'completed_at': self.completed_at.isoformat() if self.completed_at else None,
  69. 'credential_ids': self.credential_ids,
  70. 'regions': self.regions,
  71. 'project_metadata': self.project_metadata
  72. }
  73. def __repr__(self):
  74. return f'<Task {self.id} {self.name} ({self.status})>'
  75. class TaskLog(db.Model):
  76. """Task log model for storing task execution logs"""
  77. __tablename__ = 'task_logs'
  78. id = db.Column(db.Integer, primary_key=True)
  79. task_id = db.Column(db.Integer, db.ForeignKey('tasks.id'), nullable=False, index=True)
  80. level = db.Column(db.Enum('info', 'warning', 'error', name='log_level'), default='info')
  81. message = db.Column(db.Text, nullable=False)
  82. details = db.Column(db.Text) # JSON for stack trace, etc.
  83. created_at = db.Column(db.DateTime, default=datetime.utcnow, index=True)
  84. # Relationships
  85. task = db.relationship('Task', back_populates='logs')
  86. def to_dict(self) -> dict:
  87. """Convert log to dictionary"""
  88. result = {
  89. 'id': self.id,
  90. 'task_id': self.task_id,
  91. 'level': self.level,
  92. 'message': self.message,
  93. 'created_at': self.created_at.isoformat() if self.created_at else None
  94. }
  95. if self.details:
  96. result['details'] = json.loads(self.details)
  97. return result
  98. def __repr__(self):
  99. return f'<TaskLog {self.id} [{self.level}] {self.message[:50]}>'