Enhance: implementing a Transaction-Based Replacement for MySQL Chunks

cmd/server/main.go

@@ -64,17 +64,6 @@ func buildRepositories(cfg config.Config) (repository.PaperRepository, repositor
 			return nil, nil, nil, err
 		}
 		return database.Papers, database.Chunks, database.Users, nil
-
-	case "json":
-		repo, err := repository.NewJSONRepository(cfg.DBPath)
-		if err != nil {
-			return nil, nil, nil, err
-		}
-		return repo, repository.NewJSONChunkRepository(repo), nil, nil
-
-	case "memory":
-		return repository.NewMemoryPaperRepository(), repository.NewMemoryChunkRepository(), nil, nil
-
 	default:
 		return nil, nil, nil, fmt.Errorf("unsupported CERULEAN_DB_DRIVER=%q; supported: mysql, json, memory", cfg.DBDriver)
 	}

internal/dao/chunk.go

@@ -3,6 +3,7 @@ package dao
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/CeruleanFlow/cerulean/internal/domain"
 	"github.com/CeruleanFlow/cerulean/internal/entity"
@@ -135,3 +136,56 @@ func fromJSONMap(values datatypes.JSONMap) map[string]string {
 
 	return out
 }
+
+func (d *ChunkDAO) ReplaceByPaperID(ctx context.Context, paperID string, chunks []domain.Chunk) error {
+	if strings.TrimSpace(paperID) == "" {
+		return fmt.Errorf("invalid paper id")
+	}
+
+	return d.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		if err := tx.Where("paper_id = ?", paperID).Delete(&entity.Chunk{}).Error; err != nil {
+			return fmt.Errorf("delete old chunks: %w", err)
+		}
+		if len(chunks) == 0 {
+			return nil
+		}
+
+		entities := make([]entity.Chunk, 0, len(chunks))
+
+		for _, chunk := range chunks {
+			if chunk.PaperID != paperID {
+				return fmt.Errorf("invalid paper id")
+			}
+			if chunk.PaperID == "" {
+				chunk.PaperID = paperID
+			}
+
+			entities = append(entities, domainChunkToEntity(chunk))
+		}
+
+		if err := tx.Create(&entities).Error; err != nil {
+			return fmt.Errorf("insert new chunks: %w", err)
+		}
+		return nil
+	})
+}
+
+func domainChunkToEntity(chunk domain.Chunk) entity.Chunk {
+	metadata := map[string]any{}
+	for k, v := range chunk.Metadata {
+		metadata[k] = v
+	}
+
+	return entity.Chunk{
+		ID:         chunk.ID,
+		PaperID:    chunk.PaperID,
+		PageNo:     chunk.PageNo,
+		ChunkIndex: chunk.Index,
+		Text:       chunk.Text,
+		ObjectKey:  chunk.ObjectKey,
+		VectorID:   chunk.VectorID,
+		Metadata:   metadata,
+		CreatedAt:  chunk.CreatedAt,
+		UpdatedAt:  chunk.UpdatedAt,
+	}
+}

internal/ingest/service.go

@@ -140,22 +140,15 @@ func (s *Service) runPDFTextIngest(ctx context.Context, job task.Task, paper dom
 		return fmt.Errorf("pdf text parser produced no chunks")
 	}
 
-	// delete old chunks in mysql
-	if err := s.chunks.DeleteByPaperID(ctx, paper.ID); err != nil {
-		return fmt.Errorf("delete old chunks from mysql: %w", err)
+	if err := s.chunks.ReplaceByPaperID(ctx, paper.ID, chunks); err != nil {
+		return fmt.Errorf("replace chunks in mysql: %w", err)
 	}
-	// delete old chunk index in es
+
 	if s.search != nil {
 		if err := s.search.DeleteByPaperID(ctx, paper.ID); err != nil {
 			return fmt.Errorf("delete old chunks from elasticsearch: %w", err)
 		}
-	}
-	// save new chunks in mysql
-	if err := s.chunks.UpsertMany(ctx, chunks); err != nil {
-		return fmt.Errorf("save chunks to mysql: %w", err)
-	}
-	// save new chunks index in es
-	if s.search != nil {
+
 		if err := s.search.IndexChunks(ctx, chunks); err != nil {
 			return fmt.Errorf("index chunks to elasticsearch: %w", err)
 		}

internal/repository/chunk_repository.go

@@ -8,13 +8,6 @@ import (
 	"github.com/CeruleanFlow/cerulean/internal/domain"
 )
 
-type ChunkRepository interface {
-	UpsertMany(ctx context.Context, chunks []domain.Chunk) error
-	List(ctx context.Context, filters map[string]string) ([]domain.Chunk, error)
-	ListByPaperID(ctx context.Context, paperID string) ([]domain.Chunk, error)
-	DeleteByPaperID(ctx context.Context, paperID string) error
-}
-
 type MemoryChunkRepository struct {
 	mu     sync.RWMutex
 	chunks map[string]domain.Chunk

internal/repository/repository.go

@@ -0,0 +1,15 @@
+package repository
+
+import (
+	"context"
+
+	"github.com/CeruleanFlow/cerulean/internal/domain"
+)
+
+type ChunkRepository interface {
+	UpsertMany(ctx context.Context, chunks []domain.Chunk) error
+	List(ctx context.Context, filters map[string]string) ([]domain.Chunk, error)
+	ListByPaperID(ctx context.Context, paperID string) ([]domain.Chunk, error)
+	DeleteByPaperID(ctx context.Context, paperID string) error
+	ReplaceByPaperID(ctx context.Context, paperID string, chunks []domain.Chunk) error
+}