github.com/scraniel/migrate@v0.0.0-20230320185700-339088f36cee/database/yugabytedb/yugabytedb_test.go (about)

     1  package yugabytedb
     2  
     3  // error codes https://github.com/lib/pq/blob/master/error.go
     4  
     5  import (
     6  	"context"
     7  	"database/sql"
     8  	"fmt"
     9  	"log"
    10  	"strings"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/dhui/dktest"
    15  	"github.com/golang-migrate/migrate/v4"
    16  
    17  	_ "github.com/lib/pq"
    18  
    19  	dt "github.com/golang-migrate/migrate/v4/database/testing"
    20  	"github.com/golang-migrate/migrate/v4/dktesting"
    21  
    22  	_ "github.com/golang-migrate/migrate/v4/source/file"
    23  )
    24  
    25  const defaultPort = 5433
    26  
    27  var (
    28  	opts = dktest.Options{
    29  		Cmd:          []string{"bin/yugabyted", "start", "--daemon=false"},
    30  		PortRequired: true,
    31  		ReadyFunc:    isReady,
    32  		Timeout:      time.Duration(60) * time.Second,
    33  	}
    34  	// Released versions: https://docs.yugabyte.com/latest/releases/#current-supported-releases
    35  	specs = []dktesting.ContainerSpec{
    36  		{ImageName: "yugabytedb/yugabyte:2.6.16.0-b14", Options: opts},
    37  		{ImageName: "yugabytedb/yugabyte:2.8.4.0-b30", Options: opts},
    38  		{ImageName: "yugabytedb/yugabyte:2.12.2.0-b58", Options: opts},
    39  		{ImageName: "yugabytedb/yugabyte:2.13.0.1-b2", Options: opts},
    40  	}
    41  )
    42  
    43  func isReady(ctx context.Context, c dktest.ContainerInfo) bool {
    44  	ip, port, err := c.Port(defaultPort)
    45  	if err != nil {
    46  		log.Println("port error:", err)
    47  		return false
    48  	}
    49  
    50  	db, err := sql.Open("postgres", fmt.Sprintf("postgres://yugabyte:yugabyte@%v:%v?sslmode=disable", ip, port))
    51  	if err != nil {
    52  		log.Println("open error:", err)
    53  		return false
    54  	}
    55  	if err := db.PingContext(ctx); err != nil {
    56  		log.Println("ping error:", err)
    57  		return false
    58  	}
    59  	if err := db.Close(); err != nil {
    60  		log.Println("close error:", err)
    61  	}
    62  	return true
    63  }
    64  
    65  func createDB(t *testing.T, c dktest.ContainerInfo) {
    66  	ip, port, err := c.Port(defaultPort)
    67  	if err != nil {
    68  		t.Fatal(err)
    69  	}
    70  
    71  	db, err := sql.Open("postgres", fmt.Sprintf("postgres://yugabyte:yugabyte@%v:%v?sslmode=disable", ip, port))
    72  	if err != nil {
    73  		t.Fatal(err)
    74  	}
    75  	if err = db.Ping(); err != nil {
    76  		t.Fatal(err)
    77  	}
    78  	defer func() {
    79  		if err := db.Close(); err != nil {
    80  			t.Error(err)
    81  		}
    82  	}()
    83  
    84  	if _, err = db.Exec("CREATE DATABASE migrate"); err != nil {
    85  		t.Fatal(err)
    86  	}
    87  }
    88  
    89  func getConnectionString(ip, port string, options ...string) string {
    90  	options = append(options, "sslmode=disable")
    91  
    92  	return fmt.Sprintf("yugabyte://yugabyte:yugabyte@%v:%v/migrate?%s", ip, port, strings.Join(options, "&"))
    93  }
    94  
    95  func Test(t *testing.T) {
    96  	dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) {
    97  		createDB(t, ci)
    98  
    99  		ip, port, err := ci.Port(defaultPort)
   100  		if err != nil {
   101  			t.Fatal(err)
   102  		}
   103  
   104  		addr := getConnectionString(ip, port)
   105  		c := &YugabyteDB{}
   106  		d, err := c.Open(addr)
   107  		if err != nil {
   108  			t.Fatal(err)
   109  		}
   110  		dt.Test(t, d, []byte("SELECT 1"))
   111  	})
   112  }
   113  
   114  func TestMigrate(t *testing.T) {
   115  	dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) {
   116  		createDB(t, ci)
   117  
   118  		ip, port, err := ci.Port(defaultPort)
   119  		if err != nil {
   120  			t.Fatal(err)
   121  		}
   122  
   123  		addr := getConnectionString(ip, port)
   124  		c := &YugabyteDB{}
   125  		d, err := c.Open(addr)
   126  		if err != nil {
   127  			t.Fatal(err)
   128  		}
   129  
   130  		m, err := migrate.NewWithDatabaseInstance("file://./examples/migrations", "migrate", d)
   131  		if err != nil {
   132  			t.Fatal(err)
   133  		}
   134  		dt.TestMigrate(t, m)
   135  	})
   136  }
   137  
   138  func TestMultiStatement(t *testing.T) {
   139  	dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) {
   140  		createDB(t, ci)
   141  
   142  		ip, port, err := ci.Port(defaultPort)
   143  		if err != nil {
   144  			t.Fatal(err)
   145  		}
   146  
   147  		addr := getConnectionString(ip, port)
   148  		c := &YugabyteDB{}
   149  		d, err := c.Open(addr)
   150  		if err != nil {
   151  			t.Fatal(err)
   152  		}
   153  		if err := d.Run(strings.NewReader("CREATE TABLE foo (foo text); CREATE TABLE bar (bar text);")); err != nil {
   154  			t.Fatalf("expected err to be nil, got %v", err)
   155  		}
   156  
   157  		// make sure second table exists
   158  		var exists bool
   159  		if err := d.(*YugabyteDB).db.QueryRow("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'bar' AND table_schema = (SELECT current_schema()))").Scan(&exists); err != nil {
   160  			t.Fatal(err)
   161  		}
   162  		if !exists {
   163  			t.Fatal("expected table bar to exist")
   164  		}
   165  	})
   166  }
   167  
   168  func TestFilterCustomQuery(t *testing.T) {
   169  	dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) {
   170  		createDB(t, ci)
   171  
   172  		ip, port, err := ci.Port(defaultPort)
   173  		if err != nil {
   174  			t.Fatal(err)
   175  		}
   176  
   177  		addr := getConnectionString(ip, port, "x-custom=foobar")
   178  		c := &YugabyteDB{}
   179  		d, err := c.Open(addr)
   180  		if err != nil {
   181  			t.Fatal(err)
   182  		}
   183  		dt.Test(t, d, []byte("SELECT 1"))
   184  	})
   185  }