Add models and migration for Create/Delete/Get Tasks

Author: tanujnay112

go/pkg/common/errors.go

@@ -48,6 +48,12 @@ 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")
+
+	// 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/coordinator/model"
+	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
+	"github.com/chroma-core/chroma/go/pkg/types"
+	"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()
+	var outputCollectionID string
+
+	// 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
+	}
+
+	// Create output collection and task transactionally
+	err = s.catalog.txImpl.Transaction(ctx, func(txCtx context.Context) error {
+		// Create the output collection
+		createCollection := &model.CreateCollection{
+			ID:           types.NewUniqueID(),
+			Name:         req.OutputCollectionName,
+			TenantID:     req.TenantId,
+			DatabaseName: req.Database,
+			GetOrCreate:  false, // Don't get existing, must be new
+			Ts:           types.Timestamp(time.Now().Unix()),
+		}
+
+		collection, _, err := s.catalog.CreateCollection(txCtx, createCollection, createCollection.Ts)
+		if err != nil {
+			log.Error("CreateTask: failed to create output collection", zap.Error(err))
+			return err
+		}
+		outputCollectionID = collection.ID.String()
+
+		// Create the task model
+		now := time.Now()
+		task := &dbmodel.Task{
+			ID:                 taskID,
+			Name:               req.Name,
+			TenantID:           req.TenantId,
+			DatabaseID:         databases[0].ID,
+			InputCollectionID:  req.InputCollectionId,
+			OutputCollectionID: outputCollectionID,
+			OperatorID:         operatorID,
+			OperatorParams:     req.Params,
+			CompletionOffset:   0,
+			LastRun:            nil,
+			NextRun:            nil, // Will be scheduled by task scheduler
+			MinRecordsForTask:  int64(req.MinRecordsForTask),
+			CurrentAttempts:    0,
+			CreatedAt:          now,
+			UpdatedAt:          now,
+			NextNonce:          nextNonce,
+			OldestWrittenNonce: nil,
+		}
+
+		// Insert task into database
+		err = s.catalog.metaDomain.TaskDb(txCtx).Insert(task)
+		if err != nil {
+			log.Error("CreateTask: failed to insert task", zap.Error(err))
+			return err
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	log.Info("Task created", zap.String("task_id", taskID.String()), zap.String("name", req.Name))
+
+	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) {
+	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 &coordinatorpb.GetTaskByNameResponse{}, nil
+	}
+
+	// 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_id", task.OutputCollectionID))
+
+	// 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.OutputCollectionID),
+		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) {
+	// lookup 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("DeleteTask failed", zap.Error(err))
+		return nil, err
+	}
+	if len(databases) == 0 {
+		log.Error("DeleteTask failed: database not found")
+		return nil, common.ErrDatabaseNotFound
+	}
+	err = s.catalog.metaDomain.TaskDb(ctx).SoftDelete(req.TenantId, databases[0].ID, req.TaskName)
+	if err != nil {
+		log.Error("DeleteTask failed", zap.Error(err))
+		return nil, err
+	}
+
+	log.Info("Task deleted", zap.String("tenant", req.TenantId), zap.String("database", req.Database), 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("tenant", req.TenantId), zap.String("database", req.Database), 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,50 @@
+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) {
+			log.Error("GetOperatorByName: operator not found", zap.String("operator_name", operatorName))
+			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,71 @@
+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 {
+			log.Error("Postgres Error")
+			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(tenantID string, databaseID string, taskName string) error {
+	return s.db.Table("tasks").
+		Where("tenant_id = ?", tenantID).
+		Where("database_id = ?", databaseID).
+		Where("task_name = ?", taskName).
+		Updates(map[string]interface{}{
+			"is_deleted": true,
+		}).Error
+}