feat: split sheet eagerly

backend/api/gitops/webhook.go

@@ -19,12 +19,16 @@ import (
 	"github.com/labstack/echo/v4"
 	"github.com/pkg/errors"
 
+	"github.com/bytebase/bytebase/backend/utils"
+
 	"github.com/bytebase/bytebase/backend/common"
 	"github.com/bytebase/bytebase/backend/common/log"
 	"github.com/bytebase/bytebase/backend/component/activity"
 	api "github.com/bytebase/bytebase/backend/legacyapi"
 	"github.com/bytebase/bytebase/backend/plugin/vcs"
 
+	sc "github.com/bytebase/bytebase/backend/component/sheet"
+
 	"github.com/bytebase/bytebase/backend/store"
 	storepb "github.com/bytebase/bytebase/proto/generated-go/store"
 	v1pb "github.com/bytebase/bytebase/proto/generated-go/v1"
@@ -192,13 +196,22 @@ func (s *Service) createIssueFromPRInfo(ctx context.Context, project *store.Proj
 		creatorName = common.FormatUserEmail(user.Email)
 	}
 
+	engine, err := s.getDatabaseEngineSample(ctx, project, vcsConnector)
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to get database engine")
+	}
+
 	var sheets []int
 	for _, change := range prInfo.changes {
-		sheet, err := s.store.CreateSheet(ctx, &store.SheetMessage{
+		sheet, err := sc.CreateSheet(ctx, s.store, &store.SheetMessage{
 			CreatorID:  creatorID,
 			ProjectUID: project.UID,
 			Title:      change.path,
 			Statement:  change.content,
+
+			Payload: &storepb.SheetPayload{
+				Engine: engine,
+			},
 		})
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to create sheet for file %s", change.path)
@@ -295,6 +308,64 @@ func (s *Service) createIssueFromPRInfo(ctx context.Context, project *store.Proj
 	return issue, nil
 }
 
+func (s *Service) getDatabaseEngineSample(
+	ctx context.Context,
+	project *store.ProjectMessage,
+	vcsConnector *store.VCSConnectorMessage,
+) (storepb.Engine, error) {
+	sample, err := func() (*store.DatabaseMessage, error) {
+		if dbg := vcsConnector.Payload.GetDatabaseGroup(); dbg != "" {
+			projectID, databaseGroupID, err := common.GetProjectIDDatabaseGroupID(dbg)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to get project id and database group id from %q", dbg)
+			}
+			if projectID != project.ResourceID {
+				return nil, errors.Errorf("project id %q in databaseGroup %q does not match project id %q in plan config", projectID, dbg, project.ResourceID)
+			}
+			databaseGroup, err := s.store.GetDatabaseGroup(ctx, &store.FindDatabaseGroupMessage{ProjectUID: &project.UID, ResourceID: &databaseGroupID})
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to get database group %q", databaseGroupID)
+			}
+			if databaseGroup == nil {
+				return nil, errors.Errorf("database group %q not found", databaseGroupID)
+			}
+			allDatabases, err := s.store.ListDatabases(ctx, &store.FindDatabaseMessage{ProjectID: &project.ResourceID})
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to list databases for project %q", project.ResourceID)
+			}
+
+			matchedDatabases, _, err := utils.GetMatchedAndUnmatchedDatabasesInDatabaseGroup(ctx, databaseGroup, allDatabases)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to get matched and unmatched databases in database group %q", databaseGroupID)
+			}
+			if len(matchedDatabases) == 0 {
+				return nil, errors.Errorf("no matched databases found in database group %q", databaseGroupID)
+			}
+			return matchedDatabases[0], nil
+		}
+		allDatabases, err := s.store.ListDatabases(ctx, &store.FindDatabaseMessage{ProjectID: &project.ResourceID})
+		if err != nil {
+			return nil, errors.Wrapf(err, "failed to list databases for project %q", project.ResourceID)
+		}
+		if len(allDatabases) == 0 {
+			return nil, errors.Errorf("no database in the project %q", project.ResourceID)
+		}
+		return allDatabases[0], nil
+	}()
+	if err != nil {
+		return 0, errors.Wrapf(err, "failed to get sample database")
+	}
+
+	instance, err := s.store.GetInstanceV2(ctx, &store.FindInstanceMessage{ResourceID: &sample.InstanceID})
+	if err != nil {
+		return 0, errors.Wrapf(err, "failed to get instance")
+	}
+	if instance == nil {
+		return 0, errors.Errorf("instance not found")
+	}
+	return instance.Engine, nil
+}
+
 func (s *Service) getChangeSteps(
 	ctx context.Context,
 	project *store.ProjectMessage,

backend/api/v1/rollout_service.go

@@ -23,6 +23,7 @@ import (
 	"github.com/bytebase/bytebase/backend/component/dbfactory"
 	"github.com/bytebase/bytebase/backend/component/ghost"
 	"github.com/bytebase/bytebase/backend/component/iam"
+	sc "github.com/bytebase/bytebase/backend/component/sheet"
 	"github.com/bytebase/bytebase/backend/component/state"
 	enterprise "github.com/bytebase/bytebase/backend/enterprise/api"
 	api "github.com/bytebase/bytebase/backend/legacyapi"
@@ -1578,11 +1579,22 @@ func (s *RolloutService) createPipeline(ctx context.Context, project *store.Proj
 			// HACK: statement is present because the task came from a database group target plan spec.
 			// we need to create the sheet and update payload.SheetID.
 			if c.Statement != "" {
-				sheet, err := s.store.CreateSheet(ctx, &store.SheetMessage{
+				instance, err := s.store.GetInstanceV2(ctx, &store.FindInstanceMessage{UID: &c.InstanceID})
+				if err != nil {
+					return nil, errors.Wrapf(err, "failed to get instance %v", c.InstanceID)
+				}
+				if instance == nil {
+					return nil, errors.Errorf("instance not found, id %v", c.InstanceID)
+				}
+				sheet, err := sc.CreateSheet(ctx, s.store, &store.SheetMessage{
 					CreatorID:  api.SystemBotID,
 					ProjectUID: project.UID,
 					Title:      fmt.Sprintf("Sheet for task %v", c.Name),
 					Statement:  c.Statement,
+
+					Payload: &storepb.SheetPayload{
+						Engine: instance.Engine,
+					},
 				})
 				if err != nil {
 					return nil, errors.Wrapf(err, "failed to create sheet for task %v", c.Name)

backend/api/v1/rollout_service_task.go

@@ -15,6 +15,7 @@ import (
 	"github.com/bytebase/bytebase/backend/common/log"
 	"github.com/bytebase/bytebase/backend/component/dbfactory"
 	"github.com/bytebase/bytebase/backend/component/ghost"
+	sc "github.com/bytebase/bytebase/backend/component/sheet"
 	enterprise "github.com/bytebase/bytebase/backend/enterprise/api"
 	api "github.com/bytebase/bytebase/backend/legacyapi"
 	"github.com/bytebase/bytebase/backend/plugin/db"
@@ -347,11 +348,15 @@ func getTaskCreatesFromCreateDatabaseConfig(ctx context.Context, s *store.Store,
 		if err != nil {
 			return nil, err
 		}
-		sheet, err := s.CreateSheet(ctx, &store.SheetMessage{
+		sheet, err := sc.CreateSheet(ctx, s, &store.SheetMessage{
 			CreatorID:  api.SystemBotID,
 			ProjectUID: project.UID,
 			Title:      fmt.Sprintf("Sheet for creating database %v", databaseName),
 			Statement:  statement,
+
+			Payload: &storepb.SheetPayload{
+				Engine: instance.Engine,
+			},
 		})
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to create database creation sheet")

backend/api/v1/sheet_service.go

@@ -11,6 +11,7 @@ import (
 	"github.com/bytebase/bytebase/backend/common"
 	"github.com/bytebase/bytebase/backend/component/config"
 	"github.com/bytebase/bytebase/backend/component/iam"
+	sc "github.com/bytebase/bytebase/backend/component/sheet"
 	enterprise "github.com/bytebase/bytebase/backend/enterprise/api"
 	api "github.com/bytebase/bytebase/backend/legacyapi"
 	"github.com/bytebase/bytebase/backend/store"
@@ -110,7 +111,7 @@ func (s *SheetService) CreateSheet(ctx context.Context, request *v1pb.CreateShee
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, fmt.Sprintf("failed to convert sheet: %v", err))
 	}
-	sheet, err := s.store.CreateSheet(ctx, storeSheetCreate)
+	sheet, err := sc.CreateSheet(ctx, s.store, storeSheetCreate)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, fmt.Sprintf("failed to create sheet: %v", err))
 	}
@@ -376,8 +377,8 @@ func convertToStoreSheetMessage(projectUID int, databaseUID *int, creatorID int,
 		Statement:   string(sheet.Content),
 		Payload:     &storepb.SheetPayload{},
 	}
+	sheetMessage.Payload.Engine = convertEngine(sheet.Engine)
 	if sheet.Payload != nil {
-		sheetMessage.Payload.Engine = convertEngine(sheet.Engine)
 		sheetMessage.Payload.DatabaseConfig = convertV1DatabaseConfig(sheet.Payload.DatabaseConfig)
 		sheetMessage.Payload.BaselineDatabaseConfig = convertV1DatabaseConfig(sheet.Payload.BaselineDatabaseConfig)
 	}

backend/component/sheet/sheet.go

@@ -0,0 +1,47 @@
+package sheet
+
+import (
+	"context"
+	"log/slog"
+	"strings"
+
+	"github.com/bytebase/bytebase/backend/common"
+	"github.com/bytebase/bytebase/backend/plugin/parser/base"
+	"github.com/bytebase/bytebase/backend/store"
+	storepb "github.com/bytebase/bytebase/proto/generated-go/store"
+)
+
+func CreateSheet(ctx context.Context, s *store.Store, sheet *store.SheetMessage) (*store.SheetMessage, error) {
+	if sheet.Payload == nil {
+		sheet.Payload = &storepb.SheetPayload{}
+	}
+	sheet.Payload.Commands = getSheetCommands(sheet.Payload.Engine, sheet.Statement)
+
+	return s.CreateSheet(ctx, sheet)
+}
+
+func getSheetCommands(engine storepb.Engine, statement string) []*storepb.SheetCommand {
+	if len(statement) > common.MaxSheetCheckSize {
+		return nil
+	}
+
+	singleSQLs, err := base.SplitMultiSQL(engine, statement)
+	if err != nil {
+		if !strings.Contains(err.Error(), "not supported") {
+			slog.Warn("failed to split multi sql", "engine", engine.String(), "statement", statement)
+		}
+		return nil
+	}
+
+	var sheetCommands []*storepb.SheetCommand
+	p := 0
+	for _, s := range singleSQLs {
+		np := p + len(s.Text)
+		sheetCommands = append(sheetCommands, &storepb.SheetCommand{
+			Start: int32(p),
+			End:   int32(np),
+		})
+		p = np
+	}
+	return sheetCommands
+}

backend/runner/rollbackrun/runner.go

@@ -17,6 +17,7 @@ import (
 	"github.com/bytebase/bytebase/backend/common/log"
 	"github.com/bytebase/bytebase/backend/component/config"
 	"github.com/bytebase/bytebase/backend/component/dbfactory"
+	sc "github.com/bytebase/bytebase/backend/component/sheet"
 	"github.com/bytebase/bytebase/backend/component/state"
 	api "github.com/bytebase/bytebase/backend/legacyapi"
 	"github.com/bytebase/bytebase/backend/plugin/db"
@@ -145,11 +146,15 @@ func (r *Runner) generateOracleRollbackSQL(ctx context.Context, task *store.Task
 		rollbackStatement = statementsBuffer.String()
 	}
 
-	sheet, err := r.store.CreateSheet(ctx, &store.SheetMessage{
+	sheet, err := sc.CreateSheet(ctx, r.store, &store.SheetMessage{
 		CreatorID:  api.SystemBotID,
 		ProjectUID: project.UID,
 		Title:      fmt.Sprintf("Sheet for rolling back task %v", task.ID),
 		Statement:  rollbackStatement,
+
+		Payload: &storepb.SheetPayload{
+			Engine: instance.Engine,
+		},
 	})
 	if err != nil {
 		slog.Error("failed to create database creation sheet", log.BBError(err))
@@ -228,11 +233,15 @@ func (r *Runner) generateMySQLRollbackSQL(ctx context.Context, task *store.TaskM
 		rollbackStatement = rollbackSQL
 	}
 
-	sheet, err := r.store.CreateSheet(ctx, &store.SheetMessage{
+	sheet, err := sc.CreateSheet(ctx, r.store, &store.SheetMessage{
 		CreatorID:  api.SystemBotID,
 		ProjectUID: project.UID,
 		Title:      fmt.Sprintf("Sheet for rolling back task %d", task.ID),
 		Statement:  rollbackStatement,
+
+		Payload: &storepb.SheetPayload{
+			Engine: instance.Engine,
+		},
 	})
 	if err != nil {
 		slog.Error("failed to create database creation sheet", log.BBError(err))

backend/server/onboard.go

@@ -10,6 +10,8 @@ import (
 	"github.com/pkg/errors"
 	"google.golang.org/protobuf/encoding/protojson"
 
+	sc "github.com/bytebase/bytebase/backend/component/sheet"
+
 	"github.com/bytebase/bytebase/backend/common"
 	api "github.com/bytebase/bytebase/backend/legacyapi"
 	"github.com/bytebase/bytebase/backend/plugin/advisor"
@@ -184,27 +186,35 @@ func (s *Server) generateOnboardingData(ctx context.Context, user *store.UserMes
 	}
 
 	// Create a schema update issue and start with creating the sheet for the schema update.
-	testSheet, err := s.store.CreateSheet(ctx, &store.SheetMessage{
+	testSheet, err := sc.CreateSheet(ctx, s.store, &store.SheetMessage{
 		CreatorID: api.SystemBotID,
 
 		ProjectUID:  project.UID,
 		DatabaseUID: &testDatabase.UID,
 
 		Title:     "Alter table to test sample instance for sample issue",
 		Statement: "ALTER TABLE employee ADD COLUMN IF NOT EXISTS email TEXT DEFAULT '';",
+
+		Payload: &storepb.SheetPayload{
+			Engine: testInstance.Engine,
+		},
 	})
 	if err != nil {
 		return errors.Wrapf(err, "failed to create test sheet for sample project")
 	}
 
-	prodSheet, err := s.store.CreateSheet(ctx, &store.SheetMessage{
+	prodSheet, err := sc.CreateSheet(ctx, s.store, &store.SheetMessage{
 		CreatorID: api.SystemBotID,
 
 		ProjectUID:  project.UID,
 		DatabaseUID: &prodDatabase.UID,
 
 		Title:     "Alter table to prod sample instance for sample issue",
 		Statement: "ALTER TABLE employee ADD COLUMN IF NOT EXISTS email TEXT DEFAULT '';",
+
+		Payload: &storepb.SheetPayload{
+			Engine: prodInstance.Engine,
+		},
 	})
 	if err != nil {
 		return errors.Wrapf(err, "failed to create prod sheet for sample project")

backend/store/sheet.go

@@ -193,6 +193,8 @@ func (s *Store) listSheets(ctx context.Context, find *FindSheetMessage) ([]*Shee
 }
 
 // CreateSheet creates a new sheet.
+// You should not use this function directly to create sheets.
+// Use CreateSheet in component/sheet instead.
 func (s *Store) CreateSheet(ctx context.Context, create *SheetMessage) (*SheetMessage, error) {
 	if create.Payload == nil {
 		create.Payload = &storepb.SheetPayload{}