github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/postgres/testutil.go (about)

     1  //go:build ci && docker
     2  // +build ci,docker
     3  
     4  package postgres
     5  
     6  import (
     7  	"context"
     8  	"strings"
     9  
    10  	"github.com/jackc/pgx/v5"
    11  	"github.com/jackc/pgx/v5/pgconn"
    12  
    13  	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
    14  )
    15  
    16  func getExplanation(ctx context.Context, querier pgxcommon.Querier, sql string, args []any) (string, error) {
    17  	// Make sure Postgres stats are fully up-to-date so it selects the correct index.
    18  	_, err := querier.Exec(ctx, "ANALYZE;")
    19  	if err != nil {
    20  		return "", err
    21  	}
    22  
    23  	explainRows, err := querier.Query(ctx, "EXPLAIN ANALYZE "+sql, args...)
    24  	if err != nil {
    25  		return "", err
    26  	}
    27  
    28  	explanation := ""
    29  	for explainRows.Next() {
    30  		explanation += string(explainRows.RawValues()[0]) + "\n"
    31  	}
    32  	explainRows.Close()
    33  	if err := explainRows.Err(); err != nil {
    34  		return "", err
    35  	}
    36  	return explanation, nil
    37  }
    38  
    39  type withQueryInterceptor struct {
    40  	explanations map[string]string
    41  }
    42  
    43  func (ql *withQueryInterceptor) InterceptExec(ctx context.Context, querier pgxcommon.Querier, sql string, args ...interface{}) (pgconn.CommandTag, error) {
    44  	if strings.HasPrefix(sql, "WITH") {
    45  		// Note, we disable seqscan here to ensure we get an index scan for testing.
    46  		_, err := querier.Exec(ctx, "set enable_seqscan = off;")
    47  		if err != nil {
    48  			return pgconn.CommandTag{}, err
    49  		}
    50  
    51  		explanation, err := getExplanation(ctx, querier, sql, args)
    52  		if err != nil {
    53  			return pgconn.CommandTag{}, err
    54  		}
    55  
    56  		ql.explanations[sql] = explanation
    57  	}
    58  
    59  	return querier.Exec(ctx, sql, args...)
    60  }
    61  
    62  func (ql *withQueryInterceptor) InterceptQueryRow(ctx context.Context, querier pgxcommon.Querier, sql string, optionsAndArgs ...interface{}) pgx.Row {
    63  	return querier.QueryRow(ctx, sql, optionsAndArgs...)
    64  }
    65  
    66  func (ql *withQueryInterceptor) InterceptQuery(ctx context.Context, querier pgxcommon.Querier, sql string, args ...interface{}) (pgx.Rows, error) {
    67  	return querier.Query(ctx, sql, args...)
    68  }