Golang proto.Session类代码示例

func (sq *SqlQuery) ExecuteBatch(context *rpcproto.Context, queryList *proto.QueryList, reply *proto.QueryResultList) (err error) {
	defer handleError(&err, nil)
	if len(queryList.Queries) == 0 {
		panic(NewTabletError(FAIL, "Empty query list"))
	}
	sq.checkState(queryList.SessionId, false)
	begin_called := false
	var noOutput string
	session := proto.Session{
		TransactionId: queryList.TransactionId,
		SessionId:     queryList.SessionId,
	}
	reply.List = make([]mproto.QueryResult, 0, len(queryList.Queries))
	for _, bound := range queryList.Queries {
		trimmed := strings.ToLower(strings.Trim(bound.Sql, " \t\r\n"))
		switch trimmed {
		case "begin":
			if session.TransactionId != 0 {
				panic(NewTabletError(FAIL, "Nested transactions disallowed"))
			}
			var txInfo proto.TransactionInfo
			if err = sq.Begin(context, &session, &txInfo); err != nil {
				return err
			}
			session.TransactionId = txInfo.TransactionId
			begin_called = true
			reply.List = append(reply.List, mproto.QueryResult{})
		case "commit":
			if !begin_called {
				panic(NewTabletError(FAIL, "Cannot commit without begin"))
			}
			if err = sq.Commit(context, &session, &noOutput); err != nil {
				return err
			}
			session.TransactionId = 0
			begin_called = false
			reply.List = append(reply.List, mproto.QueryResult{})
		default:
			query := proto.Query{
				Sql:           bound.Sql,
				BindVariables: bound.BindVariables,
				TransactionId: session.TransactionId,
				SessionId:     session.SessionId,
			}
			var localReply mproto.QueryResult
			if err = sq.Execute(context, &query, &localReply); err != nil {
				if begin_called {
					sq.Rollback(context, &session, &noOutput)
				}
				return err
			}
			reply.List = append(reply.List, localReply)
		}
	}
	if begin_called {
		sq.Rollback(context, &session, &noOutput)
		panic(NewTabletError(FAIL, "begin called with no commit"))
	}
	return nil
}

// ExecuteBatch executes a group of queries and returns their results as a list.
// ExecuteBatch can be called for an existing transaction, or it can be called with
// the AsTransaction flag which will execute all statements inside an independent
// transaction. If AsTransaction is true, TransactionId must be 0.
func (tsv *TabletServer) ExecuteBatch(ctx context.Context, target *pb.Target, queryList *proto.QueryList, reply *proto.QueryResultList) (err error) {
	if len(queryList.Queries) == 0 {
		return NewTabletError(ErrFail, vtrpc.ErrorCode_BAD_INPUT, "Empty query list")
	}
	if queryList.AsTransaction && queryList.TransactionId != 0 {
		return NewTabletError(ErrFail, vtrpc.ErrorCode_BAD_INPUT, "cannot start a new transaction in the scope of an existing one")
	}

	allowShutdown := (queryList.TransactionId != 0)
	if err = tsv.startRequest(target, queryList.SessionId, allowShutdown); err != nil {
		return err
	}
	defer tsv.endRequest()
	defer handleError(&err, nil, tsv.qe.queryServiceStats)

	session := proto.Session{
		TransactionId: queryList.TransactionId,
		SessionId:     queryList.SessionId,
	}
	if queryList.AsTransaction {
		var txInfo proto.TransactionInfo
		if err = tsv.Begin(ctx, target, &session, &txInfo); err != nil {
			return err
		}
		session.TransactionId = txInfo.TransactionId
		// If transaction was not committed by the end, it means
		// that there was an error, roll it back.
		defer func() {
			if session.TransactionId != 0 {
				tsv.Rollback(ctx, target, &session)
			}
		}()
	}
	reply.List = make([]mproto.QueryResult, 0, len(queryList.Queries))
	for _, bound := range queryList.Queries {
		query := proto.Query{
			Sql:           bound.Sql,
			BindVariables: bound.BindVariables,
			TransactionId: session.TransactionId,
			SessionId:     session.SessionId,
		}
		var localReply mproto.QueryResult
		if err = tsv.Execute(ctx, target, &query, &localReply); err != nil {
			return err
		}
		reply.List = append(reply.List, localReply)
	}
	if queryList.AsTransaction {
		if err = tsv.Commit(ctx, target, &session); err != nil {
			session.TransactionId = 0
			return err
		}
		session.TransactionId = 0
	}
	return nil
}

func TestTabletServerStreamExecute(t *testing.T) {
	db := setUpTabletServerTest()
	testUtils := newTestUtils()
	// sql that will be executed in this test
	executeSQL := "select * from test_table limit 1000"
	executeSQLResult := &sqltypes.Result{
		RowsAffected: 1,
		Rows: [][]sqltypes.Value{
			[]sqltypes.Value{sqltypes.MakeString([]byte("row01"))},
		},
	}
	db.AddQuery(executeSQL, executeSQLResult)

	config := testUtils.newQueryServiceConfig()
	tsv := NewTabletServer(config)
	dbconfigs := testUtils.newDBConfigs(db)
	target := querypb.Target{TabletType: topodatapb.TabletType_MASTER}
	err := tsv.StartService(target, dbconfigs, []SchemaOverride{}, testUtils.newMysqld(&dbconfigs))
	if err != nil {
		t.Fatalf("StartService failed: %v", err)
	}
	defer tsv.StopService()
	ctx := context.Background()
	session := proto.Session{
		SessionId:     tsv.sessionID,
		TransactionId: 0,
	}
	txInfo := proto.TransactionInfo{TransactionId: 0}
	if err = tsv.Begin(ctx, nil, &session, &txInfo); err != nil {
		t.Fatalf("call TabletServer.Begin failed")
	}
	session.TransactionId = txInfo.TransactionId
	query := proto.Query{
		Sql:           executeSQL,
		BindVariables: nil,
		SessionId:     session.SessionId,
		TransactionId: session.TransactionId,
	}
	sendReply := func(*sqltypes.Result) error { return nil }
	if err := tsv.StreamExecute(ctx, nil, &query, sendReply); err == nil {
		t.Fatalf("TabletServer.StreamExecute should fail: %s", query.Sql)
	}
	if err := tsv.Rollback(ctx, nil, &session); err != nil {
		t.Fatalf("call TabletServer.Rollback failed")
	}
	query.TransactionId = 0
	if err := tsv.StreamExecute(ctx, nil, &query, sendReply); err != nil {
		t.Fatalf("TabletServer.StreamExecute should success: %s, but get error: %v",
			query.Sql, err)
	}
}

func TestSqlQueryStreamExecute(t *testing.T) {
	db := setUpSQLQueryTest()
	testUtils := newTestUtils()
	// sql that will be executed in this test
	executeSQL := "select * from test_table limit 1000"
	executeSQLResult := &mproto.QueryResult{
		RowsAffected: 1,
		Rows: [][]sqltypes.Value{
			[]sqltypes.Value{sqltypes.MakeString([]byte("row01"))},
		},
	}
	db.AddQuery(executeSQL, executeSQLResult)

	config := testUtils.newQueryServiceConfig()
	sqlQuery := NewSqlQuery(config)
	dbconfigs := testUtils.newDBConfigs()
	err := sqlQuery.allowQueries(nil, &dbconfigs, []SchemaOverride{}, testUtils.newMysqld(&dbconfigs))
	if err != nil {
		t.Fatalf("allowQueries failed: %v", err)
	}
	defer sqlQuery.disallowQueries()
	ctx := context.Background()
	session := proto.Session{
		SessionId:     sqlQuery.sessionID,
		TransactionId: 0,
	}
	txInfo := proto.TransactionInfo{TransactionId: 0}
	if err = sqlQuery.Begin(ctx, nil, &session, &txInfo); err != nil {
		t.Fatalf("call SqlQuery.Begin failed")
	}
	session.TransactionId = txInfo.TransactionId
	query := proto.Query{
		Sql:           executeSQL,
		BindVariables: nil,
		SessionId:     session.SessionId,
		TransactionId: session.TransactionId,
	}
	sendReply := func(*mproto.QueryResult) error { return nil }
	if err := sqlQuery.StreamExecute(ctx, nil, &query, sendReply); err == nil {
		t.Fatalf("SqlQuery.StreamExecute should fail: %s", query.Sql)
	}
	if err := sqlQuery.Rollback(ctx, nil, &session); err != nil {
		t.Fatalf("call SqlQuery.Rollback failed")
	}
	query.TransactionId = 0
	if err := sqlQuery.StreamExecute(ctx, nil, &query, sendReply); err != nil {
		t.Fatalf("SqlQuery.StreamExecute should success: %s, but get error: %v",
			query.Sql, err)
	}
}

func TestSqlQueryCommitTransaciton(t *testing.T) {
	db := setUpSqlQueryTest()
	testUtils := newTestUtils()
	// sql that will be executed in this test
	executeSql := "select * from test_table limit 1000"
	executeSqlResult := &mproto.QueryResult{
		RowsAffected: 1,
		Rows: [][]sqltypes.Value{
			[]sqltypes.Value{sqltypes.MakeString([]byte("row01"))},
		},
	}
	db.AddQuery(executeSql, executeSqlResult)
	config := testUtils.newQueryServiceConfig()
	sqlQuery := NewSqlQuery(config)
	dbconfigs := testUtils.newDBConfigs()
	err := sqlQuery.allowQueries(nil, &dbconfigs, []SchemaOverride{}, testUtils.newMysqld(&dbconfigs))
	if err != nil {
		t.Fatalf("allowQueries failed: %v", err)
	}
	defer sqlQuery.disallowQueries()
	ctx := context.Background()
	session := proto.Session{
		SessionId:     sqlQuery.sessionID,
		TransactionId: 0,
	}
	txInfo := proto.TransactionInfo{TransactionId: 0}
	if err = sqlQuery.Begin(ctx, nil, &session, &txInfo); err != nil {
		t.Fatalf("call SqlQuery.Begin failed")
	}
	session.TransactionId = txInfo.TransactionId
	query := proto.Query{
		Sql:           executeSql,
		BindVariables: nil,
		SessionId:     session.SessionId,
		TransactionId: session.TransactionId,
	}
	reply := mproto.QueryResult{}
	if err := sqlQuery.Execute(ctx, nil, &query, &reply); err != nil {
		t.Fatalf("failed to execute query: %s", query.Sql)
	}
	if err := sqlQuery.Commit(ctx, nil, &session); err != nil {
		t.Fatalf("call SqlQuery.Commit failed")
	}
}

func TestTabletServerRollback(t *testing.T) {
	db := setUpTabletServerTest()
	testUtils := newTestUtils()
	// sql that will be executed in this test
	executeSQL := "select * from test_table limit 1000"
	executeSQLResult := &mproto.QueryResult{
		RowsAffected: 1,
		Rows: [][]sqltypes.Value{
			[]sqltypes.Value{sqltypes.MakeString([]byte("row01"))},
		},
	}
	db.AddQuery(executeSQL, executeSQLResult)
	config := testUtils.newQueryServiceConfig()
	tsv := NewTabletServer(config)
	dbconfigs := testUtils.newDBConfigs(db)
	err := tsv.StartService(nil, &dbconfigs, []SchemaOverride{}, testUtils.newMysqld(&dbconfigs))
	if err != nil {
		t.Fatalf("StartService failed: %v", err)
	}
	defer tsv.StopService()
	ctx := context.Background()
	session := proto.Session{
		SessionId:     tsv.sessionID,
		TransactionId: 0,
	}
	txInfo := proto.TransactionInfo{TransactionId: 0}
	if err = tsv.Begin(ctx, nil, &session, &txInfo); err != nil {
		t.Fatalf("call TabletServer.Begin failed")
	}
	session.TransactionId = txInfo.TransactionId
	query := proto.Query{
		Sql:           executeSQL,
		BindVariables: nil,
		SessionId:     session.SessionId,
		TransactionId: session.TransactionId,
	}
	reply := mproto.QueryResult{}
	if err := tsv.Execute(ctx, nil, &query, &reply); err != nil {
		t.Fatalf("failed to execute query: %s", query.Sql)
	}
	if err := tsv.Rollback(ctx, nil, &session); err != nil {
		t.Fatalf("call TabletServer.Rollback failed")
	}
}

// ExecuteBatch executes a group of queries and returns their results as a list.
// ExecuteBatch can be called for an existing transaction, or it can also begin
// its own transaction, in which case it's expected to commit it also.
func (sq *SqlQuery) ExecuteBatch(ctx context.Context, queryList *proto.QueryList, reply *proto.QueryResultList) (err error) {
	if len(queryList.Queries) == 0 {
		return NewTabletError(ErrFail, "Empty query list")
	}

	allowShutdown := (queryList.TransactionId != 0)
	if err = sq.startRequest(queryList.SessionId, false, allowShutdown); err != nil {
		return err
	}
	defer sq.endRequest()
	defer handleError(&err, nil, sq.qe.queryServiceStats)

	beginCalled := false
	session := proto.Session{
		TransactionId: queryList.TransactionId,
		SessionId:     queryList.SessionId,
	}
	reply.List = make([]mproto.QueryResult, 0, len(queryList.Queries))
	for _, bound := range queryList.Queries {
		trimmed := strings.ToLower(strings.Trim(bound.Sql, " \t\r\n"))
		switch trimmed {
		case "begin":
			if session.TransactionId != 0 {
				panic(NewTabletError(ErrFail, "Nested transactions disallowed"))
			}
			var txInfo proto.TransactionInfo
			if err = sq.Begin(ctx, &session, &txInfo); err != nil {
				return err
			}
			session.TransactionId = txInfo.TransactionId
			beginCalled = true
			reply.List = append(reply.List, mproto.QueryResult{})
		case "commit":
			if !beginCalled {
				panic(NewTabletError(ErrFail, "Cannot commit without begin"))
			}
			if err = sq.Commit(ctx, &session); err != nil {
				return err
			}
			session.TransactionId = 0
			beginCalled = false
			reply.List = append(reply.List, mproto.QueryResult{})
		default:
			query := proto.Query{
				Sql:           bound.Sql,
				BindVariables: bound.BindVariables,
				TransactionId: session.TransactionId,
				SessionId:     session.SessionId,
			}
			var localReply mproto.QueryResult
			if err = sq.Execute(ctx, &query, &localReply); err != nil {
				if beginCalled {
					sq.Rollback(ctx, &session)
				}
				return err
			}
			reply.List = append(reply.List, localReply)
		}
	}
	if beginCalled {
		sq.Rollback(ctx, &session)
		panic(NewTabletError(ErrFail, "begin called with no commit"))
	}
	return nil
}