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  }