Add models and migration for Create/Delete/Get Tasks

Author: tanujnay112

go/pkg/common/errors.go

@@ -48,6 +48,14 @@ var (
 	// Segment metadata errors
 	ErrUnknownSegmentMetadataType = errors.New("segment metadata value type not supported")
 
+	// Task errors
+	ErrTaskAlreadyExists = errors.New("the task that was being created already exists for this collection")
+	ErrTaskNotFound      = errors.New("the requested task was not found")
+	ErrInvalidTaskName   = errors.New("task name cannot start with reserved prefix '_deleted_'")
+
+	// Operator errors
+	ErrOperatorNotFound = errors.New("operator not found")
+
 	// Others
 	ErrCompactionOffsetSomehowAhead = errors.New("system invariant was violated. Compaction offset in sysdb should always be behind or equal to offset in log")
 )

go/pkg/sysdb/coordinator/task.go

@@ -0,0 +1,210 @@
+package coordinator
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	"github.com/chroma-core/chroma/go/pkg/common"
+	"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
+	"github.com/google/uuid"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
+	"google.golang.org/protobuf/proto"
+)
+
+// CreateTask creates a new task in the database
+func (s *Coordinator) CreateTask(ctx context.Context, req *coordinatorpb.CreateTaskRequest) (*coordinatorpb.CreateTaskResponse, error) {
+	// Validate task name doesn't start with soft-deletion reserved prefix
+	if strings.HasPrefix(req.Name, "_deleted_") {
+		log.Error("CreateTask: task name cannot start with _deleted_")
+		return nil, common.ErrInvalidTaskName
+	}
+
+	// Generate new task UUID
+	taskID := uuid.New()
+	outputCollectionName := req.OutputCollectionName
+
+	// Look up database_id from databases table using database name and tenant
+	databases, err := s.catalog.metaDomain.DatabaseDb(ctx).GetDatabases(req.TenantId, req.Database)
+	if err != nil {
+		log.Error("CreateTask: failed to get database", zap.Error(err))
+		return nil, err
+	}
+	if len(databases) == 0 {
+		log.Error("CreateTask: database not found")
+		return nil, common.ErrDatabaseNotFound
+	}
+
+	// Look up operator by name from the operators table
+	operator, err := s.catalog.metaDomain.OperatorDb(ctx).GetByName(req.OperatorName)
+	if err != nil {
+		log.Error("CreateTask: failed to get operator", zap.Error(err))
+		return nil, err
+	}
+	if operator == nil {
+		log.Error("CreateTask: operator not found", zap.String("operator_name", req.OperatorName))
+		return nil, common.ErrOperatorNotFound
+	}
+	operatorID := operator.OperatorID
+
+	// Generate UUIDv7 for time-ordered nonce
+	nextNonce, err := uuid.NewV7()
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO(tanujnay112): Can combine the two collection checks into one
+	// Check if input collection exists
+	collections, err := s.catalog.metaDomain.CollectionDb(ctx).GetCollections([]string{req.InputCollectionId}, nil, req.TenantId, req.Database, nil, nil, false)
+	if err != nil {
+		log.Error("CreateTask: failed to get input collection", zap.Error(err))
+		return nil, err
+	}
+	if len(collections) == 0 {
+		log.Error("CreateTask: input collection not found")
+		return nil, common.ErrCollectionNotFound
+	}
+
+	// Check if output collection already exists
+	existingOutputCollections, err := s.catalog.metaDomain.CollectionDb(ctx).GetCollections(nil, &outputCollectionName, req.TenantId, req.Database, nil, nil, false)
+	if err != nil {
+		log.Error("CreateTask: failed to check output collection", zap.Error(err))
+		return nil, err
+	}
+	if len(existingOutputCollections) > 0 {
+		log.Error("CreateTask: output collection already exists")
+		return nil, common.ErrCollectionUniqueConstraintViolation
+	}
+
+	// Check if task already exists
+	existingTask, err := s.catalog.metaDomain.TaskDb(ctx).GetByName(req.InputCollectionId, req.Name)
+	if err != nil {
+		log.Error("CreateTask: failed to check task", zap.Error(err))
+		return nil, err
+	}
+	if existingTask != nil {
+		log.Info("CreateTask: task already exists, returning existing")
+		return &coordinatorpb.CreateTaskResponse{
+			TaskId: existingTask.ID.String(),
+		}, nil
+	}
+
+	now := time.Now()
+	task := &dbmodel.Task{
+		ID:                   taskID,
+		Name:                 req.Name,
+		TenantID:             req.TenantId,
+		DatabaseID:           databases[0].ID,
+		InputCollectionID:    req.InputCollectionId,
+		OutputCollectionName: req.OutputCollectionName,
+		OperatorID:           operatorID,
+		OperatorParams:       req.Params,
+		CompletionOffset:     0,
+		LastRun:              nil,
+		NextRun:              nil, // Will be set to zero initially, scheduled by task scheduler
+		MinRecordsForTask:    int64(req.MinRecordsForTask),
+		CurrentAttempts:      0,
+		CreatedAt:            now,
+		UpdatedAt:            now,
+		NextNonce:            nextNonce,
+		OldestWrittenNonce:   nil,
+	}
+
+	// Try to insert task into database
+	err = s.catalog.metaDomain.TaskDb(ctx).Insert(task)
+	if err != nil {
+		// Check if it's a unique constraint violation (concurrent creation)
+		if err == common.ErrTaskAlreadyExists {
+			log.Error("CreateTask: task already exists")
+			return nil, common.ErrTaskAlreadyExists
+		}
+		log.Error("CreateTask: failed to insert task", zap.Error(err))
+		return nil, err
+	}
+
+	log.Info("Task created successfully", zap.String("task_id", taskID.String()), zap.String("name", req.Name), zap.String("output_collection_name", outputCollectionName))
+	return &coordinatorpb.CreateTaskResponse{
+		TaskId: taskID.String(),
+	}, nil
+}
+
+// GetTaskByName retrieves a task by name from the database
+func (s *Coordinator) GetTaskByName(ctx context.Context, req *coordinatorpb.GetTaskByNameRequest) (*coordinatorpb.GetTaskByNameResponse, error) {
+	// Can do both calls with a JOIN
+	task, err := s.catalog.metaDomain.TaskDb(ctx).GetByName(req.InputCollectionId, req.TaskName)
+	if err != nil {
+		return nil, err
+	}
+
+	// If task not found, return empty response
+	if task == nil {
+		return nil, common.ErrTaskNotFound
+	}
+
+	// Look up operator name from operators table
+	operator, err := s.catalog.metaDomain.OperatorDb(ctx).GetByID(task.OperatorID)
+	if err != nil {
+		log.Error("GetTaskByName: failed to get operator", zap.Error(err))
+		return nil, err
+	}
+	if operator == nil {
+		log.Error("GetTaskByName: operator not found", zap.String("operator_id", task.OperatorID.String()))
+		return nil, common.ErrOperatorNotFound
+	}
+
+	// Debug logging
+	log.Info("Found task", zap.String("task_id", task.ID.String()), zap.String("name", task.Name), zap.String("input_collection_id", task.InputCollectionID), zap.String("output_collection_name", task.OutputCollectionName))
+
+	// Convert task to response
+	return &coordinatorpb.GetTaskByNameResponse{
+		TaskId:               proto.String(task.ID.String()),
+		Name:                 proto.String(task.Name),
+		OperatorName:         proto.String(operator.OperatorName),
+		InputCollectionId:    proto.String(task.InputCollectionID),
+		OutputCollectionName: proto.String(task.OutputCollectionName),
+		Params:               proto.String(task.OperatorParams),
+		CompletionOffset:     proto.Int64(task.CompletionOffset),
+		MinRecordsForTask:    proto.Uint64(uint64(task.MinRecordsForTask)),
+	}, nil
+}
+
+// DeleteTask soft deletes a task by name
+func (s *Coordinator) DeleteTask(ctx context.Context, req *coordinatorpb.DeleteTaskRequest) (*coordinatorpb.DeleteTaskResponse, error) {
+	err := s.catalog.metaDomain.TaskDb(ctx).SoftDelete(req.InputCollectionId, req.TaskName)
+	if err != nil {
+		log.Error("DeleteTask failed", zap.Error(err))
+		return nil, err
+	}
+
+	log.Info("Task deleted", zap.String("input_collection_id", req.InputCollectionId), zap.String("task_name", req.TaskName))
+
+	return &coordinatorpb.DeleteTaskResponse{
+		Success: true,
+	}, nil
+}
+
+// GetOperators retrieves all operators from the database
+func (s *Coordinator) GetOperators(ctx context.Context, req *coordinatorpb.GetOperatorsRequest) (*coordinatorpb.GetOperatorsResponse, error) {
+	operators, err := s.catalog.metaDomain.OperatorDb(ctx).GetAll()
+	if err != nil {
+		log.Error("GetOperators failed", zap.Error(err))
+		return nil, err
+	}
+
+	// Convert to proto response
+	protoOperators := make([]*coordinatorpb.Operator, len(operators))
+	for i, op := range operators {
+		protoOperators[i] = &coordinatorpb.Operator{
+			Id:   op.OperatorID.String(),
+			Name: op.OperatorName,
+		}
+	}
+
+	log.Info("GetOperators succeeded", zap.Int("count", len(operators)))
+
+	return &coordinatorpb.GetOperatorsResponse{
+		Operators: protoOperators,
+	}, nil
+}

go/pkg/sysdb/grpc/task_service.go

@@ -0,0 +1,57 @@
+package grpc
+
+import (
+	"context"
+
+	"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
+)
+
+func (s *Server) CreateTask(ctx context.Context, req *coordinatorpb.CreateTaskRequest) (*coordinatorpb.CreateTaskResponse, error) {
+	log.Info("CreateTask", zap.String("name", req.Name), zap.String("operator_name", req.OperatorName))
+
+	res, err := s.coordinator.CreateTask(ctx, req)
+	if err != nil {
+		log.Error("CreateTask failed", zap.Error(err))
+		return nil, err
+	}
+
+	return res, nil
+}
+
+func (s *Server) GetTaskByName(ctx context.Context, req *coordinatorpb.GetTaskByNameRequest) (*coordinatorpb.GetTaskByNameResponse, error) {
+	log.Info("GetTaskByName", zap.String("input_collection_id", req.InputCollectionId), zap.String("task_name", req.TaskName))
+
+	res, err := s.coordinator.GetTaskByName(ctx, req)
+	if err != nil {
+		log.Error("GetTaskByName failed", zap.Error(err))
+		return nil, err
+	}
+
+	return res, nil
+}
+
+func (s *Server) DeleteTask(ctx context.Context, req *coordinatorpb.DeleteTaskRequest) (*coordinatorpb.DeleteTaskResponse, error) {
+	log.Info("DeleteTask", zap.String("input_collection_id", req.InputCollectionId), zap.String("task_name", req.TaskName))
+
+	res, err := s.coordinator.DeleteTask(ctx, req)
+	if err != nil {
+		log.Error("DeleteTask failed", zap.Error(err))
+		return nil, err
+	}
+
+	return res, nil
+}
+
+func (s *Server) GetOperators(ctx context.Context, req *coordinatorpb.GetOperatorsRequest) (*coordinatorpb.GetOperatorsResponse, error) {
+	log.Info("GetOperators")
+
+	res, err := s.coordinator.GetOperators(ctx, req)
+	if err != nil {
+		log.Error("GetOperators failed", zap.Error(err))
+		return nil, err
+	}
+
+	return res, nil
+}

go/pkg/sysdb/metastore/db/dao/common.go

@@ -36,3 +36,11 @@ func (*MetaDomain) SegmentDb(ctx context.Context) dbmodel.ISegmentDb {
 func (*MetaDomain) SegmentMetadataDb(ctx context.Context) dbmodel.ISegmentMetadataDb {
 	return &segmentMetadataDb{dbcore.GetDB(ctx)}
 }
+
+func (*MetaDomain) TaskDb(ctx context.Context) dbmodel.ITaskDb {
+	return &taskDb{dbcore.GetDB(ctx)}
+}
+
+func (*MetaDomain) OperatorDb(ctx context.Context) dbmodel.IOperatorDb {
+	return &operatorDb{dbcore.GetDB(ctx)}
+}

go/pkg/sysdb/metastore/db/dao/operator.go

@@ -0,0 +1,60 @@
+package dao
+
+import (
+	"errors"
+
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
+	"github.com/google/uuid"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+)
+
+type operatorDb struct {
+	db *gorm.DB
+}
+
+var _ dbmodel.IOperatorDb = &operatorDb{}
+
+func (s *operatorDb) GetByName(operatorName string) (*dbmodel.Operator, error) {
+	var operator dbmodel.Operator
+	err := s.db.
+		Where("operator_name = ?", operatorName).
+		First(&operator).Error
+
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, nil
+		}
+		log.Error("GetOperatorByName failed", zap.Error(err))
+		return nil, err
+	}
+	return &operator, nil
+}
+
+func (s *operatorDb) GetByID(operatorID uuid.UUID) (*dbmodel.Operator, error) {
+	var operator dbmodel.Operator
+	err := s.db.
+		Where("operator_id = ?", operatorID).
+		First(&operator).Error
+
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, nil
+		}
+		log.Error("GetOperatorByID failed", zap.Error(err))
+		return nil, err
+	}
+	return &operator, nil
+}
+
+func (s *operatorDb) GetAll() ([]*dbmodel.Operator, error) {
+	var operators []*dbmodel.Operator
+	err := s.db.Find(&operators).Error
+
+	if err != nil {
+		log.Error("GetAllOperators failed", zap.Error(err))
+		return nil, err
+	}
+	return operators, nil
+}