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

     1  //go:build docker
     2  // +build docker
     3  
     4  package datastore
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"testing"
    10  
    11  	"github.com/google/uuid"
    12  	"github.com/jackc/pgx/v5"
    13  	"github.com/ory/dockertest/v3"
    14  	"github.com/stretchr/testify/require"
    15  
    16  	crdbmigrations "github.com/authzed/spicedb/internal/datastore/crdb/migrations"
    17  	"github.com/authzed/spicedb/pkg/datastore"
    18  	"github.com/authzed/spicedb/pkg/migrate"
    19  	"github.com/authzed/spicedb/pkg/secrets"
    20  )
    21  
    22  const (
    23  	CRDBTestVersionTag = "v23.1.16"
    24  
    25  	enableRangefeeds = `SET CLUSTER SETTING kv.rangefeed.enabled = true;`
    26  )
    27  
    28  type crdbTester struct {
    29  	conn     *pgx.Conn
    30  	hostname string
    31  	creds    string
    32  	port     string
    33  }
    34  
    35  // RunCRDBForTesting returns a RunningEngineForTest for CRDB
    36  func RunCRDBForTesting(t testing.TB, bridgeNetworkName string) RunningEngineForTest {
    37  	pool, err := dockertest.NewPool("")
    38  	require.NoError(t, err)
    39  
    40  	name := fmt.Sprintf("crds-%s", uuid.New().String())
    41  	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
    42  		Name:       name,
    43  		Repository: "mirror.gcr.io/cockroachdb/cockroach",
    44  		Tag:        CRDBTestVersionTag,
    45  		Cmd:        []string{"start-single-node", "--insecure", "--max-offset=50ms"},
    46  		NetworkID:  bridgeNetworkName,
    47  	})
    48  	require.NoError(t, err)
    49  
    50  	builder := &crdbTester{
    51  		hostname: "localhost",
    52  		creds:    "root:fake",
    53  	}
    54  	t.Cleanup(func() {
    55  		require.NoError(t, pool.Purge(resource))
    56  	})
    57  
    58  	port := resource.GetPort(fmt.Sprintf("%d/tcp", 26257))
    59  	if bridgeNetworkName != "" {
    60  		builder.hostname = name
    61  		builder.port = "26257"
    62  	} else {
    63  		builder.port = port
    64  	}
    65  
    66  	uri := fmt.Sprintf("postgres://%s@localhost:%s/defaultdb?sslmode=disable", builder.creds, port)
    67  	require.NoError(t, pool.Retry(func() error {
    68  		var err error
    69  		ctx, cancelConnect := context.WithTimeout(context.Background(), dockerBootTimeout)
    70  		defer cancelConnect()
    71  		builder.conn, err = pgx.Connect(ctx, uri)
    72  		if err != nil {
    73  			return err
    74  		}
    75  		ctx, cancelRangeFeeds := context.WithTimeout(context.Background(), dockerBootTimeout)
    76  		defer cancelRangeFeeds()
    77  		_, err = builder.conn.Exec(ctx, enableRangefeeds)
    78  		return err
    79  	}))
    80  
    81  	return builder
    82  }
    83  
    84  func (r *crdbTester) NewDatabase(t testing.TB) string {
    85  	uniquePortion, err := secrets.TokenHex(4)
    86  	require.NoError(t, err)
    87  
    88  	newDBName := "db" + uniquePortion
    89  
    90  	_, err = r.conn.Exec(context.Background(), "CREATE DATABASE "+newDBName)
    91  	require.NoError(t, err)
    92  
    93  	connectStr := fmt.Sprintf(
    94  		"postgres://%s@%s:%s/%s?sslmode=disable",
    95  		r.creds,
    96  		r.hostname,
    97  		r.port,
    98  		newDBName,
    99  	)
   100  	return connectStr
   101  }
   102  
   103  func (r *crdbTester) NewDatastore(t testing.TB, initFunc InitFunc) datastore.Datastore {
   104  	connectStr := r.NewDatabase(t)
   105  
   106  	migrationDriver, err := crdbmigrations.NewCRDBDriver(connectStr)
   107  	require.NoError(t, err)
   108  	require.NoError(t, crdbmigrations.CRDBMigrations.Run(context.Background(), migrationDriver, migrate.Head, migrate.LiveRun))
   109  
   110  	return initFunc("cockroachdb", connectStr)
   111  }