WIP

Signed-off-by: Tim Vaillancourt <[email protected]>

go/base/context.go

@@ -209,6 +209,7 @@ type MigrationContext struct {
 	CutOverCompleteFlag                    int64
 	InCutOverCriticalSectionFlag           int64
 	PanicAbort                             chan error
+	OptimizerHints                         sql.OptimizerHints
 
 	OriginalTableColumnsOnApplier    *sql.ColumnList
 	OriginalTableColumns             *sql.ColumnList

go/cmd/gh-ost/main.go

@@ -69,6 +69,7 @@ func main() {
 	flag.StringVar(&migrationContext.AlterStatement, "alter", "", "alter statement (mandatory)")
 	flag.BoolVar(&migrationContext.AttemptInstantDDL, "attempt-instant-ddl", false, "Attempt to use instant DDL for this migration first")
 	storageEngine := flag.String("storage-engine", "innodb", "Specify table storage engine (default: 'innodb'). When 'rocksdb': the session transaction isolation level is changed from REPEATABLE_READ to READ_COMMITTED.")
+	flag.StringVar(&migrationContext.OptimizerHints.ResourceGroup, "resource-group", "", "Optional MySQL Resource Group to be used for certain, heavy MySQL queries, eg: exact-rowcount SELECT query, etc")
 
 	flag.BoolVar(&migrationContext.CountTableRows, "exact-rowcount", false, "actually count table rows as opposed to estimate them (results in more accurate progress estimation)")
 	flag.BoolVar(&migrationContext.ConcurrentCountTableRows, "concurrent-rowcount", true, "(with --exact-rowcount), when true (default): count rows after row-copy begins, concurrently, and adjust row estimate later on; when false: first count rows, then start row copy")

go/logic/inspect.go

@@ -554,7 +554,11 @@ func (this *Inspector) CountTableRows(ctx context.Context) error {
 		return err
 	}
 
-	query := fmt.Sprintf(`select /* gh-ost */ count(*) as count_rows from %s.%s`, sql.EscapeName(this.migrationContext.DatabaseName), sql.EscapeName(this.migrationContext.OriginalTableName))
+	query := sql.BuildTableCountQuery(
+		sql.EscapeName(this.migrationContext.DatabaseName),
+		sql.EscapeName(this.migrationContext.OriginalTableName),
+		this.migrationContext.OptimizerHints,
+	)
 	var rowsEstimate int64
 	if err := conn.QueryRowContext(ctx, query).Scan(&rowsEstimate); err != nil {
 		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {

go/sql/builder.go

@@ -387,6 +387,7 @@ func buildUniqueKeyMinMaxValuesPreparedQuery(databaseName, tableName string, uni
 			uniqueKeyColumnOrder[i] = fmt.Sprintf("%s %s", uniqueKeyColumnNames[i], order)
 		}
 	}
+
 	query := fmt.Sprintf(`
 		select /* gh-ost %s.%s */ %s
 		from
@@ -395,7 +396,8 @@ func buildUniqueKeyMinMaxValuesPreparedQuery(databaseName, tableName string, uni
 		order by
 			%s
 		limit 1`,
-		databaseName, tableName, strings.Join(uniqueKeyColumnNames, ", "),
+		databaseName, tableName,
+		strings.Join(uniqueKeyColumnNames, ", "),
 		databaseName, tableName, uniqueKey.Name,
 		strings.Join(uniqueKeyColumnOrder, ", "),
 	)
@@ -530,3 +532,11 @@ func BuildDMLUpdateQuery(databaseName, tableName string, tableColumns, sharedCol
 	)
 	return result, sharedArgs, uniqueKeyArgs, nil
 }
+
+func BuildTableCountQuery(databaseName, originalTableName string, optimizerHints OptimizerHints) string {
+	return fmt.Sprintf(`select %s /* gh-ost */ count(*) as count_rows from %s.%s`,
+		optimizerHints,
+		databaseName,
+		originalTableName,
+	)
+}

go/sql/builder_test.go

@@ -726,3 +726,21 @@ func TestBuildDMLUpdateQuerySignedUnsigned(t *testing.T) {
 		test.S(t).ExpectTrue(reflect.DeepEqual(uniqueKeyArgs, []interface{}{uint8(253)}))
 	}
 }
+
+func TestBuildTableCountQuery(t *testing.T) {
+	databaseName := "test"
+	optimizerHints := OptimizerHints{}
+	{
+		test.S(t).ExpectEquals(
+			BuildTableCountQuery(databaseName, t.Name(), optimizerHints),
+			`select  /* gh-ost */ count(*) as count_rows from `+databaseName+`.`+t.Name(),
+		)
+	}
+	{
+		optimizerHints.ResourceGroup = "gh-ost"
+		test.S(t).ExpectEquals(
+			BuildTableCountQuery(databaseName, t.Name(), optimizerHints),
+			`select /*+ RESOURCE_GROUP(gh-ost) */ /* gh-ost */ count(*) as count_rows from `+databaseName+`.`+t.Name(),
+		)
+	}
+}

go/sql/hints.go

@@ -0,0 +1,39 @@
+package sql
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+)
+
+type OptimizerHints struct {
+	ResourceGroup    string `hint:"RESOURCE_GROUP"`
+	MaxExecutionTime int    `hint:"MAX_EXECUTION_TIME"`
+}
+
+func (oh OptimizerHints) String() (comment string) {
+	ts := reflect.TypeOf(oh)
+	if ts.NumField() == 0 {
+		return comment
+	}
+
+	vs := reflect.ValueOf(oh)
+	hintSlice := make([]string, 0, ts.NumField())
+	for i := 0; i < ts.NumField(); i++ {
+		fieldTag := ts.Field(i).Tag
+		value := vs.Field(i)
+		if hint, ok := fieldTag.Lookup("hint"); ok && !value.IsZero() {
+			switch value.Kind() {
+			case reflect.String:
+				hintSlice = append(hintSlice, fmt.Sprintf(`%s(%s)`, hint, value.String()))
+			default:
+				hintSlice = append(hintSlice, fmt.Sprintf(`%s(%d)`, hint, value.Int()))
+			}
+		}
+	}
+
+	if len(hintSlice) > 0 {
+		comment = fmt.Sprintf(`/*+ %s */`, strings.Join(hintSlice, " "))
+	}
+	return comment
+}

go/sql/hints_test.go

@@ -0,0 +1,15 @@
+package sql
+
+import (
+	"testing"
+
+	test "github.com/openark/golib/tests"
+)
+
+func TestOptimizerHintsString(t *testing.T) {
+	hints := OptimizerHints{
+		ResourceGroup:    t.Name(),
+		MaxExecutionTime: 1000,
+	}
+	test.S(t).ExpectEquals(hints.String(), `/*+ RESOURCE_GROUP(`+t.Name()+`) MAX_EXECUTION_TIME(1000) */`)
+}