feat(kol): repository layer for all KOL tables

2026-04-17 08:17:22 +08:00
parent 2c58fc591e
commit 3d4842e983
8 changed files with 1341 additions and 0 deletions
@@ -0,0 +1,118 @@
+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,
+	})
+}