github.com/blend/go-sdk@v1.20220411.3/db/context.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package db
     9  
    10  import (
    11  	"context"
    12  )
    13  
    14  type connectionKey struct{}
    15  
    16  // WithConnection adds a given connection to the context.
    17  func WithConnection(ctx context.Context, conn *Connection) context.Context {
    18  	return context.WithValue(ctx, connectionKey{}, conn)
    19  }
    20  
    21  // GetConnection adds a given connection to the context.
    22  func GetConnection(ctx context.Context) *Connection {
    23  	if value := ctx.Value(connectionKey{}); value != nil {
    24  		if typed, ok := value.(*Connection); ok {
    25  			return typed
    26  		}
    27  	}
    28  	return nil
    29  }
    30  
    31  type skipQueryLogging struct{}
    32  
    33  // WithSkipQueryLogging sets the context to skip logger listener triggers.
    34  func WithSkipQueryLogging(ctx context.Context) context.Context {
    35  	return context.WithValue(ctx, skipQueryLogging{}, true)
    36  }
    37  
    38  // IsSkipQueryLogging returns if we should skip triggering logger listeners for a context.
    39  func IsSkipQueryLogging(ctx context.Context) bool {
    40  	if v := ctx.Value(skipQueryLogging{}); v != nil {
    41  		return true
    42  	}
    43  	return false
    44  }