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"
)

// 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
}

func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
	tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset {
	return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk}
}

func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
	now := uc.clk.Now()

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

	// 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.SplitTSVLimit(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, in.Parameters, 0, now)
		if err != nil {
			return err
		}
		shards = append(shards, art)
		tasks = append(tasks, task)
		return nil
	})
	_ = rc.Close()
	if splitErr != nil {
		cleanup()
		return SubmitDatasetResult{}, splitErr
	}

	// 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
}

// 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
}