SciMesh/coordinator/internal/usecase/upload.go

package usecase

import (
	"context"
	"fmt"
	"io"

	"github.com/google/uuid"

	"github.com/emil28092005/SciMesh/coordinator/internal/chunk"
	"github.com/emil28092005/SciMesh/coordinator/internal/domain"
	"github.com/emil28092005/SciMesh/coordinator/internal/workloads"
)

// SubmitDataset accepts an uploaded dataset, splits it into shard artifacts, and
// creates the job with one task per shard — the coordinator-side counterpart of
// a client submitting pre-chunked URIs.
type SubmitDataset struct {
	blobs       BlobStore
	artifacts   ArtifactRepository
	jobs        JobRepository
	tasks       TaskRepository
	tx          TxManager
	clk         Clock
	maxAttempts int
	catalog     *workloads.Catalog
	settings    WorkloadSettingsRepository
}

func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
	tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int, catalog *workloads.Catalog, settings WorkloadSettingsRepository) *SubmitDataset {
	return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts, catalog: catalog, settings: settings}
}

func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
	if err := validateUploadedWorkload(uc.catalog, in.Workload, in.Parameters); err != nil {
		return SubmitDatasetResult{}, err
	}
	if uc.settings != nil {
		enabled, err := uc.settings.GetEnabled(ctx, in.Workload)
		if err != nil {
			return SubmitDatasetResult{}, err
		}
		if !enabled {
			return SubmitDatasetResult{}, domain.ErrWorkloadDisabled
		}
	}
	if uc.maxAttempts < 1 {
		return SubmitDatasetResult{}, domain.ErrInvalidInput
	}
	now := uc.clk.Now()

	job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now)
	if err != nil {
		return SubmitDatasetResult{}, err
	}
	job.OwnerID = ownerFromContext(ctx)

	// Everything written to blob storage, so a failed transaction can undo it.
	var putKeys []string
	cleanup := func() {
		for _, k := range putKeys {
			_ = uc.blobs.Delete(ctx, k)
		}
	}

	// 1. Stream the upload into the input artifact; we measure size and sha256.
	input, err := domain.NewArtifact(job.ID, nil, domain.ArtifactInput, in.Filename, in.ContentType, now)
	if err != nil {
		return SubmitDatasetResult{}, err
	}
	sum, size, err := uc.blobs.Put(ctx, input.StorageKey, in.Body)
	if err != nil {
		return SubmitDatasetResult{}, err
	}
	putKeys = append(putKeys, input.StorageKey)
	input.SetContent(sum, size)

	// 2. Re-open the stored input and split it into shard artifacts + tasks.
	shards := []*domain.Artifact{}
	tasks := []*domain.Task{}
	rc, err := uc.blobs.Open(ctx, input.StorageKey)
	if err != nil {
		cleanup()
		return SubmitDatasetResult{}, err
	}
	splitErr := chunk.SplitChEMBLTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
		art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
			fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
		if err != nil {
			return err
		}
		ssum, ssize, err := uc.blobs.Put(ctx, art.StorageKey, shard)
		if err != nil {
			return err
		}
		putKeys = append(putKeys, art.StorageKey)
		art.SetContent(ssum, ssize)

		task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum,
			taskParameterSubset(in.Parameters), uc.maxAttempts, now)
		if err != nil {
			return err
		}
		shards = append(shards, art)
		tasks = append(tasks, task)
		return nil
	})
	_ = rc.Close()
	if splitErr != nil {
		cleanup()
		// Dataset shape is caller input, not an internal coordinator failure.
		return SubmitDatasetResult{}, domain.ErrInvalidInput
	}

	// 3. Persist job + all artifacts + all tasks atomically.
	err = uc.tx.WithinTx(ctx, func(ctx context.Context) error {
		if err := uc.jobs.Insert(ctx, job); err != nil {
			return err
		}
		if err := uc.artifacts.Insert(ctx, input); err != nil {
			return err
		}
		for _, a := range shards {
			if err := uc.artifacts.Insert(ctx, a); err != nil {
				return err
			}
		}
		return uc.tasks.InsertBatch(ctx, tasks)
	})
	if err != nil {
		cleanup()
		return SubmitDatasetResult{}, err
	}

	return SubmitDatasetResult{
		JobID:           job.ID,
		TaskCount:       len(tasks),
		InputArtifactID: input.ID,
	}, nil
}

// validateUploadedWorkload checks the submitted workload against the embedded
// catalog: it must be an enabled, upload-ready workload whose parameters
// satisfy the declared JSON schema. Workloads that need planner-produced
// inputs (such as the graph block-pair shards) declare upload_ready=false and
// cannot be driven from a single uploaded dataset.
func validateUploadedWorkload(catalog *workloads.Catalog, workload string, parameters map[string]any) error {
	if catalog == nil {
		return domain.ErrInvalidInput
	}
	if err := catalog.ValidateParameters(workload, parameters); err != nil {
		return domain.ErrInvalidInput
	}
	if !catalog.UploadReady(workload) {
		return domain.ErrInvalidInput
	}
	return nil
}

// taskParameterSubset drops coordinator-level keys from the parameters that
// are handed to workers. max_rows is a plan-time bound applied by the chunker
// here; a worker would reject it as outside its stage projection.
func taskParameterSubset(parameters map[string]any) map[string]any {
	if _, present := parameters["max_rows"]; !present {
		return parameters
	}
	subset := make(map[string]any, len(parameters)-1)
	for key, value := range parameters {
		if key != "max_rows" {
			subset[key] = value
		}
	}
	return subset
}

// GetTaskInput resolves a task's input shard and opens it for streaming. The
// caller closes the reader.
type GetTaskInput struct {
	tasks     TaskRepository
	artifacts ArtifactRepository
	blobs     BlobStore
}

func NewGetTaskInput(tasks TaskRepository, artifacts ArtifactRepository, blobs BlobStore) *GetTaskInput {
	return &GetTaskInput{tasks: tasks, artifacts: artifacts, blobs: blobs}
}

func (uc *GetTaskInput) Execute(ctx context.Context, taskID uuid.UUID) (*domain.Artifact, io.ReadCloser, error) {
	task, err := uc.tasks.Get(ctx, taskID)
	if err != nil {
		return nil, nil, err
	}
	if task.InputArtifactID == nil {
		// A URI-based task keeps its input outside the coordinator.
		return nil, nil, domain.ErrArtifactNotFound
	}
	art, err := uc.artifacts.Get(ctx, *task.InputArtifactID)
	if err != nil {
		return nil, nil, err
	}
	rc, err := uc.blobs.Open(ctx, art.StorageKey)
	if err != nil {
		return nil, nil, err
	}
	return art, rc, nil
}