logging_dialer.go

	"log"
	"net"
	"time"
)

func loadBytes(file *zip.File) ([]byte, error) {
	r, err := file.Open()
	if err != nil {
		return nil, err
	}
	defer r.Close()
	return ioutil.ReadAll(r)
}

func NewClusterFromBundle(bundlePath string) (*gocql.ClusterConfig, error) {
	reader, err := zip.OpenReader(bundlePath)
	if err != nil {
		return nil, err
	}

	defer reader.Close()

	contents := make(map[string][]byte)

	for _, file := range reader.File {
		switch file.Name {
		case "config.json", "cert", "key", "ca.crt":
			bytes, err := loadBytes(file)
			if err != nil {
				return nil, err
			}
			contents[file.Name] = bytes
		}
	}

	config := struct {
		Host     string `json:"host"`
		Port     int    `json:"cql_port"`
		Keyspace string `json:"keyspace"`
	}{}

	err = json.Unmarshal(contents["config.json"], &config)
	if err != nil {
		return nil, err
	}

	cert, err := tls.X509KeyPair(contents["cert"], contents["key"])
	if err != nil {
		return nil, err
	}
	caCerts := x509.NewCertPool()
	if ok := caCerts.AppendCertsFromPEM(contents["ca.crt"]); !ok {
		return nil, fmt.Errorf("unable to load CA certificate")
	}
	tlsConfig := tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      caCerts,
	}

	cluster := gocql.NewCluster(config.Host)
	cluster.Port = config.Port
	cluster.Keyspace = config.Keyspace
	cluster.SslOpts = &gocql.SslOptions{
		Config: &tlsConfig,
	}

	return cluster, nil
}

var bundle = flag.String("bundle", "", "")
var username = flag.String("username", "", "")
var password = flag.String("password", "", "")

type LoggingDialer struct {
	dialer *net.Dialer
}

func NewLoggingDialer(connectTimeout, keepalive time.Duration) *LoggingDialer {
	dialer := &net.Dialer{
		Timeout: connectTimeout,
	}
	if keepalive > 0 {
		dialer.KeepAlive = keepalive
	}
	return &LoggingDialer{dialer}
}

func (l *LoggingDialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	log.Printf("Connecting to %s with timeout %d\n", addr, l.dialer.Timeout)
	conn, err := l.dialer.DialContext(ctx, network, addr)
	if err != nil {
		log.Printf("Failed to connect to %s with error: %v\n", addr, err)
	}
	return conn, err
}

func main() {
	flag.Parse()

	if bundle == nil || *bundle == "" {
		log.Fatal("-bundle argument required")
	} else if username == nil || *username == "" {
		log.Fatal("-username argument required")
	} else if password == nil || *password == "" {
		log.Fatal("-password argument required")
	}

	cluster, err := NewClusterFromBundle(*bundle)
	if err != nil {
		log.Fatalf("unable to create cluster from bundle: %v", err)
	}

	cluster.Authenticator = gocql.PasswordAuthenticator{
		Username: *username,
		Password: *password,
	}
	cluster.ConnectTimeout = 2 * time.Second
	cluster.Dialer = NewLoggingDialer(cluster.ConnectTimeout, 0)

	session, err := gocql.NewSession(*cluster)
	if err != nil {
		log.Fatalf("unable to connect session: %v", err)
	}

	iter := session.Query("SELECT release_version FROM system.local").Iter()

	var version string
	for iter.Scan(&version) {
		fmt.Println(version)
	}

	if err := iter.Close(); err != nil {
		log.Fatalf("error running query: %v", err)
	}
}