Multiple import data COPY command error handling improvements. (#220)

        - Ignore Unique constraint violation in case of transactional mode
        - Add ROWS_PER_TRANSACTION in generated COPY command
        - Added extra layer of check before enabling upsert and transaction mode
          to findout if it is supported on that server version.
        - Improved retrying logic to retry for based on retryable vs
          non-retryable errors
        - Improved logging for automation tests in case of failure
        - Enabled upsert mode by default

Author: sanyamsinghal

.github/workflows/mysql-migtests.yml

@@ -16,6 +16,8 @@ jobs:
     - name: Run installer script
       run: |
         yes | ./installer_scripts/install-yb-voyager local
+      env:
+        ON_INSTALLER_ERROR_OUTPUT_LOG: Y
 
     - name: Start MySQL
       run: |

.github/workflows/pg-migtests.yml

@@ -32,6 +32,8 @@ jobs:
       run: |
         cd installer_scripts
         yes | ./install-yb-voyager local
+      env:
+        ON_INSTALLER_ERROR_OUTPUT_LOG: Y
 
     - name: Test PostgreSQL Connection
       run: |

installer_scripts/install-yb-voyager

@@ -45,6 +45,10 @@ on_exit() {
 		echo "Done!"
 	else
 		echo "Script failed. Check log file ${LOG_FILE} ."
+		if [[ ${ON_INSTALLER_ERROR_OUTPUT_LOG:-N} = "Y" ]]
+		then
+			sudo cat ${LOG_FILE}
+		fi
 	fi
 }
 
@@ -319,13 +323,8 @@ update_bashrc() {
 
 create_yb_session_vars_file() {
 	vars_file_name="/etc/yb-voyager/ybSessionVariables.sql"
-	if [ -f "$vars_file_name" ]
-	then
-		output "No need to create ybSessionVariables.sql again. Skipping."
-		return
-	fi
 	sudo mkdir -p /etc/yb-voyager
-	sudo wget -qO $vars_file_name https://github.com/yugabyte/yb-voyager/raw/main/yb-voyager/files/ybSessionVariables.sql
+	sudo touch $vars_file_name
 }
 
 create_base_ora2pg_conf_file() {

yb-voyager/cmd/constants.go

@@ -33,23 +33,27 @@ const (
 	LAST_SPLIT_NUM              = 0
 	SPLIT_INFO_PATTERN          = "[0-9]*.[0-9]*.[0-9]*.[0-9]*"
 	LAST_SPLIT_PATTERN          = "0.[0-9]*.[0-9]*.[0-9]*"
-	COPY_MAX_RETRY_COUNT        = 5
-	MAX_SLEEP_SECOND            = 10
+	COPY_MAX_RETRY_COUNT        = 10
+	MAX_SLEEP_SECOND            = 60
 )
 
-var IMPORT_SESSION_SETTERS = []string{
-	"SET client_encoding TO 'UTF8';",
-	// Disable transactions to improve ingestion throughput.
-	"SET yb_disable_transactional_writes to true;",
-	//Disable triggers or fkeys constraint checks.
-	"SET session_replication_role TO replica;",
-	// Enable UPSERT mode instead of normal inserts into a table.
-	"SET yb_enable_upsert_mode to true;",
-}
+// import session parameters
+const (
+	SET_CLIENT_ENCODING_TO_UTF8           = "SET client_encoding TO 'UTF8'"
+	SET_SESSION_REPLICATE_ROLE_TO_REPLICA = "SET session_replication_role TO replica" //Disable triggers or fkeys constraint checks.
+	SET_YB_ENABLE_UPSERT_MODE             = "SET yb_enable_upsert_mode to true"
+	SET_YB_DISABLE_TRANSACTIONAL_WRITES   = "SET yb_disable_transactional_writes to true" // Disable transactions to improve ingestion throughput.
+)
 
 var supportedSourceDBTypes = []string{ORACLE, MYSQL, POSTGRESQL}
 
 var validSSLModes = map[string][]string{
 	"mysql":      {"disable", "prefer", "require", "verify-ca", "verify-full"},
 	"postgresql": {"disable", "allow", "prefer", "require", "verify-ca", "verify-full"},
 }
+
+var NonRetryCopyErrors = []string{
+	"Sending too long RPC message",
+	"invalid input syntax",
+	"violates unique constraint",
+}

yb-voyager/cmd/import.go

@@ -147,8 +147,8 @@ func registerCommonImportFlags(cmd *cobra.Command) {
 			"For example: \"host1:port1,host2:port2\" or \"host1,host2\"\n"+
 			"Note: use-public-ip flag will be ignored if this is used.")
 
-	cmd.Flags().BoolVar(&enableUpsert, "enable-upsert", false,
-		"true - to enable upsert for insert in target tables (default false)")
+	cmd.Flags().BoolVar(&enableUpsert, "enable-upsert", true,
+		"true - to enable upsert for insert in target tables")
 
 	// flag existence depends on fix of this gh issue: https://github.com/yugabyte/yugabyte-db/issues/12464
 	cmd.Flags().BoolVar(&disableTransactionalWrites, "disable-transactional-writes", false,

yb-voyager/cmd/importData.go

@@ -838,10 +838,12 @@ func doOneImport(task *SplitFileImportTask, connPool *tgtdb.ConnectionPool) {
 
 		// copyCommand is empty when there are no rows for that table
 		if copyCommand != "" {
-			var rowsCount int64
+			copyCommand = fmt.Sprintf(copyCommand, (task.OffsetEnd - task.OffsetStart))
+			log.Infof("COPY command: %s", copyCommand)
+			var rowsAffected int64
 			var copyErr error
 			// if retry=n, total try call will be n+1
-			copyRetryCount := COPY_MAX_RETRY_COUNT + 1
+			remainingRetries := COPY_MAX_RETRY_COUNT + 1
 
 			copyErr = connPool.WithConn(func(conn *pgx.Conn) (bool, error) {
 				// reset the reader to begin for every call
@@ -855,39 +857,44 @@ func doOneImport(task *SplitFileImportTask, connPool *tgtdb.ConnectionPool) {
 					}
 				}
 				res, err := conn.PgConn().CopyFrom(context.Background(), reader, copyCommand)
-				rowsCount = res.RowsAffected()
+				rowsAffected = res.RowsAffected()
 
+				if err != nil && utils.InsensitiveSliceContains(NonRetryCopyErrors, err.Error()) {
+					return false, err
+				}
+
+				/*
+					Note: If a user retries after deleting some row(s) from a batch,
+					yb-voyager will never be able to mark the batch as completed
+					github issue: https://github.com/yugabyte/yb-voyager/issues/223
+				*/
 				if err != nil {
 					log.Warnf("COPY FROM file %q: %s", inProgressFilePath, err)
-					if !strings.Contains(err.Error(), "violates unique constraint") {
-						log.Errorf("RETRYING.. COPY %q FROM file %q due to encountered error: %v ", task.TableName, inProgressFilePath, err)
-						duration := time.Duration(math.Min(MAX_SLEEP_SECOND, math.Pow(2, float64(COPY_MAX_RETRY_COUNT+1-copyRetryCount))))
-						log.Infof("sleep for duration %d before retrying...", duration)
-						time.Sleep(time.Second * duration) // delay for 1 sec before retrying
-						copyRetryCount--
-						return copyRetryCount > 0, err
+					log.Errorf("RETRYING.. COPY %q FROM file %q due to encountered error: %v ", task.TableName, inProgressFilePath, err)
+
+					remainingRetries--
+					if remainingRetries > 0 {
+						retryNum := COPY_MAX_RETRY_COUNT + 1 - remainingRetries
+						duration := time.Duration(math.Min(MAX_SLEEP_SECOND, 10*float64(retryNum)))
+						log.Infof("sleep for duration %d before retrying the file %s for %d time...",
+							duration, inProgressFilePath, retryNum)
+						time.Sleep(time.Second * duration)
 					}
+					return remainingRetries > 0, err
 				}
 
 				return false, err
 			})
 
-			log.Infof("%q => %d rows affected", copyCommand, rowsCount)
+			log.Infof("%q => %d rows affected", copyCommand, rowsAffected)
 			if copyErr != nil {
-				log.Warnf("COPY FROM file %q: %s", inProgressFilePath, copyErr)
-				if !strings.Contains(copyErr.Error(), "violates unique constraint") {
+				if !disableTransactionalWrites && strings.Contains(copyErr.Error(), "violates unique constraint") {
+					log.Infof("Ignoring encountered Error: %v, Assuming batch is already imported due to transactional mode", copyErr)
+				} else {
 					utils.ErrExit("COPY %q FROM file %q: %s", task.TableName, inProgressFilePath, copyErr)
-				} else { //in case of unique key violation error take row count from the split task
-					rowsCount = task.OffsetEnd - task.OffsetStart
-					log.Infof("got error:%v, assuming affected rows count %v for %q", copyErr, rowsCount, task.TableName)
 				}
 			}
 
-			if rowsCount != task.OffsetEnd-task.OffsetStart {
-				// TODO: print info/details about missed rows on the screen after progress bar is complete
-				log.Warnf("Expected to import %v records from %s. Imported %v.",
-					task.OffsetEnd-task.OffsetStart, inProgressFilePath, rowsCount)
-			}
 			incrementImportProgressBar(task.TableName, inProgressFilePath)
 		}
 		doneFilePath := getDoneFilePath(task)
@@ -1001,6 +1008,7 @@ func extractCopyStmtForTable(table string, fileToSearchIn string) {
 			utils.ErrExit("error while readline for extraction of copy stmt from file %q: %v", fileToSearchIn, err)
 		}
 		if copyCommandRegex.MatchString(line) {
+			line = strings.Trim(line, ";") + ` WITH (ROWS_PER_TRANSACTION %v)`
 			copyTableFromCommands[table] = line
 			log.Infof("copyTableFromCommand for table %q is %q", table, line)
 			return
@@ -1038,43 +1046,75 @@ func getProgressAmount(filePath string) int64 {
 }
 
 func getYBSessionInitScript() []string {
-	sessionVarsPath := "/etc/yb-voyager/ybSessionVariables.sql"
 	var sessionVars []string
-	disableTransactionalWritesCmd := fmt.Sprintf("SET yb_disable_transactional_writes to %v", disableTransactionalWrites)
-	enableUpsertCmd := fmt.Sprintf("SET yb_enable_upsert_mode to %v", enableUpsert)
-	defaultSessionVars := []string{
-		"SET client_encoding to 'UTF-8'",
-		"SET session_replication_role to replica",
-		disableTransactionalWritesCmd,
-		enableUpsertCmd,
+	if checkSessionVariableSupport(SET_CLIENT_ENCODING_TO_UTF8) {
+		sessionVars = append(sessionVars, SET_CLIENT_ENCODING_TO_UTF8)
+	}
+	if checkSessionVariableSupport(SET_SESSION_REPLICATE_ROLE_TO_REPLICA) {
+		sessionVars = append(sessionVars, SET_SESSION_REPLICATE_ROLE_TO_REPLICA)
 	}
 
+	if enableUpsert {
+		// upsert_mode parameters was introduced later than yb_disable_transactional writes in yb releases
+		// hence if upsert_mode is supported then its safe to assume yb_disable_transactional_writes is already there
+		if checkSessionVariableSupport(SET_YB_ENABLE_UPSERT_MODE) {
+			sessionVars = append(sessionVars, SET_YB_ENABLE_UPSERT_MODE)
+			// 	SET_YB_DISABLE_TRANSACTIONAL_WRITES is used only with & if upsert_mode is supported
+			if disableTransactionalWrites {
+				if checkSessionVariableSupport(SET_YB_DISABLE_TRANSACTIONAL_WRITES) {
+					sessionVars = append(sessionVars, SET_YB_DISABLE_TRANSACTIONAL_WRITES)
+				} else {
+					disableTransactionalWrites = false
+				}
+			}
+		} else {
+			log.Infof("Falling back to transactional inserts of batches during data import")
+		}
+	}
+
+	sessionVarsPath := "/etc/yb-voyager/ybSessionVariables.sql"
 	if !utils.FileOrFolderExists(sessionVarsPath) {
-		return defaultSessionVars
+		log.Infof("YBSessionInitScript: %v\n", sessionVars)
+		return sessionVars
 	}
 
 	varsFile, err := os.Open(sessionVarsPath)
 	if err != nil {
 		utils.PrintAndLog("Unable to open %s : %v. Using default values.", sessionVarsPath, err)
-		return defaultSessionVars
+		log.Infof("YBSessionInitScript: %v\n", sessionVars)
+		return sessionVars
 	}
 	defer varsFile.Close()
 	fileScanner := bufio.NewScanner(varsFile)
 
 	var curLine string
 	for fileScanner.Scan() {
 		curLine = strings.TrimSpace(fileScanner.Text())
-		sessionVars = append(sessionVars, curLine)
+		if curLine != "" && checkSessionVariableSupport(curLine) {
+			sessionVars = append(sessionVars, curLine)
+		}
 	}
+	log.Infof("YBSessionInitScript: %v\n", sessionVars)
+	return sessionVars
+}
 
-	//Only override the file if the flags are explicitly true (default false)
-	if enableUpsert {
-		sessionVars = append(sessionVars, enableUpsertCmd)
+func checkSessionVariableSupport(sqlStmt string) bool {
+	conn, err := pgx.Connect(context.Background(), target.GetConnectionUri())
+	if err != nil {
+		utils.ErrExit("error while creating connection for checking session parameter(%q) support: %v", sqlStmt, err)
 	}
-	if disableTransactionalWrites {
-		sessionVars = append(sessionVars, disableTransactionalWritesCmd)
+	defer conn.Close(context.Background())
+
+	_, err = conn.Exec(context.Background(), sqlStmt)
+	if err != nil {
+		if !strings.Contains(err.Error(), "unrecognized configuration parameter") {
+			utils.ErrExit("error while executing sqlStatement=%q: %v", sqlStmt, err)
+		} else {
+			log.Warnf("Warning: %q is not supported: %v", sqlStmt, err)
+		}
 	}
-	return sessionVars
+
+	return err == nil
 }
 
 func removeExcludeTables(tableList []string, excludeTableList []string) []string {

yb-voyager/cmd/importDataFileCommand.go

@@ -110,17 +110,18 @@ func prepareCopyCommands() {
 				if err != nil {
 					utils.ErrExit("opening datafile %q to prepare copy command: %v", err)
 				}
-				copyTableFromCommands[table] = fmt.Sprintf(`COPY %s(%s) FROM STDIN WITH (FORMAT %s, DELIMITER '%c', ESCAPE '%s', QUOTE '%s', HEADER)`,
+				copyTableFromCommands[table] = fmt.Sprintf(`COPY %s(%s) FROM STDIN WITH (FORMAT %s, DELIMITER '%c', ESCAPE '%s', QUOTE '%s', HEADER,`,
 					table, df.GetHeader(), fileFormat, []rune(delimiter)[0], fileOptsMap["escape_char"], fileOptsMap["quote_char"])
 			} else {
-				copyTableFromCommands[table] = fmt.Sprintf(`COPY %s FROM STDIN WITH (FORMAT %s, DELIMITER '%c', ESCAPE '%s', QUOTE '%s')`,
+				copyTableFromCommands[table] = fmt.Sprintf(`COPY %s FROM STDIN WITH (FORMAT %s, DELIMITER '%c', ESCAPE '%s', QUOTE '%s', `,
 					table, fileFormat, []rune(delimiter)[0], fileOptsMap["escape_char"], fileOptsMap["quote_char"])
 			}
 		} else if fileFormat == datafile.TEXT {
-			copyTableFromCommands[table] = fmt.Sprintf(`COPY %s FROM STDIN WITH (FORMAT %s, DELIMITER '%c')`, table, fileFormat, []rune(delimiter)[0])
+			copyTableFromCommands[table] = fmt.Sprintf(`COPY %s FROM STDIN WITH (FORMAT %s, DELIMITER '%c', `, table, fileFormat, []rune(delimiter)[0])
 		} else {
 			panic(fmt.Sprintf("File Type %q not implemented\n", fileFormat))
 		}
+		copyTableFromCommands[table] += ` ROWS_PER_TRANSACTION %v)`
 	}
 
 	log.Infof("copyTableFromCommands map: %+v", copyTableFromCommands)

yb-voyager/cmd/importSchemaYugabyteDB.go

@@ -72,11 +72,10 @@ func YugabyteDBImportSchema(target *tgtdb.Target, exportDir string) {
 				utils.ErrExit("Failed to run %q on target DB: %s", setSchemaQuery, err)
 			}
 
-			setClientEncQuery := IMPORT_SESSION_SETTERS[0]
-			log.Infof("Running query %q on the target DB", setClientEncQuery)
-			_, err = conn.Exec(context.Background(), setClientEncQuery)
+			log.Infof("Running query %q on the target DB", SET_CLIENT_ENCODING_TO_UTF8)
+			_, err = conn.Exec(context.Background(), SET_CLIENT_ENCODING_TO_UTF8)
 			if err != nil {
-				utils.ErrExit("Failed to run %q on target DB: %s", setClientEncQuery, err)
+				utils.ErrExit("Failed to run %q on target DB: %s", SET_CLIENT_ENCODING_TO_UTF8, err)
 			}
 		}
 

yb-voyager/files/ybSessionVariables.sql

@@ -135,7 +135,7 @@ func (pg *PostgreSQL) ExportDataPostProcessing(exportDir string, tablesProgressM
 	renameDataFiles(tablesProgressMetadata)
 	exportedRowCount := getExportedRowCount(tablesProgressMetadata)
 	dfd := datafile.Descriptor{
-		FileFormat:    datafile.CSV,
+		FileFormat:    datafile.TEXT,
 		TableRowCount: exportedRowCount,
 		Delimiter:     "\t",
 		HasHeader:     false,

yb-voyager/src/srcdb/postgres.go

@@ -3,7 +3,6 @@ package tgtdb
 import (
 	"context"
 	"math/rand"
-	"strings"
 	"sync"
 
 	"github.com/jackc/pgx/v4"
@@ -121,7 +120,7 @@ func (pool *ConnectionPool) getNextUriIndex() int {
 func (pool *ConnectionPool) initSession(conn *pgx.Conn) error {
 	for _, v := range pool.params.SessionInitScript {
 		_, err := conn.Exec(context.Background(), v)
-		if err != nil && !strings.Contains(err.Error(), "unrecognized configuration parameter") {
+		if err != nil {
 			return err
 		}
 	}

yb-voyager/src/tgtdb/conn_pool.go

@@ -278,3 +278,14 @@ func LookupIP(name string) []string {
 	}
 	return result
 }
+
+func InsensitiveSliceContains(slice []string, s string) bool {
+	for i := 0; i < len(slice); i++ {
+		if strings.Contains(strings.ToLower(s), strings.ToLower(slice[i])) {
+			log.Infof("string s=%q contains slice[i]=%q", s, slice[i])
+			return true
+		}
+	}
+	log.Infof("string s=%q did not match with any string in %v", s, slice)
+	return false
+}