root
b31d8d0096
- Added migration to harden task audit tracking by modifying audit_logs and related tables. - Introduced operator_id to several tables for better tracking of actions. - Updated article_templates with new prompt templates for various article types, enhancing content generation. - Created prompt_rules and schedule_tasks tables to manage content generation rules and scheduling. - Added foreign key constraints to articles for better data integrity.
104 lines
3.6 KiB
Go
104 lines
3.6 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
|
|
"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
|
|
)
|
|
|
|
var ErrTemplateAssistTaskNotFound = errors.New("template assist task not found")
|
|
|
|
type TemplateAssistTaskRecord struct {
|
|
ID string
|
|
Status string
|
|
ResultJSON []byte
|
|
ErrorMessage *string
|
|
}
|
|
|
|
type CreateTemplateAssistTaskInput struct {
|
|
ID string
|
|
TaskType string
|
|
TenantID int64
|
|
OperatorID int64
|
|
TemplateID int64
|
|
RequestJSON []byte
|
|
}
|
|
|
|
type TemplateAssistTaskRepository interface {
|
|
CreateTask(ctx context.Context, input CreateTemplateAssistTaskInput) error
|
|
GetTask(ctx context.Context, id, taskType string, tenantID, templateID int64) (*TemplateAssistTaskRecord, error)
|
|
MarkRunning(ctx context.Context, id, taskType string, tenantID int64, startedAt time.Time) error
|
|
CompleteTask(ctx context.Context, id, taskType string, tenantID int64, resultJSON []byte, completedAt time.Time) error
|
|
FailTask(ctx context.Context, id, taskType string, tenantID int64, errorMessage string, completedAt time.Time) error
|
|
}
|
|
|
|
type templateAssistTaskRepository struct {
|
|
db generated.DBTX
|
|
}
|
|
|
|
func NewTemplateAssistTaskRepository(db generated.DBTX) TemplateAssistTaskRepository {
|
|
return &templateAssistTaskRepository{db: db}
|
|
}
|
|
|
|
func (r *templateAssistTaskRepository) CreateTask(ctx context.Context, input CreateTemplateAssistTaskInput) error {
|
|
_, err := r.db.Exec(ctx, `
|
|
INSERT INTO template_assist_tasks (id, task_type, tenant_id, operator_id, template_id, request_json, status)
|
|
VALUES ($1, $2, $3, $4, $5, $6, 'queued')
|
|
`, input.ID, input.TaskType, input.TenantID, input.OperatorID, input.TemplateID, input.RequestJSON)
|
|
return err
|
|
}
|
|
|
|
func (r *templateAssistTaskRepository) GetTask(ctx context.Context, id, taskType string, tenantID, templateID int64) (*TemplateAssistTaskRecord, error) {
|
|
var record TemplateAssistTaskRecord
|
|
err := r.db.QueryRow(ctx, `
|
|
SELECT id, status, result_json, error_message
|
|
FROM template_assist_tasks
|
|
WHERE id = $1 AND task_type = $2 AND tenant_id = $3 AND template_id = $4
|
|
`, id, taskType, tenantID, templateID).Scan(&record.ID, &record.Status, &record.ResultJSON, &record.ErrorMessage)
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return nil, ErrTemplateAssistTaskNotFound
|
|
}
|
|
return nil, err
|
|
}
|
|
return &record, nil
|
|
}
|
|
|
|
func (r *templateAssistTaskRepository) MarkRunning(ctx context.Context, id, taskType string, tenantID int64, startedAt time.Time) error {
|
|
_, err := r.db.Exec(ctx, `
|
|
UPDATE template_assist_tasks
|
|
SET status = 'running', started_at = $1, updated_at = NOW()
|
|
WHERE id = $2 AND task_type = $3 AND tenant_id = $4
|
|
`, startedAt, id, taskType, tenantID)
|
|
return err
|
|
}
|
|
|
|
func (r *templateAssistTaskRepository) CompleteTask(ctx context.Context, id, taskType string, tenantID int64, resultJSON []byte, completedAt time.Time) error {
|
|
_, err := r.db.Exec(ctx, `
|
|
UPDATE template_assist_tasks
|
|
SET status = 'completed',
|
|
result_json = $1,
|
|
error_message = NULL,
|
|
completed_at = $2,
|
|
updated_at = NOW()
|
|
WHERE id = $3 AND task_type = $4 AND tenant_id = $5
|
|
`, resultJSON, completedAt, id, taskType, tenantID)
|
|
return err
|
|
}
|
|
|
|
func (r *templateAssistTaskRepository) FailTask(ctx context.Context, id, taskType string, tenantID int64, errorMessage string, completedAt time.Time) error {
|
|
_, err := r.db.Exec(ctx, `
|
|
UPDATE template_assist_tasks
|
|
SET status = 'failed',
|
|
error_message = $1,
|
|
completed_at = $2,
|
|
updated_at = NOW()
|
|
WHERE id = $3 AND task_type = $4 AND tenant_id = $5
|
|
`, errorMessage, completedAt, id, taskType, tenantID)
|
|
return err
|
|
}
|