119 lines
3.3 KiB
Go
119 lines
3.3 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
|
|
"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
|
|
)
|
|
|
|
type KolAssistTask struct {
|
|
ID string
|
|
TenantID int64
|
|
KolProfileID int64
|
|
PromptID *int64
|
|
OperatorID int64
|
|
TaskType string
|
|
Status string
|
|
RequestJSON []byte
|
|
ResultJSON []byte
|
|
ErrorMessage *string
|
|
StartedAt *time.Time
|
|
CompletedAt *time.Time
|
|
CreatedAt time.Time
|
|
UpdatedAt time.Time
|
|
}
|
|
|
|
type CreateKolAssistInput struct {
|
|
ID string
|
|
TenantID int64
|
|
KolProfileID int64
|
|
PromptID *int64
|
|
OperatorID int64
|
|
TaskType string
|
|
RequestJSON []byte
|
|
}
|
|
|
|
type KolAssistRepository interface {
|
|
Create(ctx context.Context, input CreateKolAssistInput) error
|
|
Get(ctx context.Context, tenantID int64, id string) (*KolAssistTask, error)
|
|
MarkStarted(ctx context.Context, tenantID int64, id string) error
|
|
MarkCompleted(ctx context.Context, tenantID int64, id string, resultJSON []byte) error
|
|
MarkFailed(ctx context.Context, tenantID int64, id string, errorMsg string) error
|
|
}
|
|
|
|
type kolAssistRepository struct {
|
|
q generated.Querier
|
|
}
|
|
|
|
func NewKolAssistRepository(db generated.DBTX) KolAssistRepository {
|
|
return &kolAssistRepository{q: newQuerier(db)}
|
|
}
|
|
|
|
func (r *kolAssistRepository) Create(ctx context.Context, input CreateKolAssistInput) error {
|
|
return r.q.CreateKolAssistTask(ctx, generated.CreateKolAssistTaskParams{
|
|
ID: input.ID,
|
|
TenantID: input.TenantID,
|
|
KolProfileID: input.KolProfileID,
|
|
PromptID: pgInt8(input.PromptID),
|
|
OperatorID: input.OperatorID,
|
|
TaskType: input.TaskType,
|
|
RequestJson: input.RequestJSON,
|
|
})
|
|
}
|
|
|
|
func (r *kolAssistRepository) Get(ctx context.Context, tenantID int64, id string) (*KolAssistTask, error) {
|
|
row, err := r.q.GetKolAssistTask(ctx, generated.GetKolAssistTaskParams{
|
|
ID: id,
|
|
TenantID: tenantID,
|
|
})
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
return &KolAssistTask{
|
|
ID: row.ID,
|
|
TenantID: row.TenantID,
|
|
KolProfileID: row.KolProfileID,
|
|
PromptID: nullableInt64(row.PromptID),
|
|
OperatorID: row.OperatorID,
|
|
TaskType: row.TaskType,
|
|
Status: row.Status,
|
|
RequestJSON: row.RequestJson,
|
|
ResultJSON: row.ResultJson,
|
|
ErrorMessage: nullableText(row.ErrorMessage),
|
|
StartedAt: optionalTime(row.StartedAt),
|
|
CompletedAt: optionalTime(row.CompletedAt),
|
|
CreatedAt: timeFromTimestamp(row.CreatedAt),
|
|
UpdatedAt: timeFromTimestamp(row.UpdatedAt),
|
|
}, nil
|
|
}
|
|
|
|
func (r *kolAssistRepository) MarkStarted(ctx context.Context, tenantID int64, id string) error {
|
|
return r.q.MarkKolAssistStarted(ctx, generated.MarkKolAssistStartedParams{
|
|
ID: id,
|
|
TenantID: tenantID,
|
|
})
|
|
}
|
|
|
|
func (r *kolAssistRepository) MarkCompleted(ctx context.Context, tenantID int64, id string, resultJSON []byte) error {
|
|
return r.q.MarkKolAssistCompleted(ctx, generated.MarkKolAssistCompletedParams{
|
|
ResultJson: resultJSON,
|
|
ID: id,
|
|
TenantID: tenantID,
|
|
})
|
|
}
|
|
|
|
func (r *kolAssistRepository) MarkFailed(ctx context.Context, tenantID int64, id string, errorMsg string) error {
|
|
return r.q.MarkKolAssistFailed(ctx, generated.MarkKolAssistFailedParams{
|
|
ErrorMessage: pgText(&errorMsg),
|
|
ID: id,
|
|
TenantID: tenantID,
|
|
})
|
|
}
|