Mazyar
1ad0206e61
continuous-integration/drone/push Build is passing
- Renamed and refactored the AI pipeline auto run functionality to a daily run, enhancing clarity and purpose. - Introduced a new `AIPipelineDailyWorker` to manage the daily execution of the AI pipeline, replacing the previous auto run implementation. - Updated configuration fields and logging messages to reflect the change from auto to daily run, ensuring consistent terminology throughout the codebase. - Removed the obsolete `ai_pipeline_auto.go` file to streamline the worker structure. This update improves the maintainability and readability of the AI pipeline management by clearly distinguishing between auto and daily run functionalities.
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"
|
|
AIPipelineDailyJobName = "ai_pipeline_daily"
|
|
)
|
|
|
|
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
|
|
}
|