Adding 5K tables integration test for schema migration (#1324)

Author: sm745052

.github/workflows/integration-tests-against-omni.yaml

@@ -40,6 +40,7 @@ jobs:
       MYSQLDB_FKACTION: test_foreign_key_action_data
       MYSQLDB_CHECK_CONSTRAINT: test_mysql_checkconstraint
       MYSQLDB_DATA_TYPES: test_mysql_data_types
+      MYSQLDB_5000_TABLES: test_mysql_5000_tables
       MYSQLPWD: root
 
       # sql server envs
@@ -109,6 +110,8 @@ jobs:
       - run: mysql -v -P 3306 --protocol=tcp -u root -proot < test_data/mysql_foreignkeyaction_dump.test.out
       - run: mysql -v -P 3306 --protocol=tcp -u root -proot < test_data/mysql_checkconstraint_dump.test.out
       - run: mysql -v -P 3306 --protocol=tcp -u root -proot < test_data/mysql_data_types_dump.test.out
+      - run: mysql -v -P 3306 --protocol=tcp -u root -proot -e "CREATE DATABASE test_mysql_5000_tables;"
+      - run: mysql -v -P 3306 --protocol=tcp -u root -proot test_mysql_5000_tables < test_data/mysql_5000_tables.sql
 
       # init sql server with test_data
       # since we use ubuntu-latest container, we should ensure that the path matches the latest from  https://packages.microsoft.com/config/ubuntu/

import_file/csv_data_test.go

@@ -136,7 +136,7 @@ func getCommonInfoSchemaMock(tableCount int) *MockInfoSchemaInterface {
 		MockProcessData: func(conv *internal.Conv, infoSchema common.InfoSchema, additionalAttributes internal.AdditionalDataAttributes) {
 		},
 		MockSetRowStats: func(conv *internal.Conv, infoSchema common.InfoSchema) {},
-		MockProcessTable: func(conv *internal.Conv, table common.SchemaAndName, infoSchema common.InfoSchema) (schema.Table, error) {
+		MockProcessTable: func(conv *internal.Conv, table common.SchemaAndName, infoSchema common.StandardInfoSchema) (schema.Table, error) {
 			return schema.Table{}, nil
 		},
 		MockGetIncludedSrcTablesFromConv: func(conv *internal.Conv) (schemaToTablesMap map[string]internal.SchemaDetails, err error) {
@@ -149,7 +149,7 @@ type MockInfoSchemaInterface struct {
 	MockGenerateSrcSchema            func(conv *internal.Conv, infoSchema common.InfoSchema, numWorkers int) (int, error)
 	MockProcessData                  func(conv *internal.Conv, infoSchema common.InfoSchema, additionalAttributes internal.AdditionalDataAttributes)
 	MockSetRowStats                  func(conv *internal.Conv, infoSchema common.InfoSchema)
-	MockProcessTable                 func(conv *internal.Conv, table common.SchemaAndName, infoSchema common.InfoSchema) (schema.Table, error)
+	MockProcessTable                 func(conv *internal.Conv, table common.SchemaAndName, infoSchema common.StandardInfoSchema) (schema.Table, error)
 	MockGetIncludedSrcTablesFromConv func(conv *internal.Conv) (schemaToTablesMap map[string]internal.SchemaDetails, err error)
 }
 
@@ -182,7 +182,7 @@ func (m *MockInfoSchemaInterface) ProcessData(conv *internal.Conv, infoSchema co
 func (m *MockInfoSchemaInterface) SetRowStats(conv *internal.Conv, infoSchema common.InfoSchema) {
 }
 
-func (m *MockInfoSchemaInterface) ProcessTable(conv *internal.Conv, table common.SchemaAndName, infoSchema common.InfoSchema) (schema.Table, error) {
+func (m *MockInfoSchemaInterface) ProcessTable(conv *internal.Conv, table common.SchemaAndName, infoSchema common.StandardInfoSchema) (schema.Table, error) {
 	return m.MockProcessTable(conv, table, infoSchema)
 }
 

internal/test_util.go

@@ -157,6 +157,10 @@ func assertSrcColDef(t *testing.T, conv *Conv, tableId string, expectedColDef, a
 		assert.NotEqual(t, colId, "")
 		actualCol := actualColDef[colId]
 		actualCol.Id = ""
+		col.DefaultValue.Value.ExpressionId = ""
+		col.GeneratedColumn.Value.ExpressionId = ""
+		actualCol.DefaultValue.Value.ExpressionId = ""
+		actualCol.GeneratedColumn.Value.ExpressionId = ""
 		assert.Equal(t, col, actualCol)
 	}
 }

sources/common/infoschema.go

@@ -34,14 +34,58 @@ type InfoSchema interface {
 	GetToDdl() ToDdl
 	GetTableName(schema string, tableName string) string
 	GetTables() ([]SchemaAndName, error)
-	GetColumns(conv *internal.Conv, table SchemaAndName, constraints map[string][]string, primaryKeys []string) (map[string]schema.Column, []string, error)
 	GetRowsFromTable(conv *internal.Conv, srcTable string) (interface{}, error)
 	GetRowCount(table SchemaAndName) (int64, error)
+	ProcessData(conv *internal.Conv, tableId string, srcSchema schema.Table, spCols []string, spSchema ddl.CreateTable, additionalAttributes internal.AdditionalDataAttributes) error
+
+}
+
+// StandardInfoSchema supports per-table fetching of metadata.
+type StandardInfoSchema interface {
+	InfoSchema
+	// GetColumns fetches column definitions for a single table.
+	// Returns a map of column ID to schema.Column, a list of column IDs in ordinal order, and any error.
+	GetColumns(conv *internal.Conv, table SchemaAndName, constraints map[string][]string, primaryKeys []string) (map[string]schema.Column, []string, error)
+	// GetConstraints fetches constraints (primary keys, check constraints, etc.) for a single table.
+	// Returns a list of primary key column names, a list of check constraints, a map of column name to list of constraint types, and any error.
 	GetConstraints(conv *internal.Conv, table SchemaAndName) ([]string, []schema.CheckConstraint, map[string][]string, error)
+	// GetForeignKeys fetches foreign key constraints for a single table.
+	// Returns a list of schema.ForeignKey and any error.
 	GetForeignKeys(conv *internal.Conv, table SchemaAndName) (foreignKeys []schema.ForeignKey, err error)
+	// GetIndexes fetches indexes for a single table.
+	// Returns a list of schema.Index and any error.
 	GetIndexes(conv *internal.Conv, table SchemaAndName, colNameIdMp map[string]string) ([]schema.Index, error)
-	ProcessData(conv *internal.Conv, tableId string, srcSchema schema.Table, spCols []string, spSchema ddl.CreateTable, additionalAttributes internal.AdditionalDataAttributes) error
+}
 
+// TableColumns contains column definitions and IDs for a table.
+type TableColumns struct {
+	ColDefs map[string]schema.Column
+	ColIds  []string
+}
+
+// TableConstraints contains constraints for a table.
+type TableConstraints struct {
+	PrimaryKeys       []string
+	CheckConstraints  []schema.CheckConstraint
+	ColumnConstraints map[string][]string // colName -> constraintTypes
+}
+
+// TableForeignKeys contains foreign keys for a table.
+type TableForeignKeys struct {
+	ForeignKeys []schema.ForeignKey
+}
+
+// BatchedInfoSchema supports bulk fetching of metadata for multiple tables.
+type BatchedInfoSchema interface {
+	InfoSchema
+	// GetColumnsBatch fetches column definitions for a batch of tables.
+	GetColumnsBatch(conv *internal.Conv, tables []SchemaAndName) (map[string]TableColumns, error)
+	// GetConstraintsBatch fetches constraints for a batch of tables.
+	GetConstraintsBatch(conv *internal.Conv, tables []SchemaAndName) (map[string]TableConstraints, error)
+	// GetForeignKeysBatch fetches foreign key constraints for a batch of tables.
+	GetForeignKeysBatch(conv *internal.Conv, tables []SchemaAndName) (map[string]TableForeignKeys, error)
+	// GetIndexesBatch fetches indexes for a batch of tables.
+	GetIndexesBatch(conv *internal.Conv, tables []SchemaAndName, colDefs map[string]TableColumns) (map[string][]schema.Index, error)
 }
 
 // SchemaAndName contains the schema and name for a table
@@ -65,7 +109,7 @@ type InfoSchemaInterface interface {
 	GenerateSrcSchema(conv *internal.Conv, infoSchema InfoSchema, numWorkers int) (int, error)
 	ProcessData(conv *internal.Conv, infoSchema InfoSchema, additionalAttributes internal.AdditionalDataAttributes)
 	SetRowStats(conv *internal.Conv, infoSchema InfoSchema)
-	ProcessTable(conv *internal.Conv, table SchemaAndName, infoSchema InfoSchema) (schema.Table, error)
+	ProcessTable(conv *internal.Conv, table SchemaAndName, infoSchema StandardInfoSchema) (schema.Table, error)
 	GetIncludedSrcTablesFromConv(conv *internal.Conv) (schemaToTablesMap map[string]internal.SchemaDetails, err error)
 }
 type InfoSchemaImpl struct{}
@@ -110,8 +154,81 @@ func (is *InfoSchemaImpl) GenerateSrcSchema(conv *internal.Conv, infoSchema Info
 		numWorkers = DefaultWorkers
 	}
 
+	var tableCount int
+	var e error
+
+	if bis, ok := infoSchema.(BatchedInfoSchema); ok {
+		tableCount, e = is.generateSrcSchemaBatched(conv, bis, tables, numWorkers)
+	} else if sis, ok := infoSchema.(StandardInfoSchema); ok {
+		tableCount, e = is.generateSrcSchemaStandard(conv, sis, tables, numWorkers)
+	} else {
+		return 0, fmt.Errorf("source does not support schema fetching")
+	}
+
+	if e != nil {
+		return 0, e
+	}
+
+	internal.ResolveForeignKeyIds(conv.SrcSchema)
+	return tableCount, nil
+}
+
+func (is *InfoSchemaImpl) generateSrcSchemaBatched(conv *internal.Conv, bis BatchedInfoSchema, tables []SchemaAndName, numWorkers int) (int, error) {
+	batchSize := 50
+	var batches [][]SchemaAndName
+	for i := 0; i < len(tables); i += batchSize {
+		end := i + batchSize
+		if end > len(tables) {
+			end = len(tables)
+		}
+		batches = append(batches, tables[i:end])
+	}
+
+	asyncProcessBatch := func(batch []SchemaAndName, mutex *sync.Mutex) task.TaskResult[[]SchemaAndName] {
+		tableCols, err := bis.GetColumnsBatch(conv, batch)
+		if err != nil {
+			return task.TaskResult[[]SchemaAndName]{Result: batch, Err: err}
+		}
+		tableConstraints, err := bis.GetConstraintsBatch(conv, batch)
+		if err != nil {
+			return task.TaskResult[[]SchemaAndName]{Result: batch, Err: err}
+		}
+		tableForeignKeys, err := bis.GetForeignKeysBatch(conv, batch)
+		if err != nil {
+			return task.TaskResult[[]SchemaAndName]{Result: batch, Err: err}
+		}
+		indexes, err := bis.GetIndexesBatch(conv, batch, tableCols)
+		if err != nil {
+			return task.TaskResult[[]SchemaAndName]{Result: batch, Err: err}
+		}
+
+		mutex.Lock()
+		for _, t := range batch {
+			name := bis.GetTableName(t.Schema, t.Name)
+			cols := tableCols[t.Name]
+			constraints := tableConstraints[t.Name]
+			fks := tableForeignKeys[t.Name]
+			
+			tableObj := BuildSchemaTable(t, name, cols.ColDefs, cols.ColIds, constraints.PrimaryKeys, constraints.CheckConstraints, indexes[t.Name], fks.ForeignKeys)
+			conv.SrcSchema[tableObj.Id] = tableObj
+		}
+		mutex.Unlock()
+
+		return task.TaskResult[[]SchemaAndName]{Result: batch, Err: nil}
+	}
+
+	r := task.RunParallelTasksImpl[[]SchemaAndName, []SchemaAndName]{}
+	_, e := r.RunParallelTasks(batches, numWorkers, asyncProcessBatch, true)
+
+	if e != nil {
+		return 0, e
+	}
+	return len(tables), nil
+}
+
+func (is *InfoSchemaImpl) generateSrcSchemaStandard(conv *internal.Conv, sis StandardInfoSchema, tables []SchemaAndName, numWorkers int) (int, error) {
 	asyncProcessTable := func(t SchemaAndName, mutex *sync.Mutex) task.TaskResult[SchemaAndName] {
-		table, e := is.ProcessTable(conv, t, infoSchema)
+		table, e := is.ProcessTable(conv, t, sis)
 		mutex.Lock()
 		conv.SrcSchema[table.Id] = table
 		mutex.Unlock()
@@ -120,13 +237,10 @@ func (is *InfoSchemaImpl) GenerateSrcSchema(conv *internal.Conv, infoSchema Info
 	}
 
 	r := task.RunParallelTasksImpl[SchemaAndName, SchemaAndName]{}
-	res, e := r.RunParallelTasks(tables, numWorkers, asyncProcessTable, false)
+	_, e := r.RunParallelTasks(tables, numWorkers, asyncProcessTable, false)
 	if e != nil {
-		logger.Log.Info(fmt.Sprintf("exiting due to error: %s , while processing schema for table %s\n", e, res))
 		return 0, e
 	}
-
-	internal.ResolveForeignKeyIds(conv.SrcSchema)
 	return len(tables), nil
 }
 
@@ -180,10 +294,9 @@ func (is *InfoSchemaImpl) SetRowStats(conv *internal.Conv, infoSchema InfoSchema
 	}
 }
 
-func (is *InfoSchemaImpl) ProcessTable(conv *internal.Conv, table SchemaAndName, infoSchema InfoSchema) (schema.Table, error) {
+func (is *InfoSchemaImpl) ProcessTable(conv *internal.Conv, table SchemaAndName, infoSchema StandardInfoSchema) (schema.Table, error) {
 	var t schema.Table
 	logger.Log.Info(fmt.Sprintf("processing schema for table %s", table))
-	tblId := internal.GenerateTableId()
 	primaryKeys, checkConstraints, constraints, err := infoSchema.GetConstraints(conv, table)
 	if err != nil {
 		return t, fmt.Errorf("couldn't get constraints for table %s.%s: %s", table.Schema, table.Name, err)
@@ -193,7 +306,6 @@ func (is *InfoSchemaImpl) ProcessTable(conv *internal.Conv, table SchemaAndName,
 		return t, fmt.Errorf("couldn't get foreign key constraints for table %s.%s: %s", table.Schema, table.Name, err)
 	}
 
-	table.Id = tblId
 	colDefs, colIds, err := infoSchema.GetColumns(conv, table, constraints, primaryKeys)
 	if err != nil {
 		return t, fmt.Errorf("couldn't get schema for table %s.%s: %s", table.Schema, table.Name, err)
@@ -209,11 +321,21 @@ func (is *InfoSchemaImpl) ProcessTable(conv *internal.Conv, table SchemaAndName,
 	}
 
 	name := infoSchema.GetTableName(table.Schema, table.Name)
+	return BuildSchemaTable(table, name, colDefs, colIds, primaryKeys, checkConstraints, indexes, foreignKeys), nil
+}
+
+// BuildSchemaTable constructs a schema.Table struct from fetched metadata.
+func BuildSchemaTable(table SchemaAndName, name string, colDefs map[string]schema.Column, colIds []string, primaryKeys []string, checkConstraints []schema.CheckConstraint, indexes []schema.Index, foreignKeys []schema.ForeignKey) schema.Table {
+	tblId := internal.GenerateTableId()
+	colNameIdMap := make(map[string]string)
+	for k, v := range colDefs {
+		colNameIdMap[v.Name] = k
+	}
 	var schemaPKeys []schema.Key
 	for _, k := range primaryKeys {
 		schemaPKeys = append(schemaPKeys, schema.Key{ColId: colNameIdMap[k]})
 	}
-	t = schema.Table{
+	return schema.Table{
 		Id:               tblId,
 		Name:             name,
 		Schema:           table.Schema,
@@ -224,9 +346,9 @@ func (is *InfoSchemaImpl) ProcessTable(conv *internal.Conv, table SchemaAndName,
 		CheckConstraints: checkConstraints,
 		Indexes:          indexes,
 		ForeignKeys:      foreignKeys}
-	return t, nil
 }
 
+
 // getIncludedSrcTablesFromConv fetches the list of tables
 // from the source database that need to be migrated.
 func (is *InfoSchemaImpl) GetIncludedSrcTablesFromConv(conv *internal.Conv) (schemaToTablesMap map[string]internal.SchemaDetails, err error) {

sources/common/mocks.go

@@ -35,7 +35,7 @@ func (mis *MockInfoSchema) GenerateSrcSchema(conv *internal.Conv, infoSchema Inf
 func (mis *MockInfoSchema) ProcessData(conv *internal.Conv, infoSchema InfoSchema, additionalAttributes internal.AdditionalDataAttributes) {
 }
 func (mis *MockInfoSchema) SetRowStats(conv *internal.Conv, infoSchema InfoSchema) {}
-func (mis *MockInfoSchema) processTable(conv *internal.Conv, table SchemaAndName, infoSchema InfoSchema) (schema.Table, error) {
+func (mis *MockInfoSchema) ProcessTable(conv *internal.Conv, table SchemaAndName, infoSchema StandardInfoSchema) (schema.Table, error) {
 	args := mis.Called(conv, table, infoSchema)
 	return args.Get(0).(schema.Table), args.Error(1)
 }