Function Query

backend/plugin/advisor/utils.go:23–138

Query runs the EXPLAIN or SELECT statements for advisors.

(ctx context.Context, qCtx QueryContext, connection *sql.DB, engine storepb.Engine, statement string)

// Query runs the EXPLAIN or SELECT statements for advisors.
func Query(ctx context.Context, qCtx QueryContext, connection *sql.DB, engine storepb.Engine, statement string) ([]any, error) {
	tx, err := connection.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	if engine == storepb.Engine_POSTGRES && qCtx.TenantMode {
		const query = `
		SELECT
			u.rolname
		FROM
			pg_roles AS u JOIN pg_database AS d ON (d.datdba = u.oid)
		WHERE
			d.datname = current_database();
		`
		var owner string
		if err := tx.QueryRowContext(ctx, query).Scan(&owner); err != nil {
			return nil, err
		}
		// Use SET SESSION ROLE to match the execution logic in backend/plugin/db/pg/pg.go
		if _, err := tx.ExecContext(ctx, fmt.Sprintf("SET SESSION ROLE '%s';", owner)); err != nil { // NOSONAR(go:S2077) owner is from pg_roles system catalog, not user input
			return nil, err
		}
	}

	for _, preExec := range qCtx.PreExecutions {
		if preExec != "" {
			if _, err := tx.ExecContext(ctx, preExec); err != nil {
				return nil, errors.Wrapf(err, "failed to execute pre-execution: %s", preExec)
			}
		}
	}

	rows, err := tx.QueryContext(ctx, statement)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	columnNames, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	columnTypes, err := rows.ColumnTypes()
	if err != nil {
		return nil, err
	}

	colCount := len(columnTypes)

	var columnTypeNames []string
	for _, v := range columnTypes {
		// DatabaseTypeName returns the database system name of the column type.
		// refer: https://pkg.go.dev/database/sql#ColumnType.DatabaseTypeName
		columnTypeNames = append(columnTypeNames, strings.ToUpper(v.DatabaseTypeName()))
	}