have added pg support

Author: matroskin13

context.go

@@ -54,5 +54,5 @@ func (c *taskContext) SetState(state any) error {
 		return err
 	}
 
-	return c.taskEngine.SetState(c.ctx, c.task.ID, b)
+	return c.taskEngine.SetState(c.ctx, c.task, b)
 }

engine.go

@@ -11,20 +11,21 @@ type Engine interface {
 }
 
 type TaskEngine interface {
-	GetRelatedTask(ctx context.Context, task string, id string) (*Task, error)
+	GetRelatedTask(ctx context.Context, task *Task) (*Task, error)
 	FindNextTask(ctx context.Context, statuses []string) (*Task, error)
-	ReleaseTask(ctx context.Context, id string) error
-	WaitTaskForSubtasks(ctx context.Context, id string) error
+	ReleaseTask(ctx context.Context, task *Task) error
+	WaitTaskForSubtasks(ctx context.Context, task *Task) error
 	FailTask(ctx context.Context, task *Task, err error, timeout time.Duration) error
 	CreateTask(ctx context.Context, task *Task) error
-	GetUnreleasedTaskChildren(ctx context.Context, id string) (*Task, error)
-	SetState(ctx context.Context, id string, state []byte) error
+	GetUnreleasedTaskChildren(ctx context.Context, task *Task) (*Task, error)
+	SetState(ctx context.Context, task *Task, state []byte) error
 }
 
 type JobEngine interface {
 	FindNextJob(ctx context.Context, statuses []string) (*Job, error)
-	GetUnreleasedJobChildren(ctx context.Context, jobName string) (*Task, error)
-	Release(ctx context.Context, jobName string, nextLaunchAt time.Time) error
-	WaitForSubtasks(ctx context.Context, jobName string) error
+	GetUnreleasedJobChildren(ctx context.Context, name string) (*Task, error)
+	Release(ctx context.Context, job *Job, nextLaunchAt time.Time) error
+	WaitForSubtasks(ctx context.Context, job *Job) error
 	RegisterJob(ctx context.Context, cfg *JobConfig) error
+	Init(ctx context.Context) error
 }

engines/mongo/mongo.go

@@ -62,8 +62,8 @@ func (m *Mongo) CreateTask(ctx context.Context, task *stepper.Task) error {
 	return err
 }
 
-func (m *Mongo) SetState(ctx context.Context, id string, state []byte) error {
-	query := bson.M{"id": id}
+func (m *Mongo) SetState(ctx context.Context, task *stepper.Task, state []byte) error {
+	query := bson.M{"id": task.ID}
 	update := bson.M{"$set": bson.M{"state": state}}
 
 	if err := m.tasks.FindOneAndUpdate(ctx, query, update).Err(); err != nil {
@@ -164,12 +164,12 @@ func (m *Mongo) GetUnreleasedJobChildren(ctx context.Context, jobId string) (*st
 	return task.ToModel(), nil
 }
 
-func (m *Mongo) GetUnreleasedTaskChildren(ctx context.Context, id string) (*stepper.Task, error) {
+func (m *Mongo) GetUnreleasedTaskChildren(ctx context.Context, forTask *stepper.Task) (*stepper.Task, error) {
 	var task Task
 
 	query := bson.M{
 		"status": bson.M{"$in": []string{"created", "in_progress"}},
-		"parent": id,
+		"parent": forTask.ID,
 	}
 
 	if err := m.tasks.FindOne(ctx, query).Decode(&task); err != nil {
@@ -183,8 +183,8 @@ func (m *Mongo) GetUnreleasedTaskChildren(ctx context.Context, id string) (*step
 	return task.ToModel(), nil
 }
 
-func (m *Mongo) GetRelatedTask(ctx context.Context, task string, id string) (*stepper.Task, error) {
-	query := bson.M{"custom_id": id, "name": task, "status": bson.M{"$ne": "released"}}
+func (m *Mongo) GetRelatedTask(ctx context.Context, task *stepper.Task) (*stepper.Task, error) {
+	query := bson.M{"custom_id": task.ID, "name": task.Name, "status": bson.M{"$ne": "released"}}
 
 	var e Task
 
@@ -199,10 +199,10 @@ func (m *Mongo) GetRelatedTask(ctx context.Context, task string, id string) (*st
 	return e.ToModel(), nil
 }
 
-func (m *Mongo) Release(ctx context.Context, name string, nextTimeLaunch time.Time) error {
+func (m *Mongo) Release(ctx context.Context, job *stepper.Job, nextTimeLaunch time.Time) error {
 	return m.jobs.FindOneAndUpdate(
 		ctx,
-		bson.M{"name": name},
+		bson.M{"name": job.Name},
 		bson.M{"$set": bson.M{
 			"lock_at":      nil,
 			"status":       "released",
@@ -231,21 +231,21 @@ func (m *Mongo) FailTask(ctx context.Context, task *stepper.Task, handlerErr err
 	).Err()
 }
 
-func (m *Mongo) ReleaseTask(ctx context.Context, id string) error {
+func (m *Mongo) ReleaseTask(ctx context.Context, task *stepper.Task) error {
 	return m.tasks.FindOneAndUpdate(
 		ctx,
-		bson.M{"id": id},
+		bson.M{"id": task.ID},
 		bson.M{"$set": bson.M{
 			"lock_at": nil,
 			"status":  "released",
 		}},
 	).Err()
 }
 
-func (m *Mongo) WaitForSubtasks(ctx context.Context, name string) error {
+func (m *Mongo) WaitForSubtasks(ctx context.Context, job *stepper.Job) error {
 	return m.jobs.FindOneAndUpdate(
 		ctx,
-		bson.M{"name": name},
+		bson.M{"name": job.Name},
 		bson.M{"$set": bson.M{
 			"lock_at":      nil,
 			"status":       "waiting",
@@ -254,14 +254,19 @@ func (m *Mongo) WaitForSubtasks(ctx context.Context, name string) error {
 	).Err()
 }
 
-func (m *Mongo) WaitTaskForSubtasks(ctx context.Context, id string) error {
+func (m *Mongo) WaitTaskForSubtasks(ctx context.Context, task *stepper.Task) error {
 	return m.tasks.FindOneAndUpdate(
 		ctx,
-		bson.M{"id": id},
+		bson.M{"id": task.ID},
 		bson.M{"$set": bson.M{
 			"lock_at":  nil,
 			"status":   "waiting",
 			"launchAt": time.Now().Add(time.Second * 1),
 		}},
 	).Err()
 }
+
+// TODO add indexes
+func (m *Mongo) Init(ctx context.Context) error {
+	return nil
+}

engines/pg/job.go

@@ -0,0 +1,8 @@
+package pg
+
+type Job struct {
+	Status       string `json:"status"`
+	Name         string `json:"name"`
+	Pattern      string `json:"pattern"`
+	NextLaunchAt int64  `json:"naxtLaunchAt"`
+}