Add models and migration for Create/Delete/Get Tasks

Author: tanujnay112

go/pkg/common/errors.go

@@ -48,6 +48,13 @@ 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")
+
+	// 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,177 @@
+package coordinator
+
+import (
+	"context"
+	"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) {
+	// 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.Error("CreateTask: task already exists")
+		return nil, common.ErrTaskAlreadyExists
+	}
+
+	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
+}

go/pkg/sysdb/grpc/task_service.go

@@ -0,0 +1,45 @@
+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
+}

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,49 @@
+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
+}

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

@@ -0,0 +1,85 @@
+package dao
+
+import (
+	"errors"
+
+	"github.com/chroma-core/chroma/go/pkg/common"
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/pingcap/log"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+)
+
+type taskDb struct {
+	db *gorm.DB
+}
+
+var _ dbmodel.ITaskDb = &taskDb{}
+
+func (s *taskDb) DeleteAll() error {
+	return s.db.Where("1 = 1").Delete(&dbmodel.Task{}).Error
+}
+
+func (s *taskDb) Insert(task *dbmodel.Task) error {
+	err := s.db.Create(task).Error
+	if err != nil {
+		log.Error("insert task failed", zap.Error(err))
+		var pgErr *pgconn.PgError
+		ok := errors.As(err, &pgErr)
+		if ok {
+			switch pgErr.Code {
+			case "23505":
+				log.Error("task already exists")
+				return common.ErrTaskAlreadyExists
+			default:
+				return err
+			}
+		}
+		return err
+	}
+	return nil
+}
+
+func (s *taskDb) GetByName(inputCollectionID string, taskName string) (*dbmodel.Task, error) {
+	var task dbmodel.Task
+	err := s.db.
+		Where("input_collection_id = ?", inputCollectionID).
+		Where("task_name = ?", taskName).
+		Where("is_deleted = ?", false).
+		First(&task).Error
+
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, nil
+		}
+		log.Error("GetTaskByName failed", zap.Error(err))
+		return nil, err
+	}
+	return &task, nil
+}
+
+func (s *taskDb) SoftDelete(inputCollectionID string, taskName string) error {
+	// Update task name and is_deleted in a single query
+	// Format: _deleted_<original_name>_<input_collection_id>_<task_id>
+	result := s.db.Exec(`
+		UPDATE tasks
+		SET task_name = CONCAT('_deleted_', task_name, '_', input_collection_id, '_', task_id::text),
+			is_deleted = true
+		WHERE input_collection_id = ?
+			AND task_name = ?
+			AND is_deleted = false
+	`, inputCollectionID, taskName)
+
+	if result.Error != nil {
+		log.Error("SoftDelete failed", zap.Error(result.Error))
+		return result.Error
+	}
+
+	// If no rows were affected, task was not found (or already deleted)
+	if result.RowsAffected == 0 {
+		return nil // Idempotent - no error if already deleted or not found
+	}
+
+	return nil
+}