github.com/pingcap/ticdc@v0.0.0-20220526033649-485a10ef2652/pkg/util/ctx.go (about)

     1  // Copyright 2020 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package util
    15  
    16  import (
    17  	"context"
    18  	"time"
    19  
    20  	"github.com/pingcap/errors"
    21  	"github.com/pingcap/tidb/kv"
    22  	"go.uber.org/zap"
    23  )
    24  
    25  type ctxKey string
    26  
    27  const (
    28  	ctxKeyTableID      = ctxKey("tableID")
    29  	ctxKeyCaptureAddr  = ctxKey("captureAddr")
    30  	ctxKeyChangefeedID = ctxKey("changefeedID")
    31  	ctxKeyIsOwner      = ctxKey("isOwner")
    32  	ctxKeyTimezone     = ctxKey("timezone")
    33  	ctxKeyKVStorage    = ctxKey("kvStorage")
    34  )
    35  
    36  // CaptureAddrFromCtx returns a capture ID stored in the specified context.
    37  // It returns an empty string if there's no valid capture ID found.
    38  func CaptureAddrFromCtx(ctx context.Context) string {
    39  	captureAddr, ok := ctx.Value(ctxKeyCaptureAddr).(string)
    40  	if !ok {
    41  		return ""
    42  	}
    43  	return captureAddr
    44  }
    45  
    46  // PutCaptureAddrInCtx returns a new child context with the specified capture ID stored.
    47  func PutCaptureAddrInCtx(ctx context.Context, captureAddr string) context.Context {
    48  	return context.WithValue(ctx, ctxKeyCaptureAddr, captureAddr)
    49  }
    50  
    51  // PutTimezoneInCtx returns a new child context with the given timezone
    52  func PutTimezoneInCtx(ctx context.Context, timezone *time.Location) context.Context {
    53  	return context.WithValue(ctx, ctxKeyTimezone, timezone)
    54  }
    55  
    56  // PutKVStorageInCtx returns a new child context with the given tikv store
    57  func PutKVStorageInCtx(ctx context.Context, store kv.Storage) context.Context {
    58  	return context.WithValue(ctx, ctxKeyKVStorage, store)
    59  }
    60  
    61  type tableinfo struct {
    62  	id   int64
    63  	name string
    64  }
    65  
    66  // PutTableInfoInCtx returns a new child context with the specified table ID and name stored.
    67  func PutTableInfoInCtx(ctx context.Context, tableID int64, tableName string) context.Context {
    68  	return context.WithValue(ctx, ctxKeyTableID, tableinfo{id: tableID, name: tableName})
    69  }
    70  
    71  // TableIDFromCtx returns a table ID
    72  func TableIDFromCtx(ctx context.Context) (int64, string) {
    73  	info, ok := ctx.Value(ctxKeyTableID).(tableinfo)
    74  	if !ok {
    75  		return 0, ""
    76  	}
    77  	return info.id, info.name
    78  }
    79  
    80  // TimezoneFromCtx returns a timezone
    81  func TimezoneFromCtx(ctx context.Context) *time.Location {
    82  	tz, ok := ctx.Value(ctxKeyTimezone).(*time.Location)
    83  	if !ok {
    84  		return nil
    85  	}
    86  	return tz
    87  }
    88  
    89  // KVStorageFromCtx returns a tikv store
    90  func KVStorageFromCtx(ctx context.Context) (kv.Storage, error) {
    91  	store, ok := ctx.Value(ctxKeyKVStorage).(kv.Storage)
    92  	if !ok {
    93  		return nil, errors.Errorf("context can not find the value associated with key: %s", ctxKeyKVStorage)
    94  	}
    95  	return store, nil
    96  }
    97  
    98  // SetOwnerInCtx returns a new child context with the owner flag set.
    99  func SetOwnerInCtx(ctx context.Context) context.Context {
   100  	return context.WithValue(ctx, ctxKeyIsOwner, true)
   101  }
   102  
   103  // IsOwnerFromCtx returns true if this capture is owner
   104  func IsOwnerFromCtx(ctx context.Context) bool {
   105  	isOwner := ctx.Value(ctxKeyIsOwner)
   106  	return isOwner != nil && isOwner.(bool)
   107  }
   108  
   109  // ChangefeedIDFromCtx returns a changefeedID stored in the specified context.
   110  // It returns an empty string if there's no valid changefeed ID found.
   111  func ChangefeedIDFromCtx(ctx context.Context) string {
   112  	changefeedID, ok := ctx.Value(ctxKeyChangefeedID).(string)
   113  	if !ok {
   114  		return ""
   115  	}
   116  	return changefeedID
   117  }
   118  
   119  // PutChangefeedIDInCtx returns a new child context with the specified changefeed ID stored.
   120  func PutChangefeedIDInCtx(ctx context.Context, changefeedID string) context.Context {
   121  	return context.WithValue(ctx, ctxKeyChangefeedID, changefeedID)
   122  }
   123  
   124  // ZapFieldCapture returns a zap field containing capture address
   125  // TODO: log redact for capture address
   126  func ZapFieldCapture(ctx context.Context) zap.Field {
   127  	return zap.String("capture", CaptureAddrFromCtx(ctx))
   128  }
   129  
   130  // ZapFieldChangefeed returns a zap field containing changefeed id
   131  func ZapFieldChangefeed(ctx context.Context) zap.Field {
   132  	return zap.String("changefeed", ChangefeedIDFromCtx(ctx))
   133  }