导出系统

internal/store/postgres/export_task_store.go

@@ -5,7 +5,9 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/bytedance/sonic"
 	"github.com/redis/go-redis/v9"
+	"gorm.io/datatypes"
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 
@@ -124,7 +126,7 @@ func (s *ExportTaskStore) TryMarkProcessingIfPending(ctx context.Context, taskID
 }
 
 // UpdateDispatchPlan 更新 dispatch 阶段计算出的总量信息。
-func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, totalRows, totalShards uint, updater uint) error {
+func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, totalRows, totalShards int, updater uint) error {
 	now := time.Now()
 	return s.db.WithContext(ctx).
 		Model(&model.ExportTask{}).
@@ -138,6 +140,40 @@ func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, t
 		}).Error
 }
 
+// UpdateResolvedHeaders 固定导出任务运行期表头。
+func (s *ExportTaskStore) UpdateResolvedHeaders(ctx context.Context, taskID uint, headers []string, updater uint) error {
+	var task model.ExportTask
+	if err := s.db.WithContext(ctx).
+		Select("id", "query_json").
+		Where("id = ?", taskID).
+		First(&task).Error; err != nil {
+		return err
+	}
+
+	raw := make(map[string]any)
+	if len(task.QueryJSON) > 0 {
+		if err := sonic.Unmarshal(task.QueryJSON, &raw); err != nil {
+			return err
+		}
+	}
+	raw["resolved_headers"] = headers
+
+	bytes, err := sonic.Marshal(raw)
+	if err != nil {
+		return err
+	}
+
+	now := time.Now()
+	return s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ?", taskID).
+		Updates(map[string]any{
+			"query_json": datatypes.JSON(bytes),
+			"updated_at": now,
+			"updater":    updater,
+		}).Error
+}
+
 // IsCancelRequested 查询是否已请求取消。
 func (s *ExportTaskStore) IsCancelRequested(ctx context.Context, taskID uint) (bool, error) {
 	var value bool