vitessdriver: add support for DistributedTx

go/vt/vitessdriver/driver.go

@@ -20,10 +20,15 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 
 	"google.golang.org/grpc"
+	"google.golang.org/protobuf/proto"
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
 
 	"vitess.io/vitess/go/vt/vtgate/grpcvtgateconn"
 	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
@@ -115,6 +120,13 @@ func (d drv) Open(name string) (driver.Conn, error) {
 
 	c.setDefaults()
 
+	if c.Configuration.SessionToken != "" {
+		c.Configuration.sessionFromToken, err = sessionTokenToSession(c.Configuration.SessionToken)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	if c.convert, err = newConverter(&c.Configuration); err != nil {
 		return nil, err
 	}
@@ -167,6 +179,10 @@ type Configuration struct {
 	//
 	// Default: "vitess"
 	DriverName string `json:"-"`
+
+	// allows for some magic
+	SessionToken     string
+	sessionFromToken *vtgatepb.Session
 }
 
 // toJSON converts Configuration to the JSON string which is required by the
@@ -205,7 +221,11 @@ func (c *conn) dial() error {
 	if err != nil {
 		return err
 	}
-	c.session = c.conn.Session(c.Target, nil)
+	if c.Configuration.sessionFromToken != nil {
+		c.session = c.conn.SessionFromPb(c.Configuration.sessionFromToken)
+	} else {
+		c.session = c.conn.Session(c.Target, nil)
+	}
 	return nil
 }
 
@@ -231,7 +251,88 @@ func (c *conn) Close() error {
 	return nil
 }
 
+// DistributedTxFromSessionToken allows users to send serialized sessions over the wire and
+// reconnect to an existing transaction. Setting the sessionToken and address on the
+// supplied configuration is the minimum required
+func DistributedTxFromSessionToken(ctx context.Context, c Configuration) (*sql.Tx, error) {
+	if c.SessionToken == "" {
+		return nil, errors.New("c.SessionToken is required")
+	}
+	if c.Address == "" {
+		return nil, errors.New("c.Address is required")
+	}
+
+	db, err := OpenWithConfiguration(c)
+	if err != nil {
+		return nil, err
+	}
+
+	// this should return the only connection associated with the db
+	return db.BeginTx(ctx, nil)
+}
+
+// SessionTokenFromTx serializes the sessionFromToken on the tx, which can be reconstituted
+// into a *sql.Tx using DistributedTxFromSessionToken
+func SessionTokenFromTx(ctx context.Context, tx *sql.Tx) (string, error) {
+	var sessionToken string
+
+	err := tx.QueryRowContext(ctx, "vt_session_token").Scan(&sessionToken)
+	if err != nil {
+		return "", err
+	}
+
+	return sessionToken, nil
+}
+
+func newSessionTokenRow(session *vtgatepb.Session, c *converter) (driver.Rows, error) {
+	sessionToken, err := sessionToSessionToken(session)
+	if err != nil {
+		return nil, err
+	}
+
+	qr := sqltypes.Result{
+		Fields: []*querypb.Field{{
+			Name: "vt_session_token",
+			Type: sqltypes.VarBinary,
+		}},
+		Rows: [][]sqltypes.Value{{
+			sqltypes.NewVarBinary(sessionToken),
+		}},
+	}
+
+	return newRows(&qr, c), nil
+}
+
+func sessionToSessionToken(session *vtgatepb.Session) (string, error) {
+	b, err := proto.Marshal(session)
+	if err != nil {
+		return "", err
+	}
+
+	return base64.StdEncoding.EncodeToString(b), nil
+}
+
+func sessionTokenToSession(sessionToken string) (*vtgatepb.Session, error) {
+	b, err := base64.StdEncoding.DecodeString(sessionToken)
+	if err != nil {
+		return nil, err
+	}
+
+	session := &vtgatepb.Session{}
+	err = proto.Unmarshal(b, session)
+	if err != nil {
+		return nil, err
+	}
+
+	return session, nil
+}
+
 func (c *conn) Begin() (driver.Tx, error) {
+	// if we're loading from an existing session, we need to avoid starting a new transaction
+	if c.Configuration.SessionToken != "" {
+		return c, nil
+	}
+
 	if _, err := c.Exec("begin", nil); err != nil {
 		return nil, err
 	}
@@ -248,11 +349,25 @@ func (c *conn) BeginTx(_ context.Context, opts driver.TxOptions) (driver.Tx, err
 }
 
 func (c *conn) Commit() error {
+	// if we're loading from an existing session, disallow committing/rolling back the transaction
+	// this isn't a technical limitation, but is enforced to prevent misuse, so that only
+	// the original creator of the transaction can commit/rollback
+	if c.Configuration.SessionToken != "" {
+		return errors.New("calling Commit from a distributed tx is not allowed")
+	}
+
 	_, err := c.Exec("commit", nil)
 	return err
 }
 
 func (c *conn) Rollback() error {
+	// if we're loading from an existing session, disallow committing/rolling back the transaction
+	// this isn't a technical limitation, but is enforced to prevent misuse, so that only
+	// the original creator of the transaction can commit/rollback
+	if c.Configuration.SessionToken != "" {
+		return errors.New("calling Rollback from a distributed tx is not allowed")
+	}
+
 	_, err := c.Exec("rollback", nil)
 	return err
 }
@@ -314,6 +429,11 @@ func (c *conn) Query(query string, args []driver.Value) (driver.Rows, error) {
 }
 
 func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	// special case for serializing the current sessionFromToken state
+	if query == "vt_session_token" {
+		return newSessionTokenRow(c.session.SessionPb(), c.convert)
+	}
+
 	bv, err := c.convert.bindVarsFromNamedValues(args)
 	if err != nil {
 		return nil, err

go/vt/vitessdriver/driver_test.go

@@ -625,3 +625,88 @@ func TestTxExecStreamingNotAllowed(t *testing.T) {
 		t.Errorf("err: %v, does not contain %s", err, want)
 	}
 }
+
+func TestSessionToken(t *testing.T) {
+	c := Configuration{
+		Protocol: "grpc",
+		Address:  testAddress,
+		Target:   "@primary",
+	}
+
+	ctx := context.Background()
+
+	db, err := OpenWithConfiguration(c)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s, err := tx.Prepare("txRequest")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = s.Exec(int64(0))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sessionToken, err := SessionTokenFromTx(ctx, tx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	distributedTxConfig := Configuration{
+		Address:      testAddress,
+		Target:       "@primary",
+		SessionToken: sessionToken,
+	}
+
+	sameTx, err := DistributedTxFromSessionToken(ctx, distributedTxConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	newS, err := sameTx.Prepare("distributedTxRequest")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = newS.Exec(int64(1))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// enforce that Rollback can't be called on the distributed tx
+	noRollbackTx, err := DistributedTxFromSessionToken(ctx, distributedTxConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = noRollbackTx.Rollback()
+	if err != nil && err.Error() != "calling Rollback from a distributed tx is not allowed" {
+		t.Fatal(err)
+	}
+
+	// enforce that Commit can't be called on the distributed tx
+	noCommitTx, err := DistributedTxFromSessionToken(ctx, distributedTxConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = noCommitTx.Commit()
+	if err != nil && err.Error() != "calling Commit from a distributed tx is not allowed" {
+		t.Fatal(err)
+	}
+
+	// finally commit the original tx
+	err = tx.Commit()
+	if err != nil {
+		t.Fatal(err)
+	}
+}

go/vt/vitessdriver/fakeserver_test.go

@@ -17,14 +17,12 @@ limitations under the License.
 package vitessdriver
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
 
-	"context"
-
 	"google.golang.org/protobuf/proto"
-
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/vtgate/vtgateservice"
 
@@ -228,6 +226,30 @@ var execMap = map[string]struct {
 		result:  &sqltypes.Result{},
 		session: session2,
 	},
+	"distributedTxRequest": {
+		execQuery: &queryExecute{
+			SQL: "distributedTxRequest",
+			BindVariables: map[string]*querypb.BindVariable{
+				"v1": sqltypes.Int64BindVariable(1),
+			},
+			Session: &vtgatepb.Session{
+				InTransaction: true,
+				ShardSessions: []*vtgatepb.Session_ShardSession{
+					{
+						Target: &querypb.Target{
+							Keyspace:   "ks",
+							Shard:      "1",
+							TabletType: topodatapb.TabletType_PRIMARY,
+						},
+						TransactionId: 1,
+					},
+				},
+				TargetString: "@rdonly",
+			},
+		},
+		result:  &sqltypes.Result{},
+		session: session2,
+	},
 	"begin": {
 		execQuery: &queryExecute{
 			SQL: "begin",

go/vt/vtgate/vtgateconn/vtgateconn.go

@@ -17,11 +17,10 @@ limitations under the License.
 package vtgateconn
 
 import (
+	"context"
 	"flag"
 	"fmt"
 
-	"context"
-
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/log"
 
@@ -55,6 +54,19 @@ func (conn *VTGateConn) Session(targetString string, options *querypb.ExecuteOpt
 	}
 }
 
+// SessionPb returns the underlying proto session.
+func (sn *VTGateSession) SessionPb() *vtgatepb.Session {
+	return sn.session
+}
+
+// SessionFromPb returns a VTGateSession based on the provided proto session.
+func (conn *VTGateConn) SessionFromPb(sn *vtgatepb.Session) *VTGateSession {
+	return &VTGateSession{
+		session: sn,
+		impl:    conn.impl,
+	}
+}
+
 // ResolveTransaction resolves the 2pc transaction.
 func (conn *VTGateConn) ResolveTransaction(ctx context.Context, dtid string) error {
 	return conn.impl.ResolveTransaction(ctx, dtid)