github.com/cloudberrydb/gpbackup@v1.0.3-0.20240118031043-5410fd45eed6/restore/data_test.go (about) 1 package restore_test 2 3 import ( 4 "fmt" 5 "regexp" 6 "sort" 7 "strings" 8 9 "github.com/DATA-DOG/go-sqlmock" 10 "github.com/cloudberrydb/gp-common-go-libs/cluster" 11 "github.com/cloudberrydb/gpbackup/backup" 12 "github.com/cloudberrydb/gpbackup/history" 13 "github.com/cloudberrydb/gpbackup/options" 14 "github.com/cloudberrydb/gpbackup/restore" 15 "github.com/cloudberrydb/gpbackup/utils" 16 "github.com/jackc/pgconn" 17 18 . "github.com/onsi/ginkgo/v2" 19 . "github.com/onsi/gomega" 20 ) 21 22 var _ = Describe("restore/data tests", func() { 23 Describe("CopyTableIn", func() { 24 BeforeEach(func() { 25 utils.SetPipeThroughProgram(utils.PipeThroughProgram{Name: "cat", OutputCommand: "cat -", InputCommand: "cat -", Extension: ""}) 26 backup.SetPluginConfig(nil) 27 _ = cmdFlags.Set(options.PLUGIN_CONFIG, "") 28 backup.SetCluster(&cluster.Cluster{ContentIDs: []int{-1, 0, 1, 2}}) 29 restore.SetBackupConfig(&history.BackupConfig{}) 30 }) 31 It("will restore a table from its own file with gzip compression", func() { 32 utils.SetPipeThroughProgram(utils.PipeThroughProgram{Name: "gzip", OutputCommand: "gzip -c -1", InputCommand: "gzip -d -c", Extension: ".gz"}) 33 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM 'cat <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456.gz | gzip -d -c' WITH CSV DELIMITER ',' ON SEGMENT") 34 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 35 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456.gz" 36 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 37 38 Expect(err).ShouldNot(HaveOccurred()) 39 }) 40 It("will restore a table from its own file with zstd compression", func() { 41 utils.SetPipeThroughProgram(utils.PipeThroughProgram{Name: "zstd", OutputCommand: "zstd --compress -1 -c", InputCommand: "zstd --decompress -c", Extension: ".zst"}) 42 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM 'cat <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456.zst | zstd --decompress -c' WITH CSV DELIMITER ',' ON SEGMENT") 43 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 44 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456.zst" 45 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 46 47 Expect(err).ShouldNot(HaveOccurred()) 48 }) 49 It("will restore a table from its own file without compression", func() { 50 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM 'cat <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456 | cat -' WITH CSV DELIMITER ',' ON SEGMENT") 51 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 52 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456" 53 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 54 55 Expect(err).ShouldNot(HaveOccurred()) 56 }) 57 It("will restore a table from a single data file", func() { 58 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM 'cat <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456 | cat -' WITH CSV DELIMITER ',' ON SEGMENT") 59 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 60 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456" 61 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, true, 0) 62 63 Expect(err).ShouldNot(HaveOccurred()) 64 }) 65 It("will restore a table from its own file with gzip compression using a plugin", func() { 66 utils.SetPipeThroughProgram(utils.PipeThroughProgram{Name: "gzip", OutputCommand: "gzip -c -1", InputCommand: "gzip -d -c", Extension: ".gz"}) 67 _ = cmdFlags.Set(options.PLUGIN_CONFIG, "/tmp/plugin_config") 68 pluginConfig := utils.PluginConfig{ExecutablePath: "/tmp/fake-plugin.sh", ConfigPath: "/tmp/plugin_config"} 69 restore.SetPluginConfig(&pluginConfig) 70 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM '/tmp/fake-plugin.sh restore_data /tmp/plugin_config <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456.gz | gzip -d -c' WITH CSV DELIMITER ',' ON SEGMENT") 71 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 72 73 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456.gz" 74 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 75 76 Expect(err).ShouldNot(HaveOccurred()) 77 }) 78 It("will restore a table from its own file with zstd compression using a plugin", func() { 79 utils.SetPipeThroughProgram(utils.PipeThroughProgram{Name: "zstd", OutputCommand: "zstd --compress -1 -c", InputCommand: "zstd --decompress -c", Extension: ".zst"}) 80 _ = cmdFlags.Set(options.PLUGIN_CONFIG, "/tmp/plugin_config") 81 pluginConfig := utils.PluginConfig{ExecutablePath: "/tmp/fake-plugin.sh", ConfigPath: "/tmp/plugin_config"} 82 restore.SetPluginConfig(&pluginConfig) 83 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM '/tmp/fake-plugin.sh restore_data /tmp/plugin_config <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456.zst | zstd --decompress -c' WITH CSV DELIMITER ',' ON SEGMENT") 84 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 85 86 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456.zst" 87 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 88 89 Expect(err).ShouldNot(HaveOccurred()) 90 }) 91 It("will restore a table from its own file without compression using a plugin", func() { 92 _ = cmdFlags.Set(options.PLUGIN_CONFIG, "/tmp/plugin_config") 93 pluginConfig := utils.PluginConfig{ExecutablePath: "/tmp/fake-plugin.sh", ConfigPath: "/tmp/plugin_config"} 94 restore.SetPluginConfig(&pluginConfig) 95 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM '/tmp/fake-plugin.sh restore_data /tmp/plugin_config <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456.gz | cat -' WITH CSV DELIMITER ',' ON SEGMENT") 96 mock.ExpectExec(execStr).WillReturnResult(sqlmock.NewResult(10, 0)) 97 98 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_pipe_3456.gz" 99 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 100 101 Expect(err).ShouldNot(HaveOccurred()) 102 }) 103 It("will output expected error string from COPY ON SEGMENT failure", func() { 104 execStr := regexp.QuoteMeta("COPY public.foo(i,j) FROM PROGRAM 'cat <SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456 | cat -' WITH CSV DELIMITER ',' ON SEGMENT") 105 pgErr := &pgconn.PgError{ 106 Severity: "ERROR", 107 Code: "22P04", 108 Message: "value of distribution key doesn't belong to segment with ID 0, it belongs to segment with ID 1", 109 Where: "COPY foo, line 1: \"5\"", 110 } 111 mock.ExpectExec(execStr).WillReturnError(pgErr) 112 filename := "<SEG_DATA_DIR>/backups/20170101/20170101010101/gpbackup_<SEGID>_20170101010101_3456" 113 _, err := restore.CopyTableIn(connectionPool, "public.foo", "(i,j)", filename, false, 0) 114 115 Expect(err).To(HaveOccurred()) 116 Expect(err.Error()).To(Equal("Error loading data into table public.foo: " + 117 "COPY foo, line 1: \"5\": " + 118 "ERROR: value of distribution key doesn't belong to segment with ID 0, it belongs to segment with ID 1 (SQLSTATE 22P04)")) 119 }) 120 }) 121 Describe("CheckRowsRestored", func() { 122 var ( 123 expectedRows int64 = 10 124 name = "public.foo" 125 ) 126 It("does nothing if the number of rows match ", func() { 127 err := restore.CheckRowsRestored(10, expectedRows, name) 128 Expect(err).ToNot(HaveOccurred()) 129 }) 130 It("returns an error if the numbers of rows do not match", func() { 131 err := restore.CheckRowsRestored(5, expectedRows, name) 132 Expect(err).To(HaveOccurred()) 133 Expect(err.Error()).To(Equal("Expected to restore 10 rows to table public.foo, but restored 5 instead")) 134 }) 135 }) 136 }) 137 138 func batchMapToString(m map[int]map[int]int) string { 139 outer := make([]string, len(m)) 140 for num, batch := range m { 141 outer[num] = fmt.Sprintf("%d: %s", num, contentMapToString(batch)) 142 } 143 return strings.Join(outer, "; ") 144 } 145 146 func contentMapToString(m map[int]int) string { 147 inner := make([]string, len(m)) 148 index := 0 149 for orig, dest := range m { 150 inner[index] = fmt.Sprintf("%d:%d", orig, dest) 151 index++ 152 } 153 sort.Strings(inner) 154 return fmt.Sprintf("{%s}", strings.Join(inner, ", ")) 155 }