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 Subscriptions 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, Subscriptions: row.Subscriptions, } } 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 }