importsdk: support multi-statement schema in size estimator

Author: GMHDBJD

pkg/importsdk/file_scanner.go

@@ -393,13 +393,13 @@ func (s *fileScanner) buildEstimateTableInfo(ctx context.Context, tblMeta *mydum
 	}
 	p := parser.New()
 	p.SetSQLMode(s.config.sqlMode)
-	stmt, err := p.ParseOneStmt(schemaSQL, "", "")
+	stmts, _, err := p.ParseSQL(schemaSQL)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	createStmt, ok := stmt.(*ast.CreateTableStmt)
-	if !ok {
-		return nil, errors.Errorf("schema file %s does not contain a CREATE TABLE statement", tblMeta.SchemaFile.FileMeta.Path)
+	createStmt, err := buildEstimateCreateTableStmt(stmts, tblMeta)
+	if err != nil {
+		return nil, err
 	}
 	tableInfo, err := ddl.BuildTableInfoFromAST(metabuild.NewContext(), createStmt)
 	if err != nil {
@@ -408,6 +408,49 @@ func (s *fileScanner) buildEstimateTableInfo(ctx context.Context, tblMeta *mydum
 	return tableInfo, nil
 }
 
+func buildEstimateCreateTableStmt(stmts []ast.StmtNode, tblMeta *mydump.MDTableMeta) (*ast.CreateTableStmt, error) {
+	var (
+		firstCreateStmt *ast.CreateTableStmt
+		createStmtCount int
+	)
+	for _, stmt := range stmts {
+		createStmt, ok := stmt.(*ast.CreateTableStmt)
+		if !ok {
+			continue
+		}
+		if firstCreateStmt == nil {
+			firstCreateStmt = createStmt
+		}
+		createStmtCount++
+		if estimateCreateTableStmtMatchesMeta(createStmt, tblMeta) {
+			return createStmt, nil
+		}
+	}
+	if createStmtCount == 1 {
+		return firstCreateStmt, nil
+	}
+	if createStmtCount == 0 {
+		return nil, errors.Errorf("schema file %s does not contain a CREATE TABLE statement", tblMeta.SchemaFile.FileMeta.Path)
+	}
+	return nil, errors.Errorf(
+		"schema file %s contains %d CREATE TABLE statements but none match table %s.%s",
+		tblMeta.SchemaFile.FileMeta.Path,
+		createStmtCount,
+		tblMeta.DB,
+		tblMeta.Name,
+	)
+}
+
+func estimateCreateTableStmtMatchesMeta(createStmt *ast.CreateTableStmt, tblMeta *mydump.MDTableMeta) bool {
+	if !strings.EqualFold(createStmt.Table.Name.String(), tblMeta.Name) {
+		return false
+	}
+	if createStmt.Table.Schema.String() == "" {
+		return true
+	}
+	return strings.EqualFold(createStmt.Table.Schema.String(), tblMeta.DB)
+}
+
 func sourceTypeToImportFormat(tp mydump.SourceType) (string, error) {
 	switch tp {
 	case mydump.SourceTypeCSV:

pkg/importsdk/file_scanner_test.go

@@ -266,6 +266,41 @@ func TestFileScanner(t *testing.T) {
 		require.Equal(t, estimate.Tables[0].SourceSize, estimate.TotalSourceSize)
 		require.Equal(t, estimate.Tables[0].TiKVSize, estimate.TotalTiKVSize)
 	})
+
+	t.Run("EstimateImportDataSizeMultiStatementSchema", func(t *testing.T) {
+		estimateDir := t.TempDir()
+		require.NoError(t, os.WriteFile(filepath.Join(estimateDir, "test_db-schema-create.sql"), []byte("CREATE DATABASE test_db;"), 0o644))
+		require.NoError(t, os.WriteFile(
+			filepath.Join(estimateDir, "test_db.users-schema.sql"),
+			[]byte(strings.Join([]string{
+				"CREATE DATABASE IF NOT EXISTS test_db;",
+				"USE test_db;",
+				"DROP TABLE IF EXISTS users;",
+				"CREATE TABLE users (id INT PRIMARY KEY, name VARCHAR(255), KEY idx_name (name));",
+			}, "\n")),
+			0o644,
+		))
+		require.NoError(t, os.WriteFile(
+			filepath.Join(estimateDir, "test_db.users.001.csv"),
+			[]byte("1,alice\n2,bob\n"),
+			0o644,
+		))
+
+		cfg := defaultSDKConfig()
+		cfg.skipInvalidFiles = true
+		estimateScanner, err := NewFileScanner(ctx, "file://"+estimateDir, db, cfg)
+		require.NoError(t, err)
+		defer estimateScanner.Close()
+
+		estimate, err := estimateScanner.EstimateImportDataSize(ctx)
+		require.NoError(t, err)
+		require.Len(t, estimate.Tables, 1)
+		require.Equal(t, "users", estimate.Tables[0].Table)
+		require.Positive(t, estimate.Tables[0].SourceSize)
+		require.Positive(t, estimate.Tables[0].TiKVSize)
+		require.Equal(t, estimate.Tables[0].SourceSize, estimate.TotalSourceSize)
+		require.Equal(t, estimate.Tables[0].TiKVSize, estimate.TotalTiKVSize)
+	})
 }
 
 func TestFileScannerWithEstimateRealSize(t *testing.T) {