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

     1  package db_local
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"sort"
     8  	"strings"
     9  
    10  	"github.com/jackc/pgx/v5/pgxpool"
    11  	"github.com/spf13/viper"
    12  	"github.com/turbot/steampipe/pkg/constants"
    13  	"github.com/turbot/steampipe/pkg/db/db_common"
    14  	"github.com/turbot/steampipe/pkg/steampipeconfig"
    15  	"github.com/turbot/steampipe/pkg/steampipeconfig/modconfig"
    16  )
    17  
    18  func SetUserSearchPath(ctx context.Context, pool *pgxpool.Pool) ([]string, error) {
    19  	var searchPath []string
    20  
    21  	// is there a user search path in the config?
    22  	// check ConfigKeyDatabaseSearchPath config (this is the value specified in the database config)
    23  	if viper.IsSet(constants.ConfigKeyServerSearchPath) {
    24  		searchPath = viper.GetStringSlice(constants.ConfigKeyServerSearchPath)
    25  		// the Internal Schema should always go at the end
    26  		searchPath = db_common.EnsureInternalSchemaSuffix(searchPath)
    27  	} else {
    28  		prefix := viper.GetStringSlice(constants.ConfigKeyServerSearchPathPrefix)
    29  		// no config set - set user search path to default
    30  		// - which is all the connection names, book-ended with public and internal
    31  		searchPath = append(prefix, getDefaultSearchPath()...)
    32  	}
    33  
    34  	// escape the schema names
    35  	escapedSearchPath := db_common.PgEscapeSearchPath(searchPath)
    36  
    37  	log.Println("[TRACE] setting user search path to", searchPath)
    38  
    39  	// get all roles which are a member of steampipe_users
    40  	conn, err := pool.Acquire(ctx)
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  	defer conn.Release()
    45  
    46  	query := fmt.Sprintf(`SELECT USENAME FROM pg_user WHERE pg_has_role(usename, '%s', 'member')`, constants.DatabaseUsersRole)
    47  	rows, err := conn.Query(ctx, query)
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  
    52  	// set the search path for all these roles
    53  	var queries = []string{
    54  		"LOCK TABLE pg_user IN SHARE ROW EXCLUSIVE MODE;",
    55  	}
    56  
    57  	for rows.Next() {
    58  		var user string
    59  		if err := rows.Scan(&user); err != nil {
    60  			return nil, err
    61  		}
    62  		if user == "root" {
    63  			continue
    64  		}
    65  		queries = append(queries, fmt.Sprintf(
    66  			"ALTER USER %s SET SEARCH_PATH TO %s;",
    67  			db_common.PgEscapeName(user),
    68  			strings.Join(escapedSearchPath, ","),
    69  		))
    70  	}
    71  
    72  	log.Printf("[TRACE] user search path sql: %v", queries)
    73  	_, err = ExecuteSqlInTransaction(ctx, conn.Conn(), queries...)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	return searchPath, nil
    78  }
    79  
    80  // GetDefaultSearchPath builds default search path from the connection schemas, book-ended with public and internal
    81  func getDefaultSearchPath() []string {
    82  	// add all connections to the seatrch path (UNLESS ImportSchema is disabled)
    83  	var searchPath []string
    84  	for connectionName, connection := range steampipeconfig.GlobalConfig.Connections {
    85  		if connection.ImportSchema == modconfig.ImportSchemaEnabled {
    86  			searchPath = append(searchPath, connectionName)
    87  		}
    88  	}
    89  
    90  	sort.Strings(searchPath)
    91  	// add the 'public' schema as the first schema in the search_path. This makes it
    92  	// easier for users to build and work with their own tables, and since it's normally
    93  	// empty, doesn't make using steampipe tables any more difficult.
    94  	searchPath = append([]string{"public"}, searchPath...)
    95  	// add 'internal' schema as last schema in the search path
    96  	searchPath = append(searchPath, constants.InternalSchema)
    97  
    98  	return searchPath
    99  }