github.com/turbot/steampipe@v1.7.0-rc.0.0.20240517123944-7cef272d4458/pkg/db/db_client/db_client_connect.go (about)

     1  package db_client
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/jackc/pgx/v5"
     8  	"github.com/jackc/pgx/v5/pgxpool"
     9  	"github.com/spf13/viper"
    10  	"github.com/turbot/go-kit/helpers"
    11  	"github.com/turbot/steampipe/pkg/constants"
    12  	"github.com/turbot/steampipe/pkg/constants/runtime"
    13  	"github.com/turbot/steampipe/pkg/db/db_common"
    14  	"github.com/turbot/steampipe/pkg/utils"
    15  )
    16  
    17  const (
    18  	MaxConnLifeTime = 10 * time.Minute
    19  	MaxConnIdleTime = 1 * time.Minute
    20  )
    21  
    22  type DbConnectionCallback func(context.Context, *pgx.Conn) error
    23  
    24  func (c *DbClient) establishConnectionPool(ctx context.Context, overrides clientConfig) error {
    25  	utils.LogTime("db_client.establishConnectionPool start")
    26  	defer utils.LogTime("db_client.establishConnectionPool end")
    27  
    28  	config, err := pgxpool.ParseConfig(c.connectionString)
    29  	if err != nil {
    30  		return err
    31  	}
    32  
    33  	locals := []string{
    34  		"127.0.0.1",
    35  		"::1",
    36  		"localhost",
    37  	}
    38  
    39  	// when connected to a service which is running a plugin compiled with SDK pre-v5, the plugin
    40  	// will not have the ability to turn off caching (feature introduced in SDKv5)
    41  	//
    42  	// the 'isLocalService' is used to set the client end cache to 'false' if caching is turned off in the local service
    43  	//
    44  	// this is a temporary workaround to make sure
    45  	// that we can turn off caching for plugins compiled with SDK pre-V5
    46  	// worst case scenario is that we don't switch off the cache for pre-V5 plugins
    47  	// refer to: https://github.com/turbot/steampipe/blob/f7f983a552a07e50e526fcadf2ccbfdb7b247cc0/pkg/db/db_client/db_client_session.go#L66
    48  	if helpers.StringSliceContains(locals, config.ConnConfig.Host) {
    49  		c.isLocalService = true
    50  	}
    51  
    52  	// MinConns should default to 0, but when not set, it actually get very high values (e.g. 80217984)
    53  	// this leads to a huge number of connections getting created
    54  	// TODO BINAEK dig into this and figure out why this is happening.
    55  	// We need to be sure that it is not an issue with service management
    56  	config.MinConns = 0
    57  	config.MaxConns = int32(db_common.MaxDbConnections())
    58  	config.MaxConnLifetime = MaxConnLifeTime
    59  	config.MaxConnIdleTime = MaxConnIdleTime
    60  	if c.onConnectionCallback != nil {
    61  		config.AfterConnect = c.onConnectionCallback
    62  	}
    63  	// set an app name so that we can track database connections from this Steampipe execution
    64  	// this is used to determine whether the database can safely be closed
    65  	config.ConnConfig.Config.RuntimeParams = map[string]string{
    66  		constants.RuntimeParamsKeyApplicationName: runtime.ClientConnectionAppName,
    67  	}
    68  
    69  	// apply any overrides
    70  	// this is used to set the pool size and lifetimes of the connections from up top
    71  	overrides.userPoolSettings.apply(config)
    72  
    73  	// this returns connection pool
    74  	dbPool, err := pgxpool.NewWithConfig(context.Background(), config)
    75  	if err != nil {
    76  		return err
    77  	}
    78  
    79  	err = db_common.WaitForPool(
    80  		ctx,
    81  		dbPool,
    82  		db_common.WithRetryInterval(constants.DBConnectionRetryBackoff),
    83  		db_common.WithTimeout(time.Duration(viper.GetInt(constants.ArgDatabaseStartTimeout))*time.Second),
    84  	)
    85  	if err != nil {
    86  		return err
    87  	}
    88  	c.userPool = dbPool
    89  
    90  	return c.establishManagementConnectionPool(ctx, config, overrides)
    91  }
    92  
    93  // establishManagementConnectionPool creates a connection pool to use to execute
    94  // system-initiated queries (loading of connection state etc.)
    95  // unlike establishConnectionPool, which is run first to create the user-query pool
    96  // this doesn't wait for the pool to completely start, as establishConnectionPool will have established and verified a connection with the service
    97  func (c *DbClient) establishManagementConnectionPool(ctx context.Context, config *pgxpool.Config, overrides clientConfig) error {
    98  	utils.LogTime("db_client.establishSystemConnectionPool start")
    99  	defer utils.LogTime("db_client.establishSystemConnectionPool end")
   100  
   101  	// create a config from the config of the user pool
   102  	copiedConfig := createManagementPoolConfig(config, overrides)
   103  
   104  	// this returns connection pool
   105  	dbPool, err := pgxpool.NewWithConfig(context.Background(), copiedConfig)
   106  	if err != nil {
   107  		return err
   108  	}
   109  	c.managementPool = dbPool
   110  	return nil
   111  }
   112  
   113  func createManagementPoolConfig(config *pgxpool.Config, overrides clientConfig) *pgxpool.Config {
   114  	// create a copy - we will be modifying this
   115  	copiedConfig := config.Copy()
   116  
   117  	// update the app name of the connection
   118  	copiedConfig.ConnConfig.Config.RuntimeParams = map[string]string{
   119  		constants.RuntimeParamsKeyApplicationName: runtime.ClientSystemConnectionAppName,
   120  	}
   121  
   122  	// remove the afterConnect hook - we don't need the session data in management connections
   123  	copiedConfig.AfterConnect = nil
   124  
   125  	overrides.managementPoolSettings.apply(copiedConfig)
   126  
   127  	return copiedConfig
   128  }