Mazyar
492f9ba3c8
continuous-integration/drone/push Build is passing
- Introduced AIPipelineAutoWorker to manage the execution of the AI pipeline auto run, including startup catch-up and scheduled tasks. - Enhanced WorkerConfig to include AIPipelineAutoEnabled and AIPipelineAutoInterval settings for better control over AI pipeline execution. - Added logging for AI pipeline auto run status, including success and error handling, to improve observability. - Updated daily job tracker to include AIPipelineAutoJobName for tracking AI pipeline job completions. This update enhances the system's capability to automate AI pipeline executions, improving efficiency and reliability in processing AI tasks.
107 lines
2.8 KiB
Go
107 lines
2.8 KiB
Go
package schedule
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"tm/pkg/logger"
|
|
"tm/pkg/mongo"
|
|
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
mongodriver "go.mongodb.org/mongo-driver/v2/mongo"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
)
|
|
|
|
const (
|
|
dailyJobRunsCollection = "daily_job_runs"
|
|
DailyJobStatusCompleted = "completed"
|
|
|
|
TEDScraperJobName = "ted_scraper"
|
|
AIPipelineAutoJobName = "ai_pipeline_auto"
|
|
)
|
|
|
|
type dailyJobRun struct {
|
|
ID string `bson:"_id"`
|
|
JobName string `bson:"job_name"`
|
|
Date string `bson:"date"`
|
|
CompletedAt int64 `bson:"completed_at"`
|
|
Status string `bson:"status"`
|
|
}
|
|
|
|
// DailyJobTracker records successful daily job runs so startup catch-up can skip
|
|
// dates that were already completed before a server restart.
|
|
type DailyJobTracker struct {
|
|
mongo *mongo.ConnectionManager
|
|
logger logger.Logger
|
|
}
|
|
|
|
func NewDailyJobTracker(mongoManager *mongo.ConnectionManager, log logger.Logger) *DailyJobTracker {
|
|
return &DailyJobTracker{
|
|
mongo: mongoManager,
|
|
logger: log,
|
|
}
|
|
}
|
|
|
|
func dailyJobRunID(jobName string, date time.Time) string {
|
|
return fmt.Sprintf("%s:%s", jobName, date.Format("2006-01-02"))
|
|
}
|
|
|
|
// IsCompleted reports whether jobName finished successfully for the given local calendar date.
|
|
func (t *DailyJobTracker) IsCompleted(ctx context.Context, jobName string, date time.Time) (bool, error) {
|
|
if jobName == "" {
|
|
return false, fmt.Errorf("daily job tracker: job name is required")
|
|
}
|
|
|
|
var doc dailyJobRun
|
|
err := t.mongo.GetCollection(dailyJobRunsCollection).FindOne(
|
|
ctx,
|
|
bson.M{"_id": dailyJobRunID(jobName, date), "status": DailyJobStatusCompleted},
|
|
).Decode(&doc)
|
|
if err != nil {
|
|
if errors.Is(err, mongodriver.ErrNoDocuments) {
|
|
return false, nil
|
|
}
|
|
return false, fmt.Errorf("daily job tracker: check completed %q: %w", jobName, err)
|
|
}
|
|
|
|
return true, nil
|
|
}
|
|
|
|
// MarkCompleted records a successful daily job run for the given local calendar date.
|
|
func (t *DailyJobTracker) MarkCompleted(ctx context.Context, jobName string, date time.Time) error {
|
|
if jobName == "" {
|
|
return fmt.Errorf("daily job tracker: job name is required")
|
|
}
|
|
|
|
dateKey := date.Format("2006-01-02")
|
|
now := time.Now().Unix()
|
|
|
|
_, err := t.mongo.GetCollection(dailyJobRunsCollection).UpdateOne(
|
|
ctx,
|
|
bson.M{"_id": dailyJobRunID(jobName, date)},
|
|
bson.M{
|
|
"$set": bson.M{
|
|
"job_name": jobName,
|
|
"date": dateKey,
|
|
"completed_at": now,
|
|
"status": DailyJobStatusCompleted,
|
|
},
|
|
"$setOnInsert": bson.M{"_id": dailyJobRunID(jobName, date)},
|
|
},
|
|
options.UpdateOne().SetUpsert(true),
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("daily job tracker: mark completed %q on %s: %w", jobName, dateKey, err)
|
|
}
|
|
|
|
t.logger.Debug("Daily job marked completed", map[string]interface{}{
|
|
"job_name": jobName,
|
|
"date": dateKey,
|
|
"completed_at": now,
|
|
})
|
|
|
|
return nil
|
|
}
|