Golang callerid.EffectiveCallerIDFromContext函數代碼示例

// BeginExecute starts a transaction and runs an Execute.
func (conn *gRPCQueryClient) BeginExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]interface{}, options *querypb.ExecuteOptions) (result *sqltypes.Result, transactionID int64, err error) {
	conn.mu.RLock()
	defer conn.mu.RUnlock()
	if conn.cc == nil {
		return nil, 0, tabletconn.ConnClosed
	}

	q, err := querytypes.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, 0, err
	}

	if *combo {
		// If combo is enabled, we combine both calls
		req := &querypb.BeginExecuteRequest{
			Target:            target,
			EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
			ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
			Query:             q,
			Options:           options,
		}
		reply, err := conn.c.BeginExecute(ctx, req)
		if err != nil {
			return nil, 0, tabletconn.TabletErrorFromGRPC(err)
		}
		if reply.Error != nil {
			return nil, reply.TransactionId, tabletconn.TabletErrorFromRPCError(reply.Error)
		}
		return sqltypes.Proto3ToResult(reply.Result), reply.TransactionId, nil
	}

	// Begin part.
	breq := &querypb.BeginRequest{
		Target:            target,
		EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
		ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
	}
	br, err := conn.c.Begin(ctx, breq)
	if err != nil {
		return nil, 0, tabletconn.TabletErrorFromGRPC(err)
	}
	transactionID = br.TransactionId

	// Execute part.
	ereq := &querypb.ExecuteRequest{
		Target:            target,
		EffectiveCallerId: breq.EffectiveCallerId,
		ImmediateCallerId: breq.ImmediateCallerId,
		Query:             q,
		TransactionId:     transactionID,
		Options:           options,
	}
	er, err := conn.c.Execute(ctx, ereq)
	if err != nil {
		return nil, transactionID, tabletconn.TabletErrorFromGRPC(err)
	}
	return sqltypes.Proto3ToResult(er.Result), transactionID, nil

}

func (conn *vtgateConn) Execute(ctx context.Context, query string, bindVars map[string]interface{}, keyspace string, tabletType topodatapb.TabletType, session interface{}) (*sqltypes.Result, interface{}, error) {
	var s *vtgatepb.Session
	if session != nil {
		s = session.(*vtgatepb.Session)
	}
	q, err := querytypes.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, session, err
	}
	request := &vtgatepb.ExecuteRequest{
		CallerId:   callerid.EffectiveCallerIDFromContext(ctx),
		Session:    s,
		Query:      q,
		Keyspace:   keyspace,
		TabletType: tabletType,
	}
	response, err := conn.c.Execute(ctx, request)
	if err != nil {
		return nil, session, vterrors.FromGRPCError(err)
	}
	if response.Error != nil {
		return nil, response.Session, vterrors.FromVtRPCError(response.Error)
	}
	return sqltypes.Proto3ToResult(response.Result), response.Session, nil
}

func (conn *vtgateConn) SplitQueryV2(
	ctx context.Context,
	keyspace string,
	query string,
	bindVars map[string]interface{},
	splitColumns []string,
	splitCount int64,
	numRowsPerQueryPart int64,
	algorithm querypb.SplitQueryRequest_Algorithm) ([]*vtgatepb.SplitQueryResponse_Part, error) {
	q, err := querytypes.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, err
	}

	request := &vtgatepb.SplitQueryRequest{
		CallerId:            callerid.EffectiveCallerIDFromContext(ctx),
		Keyspace:            keyspace,
		Query:               q,
		SplitColumn:         splitColumns,
		SplitCount:          splitCount,
		NumRowsPerQueryPart: numRowsPerQueryPart,
		Algorithm:           algorithm,
		UseSplitQueryV2:     true,
	}
	response, err := conn.c.SplitQuery(ctx, request)
	if err != nil {
		return nil, vterrors.FromGRPCError(err)
	}
	return response.Splits, nil
}

func (conn *vtgateConn) StreamExecuteKeyspaceIds(ctx context.Context, query string, keyspace string, keyspaceIds [][]byte, bindVars map[string]interface{}, tabletType topodatapb.TabletType) (sqltypes.ResultStream, error) {
	q, err := querytypes.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, err
	}
	req := &vtgatepb.StreamExecuteKeyspaceIdsRequest{
		CallerId:    callerid.EffectiveCallerIDFromContext(ctx),
		Query:       q,
		Keyspace:    keyspace,
		KeyspaceIds: keyspaceIds,
		TabletType:  tabletType,
	}
	stream, err := conn.c.StreamExecuteKeyspaceIds(ctx, req)
	if err != nil {
		return nil, vterrors.FromGRPCError(err)
	}
	return &streamExecuteAdapter{
		recv: func() (*querypb.QueryResult, error) {
			ser, err := stream.Recv()
			if err != nil {
				return nil, err
			}
			return ser.Result, nil
		},
	}, nil
}

func (conn *vtgateConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, tabletType topodatapb.TabletType) (<-chan *sqltypes.Result, vtgateconn.ErrFunc, error) {
	q, err := tproto.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, nil, err
	}
	req := &pb.StreamExecuteRequest{
		CallerId:   callerid.EffectiveCallerIDFromContext(ctx),
		Query:      q,
		TabletType: tabletType,
	}
	stream, err := conn.c.StreamExecute(ctx, req)
	if err != nil {
		return nil, nil, vterrors.FromGRPCError(err)
	}
	sr := make(chan *sqltypes.Result, 10)
	var finalError error
	go func() {
		for {
			ser, err := stream.Recv()
			if err != nil {
				if err != io.EOF {
					finalError = vterrors.FromGRPCError(err)
				}
				close(sr)
				return
			}
			sr <- sqltypes.Proto3ToResult(ser.Result)
		}
	}()
	return sr, func() error {
		return finalError
	}, nil
}

func (conn *vtgateConn) ExecuteEntityIds(ctx context.Context, query string, keyspace string, entityColumnName string, entityKeyspaceIDs []*pb.ExecuteEntityIdsRequest_EntityId, bindVars map[string]interface{}, tabletType topodatapb.TabletType, notInTransaction bool, session interface{}) (*sqltypes.Result, interface{}, error) {
	var s *pb.Session
	if session != nil {
		s = session.(*pb.Session)
	}
	q, err := tproto.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, session, err
	}
	request := &pb.ExecuteEntityIdsRequest{
		CallerId:          callerid.EffectiveCallerIDFromContext(ctx),
		Session:           s,
		Query:             q,
		Keyspace:          keyspace,
		EntityColumnName:  entityColumnName,
		EntityKeyspaceIds: entityKeyspaceIDs,
		TabletType:        tabletType,
		NotInTransaction:  notInTransaction,
	}
	response, err := conn.c.ExecuteEntityIds(ctx, request)
	if err != nil {
		return nil, session, vterrors.FromGRPCError(err)
	}
	if response.Error != nil {
		return nil, response.Session, vterrors.FromVtRPCError(response.Error)
	}
	return sqltypes.Proto3ToResult(response.Result), response.Session, nil
}

// Execute sends the query to VTTablet.
func (conn *gRPCQueryClient) Execute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*sqltypes.Result, error) {
	conn.mu.RLock()
	defer conn.mu.RUnlock()
	if conn.cc == nil {
		return nil, tabletconn.ConnClosed
	}

	q, err := tproto.BoundQueryToProto3(query, bindVars)
	if err != nil {
		return nil, err
	}

	req := &pb.ExecuteRequest{
		Target:            conn.target,
		EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
		ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
		Query:             q,
		TransactionId:     transactionID,
		SessionId:         conn.sessionID,
	}
	er, err := conn.c.Execute(ctx, req)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	return sqltypes.Proto3ToResult(er.Result), nil
}

// SplitQuery is the stub for TabletServer.SplitQuery RPC
func (conn *gRPCQueryClient) SplitQuery(ctx context.Context, query tproto.BoundQuery, splitColumn string, splitCount int) (queries []tproto.QuerySplit, err error) {
	conn.mu.RLock()
	defer conn.mu.RUnlock()
	if conn.cc == nil {
		err = tabletconn.ConnClosed
		return
	}

	q, err := tproto.BoundQueryToProto3(query.Sql, query.BindVariables)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	req := &pb.SplitQueryRequest{
		Target:            conn.target,
		EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
		ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
		Query:             q,
		SplitColumn:       splitColumn,
		SplitCount:        int64(splitCount),
		SessionId:         conn.sessionID,
	}
	sqr, err := conn.c.SplitQuery(ctx, req)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	split, err := tproto.Proto3ToQuerySplits(sqr.Queries)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	return split, nil
}

// SplitQuery is the stub for TabletServer.SplitQuery RPC
// TODO(erez): Remove this method and rename SplitQueryV2 to SplitQuery once
// the migration to SplitQuery V2 is done.
func (conn *gRPCQueryClient) SplitQuery(ctx context.Context, target *querypb.Target, query querytypes.BoundQuery, splitColumn string, splitCount int64) (queries []querytypes.QuerySplit, err error) {
	conn.mu.RLock()
	defer conn.mu.RUnlock()
	if conn.cc == nil {
		err = tabletconn.ConnClosed
		return
	}

	q, err := querytypes.BoundQueryToProto3(query.Sql, query.BindVariables)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	req := &querypb.SplitQueryRequest{
		Target:              target,
		EffectiveCallerId:   callerid.EffectiveCallerIDFromContext(ctx),
		ImmediateCallerId:   callerid.ImmediateCallerIDFromContext(ctx),
		Query:               q,
		SplitColumn:         []string{splitColumn},
		SplitCount:          splitCount,
		NumRowsPerQueryPart: 0,
		Algorithm:           querypb.SplitQueryRequest_EQUAL_SPLITS,
		UseSplitQueryV2:     false,
	}
	sqr, err := conn.c.SplitQuery(ctx, req)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	split, err := querytypes.Proto3ToQuerySplits(sqr.Queries)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(err)
	}
	return split, nil
}

func (conn *vtgateConn) ExecuteBatchKeyspaceIds(ctx context.Context, queries []proto.BoundKeyspaceIdQuery, tabletType pbt.TabletType, asTransaction bool, session interface{}) ([]mproto.QueryResult, interface{}, error) {
	var s *pb.Session
	if session != nil {
		s = session.(*pb.Session)
	}
	qs, err := proto.BoundKeyspaceIdQueriesToProto(queries)
	if err != nil {
		return nil, session, err
	}
	request := &pb.ExecuteBatchKeyspaceIdsRequest{
		CallerId:      callerid.EffectiveCallerIDFromContext(ctx),
		Session:       s,
		Queries:       qs,
		TabletType:    tabletType,
		AsTransaction: asTransaction,
	}
	response, err := conn.c.ExecuteBatchKeyspaceIds(ctx, request)
	if err != nil {
		return nil, session, vterrors.FromGRPCError(err)
	}
	if response.Error != nil {
		return nil, response.Session, vterrors.FromVtRPCError(response.Error)
	}
	return mproto.Proto3ToQueryResults(response.Results), response.Session, nil
}

func (conn *vtgateConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, tabletType pbt.TabletType) (<-chan *mproto.QueryResult, vtgateconn.ErrFunc, error) {
	req := &pb.StreamExecuteRequest{
		CallerId:   callerid.EffectiveCallerIDFromContext(ctx),
		Query:      tproto.BoundQueryToProto3(query, bindVars),
		TabletType: tabletType,
	}
	stream, err := conn.c.StreamExecute(ctx, req)
	if err != nil {
		return nil, nil, vterrors.FromGRPCError(err)
	}
	sr := make(chan *mproto.QueryResult, 10)
	// TODO(aaijazi): I'm pretty sure ther error handling going on down here is overkill now...
	var finalError error
	go func() {
		for {
			ser, err := stream.Recv()
			if err != nil {
				if err != io.EOF {
					finalError = err
				}
				close(sr)
				return
			}
			if ser.Error != nil {
				finalError = vterrors.FromVtRPCError(ser.Error)
				close(sr)
				return
			}
			sr <- mproto.Proto3ToQueryResult(ser.Result)
		}
	}()
	return sr, func() error {
		return finalError
	}, nil
}

func (f *FakeQueryService) checkSessionTargetCallerID(ctx context.Context, name string, target *querypb.Target, sessionID int64) {
	if f.checkTarget {
		if !reflect.DeepEqual(target, testTarget) {
			f.t.Errorf("invalid Target for %v: got %#v expected %#v", name, target, testTarget)
		}
	} else {
		if sessionID != testSessionID {
			f.t.Errorf("invalid sessionID for %v: got %v expected %v", name, sessionID, testSessionID)
		}
	}

	ef := callerid.EffectiveCallerIDFromContext(ctx)
	if ef == nil {
		f.t.Errorf("no effective caller id for %v", name)
	} else {
		if !reflect.DeepEqual(ef, testCallerID) {
			f.t.Errorf("invalid effective caller id for %v: got %v expected %v", name, ef, testCallerID)
		}
	}

	im := callerid.ImmediateCallerIDFromContext(ctx)
	if im == nil {
		f.t.Errorf("no immediate caller id for %v", name)
	} else {
		if !reflect.DeepEqual(im, testVTGateCallerID) {
			f.t.Errorf("invalid immediate caller id for %v: got %v expected %v", name, im, testVTGateCallerID)
		}
	}
}

// Begin begins a transaction, and returns the associated transaction id.
// Subsequent statements can access the connection through the transaction id.
func (axp *TxPool) Begin(ctx context.Context) int64 {
	poolCtx := ctx
	if deadline, ok := ctx.Deadline(); ok {
		var cancel func()
		poolCtx, cancel = context.WithDeadline(ctx, deadline.Add(-10*time.Millisecond))
		defer cancel()
	}
	conn, err := axp.pool.Get(poolCtx)
	if err != nil {
		switch err {
		case ErrConnPoolClosed:
			panic(err)
		case pools.ErrTimeout:
			axp.LogActive()
			panic(NewTabletError(ErrTxPoolFull, vtrpcpb.ErrorCode_RESOURCE_EXHAUSTED, "Transaction pool connection limit exceeded"))
		}
		panic(NewTabletErrorSQL(ErrFatal, vtrpcpb.ErrorCode_INTERNAL_ERROR, err))
	}
	if _, err := conn.Exec(ctx, "begin", 1, false); err != nil {
		conn.Recycle()
		panic(NewTabletErrorSQL(ErrFail, vtrpcpb.ErrorCode_UNKNOWN_ERROR, err))
	}
	transactionID := axp.lastID.Add(1)
	axp.activePool.Register(
		transactionID,
		newTxConnection(
			conn,
			transactionID,
			axp,
			callerid.ImmediateCallerIDFromContext(ctx),
			callerid.EffectiveCallerIDFromContext(ctx),
		),
	)
	return transactionID
}

func (conn *vtgateConn) StreamExecuteKeyspaceIds(ctx context.Context, query string, keyspace string, keyspaceIds [][]byte, bindVars map[string]interface{}, tabletType pbt.TabletType) (<-chan *mproto.QueryResult, vtgateconn.ErrFunc, error) {
	req := &pb.StreamExecuteKeyspaceIdsRequest{
		CallerId:    callerid.EffectiveCallerIDFromContext(ctx),
		Query:       tproto.BoundQueryToProto3(query, bindVars),
		Keyspace:    keyspace,
		KeyspaceIds: keyspaceIds,
		TabletType:  tabletType,
	}
	stream, err := conn.c.StreamExecuteKeyspaceIds(ctx, req)
	if err != nil {
		return nil, nil, vterrors.FromGRPCError(err)
	}
	sr := make(chan *mproto.QueryResult, 10)
	var finalError error
	go func() {
		for {
			ser, err := stream.Recv()
			if err != nil {
				if err != io.EOF {
					finalError = vterrors.FromGRPCError(err)
				}
				close(sr)
				return
			}
			sr <- mproto.Proto3ToQueryResult(ser.Result)
		}
	}()
	return sr, func() error {
		return finalError
	}, nil
}

// ExecuteBatch sends a batch query to VTTablet.
func (conn *gRPCQueryClient) ExecuteBatch(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
	conn.mu.RLock()
	defer conn.mu.RUnlock()
	if conn.cc == nil {
		return nil, tabletconn.ConnClosed
	}

	req := &pb.ExecuteBatchRequest{
		Target:            conn.target,
		EffectiveCallerId: callerid.EffectiveCallerIDFromContext(ctx),
		ImmediateCallerId: callerid.ImmediateCallerIDFromContext(ctx),
		Queries:           make([]*pb.BoundQuery, len(queries)),
		AsTransaction:     asTransaction,
		TransactionId:     transactionID,
		SessionId:         conn.sessionID,
	}
	for i, q := range queries {
		req.Queries[i] = tproto.BoundQueryToProto3(q.Sql, q.BindVariables)
	}
	ebr, err := conn.c.ExecuteBatch(ctx, req)
	if err != nil {
		return nil, tabletErrorFromGRPC(err)
	}
	return tproto.Proto3ToQueryResultList(ebr.Results), nil
}