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,198 @@
+package repository
+
+import (
+	"context"
+	"time"
+
+	"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
+)
+
+type KolUsageLog struct {
+	ID               int64
+	TenantID         int64
+	GenerationTaskID *int64
+	Status           string
+	CreatedAt        time.Time
+}
+
+type KolUsageByTaskRecord struct {
+	ID        int64
+	TenantID  int64
+	PackageID int64
+	PromptID  int64
+	Status    string
+	CreatedAt time.Time
+}
+
+type CreateKolUsageLogInput struct {
+	TenantID             int64
+	SubscriptionID       *int64
+	SubscriptionPromptID *int64
+	PackageID            int64
+	PromptID             int64
+	PromptRevisionNo     int
+	GenerationTaskID     *int64
+	VariablesSnapshot    []byte
+	SchemaSnapshot       []byte
+	RenderedPromptHash   *string
+}
+
+type KolDashboardOverview struct {
+	TotalSubs    int64
+	MonthSubs    int64
+	TotalUsage   int64
+	MonthUsage   int64
+	FailureCount int64
+}
+
+type KolDashboardTrendPoint struct {
+	Day    time.Time
+	Usages int64
+}
+
+type KolUsagePackageStats struct {
+	PackageID int64
+	Total     int64
+	Completed int64
+	Failed    int64
+}
+
+type KolUsageRepository interface {
+	Create(ctx context.Context, input CreateKolUsageLogInput) (*KolUsageLog, error)
+	MarkCompletedByTask(ctx context.Context, tenantID, taskID, articleID int64) error
+	MarkFailedByTask(ctx context.Context, tenantID, taskID int64, errorMsg string) error
+	GetByGenerationTask(ctx context.Context, tenantID, taskID int64) (*KolUsageByTaskRecord, error)
+	DashboardOverview(ctx context.Context, packageIDs []int64) (*KolDashboardOverview, error)
+	DashboardTrend(ctx context.Context, packageIDs []int64, periodDays int) ([]KolDashboardTrendPoint, error)
+	PerPackageStats(ctx context.Context, packageIDs []int64) ([]KolUsagePackageStats, error)
+}
+
+type kolUsageRepository struct {
+	q generated.Querier
+}
+
+func NewKolUsageRepository(db generated.DBTX) KolUsageRepository {
+	return &kolUsageRepository{q: newQuerier(db)}
+}
+
+func (r *kolUsageRepository) Create(ctx context.Context, input CreateKolUsageLogInput) (*KolUsageLog, error) {
+	row, err := r.q.CreateKolUsageLog(ctx, generated.CreateKolUsageLogParams{
+		TenantID:             input.TenantID,
+		SubscriptionID:       pgInt8(input.SubscriptionID),
+		SubscriptionPromptID: pgInt8(input.SubscriptionPromptID),
+		PackageID:            input.PackageID,
+		PromptID:             input.PromptID,
+		PromptRevisionNo:     int32(input.PromptRevisionNo),
+		GenerationTaskID:     pgInt8(input.GenerationTaskID),
+		VariablesSnapshot:    input.VariablesSnapshot,
+		SchemaSnapshot:       input.SchemaSnapshot,
+		RenderedPromptHash:   pgText(input.RenderedPromptHash),
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &KolUsageLog{
+		ID:               row.ID,
+		TenantID:         row.TenantID,
+		GenerationTaskID: nullableInt64(row.GenerationTaskID),
+		Status:           row.Status,
+		CreatedAt:        timeFromTimestamp(row.CreatedAt),
+	}, nil
+}
+
+func (r *kolUsageRepository) MarkCompletedByTask(ctx context.Context, tenantID, taskID, articleID int64) error {
+	return r.q.MarkKolUsageCompletedByTask(ctx, generated.MarkKolUsageCompletedByTaskParams{
+		ArticleID:        pgInt8(&articleID),
+		GenerationTaskID: pgInt8(&taskID),
+		TenantID:         tenantID,
+	})
+}
+
+func (r *kolUsageRepository) MarkFailedByTask(ctx context.Context, tenantID, taskID int64, errorMsg string) error {
+	return r.q.MarkKolUsageFailedByTask(ctx, generated.MarkKolUsageFailedByTaskParams{
+		ErrorMessage:     pgText(&errorMsg),
+		GenerationTaskID: pgInt8(&taskID),
+		TenantID:         tenantID,
+	})
+}
+
+func (r *kolUsageRepository) GetByGenerationTask(ctx context.Context, tenantID, taskID int64) (*KolUsageByTaskRecord, error) {
+	row, err := r.q.GetKolUsageByGenerationTask(ctx, generated.GetKolUsageByGenerationTaskParams{
+		GenerationTaskID: pgInt8(&taskID),
+		TenantID:         tenantID,
+	})
+	if err != nil {
+		// Callers use errors.Is(err, pgx.ErrNoRows) for "not found"; preserve err semantics.
+		return nil, err
+	}
+	return &KolUsageByTaskRecord{
+		ID:        row.ID,
+		TenantID:  row.TenantID,
+		PackageID: row.PackageID,
+		PromptID:  row.PromptID,
+		Status:    row.Status,
+		CreatedAt: timeFromTimestamp(row.CreatedAt),
+	}, nil
+}
+
+func (r *kolUsageRepository) DashboardOverview(ctx context.Context, packageIDs []int64) (*KolDashboardOverview, error) {
+	if len(packageIDs) == 0 {
+		return &KolDashboardOverview{}, nil
+	}
+	row, err := r.q.KolDashboardOverview(ctx, packageIDs)
+	if err != nil {
+		return nil, err
+	}
+	return &KolDashboardOverview{
+		TotalSubs:    row.TotalSubs,
+		MonthSubs:    row.MonthSubs,
+		TotalUsage:   row.TotalUsage,
+		MonthUsage:   row.MonthUsage,
+		FailureCount: row.FailureCount,
+	}, nil
+}
+
+func (r *kolUsageRepository) DashboardTrend(ctx context.Context, packageIDs []int64, periodDays int) ([]KolDashboardTrendPoint, error) {
+	if len(packageIDs) == 0 {
+		return []KolDashboardTrendPoint{}, nil
+	}
+	rows, err := r.q.KolDashboardTrend(ctx, generated.KolDashboardTrendParams{
+		PackageIds: packageIDs,
+		PeriodDays: int32(periodDays),
+	})
+	if err != nil {
+		return nil, err
+	}
+	result := make([]KolDashboardTrendPoint, len(rows))
+	for i, row := range rows {
+		var day time.Time
+		if row.D.Valid {
+			day = row.D.Time
+		}
+		result[i] = KolDashboardTrendPoint{
+			Day:    day,
+			Usages: row.Usages,
+		}
+	}
+	return result, nil
+}
+
+func (r *kolUsageRepository) PerPackageStats(ctx context.Context, packageIDs []int64) ([]KolUsagePackageStats, error) {
+	if len(packageIDs) == 0 {
+		return []KolUsagePackageStats{}, nil
+	}
+	rows, err := r.q.KolUsageCountByPackage(ctx, packageIDs)
+	if err != nil {
+		return nil, err
+	}
+	result := make([]KolUsagePackageStats, len(rows))
+	for i, row := range rows {
+		result[i] = KolUsagePackageStats{
+			PackageID: row.PackageID,
+			Total:     row.Total,
+			Completed: row.Completed,
+			Failed:    row.Failed,
+		}
+	}
+	return result, nil
+}