func testStreamExecuteErrors(t *testing.T, conn *vtgateconn.VTGateConn) {
ctx := context.Background()
checkStreamExecuteErrors(t, func(query string) error {
return getStreamError(conn.StreamExecute(ctx, query, bindVars, tabletType))
})
checkStreamExecuteErrors(t, func(query string) error {
return getStreamError(conn.StreamExecuteShards(ctx, query, keyspace, shards, bindVars, tabletType))
})
checkStreamExecuteErrors(t, func(query string) error {
return getStreamError(conn.StreamExecuteKeyspaceIds(ctx, query, keyspace, keyspaceIDs, bindVars, tabletType))
})
checkStreamExecuteErrors(t, func(query string) error {
return getStreamError(conn.StreamExecuteKeyRanges(ctx, query, keyspace, keyRanges, bindVars, tabletType))
})
}
func testEchoStreamExecute(t *testing.T, conn *vtgateconn.VTGateConn) {
var qrc <-chan *mproto.QueryResult
var err error
ctx := callerid.NewContext(context.Background(), callerID, nil)
qrc, _, err = conn.StreamExecute(ctx, echoPrefix+query, bindVars, tabletType)
checkEcho(t, "StreamExecute", <-qrc, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
qrc, _, err = conn.StreamExecuteShards(ctx, echoPrefix+query, keyspace, shards, bindVars, tabletType)
checkEcho(t, "StreamExecuteShards", <-qrc, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"keyspace": keyspace,
"shards": shardsEcho,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
qrc, _, err = conn.StreamExecuteKeyspaceIds(ctx, echoPrefix+query, keyspace, keyspaceIDs, bindVars, tabletType)
checkEcho(t, "StreamExecuteKeyspaceIds", <-qrc, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"keyspace": keyspace,
"keyspaceIds": keyspaceIDsEcho,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
qrc, _, err = conn.StreamExecuteKeyRanges(ctx, echoPrefix+query, keyspace, keyRanges, bindVars, tabletType)
checkEcho(t, "StreamExecuteKeyRanges", <-qrc, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"keyspace": keyspace,
"keyRanges": keyRangesEcho,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
}
func testEchoStreamExecute(t *testing.T, conn *vtgateconn.VTGateConn) {
var stream sqltypes.ResultStream
var err error
var qr *sqltypes.Result
ctx := callerid.NewContext(context.Background(), callerID, nil)
stream, err = conn.StreamExecute(ctx, echoPrefix+query, bindVars, tabletType)
if err != nil {
t.Fatal(err)
}
qr, err = stream.Recv()
checkEcho(t, "StreamExecute", qr, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
stream, err = conn.StreamExecuteShards(ctx, echoPrefix+query, keyspace, shards, bindVars, tabletType)
if err != nil {
t.Fatal(err)
}
qr, err = stream.Recv()
checkEcho(t, "StreamExecuteShards", qr, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"keyspace": keyspace,
"shards": shardsEcho,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
stream, err = conn.StreamExecuteKeyspaceIds(ctx, echoPrefix+query, keyspace, keyspaceIDs, bindVars, tabletType)
if err != nil {
t.Fatal(err)
}
qr, err = stream.Recv()
checkEcho(t, "StreamExecuteKeyspaceIds", qr, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"keyspace": keyspace,
"keyspaceIds": keyspaceIDsEcho,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
stream, err = conn.StreamExecuteKeyRanges(ctx, echoPrefix+query, keyspace, keyRanges, bindVars, tabletType)
if err != nil {
t.Fatal(err)
}
qr, err = stream.Recv()
checkEcho(t, "StreamExecuteKeyRanges", qr, err, map[string]string{
"callerId": callerIDEcho,
"query": echoPrefix + query,
"keyspace": keyspace,
"keyRanges": keyRangesEcho,
"bindVars": bindVarsEcho,
"tabletType": tabletTypeEcho,
})
}
// testCallerID adds a caller ID to a context, and makes sure the server
// gets it.
func testCallerID(t *testing.T, conn *vtgateconn.VTGateConn) {
t.Log("testCallerID")
ctx := context.Background()
callerID := callerid.NewEffectiveCallerID("test_principal", "test_component", "test_subcomponent")
ctx = callerid.NewContext(ctx, callerID, nil)
data, err := json.Marshal(callerID)
if err != nil {
t.Errorf("failed to marshal callerid: %v", err)
return
}
query := services.CallerIDPrefix + string(data)
// test Execute calls forward the callerID
_, err = conn.Execute(ctx, query, nil, topodatapb.TabletType_MASTER, nil)
checkCallerIDError(t, "Execute", err)
_, err = conn.ExecuteShards(ctx, query, "", nil, nil, topodatapb.TabletType_MASTER, nil)
checkCallerIDError(t, "ExecuteShards", err)
_, err = conn.ExecuteKeyspaceIds(ctx, query, "", nil, nil, topodatapb.TabletType_MASTER, nil)
checkCallerIDError(t, "ExecuteKeyspaceIds", err)
_, err = conn.ExecuteKeyRanges(ctx, query, "", nil, nil, topodatapb.TabletType_MASTER, nil)
checkCallerIDError(t, "ExecuteKeyRanges", err)
_, err = conn.ExecuteEntityIds(ctx, query, "", "", nil, nil, topodatapb.TabletType_MASTER, nil)
checkCallerIDError(t, "ExecuteEntityIds", err)
// test ExecuteBatch calls forward the callerID
_, err = conn.ExecuteBatchShards(ctx, []*vtgatepb.BoundShardQuery{
{
Query: &querypb.BoundQuery{
Sql: query,
},
},
}, topodatapb.TabletType_MASTER, false, nil)
checkCallerIDError(t, "ExecuteBatchShards", err)
_, err = conn.ExecuteBatchKeyspaceIds(ctx, []*vtgatepb.BoundKeyspaceIdQuery{
{
Query: &querypb.BoundQuery{
Sql: query,
},
},
}, topodatapb.TabletType_MASTER, false, nil)
checkCallerIDError(t, "ExecuteBatchKeyspaceIds", err)
// test StreamExecute calls forward the callerID
err = getStreamError(conn.StreamExecute(ctx, query, nil, topodatapb.TabletType_MASTER, nil))
checkCallerIDError(t, "StreamExecute", err)
err = getStreamError(conn.StreamExecuteShards(ctx, query, "", nil, nil, topodatapb.TabletType_MASTER, nil))
checkCallerIDError(t, "StreamExecuteShards", err)
err = getStreamError(conn.StreamExecuteKeyspaceIds(ctx, query, "", nil, nil, topodatapb.TabletType_MASTER, nil))
checkCallerIDError(t, "StreamExecuteKeyspaceIds", err)
err = getStreamError(conn.StreamExecuteKeyRanges(ctx, query, "", nil, nil, topodatapb.TabletType_MASTER, nil))
checkCallerIDError(t, "StreamExecuteKeyRanges", err)
// test UpdateStream forwards the callerID
err = getUpdateStreamError(conn.UpdateStream(ctx, query, nil, topodatapb.TabletType_MASTER, 0, nil))
checkCallerIDError(t, "UpdateStream", err)
}