fix(kbfile): reprocessing fixes (#246)

Because

- Old failed files can't be reprocessed because the previous pipeline
  version is used
- Even if reprocessing succeeds, failure reason isn't cleaned up from
  previous failure
- Console passes an empty string as the conversion pipeline when setting
  it as the default pipeline

This commit

- Adds a migration so failed files use the new pipeline
- Clears failure reason when reprocessing a file
- Sanitizes conversion pipeline input on catalog creation / update

Author: jvallesm

cmd/migration/main.go

@@ -1,17 +1,22 @@
 package main
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"os"
 
 	"github.com/golang-migrate/migrate/v4"
+	"go.uber.org/zap"
 
 	_ "github.com/golang-migrate/migrate/v4/database/postgres"
 	_ "github.com/golang-migrate/migrate/v4/source/file"
 
 	"github.com/instill-ai/artifact-backend/config"
 	"github.com/instill-ai/artifact-backend/pkg/db/migration"
+
+	database "github.com/instill-ai/artifact-backend/pkg/db"
+	logx "github.com/instill-ai/x/log"
 )
 
 func dbExistsOrCreate(databaseConfig config.DatabaseConfig) error {
@@ -54,15 +59,25 @@ func dbExistsOrCreate(databaseConfig config.DatabaseConfig) error {
 
 	return nil
 }
-func main() {
 
+func main() {
 	if err := config.Init(config.ParseConfigFlag()); err != nil {
 		panic(err)
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	logx.Debug = config.Config.Server.Debug
+	logger, _ := logx.GetZapLogger(ctx)
+	defer func() {
+		// can't handle the error due to https://github.com/uber-go/zap/issues/880
+		_ = logger.Sync()
+	}()
+
 	databaseConfig := config.Config.Database
 	if err := dbExistsOrCreate(databaseConfig); err != nil {
-		panic(err)
+		logger.Fatal("Checking database existence", zap.Error(err))
 	}
 
 	dsn := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?%s",
@@ -74,39 +89,70 @@ func main() {
 		"sslmode=disable",
 	)
 
-	migrateFolder, _ := os.Getwd()
+	db := database.GetSharedConnection().WithContext(ctx)
+	codeMigrator := &migration.CodeMigrator{
+		Logger: logger,
+		DB:     db,
+	}
+
+	defer func() { database.Close(db) }()
+
+	if err := runMigration(dsn, migration.TargetSchemaVersion, codeMigrator.Migrate, logger); err != nil {
+		logger.Fatal("Running migration", zap.Error(err))
+	}
+}
+
+func runMigration(
+	dsn string,
+	expectedVersion uint,
+	execCode func(version uint) error,
+	logger *zap.Logger,
+) error {
+	migrateFolder, err := os.Getwd()
+	if err != nil {
+		return fmt.Errorf("accessing base path: %w", err)
+	}
+
 	m, err := migrate.New(fmt.Sprintf("file:///%s/pkg/db/migration", migrateFolder), dsn)
 	if err != nil {
-		panic(err)
+		return fmt.Errorf("creating migration: %w", err)
 	}
 
-	expectedVersion := migration.TargetSchemaVersion
 	curVersion, dirty, err := m.Version()
 	if err != nil && curVersion != 0 {
-		panic(err)
+		return fmt.Errorf("getting current version: %w", err)
 	}
 
-	fmt.Printf("Expected migration version is %d\n", expectedVersion)
-	fmt.Printf("The current schema version is %d, and dirty flag is %t\n", curVersion, dirty)
+	logger.Info("Running migration",
+		zap.Uint("expectedVersion", expectedVersion),
+		zap.Uint("currentVersion", curVersion),
+		zap.Bool("dirty", dirty),
+	)
+
 	if dirty {
-		panic("the database's dirty flag is set, please fix it")
+		return fmt.Errorf("database is dirty, please fix it")
 	}
 
 	step := curVersion
 	for {
 		if expectedVersion <= step {
-			fmt.Printf("Migration to version %d complete\n", expectedVersion)
+			logger.Info("Migration completed", zap.Uint("expectedVersion", expectedVersion))
 			break
 		}
 
-		fmt.Printf("Step up to version %d\n", step+1)
+		logger.Info("Step up", zap.Uint("step", step+1))
 		if err := m.Steps(1); err != nil {
-			panic(err)
+			return fmt.Errorf("stepping up: %w", err)
+		}
+
+		if step, _, err = m.Version(); err != nil {
+			return fmt.Errorf("getting new version: %w", err)
 		}
 
-		step, _, err = m.Version()
-		if err != nil {
-			panic(err)
+		if err := execCode(step); err != nil {
+			return fmt.Errorf("running associated code: %w", err)
 		}
 	}
+
+	return nil
 }

pkg/db/migration/000025_bump_conversion_pipeline_on_failed_files.down.sql

@@ -0,0 +1,49 @@
+package convert000025
+
+import (
+	"fmt"
+
+	"gorm.io/gorm"
+
+	"github.com/instill-ai/artifact-backend/pkg/db/migration/convert"
+	"github.com/instill-ai/artifact-backend/pkg/repository"
+
+	artifactpb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha"
+)
+
+const batchSize = 100
+
+type BumpConversionPipeline struct {
+	convert.Basic
+}
+
+const (
+	oldConvertingPipe = "preset/indexing-advanced-convert-doc@v1.3.1"
+	newConvertingPipe = "preset/indexing-advanced-convert-doc@v1.3.2"
+)
+
+func (c *BumpConversionPipeline) Migrate() error {
+	files := make([]*repository.KnowledgeBaseFile, 0, batchSize)
+	q := c.DB.Select("uid").
+		Where("process_status = ?", artifactpb.FileProcessStatus_FILE_PROCESS_STATUS_FAILED.String()).
+		Where("extra_meta_data->>'converting_pipe' = ?", oldConvertingPipe)
+
+	return q.FindInBatches(&files, batchSize, func(tx *gorm.DB, _ int) error {
+		for _, f := range files {
+			// Update the process status of the files
+			updates := map[string]any{
+				// Clear previous failure reason
+				"extra_meta_data": gorm.Expr(
+					"COALESCE(extra_meta_data, '{}'::jsonb) || ?::jsonb",
+					fmt.Sprintf(`{"converting_pipe": "%s"}`, newConvertingPipe),
+				),
+			}
+
+			if err := tx.Model(f).Where("uid = ?", f.UID).Updates(updates).Error; err != nil {
+				return fmt.Errorf("updating record %s: %w", f.UID.String(), err)
+			}
+		}
+
+		return nil
+	}).Error
+}

pkg/db/migration/000025_bump_conversion_pipeline_on_failed_files.up.sql

@@ -0,0 +1,12 @@
+package convert
+
+import (
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+)
+
+// Basic contains the basic elements to execute a conversion migration.
+type Basic struct {
+	DB     *gorm.DB
+	Logger *zap.Logger
+}

pkg/db/migration/convert/convert000025/convert.go

@@ -1,4 +1,49 @@
 package migration
 
+import (
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+
+	"github.com/instill-ai/artifact-backend/pkg/db/migration/convert"
+	"github.com/instill-ai/artifact-backend/pkg/db/migration/convert/convert000025"
+)
+
 // TargetSchemaVersion determines the database schema version.
-const TargetSchemaVersion uint = 24
+const TargetSchemaVersion uint = 25
+
+type migration interface {
+	Migrate() error
+}
+
+// CodeMigrator orchestrates the execution of the code associated with the
+// different database migrations and holds their dependencies.
+type CodeMigrator struct {
+	Logger *zap.Logger
+
+	DB *gorm.DB
+}
+
+// Migrate executes custom code as part of a database migration. This code is
+// intended to be run only once and typically goes along a change
+// in the database schemas. Some use cases might be backfilling a table or
+// updating some existing records according to the schema changes.
+//
+// Note that the changes in the database schemas shouldn't be run here, only
+// code accompanying them.
+func (cm *CodeMigrator) Migrate(version uint) error {
+	var m migration
+
+	bc := convert.Basic{
+		DB:     cm.DB,
+		Logger: cm.Logger,
+	}
+
+	switch version {
+	case 25:
+		m = &convert000025.BumpConversionPipeline{Basic: bc}
+	default:
+		return nil
+	}
+
+	return m.Migrate()
+}

pkg/db/migration/convert/converter.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math/rand"
 	"regexp"
+	"slices"
 	"strings"
 	"time"
 
@@ -101,25 +102,11 @@ func (ph *PublicHandler) CreateCatalog(ctx context.Context, req *artifactpb.Crea
 	}
 
 	// Read conversion pipelines from request.
-	// TODO jvallesm: validate existence, permissions & recipe of provided
-	// pipelines.
-	convertingPipelines := req.GetConvertingPipelines()
-	for _, pipelineName := range convertingPipelines {
-		if _, err := service.PipelineReleaseFromName(pipelineName); err != nil {
-			err = fmt.Errorf("%w: invalid conversion pipeline format: %w", errorsx.ErrInvalidArgument, err)
-			return nil, errorsx.AddMessage(
-				err,
-				`Conversion pipeline must have the format "{namespaceID}/{pipelineID}@{version}"`,
-			)
-		}
+	convertingPipelines, err := sanitizeConvertingPipelines(req.GetConvertingPipelines())
+	if err != nil {
+		return nil, err
 	}
 
-	/*
-		if err := ph.service.ValidateConvertingPipelines(convertingPipelines); err != nil {
-			return nil, fmt.Errorf("validating pipelines: %w")
-		}
-	*/
-
 	// create catalog
 	dbData, err := ph.service.Repository().CreateKnowledgeBase(
 		ctx,
@@ -296,18 +283,20 @@ func (ph *PublicHandler) UpdateCatalog(ctx context.Context, req *artifactpb.Upda
 		return nil, fmt.Errorf("%w: no permission over catalog", errorsx.ErrUnauthorized)
 	}
 
+	convertingPipelines, err := sanitizeConvertingPipelines(req.GetConvertingPipelines())
+	if err != nil {
+		return nil, err
+	}
+
 	// update catalog
 	kb, err = ph.service.Repository().UpdateKnowledgeBase(
 		ctx,
 		req.GetCatalogId(),
 		ns.NsUID.String(),
 		repository.KnowledgeBase{
-			Description: req.GetDescription(),
-			Tags:        req.GetTags(),
-
-			// TODO jvallesm: validate existence, permissions & recipe of provided
-			// pipelines.
-			ConvertingPipelines: req.GetConvertingPipelines(),
+			Description:         req.GetDescription(),
+			Tags:                req.GetTags(),
+			ConvertingPipelines: convertingPipelines,
 		},
 	)
 	if err != nil {
@@ -349,12 +338,9 @@ func (ph *PublicHandler) UpdateCatalog(ctx context.Context, req *artifactpb.Upda
 		UsedStorage:    uint64(kb.Usage),
 	}
 
-	if len(kb.ConvertingPipelines) == 0 {
-		catalog.ConvertingPipelines = defaultConvertingPipelines()
-	}
-
 	return &artifactpb.UpdateCatalogResponse{Catalog: catalog}, nil
 }
+
 func (ph *PublicHandler) DeleteCatalog(ctx context.Context, req *artifactpb.DeleteCatalogRequest) (*artifactpb.DeleteCatalogResponse, error) {
 	logger, _ := logx.GetZapLogger(ctx)
 	authUID, err := getUserUIDFromContext(ctx)
@@ -518,3 +504,41 @@ func defaultConvertingPipelines() []string {
 		service.ConvertDocToMDPipeline.Name(),
 	}
 }
+
+// sanitizeConvertingPipelines validates an input array of strings that
+// represent the conversion pipelines of a catalog. It checks the string format
+// is correct.
+// TODO we also want to validate the existence of the pipelines, permissions of
+// the requester over that pipeline and the validity of its recipe.
+func sanitizeConvertingPipelines(pipelines []string) ([]string, error) {
+	validPipelines := make([]string, 0, len(pipelines))
+	for _, pipelineName := range pipelines {
+		// Console passes an empty string to reset the catalog conversion
+		// pipeline to the default one.
+		if pipelineName == "" {
+			continue
+		}
+
+		// Remove duplicates.
+		if slices.Contains(validPipelines, pipelineName) {
+			continue
+		}
+
+		if _, err := service.PipelineReleaseFromName(pipelineName); err != nil {
+			err = fmt.Errorf("%w: invalid conversion pipeline format: %w", errorsx.ErrInvalidArgument, err)
+			return nil, errorsx.AddMessage(
+				err,
+				`Conversion pipeline must have the format "{namespaceID}/{pipelineID}@{version}"`,
+			)
+		}
+
+		validPipelines = append(validPipelines, pipelineName)
+
+	}
+
+	if len(validPipelines) == 0 {
+		validPipelines = defaultConvertingPipelines()
+	}
+
+	return validPipelines, nil
+}