func testExecuteErrors(t *testing.T, conn *vtgateconn.VTGateConn) { ctx := context.Background() checkExecuteErrors(t, func(query string) error { _, err := conn.Execute(ctx, query, bindVars, tabletType) return err }) checkExecuteErrors(t, func(query string) error { _, err := conn.ExecuteShards(ctx, query, keyspace, shards, bindVars, tabletType) return err }) checkExecuteErrors(t, func(query string) error { _, err := conn.ExecuteKeyspaceIds(ctx, query, keyspace, keyspaceIDs, bindVars, tabletType) return err }) checkExecuteErrors(t, func(query string) error { _, err := conn.ExecuteKeyRanges(ctx, query, keyspace, keyRanges, bindVars, tabletType) return err }) checkExecuteErrors(t, func(query string) error { _, err := conn.ExecuteEntityIds(ctx, query, keyspace, "column1", entityKeyspaceIDs, bindVars, tabletType) return err }) checkExecuteErrors(t, func(query string) error { _, err := conn.ExecuteBatchShards(ctx, []*vtgatepb.BoundShardQuery{ { Query: &querypb.BoundQuery{ Sql: query, BindVariables: bindVarsP3, }, Keyspace: keyspace, Shards: shards, }, }, tabletType, true) return err }) checkExecuteErrors(t, func(query string) error { _, err := conn.ExecuteBatchKeyspaceIds(ctx, []*vtgatepb.BoundKeyspaceIdQuery{ { Query: &querypb.BoundQuery{ Sql: query, BindVariables: bindVarsP3, }, Keyspace: keyspace, KeyspaceIds: keyspaceIDs, }, }, tabletType, true) return err }) }
func testEchoExecute(t *testing.T, conn *vtgateconn.VTGateConn) { var qr *mproto.QueryResult var err error ctx := callerid.NewContext(context.Background(), callerID, nil) qr, err = conn.Execute(ctx, echoPrefix+query, bindVars, tabletType) checkEcho(t, "Execute", qr, err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, }) qr, err = conn.ExecuteShards(ctx, echoPrefix+query, keyspace, shards, bindVars, tabletType) checkEcho(t, "ExecuteShards", qr, err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "keyspace": keyspace, "shards": shardsEcho, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, }) qr, err = conn.ExecuteKeyspaceIds(ctx, echoPrefix+query, keyspace, keyspaceIDs, bindVars, tabletType) checkEcho(t, "ExecuteKeyspaceIds", qr, err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "keyspace": keyspace, "keyspaceIds": keyspaceIDsEcho, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, }) qr, err = conn.ExecuteKeyRanges(ctx, echoPrefix+query, keyspace, keyRanges, bindVars, tabletType) checkEcho(t, "ExecuteKeyRanges", qr, err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "keyspace": keyspace, "keyRanges": keyRangesEcho, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, }) qr, err = conn.ExecuteEntityIds(ctx, echoPrefix+query, keyspace, "column1", entityKeyspaceIDs, bindVars, tabletType) checkEcho(t, "ExecuteEntityIds", qr, err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "keyspace": keyspace, "entityColumnName": "column1", "entityIds": entityKeyspaceIDsEcho, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, }) var qrs []mproto.QueryResult qrs, err = conn.ExecuteBatchShards(ctx, []gproto.BoundShardQuery{ gproto.BoundShardQuery{ Sql: echoPrefix + query, Keyspace: keyspace, Shards: shards, BindVariables: bindVars, }, }, tabletType, true) checkEcho(t, "ExecuteBatchShards", &qrs[0], err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "keyspace": keyspace, "shards": shardsEcho, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, "asTransaction": "true", }) qrs, err = conn.ExecuteBatchKeyspaceIds(ctx, []gproto.BoundKeyspaceIdQuery{ gproto.BoundKeyspaceIdQuery{ Sql: echoPrefix + query, Keyspace: keyspace, KeyspaceIds: keyspaceIDs, BindVariables: bindVars, }, }, tabletType, true) checkEcho(t, "ExecuteBatchKeyspaceIds", &qrs[0], err, map[string]string{ "callerId": callerIDEcho, "query": echoPrefix + query, "keyspace": keyspace, "keyspaceIds": keyspaceIDsEcho, "bindVars": bindVarsEcho, "tabletType": tabletTypeEcho, "asTransaction": "true", }) }
// 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) }