First class schema support through RIVER_SCHEMA and -schema

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add first-class custom database schema support to the `riverui` and `riverproui` executables through the `-schema` flag and `RIVER_SCHEMA` environment variable. [PR #629](https://github.com/riverqueue/riverui/pull/629).
+
 ### Fixed
 
 - Job retry: keep the current page open and show an actionable error when a retry conflicts with another active unique job. [PR #623](https://github.com/riverqueue/riverui/pull/623).

cmd/riverui/main.go

@@ -14,8 +14,8 @@ import (
 
 func main() {
 	riveruicmd.Run(
-		func(dbPool *pgxpool.Pool) (*river.Client[pgx.Tx], error) {
-			return river.NewClient(riverpgxv5.New(dbPool), &river.Config{})
+		func(dbPool *pgxpool.Pool, opts *riveruicmd.ClientOpts) (*river.Client[pgx.Tx], error) {
+			return river.NewClient(riverpgxv5.New(dbPool), &river.Config{Schema: opts.Schema})
 		},
 		func(client *river.Client[pgx.Tx]) uiendpoints.Bundle {
 			return riverui.NewEndpoints(client, nil)

docs/README.md

@@ -47,6 +47,42 @@ See [health checks](health_checks.md).
 
 ## Configuration
 
+### Custom database schema
+
+River UI can query River tables in a non-default PostgreSQL schema without changing the connection's `search_path`. Set the schema with the `-schema` flag or the `RIVER_SCHEMA` environment variable. An explicit flag takes precedence over the environment variable. If neither is set, River UI continues to use PostgreSQL's configured `search_path`.
+
+Use the same schema when running River migrations and River UI:
+
+```sh
+$ export RIVER_SCHEMA=river
+$ river migrate-up --database-url "$DATABASE_URL" --schema "$RIVER_SCHEMA"
+$ riverui
+```
+
+The schema can also be supplied directly with `riverui -schema=river`. Both `riverui` and `riverproui` support the same setting.
+
+For the container image, pass `RIVER_SCHEMA` alongside the database URL:
+
+```sh
+$ docker run -p 8080:8080 --env DATABASE_URL --env RIVER_SCHEMA ghcr.io/riverqueue/riverui:latest
+```
+
+When embedding River UI in a Go application, configure the schema on the River client:
+
+```go
+client, err := river.NewClient(driver, &river.Config{
+	Schema: "river",
+})
+if err != nil {
+	// handle error
+}
+
+handler, err := riverui.NewHandler(&riverui.HandlerOpts{
+	Endpoints: riverui.NewEndpoints(client, nil),
+	// ...
+})
+```
+
 ### Custom path prefix
 
 Serve River UI under a URL prefix like `/ui` by setting `-prefix` (binary) or `PATH_PREFIX` (Docker). Rules: **must start with `/`**, use `/` for no prefix, and a trailing `/` is ignored.

handler_api_endpoint.go

@@ -117,6 +117,7 @@ func (a *autocompleteListEndpoint[TTx]) Execute(ctx context.Context, req *autoco
 				Exclude: req.Exclude,
 				Match:   match,
 				Max:     100,
+				Schema:  a.Client.Schema(),
 			})
 			if err != nil {
 				return nil, fmt.Errorf("error listing job kinds: %w", err)
@@ -136,6 +137,7 @@ func (a *autocompleteListEndpoint[TTx]) Execute(ctx context.Context, req *autoco
 				Exclude: req.Exclude,
 				Match:   match,
 				Max:     100,
+				Schema:  a.Client.Schema(),
 			})
 			if err != nil {
 				return nil, fmt.Errorf("error listing queue names: %w", err)

handler_api_endpoint_test.go

@@ -78,6 +78,48 @@ func setupEndpointWithOpts[TEndpoint any](ctx context.Context, t *testing.T, ini
 	}
 }
 
+func setupEndpointWithCustomSchema[TEndpoint any](ctx context.Context, t *testing.T, initFunc func(bundle apibundle.APIBundle[pgx.Tx]) *TEndpoint) (*TEndpoint, *setupEndpointTestBundle) {
+	t.Helper()
+
+	var (
+		logger = riversharedtest.Logger(t)
+		driver = riverpgxv5.New(riversharedtest.DBPool(ctx, t))
+		schema = riverdbtest.TestSchema(ctx, t, driver, nil)
+	)
+
+	exec, err := driver.GetExecutor().Begin(ctx)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, exec.Rollback(ctx)) })
+	require.NoError(t, exec.Exec(ctx, "SET LOCAL search_path TO public"))
+
+	client, err := river.NewClient(driver, &river.Config{
+		Logger: logger,
+		Schema: schema,
+	})
+	require.NoError(t, err)
+
+	endpoint := initFunc(apibundle.APIBundle[pgx.Tx]{
+		Archetype:  riversharedtest.BaseServiceArchetype(t),
+		Client:     client,
+		DB:         exec,
+		Driver:     driver,
+		Extensions: func(_ context.Context) (map[string]bool, error) { return map[string]bool{}, nil },
+		Logger:     logger,
+	})
+
+	if service, ok := any(endpoint).(startstop.Service); ok {
+		require.NoError(t, service.Start(ctx))
+		t.Cleanup(service.Stop)
+	}
+
+	return endpoint, &setupEndpointTestBundle{
+		client: client,
+		exec:   exec,
+		logger: logger,
+		tx:     driver.UnwrapTx(exec),
+	}
+}
+
 func testMountOpts(t *testing.T) *apiendpoint.MountOpts {
 	t.Helper()
 	return &apiendpoint.MountOpts{
@@ -236,6 +278,42 @@ func TestAPIHandlerAutocompleteList(t *testing.T) {
 	})
 }
 
+func TestAPIHandlerAutocompleteListCustomSchema(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	endpoint, bundle := setupEndpointWithCustomSchema(ctx, t, newAutocompleteListEndpoint)
+	schema := bundle.client.Schema()
+
+	jobParams := testfactory.Job_Build(t, &testfactory.JobOpts{Kind: ptrutil.Ptr("custom_schema_job")})
+	jobParams.Schema = schema
+	_, err := bundle.exec.JobInsertFull(ctx, jobParams)
+	require.NoError(t, err)
+
+	_, err = bundle.exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{
+		Metadata: []byte("{}"),
+		Name:     "custom_schema_queue",
+		Schema:   schema,
+	})
+	require.NoError(t, err)
+
+	jobKindResp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &autocompleteListRequest{
+		Facet: autocompleteFacetJobKind,
+		Match: ptrutil.Ptr("custom_schema_job"),
+	})
+	require.NoError(t, err)
+	require.Len(t, jobKindResp.Data, 1)
+	require.Equal(t, "custom_schema_job", *jobKindResp.Data[0])
+
+	queueNameResp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &autocompleteListRequest{
+		Facet: autocompleteFacetQueueName,
+		Match: ptrutil.Ptr("custom_schema_queue"),
+	})
+	require.NoError(t, err)
+	require.Len(t, queueNameResp.Data, 1)
+	require.Equal(t, "custom_schema_queue", *queueNameResp.Data[0])
+}
+
 func TestAPIHandlerFeaturesGet(t *testing.T) {
 	t.Parallel()
 

internal/riveruicmd/auth_middleware_test.go

@@ -41,8 +41,8 @@ func TestAuthMiddleware(t *testing.T) { //nolint:tparallel
 				logger:     riversharedtest.Logger(t),
 				pathPrefix: prefix,
 			},
-			func(dbPool *pgxpool.Pool) (*river.Client[pgx.Tx], error) {
-				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{})
+			func(dbPool *pgxpool.Pool, opts *ClientOpts) (*river.Client[pgx.Tx], error) {
+				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{Schema: opts.Schema})
 			},
 			func(client *river.Client[pgx.Tx]) uiendpoints.Bundle {
 				return riverui.NewEndpoints(client, nil)

internal/riveruicmd/riveruicmd.go

@@ -28,7 +28,11 @@ type BundleOpts struct {
 	JobListHideArgsByDefault bool
 }
 
-func Run[TClient any](createClient func(*pgxpool.Pool) (TClient, error), createBundle func(TClient) uiendpoints.Bundle) {
+type ClientOpts struct {
+	Schema string
+}
+
+func Run[TClient any](createClient func(*pgxpool.Pool, *ClientOpts) (TClient, error), createBundle func(TClient) uiendpoints.Bundle) {
 	ctx := context.Background()
 
 	logger := slog.New(getLogHandler(&slog.HandlerOptions{
@@ -38,6 +42,9 @@ func Run[TClient any](createClient func(*pgxpool.Pool) (TClient, error), createB
 	var pathPrefix string
 	flag.StringVar(&pathPrefix, "prefix", "/", "path prefix for API and UI routes (must start with '/', use '/' for no prefix)")
 
+	var schema string
+	flag.StringVar(&schema, "schema", os.Getenv("RIVER_SCHEMA"), "name of non-default database schema where River tables are located")
+
 	var healthCheckName string
 	flag.StringVar(&healthCheckName, "healthcheck", "", "the name of the health checks: minimal or complete")
 
@@ -57,6 +64,7 @@ func Run[TClient any](createClient func(*pgxpool.Pool) (TClient, error), createB
 	initRes, err := initServer(ctx, &initServerOpts{
 		logger:             logger,
 		pathPrefix:         pathPrefix,
+		schema:             schema,
 		silentHealthChecks: silentHealthChecks,
 	}, createClient, createBundle)
 	if err != nil {
@@ -141,10 +149,11 @@ type initServerResult struct {
 type initServerOpts struct {
 	logger             *slog.Logger
 	pathPrefix         string
+	schema             string
 	silentHealthChecks bool
 }
 
-func initServer[TClient any](ctx context.Context, opts *initServerOpts, createClient func(*pgxpool.Pool) (TClient, error), createBundle func(TClient) uiendpoints.Bundle) (*initServerResult, error) {
+func initServer[TClient any](ctx context.Context, opts *initServerOpts, createClient func(*pgxpool.Pool, *ClientOpts) (TClient, error), createBundle func(TClient) uiendpoints.Bundle) (*initServerResult, error) {
 	if opts == nil {
 		return nil, errors.New("opts is required")
 	}
@@ -184,7 +193,7 @@ func initServer[TClient any](ctx context.Context, opts *initServerOpts, createCl
 		return nil, fmt.Errorf("error connecting to db: %w", err)
 	}
 
-	client, err := createClient(dbPool)
+	client, err := createClient(dbPool, &ClientOpts{Schema: opts.schema})
 	if err != nil {
 		return nil, err
 	}

internal/riveruicmd/riveruicmd_test.go

@@ -42,8 +42,8 @@ func TestInitServer(t *testing.T) { //nolint:tparallel
 			logger:     riversharedtest.Logger(t),
 			pathPrefix: "/",
 		},
-			func(dbPool *pgxpool.Pool) (*river.Client[pgx.Tx], error) {
-				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{})
+			func(dbPool *pgxpool.Pool, opts *ClientOpts) (*river.Client[pgx.Tx], error) {
+				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{Schema: opts.Schema})
 			},
 			func(client *river.Client[pgx.Tx]) uiendpoints.Bundle {
 				return riverui.NewEndpoints(client, nil)
@@ -84,6 +84,30 @@ func TestInitServer(t *testing.T) { //nolint:tparallel
 		require.NoError(t, err)
 	})
 
+	t.Run("PassesClientOptions", func(t *testing.T) {
+		t.Parallel()
+
+		const schema = "river_custom"
+
+		var receivedOpts *ClientOpts
+		initRes, err := initServer(ctx, &initServerOpts{
+			logger:     riversharedtest.Logger(t),
+			pathPrefix: "/",
+			schema:     schema,
+		},
+			func(dbPool *pgxpool.Pool, opts *ClientOpts) (*river.Client[pgx.Tx], error) {
+				receivedOpts = opts
+				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{})
+			},
+			func(client *river.Client[pgx.Tx]) uiendpoints.Bundle {
+				return riverui.NewEndpoints(client, nil)
+			},
+		)
+		require.NoError(t, err)
+		t.Cleanup(initRes.dbPool.Close)
+		require.Equal(t, &ClientOpts{Schema: schema}, receivedOpts)
+	})
+
 	t.Run("JobListHideArgsByDefault", func(t *testing.T) {
 		t.Run("DefaultValueIsFalse", func(t *testing.T) {
 			t.Parallel()
@@ -182,8 +206,8 @@ func TestSilentHealthchecks_SuppressesLogs(t *testing.T) {
 			pathPrefix:         prefix,
 			silentHealthChecks: silent,
 		},
-			func(dbPool *pgxpool.Pool) (*river.Client[pgx.Tx], error) {
-				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{})
+			func(dbPool *pgxpool.Pool, opts *ClientOpts) (*river.Client[pgx.Tx], error) {
+				return river.NewClient(riverpgxv5.New(dbPool), &river.Config{Schema: opts.Schema})
 			},
 			func(client *river.Client[pgx.Tx]) uiendpoints.Bundle {
 				return riverui.NewEndpoints(client, nil)

riverproui/cmd/riverproui/main.go

@@ -4,6 +4,8 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 
+	"github.com/riverqueue/river"
+
 	"riverqueue.com/riverpro"
 	"riverqueue.com/riverpro/driver/riverpropgxv5"
 
@@ -14,8 +16,10 @@ import (
 
 func main() {
 	riveruicmd.Run(
-		func(dbPool *pgxpool.Pool) (*riverpro.Client[pgx.Tx], error) {
-			return riverpro.NewClient(riverpropgxv5.New(dbPool), &riverpro.Config{})
+		func(dbPool *pgxpool.Pool, opts *riveruicmd.ClientOpts) (*riverpro.Client[pgx.Tx], error) {
+			return riverpro.NewClient(riverpropgxv5.New(dbPool), &riverpro.Config{
+				Config: river.Config{Schema: opts.Schema},
+			})
 		},
 		func(client *riverpro.Client[pgx.Tx]) uiendpoints.Bundle {
 			return riverproui.NewEndpoints(client, nil)

riverproui/go.mod

@@ -67,4 +67,4 @@ retract (
 
 // replace riverqueue.com/riverpro/driver/riverpropgxv5 => ../../riverpro/driver/riverpropgxv5
 
-// replace riverqueue.com/riverui => ../
+replace riverqueue.com/riverui => ../

riverproui/internal/prohandler/pro_handler_api_endpoints.go

@@ -722,6 +722,7 @@ func (a *workflowListEndpoint[TTx]) Execute(ctx context.Context, req *workflowLi
 		workflows, err := a.DB.WorkflowListActive(ctx, &riverprodriver.WorkflowListParams{
 			After:           ptrutil.ValOrDefault(req.After, ""),
 			PaginationLimit: min(ptrutil.ValOrDefault(req.Limit, 100), 1000),
+			Schema:          a.Client.Schema(),
 		})
 		if err != nil {
 			return nil, fmt.Errorf("error listing workflows: %w", err)
@@ -732,6 +733,7 @@ func (a *workflowListEndpoint[TTx]) Execute(ctx context.Context, req *workflowLi
 		workflows, err := a.DB.WorkflowListInactive(ctx, &riverprodriver.WorkflowListParams{
 			After:           ptrutil.ValOrDefault(req.After, ""),
 			PaginationLimit: min(ptrutil.ValOrDefault(req.Limit, 100), 1000),
+			Schema:          a.Client.Schema(),
 		})
 		if err != nil {
 			return nil, fmt.Errorf("error listing workflows: %w", err)
@@ -741,6 +743,7 @@ func (a *workflowListEndpoint[TTx]) Execute(ctx context.Context, req *workflowLi
 		workflows, err := a.DB.WorkflowListAll(ctx, &riverprodriver.WorkflowListParams{
 			After:           ptrutil.ValOrDefault(req.After, ""),
 			PaginationLimit: min(ptrutil.ValOrDefault(req.Limit, 100), 1000),
+			Schema:          a.Client.Schema(),
 		})
 		if err != nil {
 			return nil, fmt.Errorf("error listing workflows: %w", err)

riverproui/internal/prohandler/pro_handler_api_endpoints_test.go

@@ -255,6 +255,41 @@ func TestProAPIHandlerWorkflowGet(t *testing.T) {
 	})
 }
 
+func TestProAPIHandlerWorkflowListCustomSchema(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	endpoint, bundle := setupEndpoint(ctx, t, NewWorkflowListEndpoint)
+
+	require.NoError(t, bundle.exec.WorkflowInsertMany(ctx, &driver.WorkflowInsertManyParams{
+		IDs:    []string{"active_workflow", "inactive_workflow"},
+		Names:  []string{"Active workflow", "Inactive workflow"},
+		Schema: bundle.schema,
+	}))
+
+	makeWorkflowJob(ctx, t, bundle.exec, bundle.schema, "active_workflow", "active_task", nil)
+	jobWithSchema(ctx, t, bundle.exec, bundle.schema, &testfactory.JobOpts{
+		FinalizedAt: ptrutil.Ptr(time.Now()),
+		Metadata:    workflowMetadata("inactive_workflow", "inactive_task", nil),
+		State:       ptrutil.Ptr(rivertype.JobStateCompleted),
+	})
+
+	activeResp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &workflowListRequest{State: "active"})
+	require.NoError(t, err)
+	require.Len(t, activeResp.Data, 1)
+	require.Equal(t, "active_workflow", activeResp.Data[0].ID)
+
+	inactiveResp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &workflowListRequest{State: "inactive"})
+	require.NoError(t, err)
+	require.Len(t, inactiveResp.Data, 1)
+	require.Equal(t, "inactive_workflow", inactiveResp.Data[0].ID)
+
+	allResp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &workflowListRequest{})
+	require.NoError(t, err)
+	require.Len(t, allResp.Data, 2)
+	require.ElementsMatch(t, []string{"active_workflow", "inactive_workflow"}, []string{allResp.Data[0].ID, allResp.Data[1].ID})
+}
+
 func TestProAPIHandlerWorkflowListRequiresV2Tables(t *testing.T) {
 	t.Parallel()