Add role-based filtering and imporve code

services/training/app/models/__init__.py

@@ -4,6 +4,13 @@ Training Service Models Package
 Import all models to ensure they are registered with SQLAlchemy Base.
 """
 
+# Import AuditLog model for this service
+from shared.security import create_audit_log_model
+from shared.database.base import Base
+
+# Create audit log model for this service
+AuditLog = create_audit_log_model(Base)
+
 # Import all models to register them with the Base metadata
 from .training import (
     TrainedModel,
@@ -20,4 +27,5 @@ __all__ = [
     "ModelPerformanceMetric",
     "TrainingJobQueue",
     "ModelArtifact",
+    "AuditLog",
 ]

services/training/app/models/training.py

@@ -193,4 +193,59 @@ class TrainedModel(Base):
             "training_start_date": self.training_start_date.isoformat() if self.training_start_date else None,
             "training_end_date": self.training_end_date.isoformat() if self.training_end_date else None,
             "data_quality_score": self.data_quality_score
+        }
+
+
+class TrainingPerformanceMetrics(Base):
+    """
+    Table to track historical training performance for time estimation.
+    Stores aggregated metrics from completed training jobs.
+    """
+    __tablename__ = "training_performance_metrics"
+
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    tenant_id = Column(UUID(as_uuid=True), nullable=False, index=True)
+    job_id = Column(String(255), nullable=False, index=True)
+
+    # Training job statistics
+    total_products = Column(Integer, nullable=False)
+    successful_products = Column(Integer, nullable=False)
+    failed_products = Column(Integer, nullable=False)
+
+    # Time metrics
+    total_duration_seconds = Column(Float, nullable=False)
+    avg_time_per_product = Column(Float, nullable=False)  # Key metric for estimation
+    data_analysis_time_seconds = Column(Float, nullable=True)
+    training_time_seconds = Column(Float, nullable=True)
+    finalization_time_seconds = Column(Float, nullable=True)
+
+    # Job metadata
+    completed_at = Column(DateTime(timezone=True), nullable=False, default=lambda: datetime.now(timezone.utc))
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+
+    def __repr__(self):
+        return (
+            f"<TrainingPerformanceMetrics("
+            f"tenant_id={self.tenant_id}, "
+            f"job_id={self.job_id}, "
+            f"total_products={self.total_products}, "
+            f"avg_time_per_product={self.avg_time_per_product:.2f}s"
+            f")>"
+        )
+
+    def to_dict(self):
+        return {
+            "id": str(self.id),
+            "tenant_id": str(self.tenant_id),
+            "job_id": self.job_id,
+            "total_products": self.total_products,
+            "successful_products": self.successful_products,
+            "failed_products": self.failed_products,
+            "total_duration_seconds": self.total_duration_seconds,
+            "avg_time_per_product": self.avg_time_per_product,
+            "data_analysis_time_seconds": self.data_analysis_time_seconds,
+            "training_time_seconds": self.training_time_seconds,
+            "finalization_time_seconds": self.finalization_time_seconds,
+            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
+            "created_at": self.created_at.isoformat() if self.created_at else None
         }