vitessdriver: Provide helper functions to open database.

Updated current users vtclient and client.go accordingly.

examples/local/client.go

@@ -10,14 +10,12 @@
 package main
 
 import (
-	"database/sql"
 	"flag"
 	"fmt"
 	"os"
 	"time"
 
-	// import the 'vitess' sql driver
-	_ "github.com/youtube/vitess/go/vt/vitessdriver"
+	"github.com/youtube/vitess/go/vt/vitessdriver"
 )
 
 var (
@@ -25,16 +23,14 @@ var (
 )
 
 func main() {
-	keyspace := "test_keyspace"
-	timeout := (10 * time.Second).Nanoseconds()
-	shard := "0"
-
 	flag.Parse()
 
+	keyspace := "test_keyspace"
+	shard := "0"
+	timeout := 10 * time.Second
+
 	// Connect to vtgate.
-	connStr := fmt.Sprintf(`{"protocol": "grpc", "address": "%s", "keyspace": "%s", "shard": "%s", "tablet_type": "%s", "streaming": %v, "timeout": %d}`,
-		*server, keyspace, shard, "master", false, timeout)
-	db, err := sql.Open("vitess", connStr)
+	db, err := vitessdriver.OpenShard(*server, keyspace, shard, "master", timeout)
 	if err != nil {
 		fmt.Printf("client error: %v\n", err)
 		os.Exit(1)
@@ -82,9 +78,7 @@ func main() {
 	// Note that this may be behind master due to replication lag.
 	fmt.Println("Reading from replica...")
 
-	connStr = fmt.Sprintf(`{"protocol": "grpc", "address": "%s", "keyspace": "%s", "shard": "%s", "tablet_type": "%s", "streaming": %v, "timeout": %d}`,
-		*server, keyspace, shard, "replica", false, timeout)
-	dbr, err := sql.Open("vitess", connStr)
+	dbr, err := vitessdriver.OpenShard(*server, keyspace, shard, "replica", timeout)
 	if err != nil {
 		fmt.Printf("client error: %v\n", err)
 		os.Exit(1)

go/cmd/vtclient/vtclient.go

@@ -5,7 +5,6 @@
 package main
 
 import (
-	"database/sql"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -17,8 +16,7 @@ import (
 	"github.com/youtube/vitess/go/exit"
 	"github.com/youtube/vitess/go/vt/logutil"
 
-	// import the 'vitess' sql driver
-	_ "github.com/youtube/vitess/go/vt/vitessdriver"
+	"github.com/youtube/vitess/go/vt/vitessdriver"
 )
 
 var (
@@ -106,8 +104,15 @@ func main() {
 		exit.Return(1)
 	}
 
-	connStr := fmt.Sprintf(`{"address": "%s", "keyspace": "%s", "shard": "%s", "tablet_type": "%s", "streaming": %v, "timeout": %d}`, *server, *keyspace, *shard, *tabletType, *streaming, int64(30*(*timeout)))
-	db, err := sql.Open("vitess", connStr)
+	c := vitessdriver.Configuration{
+		Address:    *server,
+		Keyspace:   *keyspace,
+		Shard:      *shard,
+		TabletType: *tabletType,
+		Timeout:    *timeout,
+		Streaming:  *streaming,
+	}
+	db, err := vitessdriver.OpenWithConfiguration(c)
 	if err != nil {
 		log.Errorf("client error: %v", err)
 		exit.Return(1)

go/vt/vitessdriver/driver.go

@@ -25,27 +25,86 @@ func init() {
 	sql.Register("vitess", drv{})
 }
 
-// TODO(mberlin): Add helper methods.
+// Open is a Vitess helper function for sql.Open().
+//
+// It opens a database connection to vtgate running at "address".
+//
+// Note that this is the vtgate v3 mode and requires a loaded VSchema.
+func Open(address, tabletType string, timeout time.Duration) (*sql.DB, error) {
+	return OpenShard(address, "" /* keyspace */, "" /* shard */, tabletType, timeout)
+}
+
+// OpenShard connects to vtgate running at "address".
+//
+// Unlike Open(), all queries will target a specific shard in a given keyspace
+// ("fallback" mode to vtgate v2).
+//
+// This mode is recommended when you want to try out Vitess initially because it
+// does not require defining a VSchema. Just replace the MySQL/MariaDB driver
+// invocation in your application with the Vitess driver.
+func OpenShard(address, keyspace, shard, tabletType string, timeout time.Duration) (*sql.DB, error) {
+	c := newDefaultConfiguration()
+	c.Address = address
+	c.Keyspace = keyspace
+	c.Shard = shard
+	c.TabletType = tabletType
+	c.Timeout = timeout
+	return OpenWithConfiguration(c)
+}
+
+// OpenForStreaming is the same as Open() but uses streaming RPCs to retrieve
+// the results.
+//
+// The streaming mode is recommended for large results.
+func OpenForStreaming(address, tabletType string, timeout time.Duration) (*sql.DB, error) {
+	return OpenShardForStreaming(address, "" /* keyspace */, "" /* shard */, tabletType, timeout)
+}
+
+// OpenShardForStreaming is the same as OpenShard() but uses streaming RPCs to
+// retrieve the results.
+//
+// The streaming mode is recommended for large results.
+func OpenShardForStreaming(address, keyspace, shard, tabletType string, timeout time.Duration) (*sql.DB, error) {
+	c := newDefaultConfiguration()
+	c.Address = address
+	c.Keyspace = keyspace
+	c.Shard = shard
+	c.TabletType = tabletType
+	c.Timeout = timeout
+	c.Streaming = true
+	return OpenWithConfiguration(c)
+}
+
+// OpenWithConfiguration is the generic Vitess helper function for sql.Open().
+//
+// It allows to pass in a Configuration struct to control all possible
+// settings of the Vitess Go SQL driver.
+func OpenWithConfiguration(c Configuration) (*sql.DB, error) {
+	jsonBytes, err := json.Marshal(c)
+	if err != nil {
+		return nil, err
+	}
+	return sql.Open("vitess", string(jsonBytes))
+}
 
 type drv struct {
 }
 
-// Open must be called with a JSON string that looks like this:
+// Open implements the database/sql/driver.Driver interface.
+//
+// For "name", the Vitess driver requires that a JSON object is passed in.
+//
+// Instead of using this call and passing in a hand-crafted JSON string, it's
+// recommended to use the public Vitess helper functions like
+// Open(), OpenShard() or OpenWithConfiguration() instead. These will generate
+// the required JSON string behind the scenes for you.
+//
+// Example for a JSON string:
 //
 //   {"protocol": "gorpc", "address": "localhost:1111", "tablet_type": "master", "timeout": 1000000000}
 //
-// protocol specifies the rpc protocol to use.
-// address specifies the address for the VTGate to connect to.
-// tablet_type represents the consistency level of your operations.
-// For example "replica" means eventually consistent reads, while
-// "master" supports transactions and gives you read-after-write consistency.
-// timeout is specified in nanoseconds. It applies for all operations.
-//
-// If you want to execute queries which are not supported by vtgate v3, you can
-// run queries against a specific keyspace and shard.
-// Therefore, add the fields "keyspace" and "shard" to the JSON string. Example:
-//
-//   {"protocol": "gorpc", "address": "localhost:1111", "keyspace": "ks1", "shard": "0", "tablet_type": "master", "timeout": 1000000000}
+// For a description of the available fields, see the Configuration struct.
+// Note: In the JSON string, timeout has to be specified in nanoseconds.
 //
 // Note that this function will always create a connection to vtgate i.e. there
 // is no need to call DB.Ping() to verify the connection.