github.com/CloudCom/goose@v0.0.0-20151110184009-e03c3249c21b/travis.setup_redshift.go (about)

     1  // +build travis
     2  
     3  package main
     4  
     5  import (
     6  	"database/sql"
     7  	"fmt"
     8  	"net/url"
     9  	"os"
    10  	"path/filepath"
    11  	"strconv"
    12  	"strings"
    13  	"time"
    14  
    15  	_ "github.com/lib/pq"
    16  )
    17  
    18  func main() {
    19  	var cmd string
    20  	if len(os.Args) == 2 {
    21  		cmd = os.Args[1]
    22  	}
    23  	switch cmd {
    24  	case "setup":
    25  		setup()
    26  	case "destroy":
    27  		destroy()
    28  	default:
    29  		fmt.Fprintf(os.Stderr, "usage: %s {setup|destroy}\n", filepath.Base(os.Args[0]))
    30  		os.Exit(1)
    31  	}
    32  }
    33  
    34  func setup() {
    35  	dsn := os.Getenv("REDSHIFT_DATABASE_DSN")
    36  
    37  	newdb := "goose-" + strconv.FormatInt(time.Now().UnixNano(), 36)
    38  	var newDSN string
    39  
    40  	if strings.HasPrefix(dsn, "postgres://") {
    41  		u, err := url.Parse(dsn)
    42  		if err != nil {
    43  			fmt.Fprintf(os.Stderr, "Could not parse $REDSHIFT_DATABASE_DSN\n")
    44  			os.Exit(1)
    45  		}
    46  		u.Path = newdb
    47  		newDSN = u.String()
    48  	} else {
    49  		foundDBName := false
    50  		for _, f := range strings.Fields(dsn) {
    51  			kv := strings.SplitN(f, "=", 2)
    52  			if kv[0] == "dbname" {
    53  				f = "dbname=" + newdb
    54  				foundDBName = true
    55  			}
    56  			newDSN = newDSN + f + " "
    57  		}
    58  		if !foundDBName {
    59  			newDSN = newDSN + "dbname=" + newdb
    60  		}
    61  	}
    62  
    63  	db, err := sql.Open("postgres", dsn)
    64  	if err != nil {
    65  		fmt.Fprintf(os.Stderr, "Could not connect to Redshift: %s\n", err)
    66  		os.Exit(1)
    67  	}
    68  	_, err = db.Exec(fmt.Sprintf("create database %q", newdb))
    69  	if err != nil {
    70  		fmt.Fprintf(os.Stderr, "Could not create database: %s\n", err)
    71  		os.Exit(1)
    72  	}
    73  
    74  	fmt.Printf("REDSHIFT_DATABASE_DSN=%s\n", newDSN)
    75  	os.Exit(0)
    76  }
    77  
    78  func destroy() {
    79  	dsn := os.Getenv("REDSHIFT_DATABASE_DSN")
    80  
    81  	var dbname string
    82  
    83  	if strings.HasPrefix(dsn, "postgres://") {
    84  		u, err := url.Parse(dsn)
    85  		if err != nil {
    86  			fmt.Fprintf(os.Stderr, "Could not parse $REDSHIFT_DATABASE_DSN\n")
    87  			os.Exit(1)
    88  		}
    89  		dbname = u.Path
    90  	} else {
    91  		for _, f := range strings.Fields(dsn) {
    92  			kv := strings.SplitN(f, "=", 2)
    93  			if len(kv) != 2 || kv[0] == "dbname" {
    94  				dbname = kv[1]
    95  				break
    96  			}
    97  		}
    98  	}
    99  	if dbname == "" {
   100  		fmt.Fprintf(os.Stderr, "Could not find db name in $REDSHIFT_DATABASE_DSN\n")
   101  		os.Exit(1)
   102  	}
   103  
   104  	db, err := sql.Open("postgres", strings.Replace(dsn, dbname, "dev", 1))
   105  	if err != nil {
   106  		fmt.Fprintf(os.Stderr, "Could not connect to Redshift: %s\n", err)
   107  		os.Exit(1)
   108  	}
   109  	_, err = db.Exec(fmt.Sprintf("drop database %q", dbname))
   110  	if err != nil {
   111  		fmt.Fprintf(os.Stderr, "Could not drop database: %s\n", err)
   112  		os.Exit(1)
   113  	}
   114  }