diff --git a/internal/database/models.go b/internal/database/models.go index d1a6448..975422a 100644 --- a/internal/database/models.go +++ b/internal/database/models.go @@ -65,6 +65,7 @@ type StepResult struct { ID string `bun:"id,pk,type:text" json:"id"` RunID int64 `bun:"run_id,notnull" json:"run_id"` + RunUUID string `bun:"run_uuid" json:"run_uuid,omitempty"` StepName string `bun:"step_name,notnull" json:"step_name"` StepType string `bun:"step_type,notnull" json:"step_type"` Status string `bun:"status,notnull" json:"status"` diff --git a/internal/database/repository/run_repo.go b/internal/database/repository/run_repo.go index 927defa..ea8b70b 100644 --- a/internal/database/repository/run_repo.go +++ b/internal/database/repository/run_repo.go @@ -42,7 +42,7 @@ func (r *RunRepository) GetByRunID(ctx context.Context, runID string) (*database scan := new(database.Run) err := r.db.NewSelect(). Model(scan). - Where("run_id = ?", runID). + Where("run_uuid = ?", runID). Scan(ctx) if err != nil { return nil, err diff --git a/internal/database/write_coordinator.go b/internal/database/write_coordinator.go index d4ec326..544547a 100644 --- a/internal/database/write_coordinator.go +++ b/internal/database/write_coordinator.go @@ -71,6 +71,7 @@ func (wc *WriteCoordinator) AddStepResult(stepName, stepType, status, command, o result := &StepResult{ ID: uuid.New().String(), RunID: wc.runID, + RunUUID: wc.runUUID, StepName: stepName, StepType: stepType, Status: status, diff --git a/internal/distributed/master.go b/internal/distributed/master.go index b784939..19298e9 100644 --- a/internal/distributed/master.go +++ b/internal/distributed/master.go @@ -449,6 +449,9 @@ func (m *Master) processWorkerData(ctx context.Context, key string, envelope *Da return } + if m.db == nil { + m.db = database.GetDB() + } if m.db == nil { m.logger.Debug("skipping data processing - no database connection", zap.String("key", key), @@ -505,10 +508,27 @@ func (m *Master) processStepData(ctx context.Context, envelope *DataEnvelope) { return } - // Insert step result - _, err := m.db.NewInsert().Model(&step).Exec(ctx) - if err != nil { - m.printer.Warning("Failed to create step result %s: %s", step.StepName, err) + // Resolve the real DB RunID using RunUUID sent by the worker + if step.RunUUID != "" { + repo := repository.NewRunRepository(m.db) + existing, err := repo.GetByRunID(ctx, step.RunUUID) + if err == nil && existing != nil { + step.RunID = existing.ID + } + } + + // Insert step result and increment progress + if step.RunID > 0 { + _, err := m.db.NewInsert().Model(&step).Exec(ctx) + if err != nil { + m.printer.Warning("Failed to create step result %s: %s", step.StepName, err) + } else { + // Increment completed_steps on the run + _, _ = m.db.NewUpdate().Model((*database.Run)(nil)). + Set("completed_steps = completed_steps + 1"). + Where("id = ?", step.RunID). + Exec(ctx) + } } } diff --git a/internal/distributed/worker.go b/internal/distributed/worker.go index 32ca7ac..11f7d5e 100644 --- a/internal/distributed/worker.go +++ b/internal/distributed/worker.go @@ -321,7 +321,10 @@ func (w *Worker) executeTask(ctx context.Context, task *Task) *TaskResult { // Create run record for distributed tracking now := time.Now() - runUUID := uuid.New().String() + runUUID := task.ScanID + if runUUID == "" { + runUUID = uuid.New().String() + } paramsInterface := make(map[string]interface{}) for k, v := range params { paramsInterface[k] = v @@ -348,6 +351,7 @@ func (w *Worker) executeTask(ctx context.Context, task *Task) *TaskResult { // Wire up executor for run tracking w.executor.SetDBRunUUID(runUUID) + w.executor.SetDBRunID(time.Now().UnixNano() % 1000000000) // Execute based on workflow kind var wfResult *core.WorkflowResult diff --git a/pkg/cli/worker.go b/pkg/cli/worker.go index 25bfcc6..71d41cf 100644 --- a/pkg/cli/worker.go +++ b/pkg/cli/worker.go @@ -65,6 +65,11 @@ var workerJoinCmd = &cobra.Command{ // Ensure external-binaries are in PATH so workflow steps can find tools ensureExternalBinariesInPath(cfg) + // Connect to database so db_import_* functions work on the worker + if _, err := database.Connect(cfg); err != nil { + fmt.Fprintf(os.Stderr, "Worker database connection failed (asset imports will be skipped): %v\n", err) + } + // Create worker worker, err := distributed.NewWorker(cfg, &distributed.WorkerOptions{ GetPublicIP: getPublicIP, diff --git a/pkg/server/handlers/runs.go b/pkg/server/handlers/runs.go index cbb735a..3b70cf2 100644 --- a/pkg/server/handlers/runs.go +++ b/pkg/server/handlers/runs.go @@ -493,6 +493,9 @@ func CreateRun(cfg *config.Config, master *distributed.Master) fiber.Handler { Target: target, Params: taskParams, } + if run != nil { + task.ScanID = run.RunUUID + } if err := master.SubmitTask(ctx, task); err != nil { lgr := logger.Get()