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

     1  package db_common
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  
     8  	"github.com/jackc/pgx/v5"
     9  	"github.com/turbot/steampipe-plugin-sdk/v5/sperr"
    10  	"github.com/turbot/steampipe/pkg/constants"
    11  	"github.com/turbot/steampipe/pkg/constants/runtime"
    12  )
    13  
    14  // SystemClientExecutor is the executor function that is called within a transaction
    15  // make sure that by the time the executor finishes execution, the connection is freed
    16  // otherwise we will get a `conn is busy` error
    17  type SystemClientExecutor func(context.Context, pgx.Tx) error
    18  
    19  // ExecuteSystemClientCall creates a transaction and sets the application_name to the
    20  // one used by the system client, executes the callback and sets the application name back to the client app name
    21  func ExecuteSystemClientCall(ctx context.Context, conn *pgx.Conn, executor SystemClientExecutor) error {
    22  	if !IsClientAppName(conn.Config().RuntimeParams[constants.RuntimeParamsKeyApplicationName]) {
    23  		// this should NEVER happen
    24  		return sperr.New("ExecuteSystemClientCall called with appname other than client: %s", conn.Config().RuntimeParams[constants.RuntimeParamsKeyApplicationName])
    25  	}
    26  
    27  	return pgx.BeginFunc(ctx, conn, func(tx pgx.Tx) (e error) {
    28  		// if the appName is the ClientAppName, we need to set it to ClientSystemAppName
    29  		// and then revert when done
    30  		_, err := tx.Exec(ctx, fmt.Sprintf("SET application_name TO '%s'", runtime.ClientSystemConnectionAppName))
    31  		if err != nil {
    32  			return sperr.WrapWithRootMessage(err, "could not set application name on connection")
    33  		}
    34  		defer func() {
    35  			// set back the original application name
    36  			_, err = tx.Exec(ctx, fmt.Sprintf("SET application_name TO '%s'", conn.Config().RuntimeParams[constants.RuntimeParamsKeyApplicationName]))
    37  			if err != nil {
    38  				log.Println("[TRACE] could not reset application_name", e)
    39  			}
    40  			// if there is not already an error, set the error
    41  			if e == nil {
    42  				e = err
    43  			}
    44  		}()
    45  
    46  		if err := executor(ctx, tx); err != nil {
    47  			return sperr.WrapWithMessage(err, "system client query execution failed")
    48  		}
    49  		return nil
    50  	})
    51  }