github.com/tuhaihe/gpbackup@v1.0.3/testutils/functions.go

package testutils

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"

	"github.com/tuhaihe/gp-common-go-libs/gplog"
	"github.com/tuhaihe/gp-common-go-libs/operating"
	"github.com/jmoiron/sqlx"
	"github.com/pkg/errors"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/tuhaihe/gp-common-go-libs/cluster"
	"github.com/tuhaihe/gp-common-go-libs/dbconn"
	"github.com/tuhaihe/gp-common-go-libs/structmatcher"
	"github.com/tuhaihe/gp-common-go-libs/testhelper"
	"github.com/tuhaihe/gpbackup/backup"
	"github.com/tuhaihe/gpbackup/filepath"
	"github.com/tuhaihe/gpbackup/restore"
	"github.com/tuhaihe/gpbackup/toc"
	"github.com/tuhaihe/gpbackup/utils"
	"github.com/sergi/go-diff/diffmatchpatch"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
)

/*
 * Functions for setting up the test environment and mocking out variables
 */

func SetupTestEnvironment() (*dbconn.DBConn, sqlmock.Sqlmock, *Buffer, *Buffer, *Buffer) {
	connectionPool, mock, testStdout, testStderr, testLogfile := testhelper.SetupTestEnvironment()

	// Default if not set is GPDB version `5.1.0`
	envTestGpdbVersion := os.Getenv("TEST_GPDB_VERSION")
	if envTestGpdbVersion != "" {
		testhelper.SetDBVersion(connectionPool, envTestGpdbVersion)
	}

	SetupTestCluster()
	backup.SetVersion("0.1.0")
	return connectionPool, mock, testStdout, testStderr, testLogfile
}

func SetupTestCluster() *cluster.Cluster {
	testCluster := SetDefaultSegmentConfiguration()
	backup.SetCluster(testCluster)
	restore.SetCluster(testCluster)
	testFPInfo := filepath.NewFilePathInfo(testCluster, "", "20170101010101", "gpseg")
	backup.SetFPInfo(testFPInfo)
	restore.SetFPInfo(testFPInfo)
	return testCluster
}

func SetupTestDbConn(dbname string) *dbconn.DBConn {
	conn := dbconn.NewDBConnFromEnvironment(dbname)
	conn.MustConnect(1)
	return conn
}

// Connects to specific segment in utility mode
func SetupTestDBConnSegment(dbname string, port int, gpVersion dbconn.GPDBVersion) *dbconn.DBConn {

	if dbname == "" {
		gplog.Fatal(errors.New("No database provided"), "")
	}
	if port == 0 {
		gplog.Fatal(errors.New("No segment port provided"), "")
	}
	username := operating.System.Getenv("PGUSER")
	if username == "" {
		currentUser, _ := operating.System.CurrentUser()
		username = currentUser.Username
	}
	host := operating.System.Getenv("PGHOST")
	if host == "" {
		host, _ = operating.System.Hostname()
	}

	conn := &dbconn.DBConn{
		ConnPool: nil,
		NumConns: 0,
		Driver:   &dbconn.GPDBDriver{},
		User:     username,
		DBName:   dbname,
		Host:     host,
		Port:     port,
		Tx:       nil,
		Version:  dbconn.GPDBVersion{},
	}

	var gpRoleGuc string
	gpRoleGuc = "gp_role"

	connStr := fmt.Sprintf("postgres://%s@%s:%d/%s?sslmode=disable&statement_cache_capacity=0&%s=utility", conn.User, conn.Host, conn.Port, conn.DBName, gpRoleGuc)

	segConn, err := conn.Driver.Connect("pgx", connStr)
	if err != nil {
		gplog.FatalOnError(err)
	}
	conn.ConnPool = make([]*sqlx.DB, 1)
	conn.ConnPool[0] = segConn

	conn.Tx = make([]*sqlx.Tx, 1)
	conn.NumConns = 1
	version, err := dbconn.InitializeVersion(conn)
	if err != nil {
		gplog.FatalOnError(err)
	}
	conn.Version = version
	return conn
}

func SetDefaultSegmentConfiguration() *cluster.Cluster {
	configCoordinator := cluster.SegConfig{ContentID: -1, Hostname: "localhost", DataDir: "gpseg-1"}
	configSegOne := cluster.SegConfig{ContentID: 0, Hostname: "localhost", DataDir: "gpseg0"}
	configSegTwo := cluster.SegConfig{ContentID: 1, Hostname: "localhost", DataDir: "gpseg1"}
	return cluster.NewCluster([]cluster.SegConfig{configCoordinator, configSegOne, configSegTwo})
}

func SetupTestFilespace(connectionPool *dbconn.DBConn, testCluster *cluster.Cluster) {
	remoteOutput := testCluster.GenerateAndExecuteCommand("Creating filespace test directory",
		cluster.ON_HOSTS|cluster.INCLUDE_COORDINATOR,
		func(contentID int) string {
			return fmt.Sprintf("mkdir -p /tmp/test_dir")
		})
	if remoteOutput.NumErrors != 0 {
		Fail("Could not create filespace test directory on 1 or more hosts")
	}
	// Construct a filespace config like the one that gpfilespace generates
	filespaceConfigQuery := `COPY (SELECT hostname || ':' || dbid || ':/tmp/test_dir/' || preferred_role || content FROM gp_segment_configuration AS subselect) TO '/tmp/temp_filespace_config';`
	testhelper.AssertQueryRuns(connectionPool, filespaceConfigQuery)
	out, err := exec.Command("bash", "-c", "echo \"filespace:test_dir\" > /tmp/filespace_config").CombinedOutput()
	if err != nil {
		Fail(fmt.Sprintf("Cannot create test filespace configuration: %s: %s", out, err.Error()))
	}
	out, err = exec.Command("bash", "-c", "cat /tmp/temp_filespace_config >> /tmp/filespace_config").CombinedOutput()
	if err != nil {
		Fail(fmt.Sprintf("Cannot finalize test filespace configuration: %s: %s", out, err.Error()))
	}
	// Create the filespace and verify it was created successfully
	out, err = exec.Command("bash", "-c", "gpfilespace --config /tmp/filespace_config").CombinedOutput()
	if err != nil {
		Fail(fmt.Sprintf("Cannot create test filespace: %s: %s", out, err.Error()))
	}
	filespaceName := dbconn.MustSelectString(connectionPool, "SELECT fsname AS string FROM pg_filespace WHERE fsname = 'test_dir';")
	if filespaceName != "test_dir" {
		Fail("Filespace test_dir was not successfully created")
	}
}

func DestroyTestFilespace(connectionPool *dbconn.DBConn) {
	filespaceName := dbconn.MustSelectString(connectionPool, "SELECT fsname AS string FROM pg_filespace WHERE fsname = 'test_dir';")
	if filespaceName != "test_dir" {
		return
	}
	testhelper.AssertQueryRuns(connectionPool, "DROP FILESPACE test_dir")
	out, err := exec.Command("bash", "-c", "rm -rf /tmp/test_dir /tmp/filespace_config /tmp/temp_filespace_config").CombinedOutput()
	if err != nil {
		Fail(fmt.Sprintf("Could not remove test filespace directory and configuration files: %s: %s", out, err.Error()))
	}
}

func DefaultMetadata(objType string, hasPrivileges bool, hasOwner bool, hasComment bool, hasSecurityLabel bool) backup.ObjectMetadata {
	privileges := make([]backup.ACL, 0)
	if hasPrivileges {
		privileges = []backup.ACL{DefaultACLForType("testrole", objType)}
	}
	owner := ""
	if hasOwner {
		owner = "testrole"
	}
	comment := ""
	if hasComment {
		n := ""
		switch objType[0] {
		case 'A', 'E', 'I', 'O', 'U':
			n = "n"
		}
		comment = fmt.Sprintf("This is a%s %s comment.", n, strings.ToLower(objType))
	}
	securityLabelProvider := ""
	securityLabel := ""
	if hasSecurityLabel {
		securityLabelProvider = "dummy"
		securityLabel = "unclassified"
	}
	switch objType {
	case "DOMAIN":
		objType = "TYPE"
	case "FOREIGN SERVER":
		objType = "SERVER"
	case "MATERIALIZED VIEW":
		objType = "RELATION"
	case "SEQUENCE":
		objType = "RELATION"
	case "TABLE":
		objType = "RELATION"
	case "VIEW":
		objType = "RELATION"
	}
	return backup.ObjectMetadata{
		Privileges:            privileges,
		ObjectType:            objType,
		Owner:                 owner,
		Comment:               comment,
		SecurityLabelProvider: securityLabelProvider,
		SecurityLabel:         securityLabel,
	}

}

// objType should be an all-caps string like TABLE, INDEX, etc.
func DefaultMetadataMap(objType string, hasPrivileges bool, hasOwner bool, hasComment bool, hasSecurityLabel bool) backup.MetadataMap {
	return backup.MetadataMap{
		backup.UniqueID{ClassID: ClassIDFromObjectName(objType), Oid: 1}: DefaultMetadata(objType, hasPrivileges, hasOwner, hasComment, hasSecurityLabel),
	}
}

var objNameToClassID = map[string]uint32{
	"AGGREGATE":                 1255,
	"CAST":                      2605,
	"COLLATION":                 3456,
	"CONSTRAINT":                2606,
	"CONVERSION":                2607,
	"DATABASE":                  1262,
	"DOMAIN":                    1247,
	"EVENT TRIGGER":             3466,
	"EXTENSION":                 3079,
	"FOREIGN DATA WRAPPER":      2328,
	"FOREIGN SERVER":            1417,
	"FUNCTION":                  1255,
	"INDEX":                     2610,
	"LANGUAGE":                  2612,
	"OPERATOR CLASS":            2616,
	"OPERATOR FAMILY":           2753,
	"OPERATOR":                  2617,
	"PROTOCOL":                  7175,
	"RESOURCE GROUP":            6436,
	"RESOURCE QUEUE":            6026,
	"ROLE":                      1260,
	"RULE":                      2618,
	"SCHEMA":                    2615,
	"SEQUENCE":                  1259,
	"TABLE":                     1259,
	"TABLESPACE":                1213,
	"TEXT SEARCH CONFIGURATION": 3602,
	"TEXT SEARCH DICTIONARY":    3600,
	"TEXT SEARCH PARSER":        3601,
	"TEXT SEARCH TEMPLATE":      3764,
	"TRIGGER":                   2620,
	"TYPE":                      1247,
	"USER MAPPING":              1418,
	"VIEW":                      1259,
	"MATERIALIZED VIEW":         1259,
}

func ClassIDFromObjectName(objName string) uint32 {
	return objNameToClassID[objName]

}
func DefaultACLForType(grantee string, objType string) backup.ACL {
	return backup.ACL{
		Grantee:    grantee,
		Select:     objType == "PROTOCOL" || objType == "SEQUENCE" || objType == "TABLE" || objType == "VIEW" || objType == "FOREIGN TABLE" || objType == "MATERIALIZED VIEW",
		Insert:     objType == "PROTOCOL" || objType == "TABLE" || objType == "VIEW" || objType == "FOREIGN TABLE" || objType == "MATERIALIZED VIEW",
		Update:     objType == "SEQUENCE" || objType == "TABLE" || objType == "VIEW" || objType == "FOREIGN TABLE" || objType == "MATERIALIZED VIEW",
		Delete:     objType == "TABLE" || objType == "VIEW" || objType == "FOREIGN TABLE" || objType == "MATERIALIZED VIEW",
		Truncate:   objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		References: objType == "TABLE" || objType == "VIEW" || objType == "FOREIGN TABLE" || objType == "MATERIALIZED VIEW",
		Trigger:    objType == "TABLE" || objType == "VIEW" || objType == "FOREIGN TABLE" || objType == "MATERIALIZED VIEW",
		Usage:      objType == "LANGUAGE" || objType == "SCHEMA" || objType == "SEQUENCE" || objType == "FOREIGN DATA WRAPPER" || objType == "FOREIGN SERVER",
		Execute:    objType == "FUNCTION" || objType == "AGGREGATE",
		Create:     objType == "DATABASE" || objType == "SCHEMA" || objType == "TABLESPACE",
		Temporary:  objType == "DATABASE",
		Connect:    objType == "DATABASE",
	}
}

func DefaultACLForTypeWithGrant(grantee string, objType string) backup.ACL {
	return backup.ACL{
		Grantee:             grantee,
		SelectWithGrant:     objType == "PROTOCOL" || objType == "SEQUENCE" || objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		InsertWithGrant:     objType == "PROTOCOL" || objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		UpdateWithGrant:     objType == "SEQUENCE" || objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		DeleteWithGrant:     objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		TruncateWithGrant:   objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		ReferencesWithGrant: objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		TriggerWithGrant:    objType == "TABLE" || objType == "VIEW" || objType == "MATERIALIZED VIEW",
		UsageWithGrant:      objType == "LANGUAGE" || objType == "SCHEMA" || objType == "SEQUENCE" || objType == "FOREIGN DATA WRAPPER" || objType == "FOREIGN SERVER",
		ExecuteWithGrant:    objType == "FUNCTION",
		CreateWithGrant:     objType == "DATABASE" || objType == "SCHEMA" || objType == "TABLESPACE",
		TemporaryWithGrant:  objType == "DATABASE",
		ConnectWithGrant:    objType == "DATABASE",
	}
}

func DefaultACLWithout(grantee string, objType string, revoke ...string) backup.ACL {
	defaultACL := DefaultACLForType(grantee, objType)
	for _, priv := range revoke {
		switch priv {
		case "SELECT":
			defaultACL.Select = false
		case "INSERT":
			defaultACL.Insert = false
		case "UPDATE":
			defaultACL.Update = false
		case "DELETE":
			defaultACL.Delete = false
		case "TRUNCATE":
			defaultACL.Truncate = false
		case "REFERENCES":
			defaultACL.References = false
		case "TRIGGER":
			defaultACL.Trigger = false
		case "EXECUTE":
			defaultACL.Execute = false
		case "USAGE":
			defaultACL.Usage = false
		case "CREATE":
			defaultACL.Create = false
		case "TEMPORARY":
			defaultACL.Temporary = false
		case "CONNECT":
			defaultACL.Connect = false
		}
	}
	return defaultACL
}

func DefaultACLWithGrantWithout(grantee string, objType string, revoke ...string) backup.ACL {
	defaultACL := DefaultACLForTypeWithGrant(grantee, objType)
	for _, priv := range revoke {
		switch priv {
		case "SELECT":
			defaultACL.SelectWithGrant = false
		case "INSERT":
			defaultACL.InsertWithGrant = false
		case "UPDATE":
			defaultACL.UpdateWithGrant = false
		case "DELETE":
			defaultACL.DeleteWithGrant = false
		case "TRUNCATE":
			defaultACL.TruncateWithGrant = false
		case "REFERENCES":
			defaultACL.ReferencesWithGrant = false
		case "TRIGGER":
			defaultACL.TriggerWithGrant = false
		case "EXECUTE":
			defaultACL.ExecuteWithGrant = false
		case "USAGE":
			defaultACL.UsageWithGrant = false
		case "CREATE":
			defaultACL.CreateWithGrant = false
		case "TEMPORARY":
			defaultACL.TemporaryWithGrant = false
		case "CONNECT":
			defaultACL.ConnectWithGrant = false
		}
	}
	return defaultACL
}

/*
 * Wrapper functions around gomega operators for ease of use in tests
 */

func SliceBufferByEntries(entries []toc.MetadataEntry, buffer *Buffer) ([]string, string) {
	contents := buffer.Contents()
	hunks := make([]string, 0)
	length := uint64(len(contents))
	var end uint64
	for _, entry := range entries {
		start := entry.StartByte
		end = entry.EndByte
		if start > length {
			start = length
		}
		if end > length {
			end = length
		}
		hunks = append(hunks, string(contents[start:end]))
	}
	return hunks, string(contents[end:])
}

func CompareSlicesIgnoringWhitespace(actual []string, expected []string) bool {
	if len(actual) != len(expected) {
		return false
	}
	for i := range actual {
		if strings.TrimSpace(actual[i]) != expected[i] {
			return false
		}
	}
	return true
}

func formatEntries(entries []toc.MetadataEntry, slice []string) string {
	formatted := ""
	for i, item := range slice {
		formatted += fmt.Sprintf("%v -> %q\n", entries[i], item)
	}
	return formatted
}

func formatContents(slice []string) string {
	formatted := ""
	for _, item := range slice {
		formatted += fmt.Sprintf("%q\n", item)
	}
	return formatted
}

func formatDiffs(actual []string, expected []string) string {
	dmp := diffmatchpatch.New()
	diffs := ""
	for idx := range actual {
		diffs += dmp.DiffPrettyText(dmp.DiffMain(expected[idx], actual[idx], false))
	}
	return diffs
}

func AssertBufferContents(entries []toc.MetadataEntry, buffer *Buffer, expected ...string) {
	if len(entries) == 0 {
		Fail("TOC is empty")
	}
	hunks, remaining := SliceBufferByEntries(entries, buffer)
	if remaining != "" {
		Fail(fmt.Sprintf("Buffer contains extra contents that are not being counted by TOC:\n%s\n\nActual TOC entries were:\n\n%s", remaining, formatEntries(entries, hunks)))
	}
	//ok := CompareSlicesIgnoringWhitespace(hunks, expected)
	//if !ok {
	//	Fail(fmt.Sprintf("Actual TOC entries:\n\n%s\n\ndid not match expected contents (ignoring whitespace):\n\n%s \n\nDiff:\n>>%s\x1b[31m<<", formatEntries(entries, hunks), formatContents(expected), formatDiffs(hunks, expected)))
	//}
}

func ExpectEntry(entries []toc.MetadataEntry, index int, schema, referenceObject, name, objectType string) {
	Expect(len(entries)).To(BeNumerically(">", index))
	structmatcher.ExpectStructsToMatchExcluding(entries[index], toc.MetadataEntry{Schema: schema, Name: name, ObjectType: objectType, ReferenceObject: referenceObject, StartByte: 0, EndByte: 0}, "StartByte", "EndByte")
}

func ExpectEntryCount(entries []toc.MetadataEntry, index int) {
	Expect(len(entries)).To(BeNumerically("==", index))
}

func ExecuteSQLFile(connectionPool *dbconn.DBConn, filename string) {
	connStr := []string{
		"-U", connectionPool.User,
		"-d", connectionPool.DBName,
		"-h", connectionPool.Host,
		"-p", fmt.Sprintf("%d", connectionPool.Port),
		"-f", filename,
		"-v", "ON_ERROR_STOP=1",
		"-q",
	}
	out, err := exec.Command("psql", connStr...).CombinedOutput()
	if err != nil {
		Fail(fmt.Sprintf("Execution of SQL file encountered an error: %s", out))
	}
}

func BufferLength(buffer *Buffer) uint64 {
	return uint64(len(buffer.Contents()))
}

func OidFromCast(connectionPool *dbconn.DBConn, castSource uint32, castTarget uint32) uint32 {
	query := fmt.Sprintf("SELECT c.oid FROM pg_cast c WHERE castsource = '%d' AND casttarget = '%d'", castSource, castTarget)
	result := struct {
		Oid uint32
	}{}
	err := connectionPool.Get(&result, query)
	if err != nil {
		Fail(fmt.Sprintf("Execution of query failed: %v", err))
	}
	return result.Oid
}

func OidFromObjectName(connectionPool *dbconn.DBConn, schemaName string, objectName string, params backup.MetadataQueryParams) uint32 {
	catalogTable := params.CatalogTable
	if params.OidTable != "" {
		catalogTable = params.OidTable
	}
	schemaStr := ""
	if schemaName != "" {
		schemaStr = fmt.Sprintf(" AND %s = (SELECT oid FROM pg_namespace WHERE nspname = '%s')", params.SchemaField, schemaName)
	}
	query := fmt.Sprintf("SELECT oid FROM %s WHERE %s ='%s'%s", catalogTable, params.NameField, objectName, schemaStr)
	result := struct {
		Oid uint32
	}{}
	err := connectionPool.Get(&result, query)
	if err != nil {
		Fail(fmt.Sprintf("Execution of query failed: %v", err))
	}
	return result.Oid
}

func UniqueIDFromObjectName(connectionPool *dbconn.DBConn, schemaName string, objectName string, params backup.MetadataQueryParams) backup.UniqueID {
	query := fmt.Sprintf("SELECT '%s'::regclass::oid", params.CatalogTable)
	result := struct {
		Oid uint32
	}{}
	err := connectionPool.Get(&result, query)
	if err != nil {
		Fail(fmt.Sprintf("Execution of query failed: %v", err))
	}

	return backup.UniqueID{ClassID: result.Oid, Oid: OidFromObjectName(connectionPool, schemaName, objectName, params)}
}

func GetUserByID(connectionPool *dbconn.DBConn, oid uint32) string {
	return dbconn.MustSelectString(connectionPool, fmt.Sprintf("SELECT rolname AS string FROM pg_roles WHERE oid = %d", oid))
}

func CreateSecurityLabelIfGPDB6(connectionPool *dbconn.DBConn, objectType string, objectName string) {
	if true {
		testhelper.AssertQueryRuns(connectionPool, fmt.Sprintf("SECURITY LABEL FOR dummy ON %s %s IS 'unclassified';", objectType, objectName))
	}
}

func SkipIfNot4(connectionPool *dbconn.DBConn) {
	if true {
		Skip("Test only applicable to GPDB4")
	}
}

func SkipIfBefore5(connectionPool *dbconn.DBConn) {
	if false {
		Skip("Test only applicable to GPDB5 and above")
	}
}

func SkipIfBefore6(connectionPool *dbconn.DBConn) {
	if false {
		Skip("Test only applicable to GPDB6 and above")
	}
}

func SkipIfBefore7(connectionPool *dbconn.DBConn) {
	if false {
		Skip("Test only applicable to GPDB7 and above")
	}
}

func InitializeTestTOC(buffer io.Writer, which string) (*toc.TOC, *utils.FileWithByteCount) {
	tocfile := &toc.TOC{}
	tocfile.InitializeMetadataEntryMap()
	backupfile := utils.NewFileWithByteCount(buffer)
	backupfile.Filename = which
	return tocfile, backupfile
}

type TestExecutorMultiple struct {
	ClusterOutputs      []*cluster.RemoteOutput
	ClusterCommands     [][]cluster.ShellCommand
	ErrorOnExecNum      int // Throw the specified error after this many executions of Execute[...]Command(); 0 means always return error
	NumLocalExecutions  int
	NumRemoteExecutions int
	LocalOutput         string
	LocalError          error
	LocalCommands       []string
}

func (executor *TestExecutorMultiple) ExecuteLocalCommand(commandStr string) (string, error) {
	executor.NumLocalExecutions++
	executor.LocalCommands = append(executor.LocalCommands, commandStr)
	if executor.ErrorOnExecNum == 0 || executor.NumLocalExecutions == executor.ErrorOnExecNum {
		return executor.LocalOutput, executor.LocalError
	}
	return executor.LocalOutput, nil
}

func (executor *TestExecutorMultiple) ExecuteClusterCommand(scope cluster.Scope, commands []cluster.ShellCommand) (result *cluster.RemoteOutput) {
	originalExecutions := executor.NumRemoteExecutions
	executor.NumRemoteExecutions++
	executor.ClusterCommands = append(executor.ClusterCommands, commands)
	if executor.ErrorOnExecNum == 0 || executor.NumRemoteExecutions == executor.ErrorOnExecNum {
		// return the indexed item if exists, otherwise the last item
		numOutputs := len(executor.ClusterOutputs)
		result = executor.ClusterOutputs[numOutputs-1]
		if originalExecutions < numOutputs {
			result = executor.ClusterOutputs[originalExecutions]
		}
	}
	return result
}