task.py 4.4 KB

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