github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/pkg/interactive/metaquery/handler_cache.go (about)

     1  package metaquery
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math"
     7  	"strconv"
     8  	"strings"
     9  	"time"
    10  
    11  	"github.com/spf13/viper"
    12  	"github.com/turbot/steampipe-plugin-sdk/v5/sperr"
    13  	"github.com/turbot/steampipe/pkg/constants"
    14  	"github.com/turbot/steampipe/pkg/db/db_common"
    15  )
    16  
    17  // controls the cache in the connected FDW
    18  func cacheControl(ctx context.Context, input *HandlerInput) error {
    19  	if len(input.args()) == 0 {
    20  		return showCache(ctx, input)
    21  	}
    22  
    23  	// just get the active session from the connection pool
    24  	// and set the cache parameters on it.
    25  	// NOTE: this works because the interactive client
    26  	// always has only one active connection due to the way it works
    27  	sessionResult := input.Client.AcquireSession(ctx)
    28  	if sessionResult.Error != nil {
    29  		return sessionResult.Error
    30  	}
    31  	defer func() {
    32  		// we need to do this in a closure, otherwise the ctx will be evaluated immediately
    33  		// and not in call-time
    34  		sessionResult.Session.Close(false)
    35  	}()
    36  
    37  	conn := sessionResult.Session.Connection.Conn()
    38  	command := strings.ToLower(input.args()[0])
    39  	switch command {
    40  	case constants.ArgOn:
    41  		serverSettings := input.Client.ServerSettings()
    42  		if serverSettings != nil && !serverSettings.CacheEnabled {
    43  			fmt.Println("Caching is disabled on the server.")
    44  		}
    45  		viper.Set(constants.ArgClientCacheEnabled, true)
    46  		return db_common.SetCacheEnabled(ctx, true, conn)
    47  	case constants.ArgOff:
    48  		viper.Set(constants.ArgClientCacheEnabled, false)
    49  		return db_common.SetCacheEnabled(ctx, false, conn)
    50  	case constants.ArgClear:
    51  		return db_common.CacheClear(ctx, conn)
    52  	}
    53  
    54  	return fmt.Errorf("invalid command")
    55  }
    56  
    57  // sets the cache TTL
    58  func cacheTTL(ctx context.Context, input *HandlerInput) error {
    59  	if len(input.args()) == 0 {
    60  		return showCacheTtl(ctx, input)
    61  	}
    62  	seconds, err := strconv.Atoi(input.args()[0])
    63  	if err != nil {
    64  		return sperr.WrapWithMessage(err, "valid value is the number of seconds")
    65  	}
    66  	if seconds <= 0 {
    67  		return sperr.New("TTL must be greater than 0")
    68  	}
    69  	if can, whyCannotSet := db_common.CanSetCacheTtl(input.Client.ServerSettings(), seconds); !can {
    70  		fmt.Println(whyCannotSet)
    71  	}
    72  	sessionResult := input.Client.AcquireSession(ctx)
    73  	if sessionResult.Error != nil {
    74  		return sessionResult.Error
    75  	}
    76  	defer func() {
    77  		// we need to do this in a closure, otherwise the ctx will be evaluated immediately
    78  		// and not in call-time
    79  		sessionResult.Session.Close(false)
    80  		viper.Set(constants.ArgCacheTtl, seconds)
    81  	}()
    82  	return db_common.SetCacheTtl(ctx, time.Duration(seconds)*time.Second, sessionResult.Session.Connection.Conn())
    83  }
    84  
    85  func showCache(_ context.Context, input *HandlerInput) error {
    86  	if input.Client.ServerSettings() != nil && !input.Client.ServerSettings().CacheEnabled {
    87  		fmt.Println("Caching is disabled on the server.")
    88  		return nil
    89  	}
    90  
    91  	currentStatusString := "off"
    92  	action := "on"
    93  
    94  	if !viper.IsSet(constants.ArgClientCacheEnabled) || viper.GetBool(constants.ArgClientCacheEnabled) {
    95  		currentStatusString = "on"
    96  		action = "off"
    97  	}
    98  
    99  	fmt.Printf(
   100  		`Caching is %s. To turn it %s, type %s`,
   101  		constants.Bold(currentStatusString),
   102  		constants.Bold(action),
   103  		constants.Bold(fmt.Sprintf(".cache %s", action)),
   104  	)
   105  
   106  	// add an empty line here so that the rendering buffer can start from the next line
   107  	fmt.Println()
   108  
   109  	return nil
   110  }
   111  
   112  func showCacheTtl(ctx context.Context, input *HandlerInput) error {
   113  	if viper.IsSet(constants.ArgCacheTtl) {
   114  		ttl := getEffectiveCacheTtl(input.Client.ServerSettings(), viper.GetInt(constants.ArgCacheTtl))
   115  		fmt.Println("Cache TTL is", ttl, "seconds.")
   116  	} else if input.Client.ServerSettings() != nil {
   117  		serverTtl := input.Client.ServerSettings().CacheMaxTtl
   118  		fmt.Println("Cache TTL is", serverTtl, "seconds.")
   119  	}
   120  	errorsAndWarnings := db_common.ValidateClientCacheTtl(input.Client)
   121  	errorsAndWarnings.ShowWarnings()
   122  	// we don't know what the setting is
   123  	return nil
   124  }
   125  
   126  // getEffectiveCacheTtl returns the lower of the server TTL and the clientTtl
   127  func getEffectiveCacheTtl(serverSettings *db_common.ServerSettings, clientTtl int) int {
   128  	if serverSettings != nil {
   129  		return int(math.Min(float64(serverSettings.CacheMaxTtl), float64(clientTtl)))
   130  	}
   131  	return clientTtl
   132  }