Proof out explicit schema injection without `search_path` #798

brandur · 2025-03-08T06:13:18Z

Here, follow up #794 to proof out the injection of an explicit schema
into SQL queries so that putting River in a custom schema doesn't
require the use of search_path. We use the newly introduced sqlc
template package to prefix table names like so:

-- name: JobCountByState :one
SELECT count(*)
FROM /* TEMPLATE: schema */river_job
WHERE state = @state;

The main advantage of this approach compared to search_path is that
connections no longer require that their search_path be set with an
explicit value to work. search_path can be set or misset, but because
table names are prefixed with the right schema name already, queries
still go through. This is especially useful in the context of PgBouncer,
where a valid search_path setting can't be guaranteed.

Notably, this change doesn't bring in enough new testing to prove that
River really works with explicit schema injection, so for the time being
this setting remains completely internal. What I'd like to try and do is
get some basic infrastructure like this in place first, then prove it
out by starting to move the test suite over to schema-specific tests. By
virtue of doing that we'd be putting the entire load of the River test
suite on the new paths, which should give us high confidence that it's
all working as intended.

brandur · 2025-03-08T06:56:18Z

riverdriver/river_driver_interface.go

+}
+
+type JobCountByStateParams struct {
+	Schema string


So I'd originally had Schema as a configuration parameter to the driver, but I realized later that I don't think that's appropriate ... not only is the UX a bit questionable, but it should really be only the client that controls its insert/work schema, and it should be possible for two clients to share a single driver and be using two separate schemas.

The slightly unfortunate part is that this approach would necessitate adding a Schema string parameter to every driver params structs where the underlying query targets a River table, but compared to the other possibilities, this might be the best/most explicit way to go.

Yep, I think that's reasonable. Explicitly passing through the schema makes the most sense and also makes it clear the driver needs to handle running its queries on an arbitrary schema.

brandur · 2025-03-08T06:57:02Z

@bgentry This one's going to be quite a bit more work to take all the way to fruition, so I just converted two queries to get your thoughts on the general approach before doing anything more.

bgentry

I definitely don't love this because it makes the queries uglier + harder to read, and it's quite the (clever) hack to work around sqlc limitations. But, it's also probably the best we can do without changing out that core piece of the stack. I think we should move forward with it.

Maybe we also want to make it so we can run a full test suite against a database that only has our migrations run in an alternate schema? Even if it mainly runs in CI it would be an effective way to ensure this all works & keeps working.

This one follows up #794 again. That change included a template cache, with the idea being that we could reuse a rendered sqlc query when all input values were expected to be stable. For example, when replacing something like a schema name, we'd presumably always be replacing the template value with the same schema over and over again, so it'd be better to save on the work. But that caching system hadn't adequately been thought through, and wouldn't really work. I realized when I was putting #798 (explicit schemas) together that if you injected two schema values from two different clients then you'd end up using the same cache value, which is wrong. Here, we tool out the cache layer a little more so that it considers input values and named args, which all must be consistent for a cached value to be returned. We also add tests to make sure of it. Building a cache key unfortunately has to rely on concatenating some strings together and presorting map keys for stability, but even with the extra work involved, a cache hit still comes out to be quite a bit faster than a miss (~2.5x faster), so it seems worth doing: $ go test ./rivershared/sqlctemplate -bench=. goos: darwin goarch: arm64 pkg: github.com/riverqueue/river/rivershared/sqlctemplate cpu: Apple M4 BenchmarkReplacer/WithCache-10 1626988 735.7 ns/op BenchmarkReplacer/WithoutCache-10 695517 1710 ns/op PASS ok github.com/riverqueue/river/rivershared/sqlctemplate 3.419s

brandur · 2025-03-09T15:15:59Z

I definitely don't love this because it makes the queries uglier + harder to read, and it's quite the (clever) hack to work around sqlc limitations. But, it's also probably the best we can do without changing out that core piece of the stack. I think we should move forward with it.

Yeah, I'm going to see if I can pass this along to Kyle. Low probability of anything happening, but maybe it'd inspire him as to what a potentially native solution for this would look like. Ideally, this lives in the sqlc layer.

Maybe we also want to make it so we can run a full test suite against a database that only has our migrations run in an alternate schema? Even if it mainly runs in CI it would be an effective way to ensure this all works & keeps working.

Yeah, we'll need something like that I think. Another one I'm considering is if we could move our test databases over to test schemas, we'd get full 100% vetting that things are working because by necessity, the test suite depends on it.

I worked on this another couple hours last night though, and even just getting the basic plumbing in is thousands of LOCs changed, and all in areas that are very likely to conflict with anything else we try to do. And I haven't even started on testing or rehabilitating the listen/notify system to be schema-aware. I'm kind of wondering if it might be smart to bring in some of the high churn changes early, but keep the new Schema config unexported for now (so schema), then follow on by doing some of the other major work streams from there, only releasing it once it's fully ready/tested.

bgentry · 2025-03-09T17:20:53Z

I'm kind of wondering if it might be smart to bring in some of the high churn changes early, but keep the new Schema config unexported for now (so schema), then follow on by doing some of the other major work streams from there, only releasing it once it's fully ready/tested.

Makes sense. It's best if we can avoid exposing it until we're confident it's ready-to-go. And given the low chance of upstream sqlc changes in a short timeframe, IMO it's also fine for us to move forward with the hacky DIY approach. We can always switch over if an official easier option becomes available.

This one follows up #794 again. That change included a template cache, with the idea being that we could reuse a rendered sqlc query when all input values were expected to be stable. For example, when replacing something like a schema name, we'd presumably always be replacing the template value with the same schema over and over again, so it'd be better to save on the work. But that caching system hadn't adequately been thought through, and wouldn't really work. I realized when I was putting #798 (explicit schemas) together that if you injected two schema values from two different clients then you'd end up using the same cache value, which is wrong. Here, we tool out the cache layer a little more so that it considers input values and named args, which all must be consistent for a cached value to be returned. We also add tests to make sure of it. Building a cache key unfortunately has to rely on concatenating some strings together and presorting map keys for stability, but even with the extra work involved, a cache hit still comes out to be quite a bit faster than a miss (~2.5x faster), so it seems worth doing: $ go test ./rivershared/sqlctemplate -bench=. goos: darwin goarch: arm64 pkg: github.com/riverqueue/river/rivershared/sqlctemplate cpu: Apple M4 BenchmarkReplacer/WithCache-10 1626988 735.7 ns/op BenchmarkReplacer/WithoutCache-10 695517 1710 ns/op PASS ok github.com/riverqueue/river/rivershared/sqlctemplate 3.419s

brandur · 2025-04-14T00:11:54Z

riverdriver/river_driver_interface.go

 	// ColumnExists checks whether a column for a particular table exists for
 	// the schema in the current search schema.
-	ColumnExists(ctx context.Context, tableName, columnName string) (bool, error)
+	ColumnExists(ctx context.Context, params *ColumnExistsParams) (bool, error)


Since everything was getting a new schema argument, I ended up just changing everything into a parameter struct so (1) we don't have some awkward functions that need to take three strings like this one, and (2) just so it's all predictable. Everything just takes a param struct.

This ended up requiring a lot of refactoring so not sure it was worth it in the end, but oh well.

brandur · 2025-04-14T00:38:11Z

@bgentry Alright, so I ended up going through and converting everything over to take a schema template param.

As mentioned above, it's unlikely this fully works given we'd have to add significant new test coverage to root out all the edges, so all the functionality is kept internal-only for the time being until the feature can be fully finished. I'm thinking next we could try to switch over to schema-based testing, and use that strategy to fully vet what's happening here as schema-based testing would imply the entire weight of the test suite on the schema injection logic here.

Thoughts?

brandur · 2025-04-14T00:43:28Z

rivermigrate/river_migrate.go

 	// or higher.
 	Logger *slog.Logger
+
+	schema string


The migrate tests do actually use schemas already, so unlike most of the rest of the new schema code, all the migrations are actually fully tested against schema injection and verified to work.

bgentry

Herculean effort 💪 I wish this wasn't all necessary but it seems it's either this brute force method or ditching sqlc altogether. Looking forward to exercising this further with schema-based tests instead of DB-based 🚀

bgentry · 2025-04-15T02:15:34Z

client.go

+	// schema is a non-standard schema where River tables are located. All table
+	// references in database queries will use this value as a prefix.
+	//
+	// Defaults to empty, which causes the client to look for tables using the
+	// setting of Postgres `search_path`.
+	schema string


I wonder, would we be better off defaulting to this explicitly setting public vs letting it default to conn string search path? Probably not, just trying to think through the full cycle of existing users migrating to this setup once we have full schema support.

I thinkk at this point we'll probably have to leave empty as default for backwards compatibility if nothing else.

Starting from scratch, yeah fair question. If we gave it a default though, we'd have to also provide a way to zero it for anyone who'd prefer to use search_path by choice.

bgentry · 2025-04-15T02:19:21Z

rivershared/riverpilot/pilot.go

 		ctx context.Context,
 		tx riverdriver.ExecutorTx,
-		params []*riverdriver.JobInsertFastParams,
+		params *riverdriver.JobInsertFastManyParams,


I'm sure you've already considered it but this will definitely require corresponding changes in Pro. Then again so will a lot of this PR :)

Yeah, I think I finished most of those up this evening thankfully. I'll send them over today or tomorrow.

bgentry · 2025-04-15T02:27:01Z

internal/jobexecutor/job_executor_test.go

+		job, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{
+			ID:     bundle.jobRow.ID,
+			Schema: "",
+		})


These all get quite a bit more verbose, unfortunately 😕

Yeah :/ I think we can clean a lot of this up again with a few more helpers for test functions that end up generating a lot of objects.

bgentry · 2025-04-15T02:27:44Z

internal/leadership/elector_test.go

 				archetype = riversharedtest.BaseServiceArchetype(t)
 				dbPool    = riverinternaltest.TestDB(ctx, t)
 				driver    = riverpgxv5.New(dbPool)
+				schema    = "" // try to make tests schema-based rather than database-based in the future


💯 I have thoughts on this, might be worth talking about possible approaches together before diving in

Soo I ended up writing a semi-functional implementation yesterday, but yeah, happy to cover more. I think this is going to be a good path though — there are a ton of advantages compared to test DBs (e.g. can create as many schemas as you need so nothing blocks waiting on DBs, full parallel testing, still get efficient reuse where possible, no need to run testdbman).

Here, follow up #794 to proof out the injection of an explicit schema into SQL queries so that putting River in a custom schema doesn't require the use of `search_path`. We use the newly introduced sqlc template package to prefix table names like so: -- name: JobCountByState :one SELECT count(*) FROM /* TEMPLATE: schema */river_job WHERE state = @State; The main advantage of this approach compared to `search_path` is that connections no longer require that their `search_path` be set with an explicit value to work. `search_path` can be set or misset, but because table names are prefixed with the right schema name already, queries still go through. This is especially useful in the context of PgBouncer, where a valid `search_path` setting can't be guaranteed. Notably, this change doesn't bring in enough new testing to prove that River really works with explicit schema injection, so for the time being this setting remains completely internal. What I'd like to try and do is get some basic infrastructure like this in place first, then prove it out by starting to move the test suite over to schema-specific tests. By virtue of doing that we'd be putting the entire load of the River test suite on the new paths, which should give us high confidence that it's all working as intended.

brandur · 2025-04-15T08:29:07Z

Thanks! Okay this is still a bit rough (schemas definitely do not work yet), but I have another branch based on this one where I've started to get large swathes of the test suite working a schema-based design. There's still a bunch of debugging that needs to be done, but I think I can get something up by EOW.

The schema stuff will be a bit of a mess in master and not usable, but as far as I know, not harmful.

Here, introduce a full version of schema injection as initially started in #798, and exposing a new `Schema` option to make it configurable even outside the test suite. This involved fixing a lot of bugs from #798, and the only way to make it possible to root them all out was to make full use of schemas across the entire test suite. The original "test DB manager" system has been replaced with a new `riverschematest.TestSchema` helper that generates a schema for use with a test case or prefers an existing idle one that was already generated for the same test run. `TestSchema` runs migrations for generated schemas which also means we don't need to use `testdbman` anymore, with tests capable of bootstrapping themselves at run time. We reuse schemas to avoid this extra work when possible, but migrating a new schema is also surprisingly fast, taking up to 50ms, but getting more down to a stable ~10ms once things are warmed up. Here's a series of sample timings: $ go test ./rivermigrate -run TestMigrator/MigrateUpDefault -test.v -count 50 | grep 'migrations ran in' river_migrate_test.go:492: migrations ran in 45.571209ms river_migrate_test.go:492: migrations ran in 24.642458ms river_migrate_test.go:492: migrations ran in 16.749708ms river_migrate_test.go:492: migrations ran in 22.970375ms river_migrate_test.go:492: migrations ran in 16.201375ms river_migrate_test.go:492: migrations ran in 15.727625ms river_migrate_test.go:492: migrations ran in 13.291333ms river_migrate_test.go:492: migrations ran in 13.680708ms river_migrate_test.go:492: migrations ran in 14.867416ms river_migrate_test.go:492: migrations ran in 15.631916ms river_migrate_test.go:492: migrations ran in 13.873791ms river_migrate_test.go:492: migrations ran in 14.8645ms river_migrate_test.go:492: migrations ran in 14.92575ms river_migrate_test.go:492: migrations ran in 12.541834ms river_migrate_test.go:492: migrations ran in 14.753875ms river_migrate_test.go:492: migrations ran in 12.694334ms river_migrate_test.go:492: migrations ran in 13.955917ms river_migrate_test.go:492: migrations ran in 12.126458ms river_migrate_test.go:492: migrations ran in 14.095958ms river_migrate_test.go:492: migrations ran in 13.273375ms river_migrate_test.go:492: migrations ran in 13.988917ms river_migrate_test.go:492: migrations ran in 13.141459ms river_migrate_test.go:492: migrations ran in 12.394417ms river_migrate_test.go:492: migrations ran in 11.539208ms river_migrate_test.go:492: migrations ran in 11.577834ms river_migrate_test.go:492: migrations ran in 10.883375ms river_migrate_test.go:492: migrations ran in 10.547417ms river_migrate_test.go:492: migrations ran in 12.330375ms river_migrate_test.go:492: migrations ran in 11.54575ms river_migrate_test.go:492: migrations ran in 11.437458ms river_migrate_test.go:492: migrations ran in 10.957ms river_migrate_test.go:492: migrations ran in 10.589083ms river_migrate_test.go:492: migrations ran in 9.758583ms Removal of the "test DB manager" system also means that we can ungate test from `-p 1` because they're all able to run in parallel now. The limiting factor I ran in is that we need to keep max pool connections within each package's tests to a relatively modest number (I found 15 seemed to maximum success) so parallel packages don't exceed the default Postgres configuration of 100 connections. Something that can be kind of annoying is that in case a schema isn't used properly somewhere in a test case (i.e. `TestSchema` is run, but then not used), inserts/operations will go the default schema, which will leave debris there, and that will interfere with test cases using `TestTx` (with test DB manager, all debris would go to a different database so you wouldn't notice). To remediate this, I've added a cleanup hook to `TestSchema` that looks for leftovers that may have been added to the default schema. This isn't perfect because those leftovers may have come from another test case running in parallel or which ran previously, but it helps to zero in on the original source of the issue.

* Full schema insertion + schema testing Here, introduce a full version of schema injection as initially started in #798, and exposing a new `Schema` option to make it configurable even outside the test suite. This involved fixing a lot of bugs from #798, and the only way to make it possible to root them all out was to make full use of schemas across the entire test suite. The original "test DB manager" system has been replaced with a new `riverschematest.TestSchema` helper that generates a schema for use with a test case or prefers an existing idle one that was already generated for the same test run. `TestSchema` runs migrations for generated schemas which also means we don't need to use `testdbman` anymore, with tests capable of bootstrapping themselves at run time. We reuse schemas to avoid this extra work when possible, but migrating a new schema is also surprisingly fast, taking up to 50ms, but getting more down to a stable ~10ms once things are warmed up. Here's a series of sample timings: $ go test ./rivermigrate -run TestMigrator/MigrateUpDefault -test.v -count 50 | grep 'migrations ran in' river_migrate_test.go:492: migrations ran in 45.571209ms river_migrate_test.go:492: migrations ran in 24.642458ms river_migrate_test.go:492: migrations ran in 16.749708ms river_migrate_test.go:492: migrations ran in 22.970375ms river_migrate_test.go:492: migrations ran in 16.201375ms river_migrate_test.go:492: migrations ran in 15.727625ms river_migrate_test.go:492: migrations ran in 13.291333ms river_migrate_test.go:492: migrations ran in 13.680708ms river_migrate_test.go:492: migrations ran in 14.867416ms river_migrate_test.go:492: migrations ran in 15.631916ms river_migrate_test.go:492: migrations ran in 13.873791ms river_migrate_test.go:492: migrations ran in 14.8645ms river_migrate_test.go:492: migrations ran in 14.92575ms river_migrate_test.go:492: migrations ran in 12.541834ms river_migrate_test.go:492: migrations ran in 14.753875ms river_migrate_test.go:492: migrations ran in 12.694334ms river_migrate_test.go:492: migrations ran in 13.955917ms river_migrate_test.go:492: migrations ran in 12.126458ms river_migrate_test.go:492: migrations ran in 14.095958ms river_migrate_test.go:492: migrations ran in 13.273375ms river_migrate_test.go:492: migrations ran in 13.988917ms river_migrate_test.go:492: migrations ran in 13.141459ms river_migrate_test.go:492: migrations ran in 12.394417ms river_migrate_test.go:492: migrations ran in 11.539208ms river_migrate_test.go:492: migrations ran in 11.577834ms river_migrate_test.go:492: migrations ran in 10.883375ms river_migrate_test.go:492: migrations ran in 10.547417ms river_migrate_test.go:492: migrations ran in 12.330375ms river_migrate_test.go:492: migrations ran in 11.54575ms river_migrate_test.go:492: migrations ran in 11.437458ms river_migrate_test.go:492: migrations ran in 10.957ms river_migrate_test.go:492: migrations ran in 10.589083ms river_migrate_test.go:492: migrations ran in 9.758583ms Removal of the "test DB manager" system also means that we can ungate test from `-p 1` because they're all able to run in parallel now. The limiting factor I ran in is that we need to keep max pool connections within each package's tests to a relatively modest number (I found 15 seemed to maximum success) so parallel packages don't exceed the default Postgres configuration of 100 connections. Something that can be kind of annoying is that in case a schema isn't used properly somewhere in a test case (i.e. `TestSchema` is run, but then not used), inserts/operations will go the default schema, which will leave debris there, and that will interfere with test cases using `TestTx` (with test DB manager, all debris would go to a different database so you wouldn't notice). To remediate this, I've added a cleanup hook to `TestSchema` that looks for leftovers that may have been added to the default schema. This isn't perfect because those leftovers may have come from another test case running in parallel or which ran previously, but it helps to zero in on the original source of the issue. * Use schema testing for `TestTx`, sharing one schema per set of migration lines Here, start using a new isolated schema for test transaction invocations rather than falling back on whatever happens to be in `river_test`. A major benefit of this approach is that it makes migration inconsistencies impossible so that when checking out a branch with a new migration, your tests will always get the new migration, but then drop it when moving back to the `master` branch. It has the additional advantage that when alternatively running River and River Pro tests locally pointing to the same `river_test` database, we get the right migration lines raised for `TestTx` without having to worry about whether Pro migrations have been made available yet or not. We attempt to keep test transactions very efficient by only raising one schema per package being tested (and per set of migration lines, although this is rarely used) so we don't need to create a schema and run migrations once per test, but rather once for _all_ tests in the package. Similar to `TestSchema`, `TestTx` picks up a driver so that it can use `GetMigrationDefaultLines` to determine what migration lines a test case is expecting to exist before running. Once again, this is mainly useful for River Pro. `TestTx` gets moved to the higher level `riverschematest` package because it needs access to high level constructs like migrations, and I rename `riverschematest` to `riverdbtest` because with test transactions in it, its use becomes more general. I would've also moved functions in `riversharedtest` like `DBPool` and `TestDatabaseURL` over, but unfortunately we still need these at a lower level so that the tests in `rivermigrate` can use them. I removed the `riverdbtest` functionality that checks for extraneous rows. The `public` schema no longer even exists in CI, and the check is fairly unnecessary now that zeroing `search_path` makes it very difficult to accidentally insert records into anything but the current test schema. * Inject schema to completer instead of each job being completed Per code review feedback, instead of sending schema into the executor and then through into each completion being made, inject it once into the completer, then have the completer add the schema as it's completing batches. In practice, this ends up removing a lot of schema noise, and the job executor no longer needs to know about schema at all.

brandur commented Mar 8, 2025

View reviewed changes

brandur requested a review from bgentry March 8, 2025 06:56

brandur mentioned this pull request Mar 8, 2025

Add sqlc templates for arbitrary text substitution #792

Closed

bgentry reviewed Mar 9, 2025

View reviewed changes

brandur mentioned this pull request Mar 9, 2025

Better thought out caching mechanism that actually works #802

Merged

brandur force-pushed the brandur-schema-injection-proof branch 3 times, most recently from a287a84 to 95a535a Compare April 14, 2025 00:10

brandur commented Apr 14, 2025

View reviewed changes

brandur requested a review from bgentry April 14, 2025 00:13

brandur force-pushed the brandur-schema-injection-proof branch from 95a535a to f86f80f Compare April 14, 2025 00:35

brandur force-pushed the brandur-schema-injection-proof branch from f86f80f to 9b22b77 Compare April 14, 2025 00:38

brandur commented Apr 14, 2025

View reviewed changes

brandur force-pushed the brandur-schema-injection-proof branch 3 times, most recently from 2f744dd to a76bfc5 Compare April 14, 2025 06:01

bgentry approved these changes Apr 15, 2025

View reviewed changes

brandur force-pushed the brandur-schema-injection-proof branch from a76bfc5 to 69dc933 Compare April 15, 2025 05:13

brandur force-pushed the brandur-schema-injection-proof branch from 69dc933 to 7afc913 Compare April 15, 2025 05:13

brandur mentioned this pull request Apr 17, 2025

Full schema insertion + schema testing #848

Merged

Proof out explicit schema injection without search_path #798

Proof out explicit schema injection without search_path #798

Uh oh!

Conversation

brandur commented Mar 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brandur commented Mar 8, 2025

Uh oh!

bgentry left a comment

Choose a reason for hiding this comment

Uh oh!

brandur commented Mar 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bgentry commented Mar 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brandur commented Apr 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bgentry left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brandur commented Apr 15, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Proof out explicit schema injection without `search_path` #798

Proof out explicit schema injection without `search_path` #798

brandur commented Mar 8, 2025 •

edited

Loading

brandur commented Mar 9, 2025 •

edited

Loading