feat: multi-shard LIMIT+OFFSET in-flight re-write in PREPARE+EXECUTE

[jkaczman]

Implements functionality that handles multi-shard queries that use LIMIT + OFFSET in PREPARE + EXECUTE to re-write the LIMIT and OFFSET values (LIMIT = LIMIT + OFFSET, OFFSET = 0) and update the Route to expect them later to handle the appropriate multi-shard response for the client.

Specifically,

• When we first receive a PreparedStatement with a SELECT query that has a LIMIT + OFFSET, it re-writes any A_Const nodes (constant numbers) to be ParamRef nodes. This allows us the ability to always be able to modify the limit/offset parameters later on, in a helper method stemming off of OffsetPlan::apply_after_parser. This later-on modification (during execution) only happens if it's multi-shard.
• When we get the matching ExecuteStatement, if the Client previously used A_Const nodes, which means we had to store them in cache, we append the relevant cached numbers to the parameters list for the ExecuteStatement (as the Client does not know we internally re-wrote the PreparedStatement in our cache). If they previously used ParamRef nodes, the (cached, cloned) OffsetPlan is updated to the newly passed in values to complete it, so that we can use it in the OffsetPlan::apply_after_parser helper method.
• Added Option<OffsetPlan> to prepared statement. Did not modify CacheKey::Simple; I modified it to use the original query (before rewrite) to store in the cache, and store the rewritten query (if any) in the cache. This allows us to later (1) which parameter is offset / which is limit, as well as (2) if the client pre-specified A_Const nodes, we must save them in our cache, as they will not be given to us in the ExecuteStatement.

Added some integration tests to test various cases: both limit/offset being ParamRefs, one being a ParamRef, both being A_Const, switching up the ordering (LIMIT $2 OFFSET $1), adding in an unrelated ParamRef (WHERE id < $2)... used ORDER BY for determinism in ensuring that the response we got is accurate.

Integration test and unit test also checks to make sure that two Queries that are the same due to a re-write won't conflict in GlobalCache.

Fixes #1383.

[codecov]

Codecov Report

❌ Patch coverage is 94.89796% with 15 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...router/parser/rewrite/statement/simple_prepared.rs	93.59%	13 Missing ⚠️
...gdog/src/frontend/prepared_statements/cache_key.rs	0.00%	1 Missing ⚠️
...frontend/router/parser/rewrite/statement/offset.rs	97.91%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[levkk]

We do this already here:

pgdog/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

Line 147 in af2edf5

let mut next_param = plan.params as i32 + 1;

and here

pgdog/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

Lines 125 to 131 in af2edf5

	if let Some(parameterized_stmt) = parameterized_stmt {
	walk::walk(parameterized_stmt, |node| {
	if let Node::ParamRef(param) = node {
	plan.params = plan.params.max(param.number as u16)
	}
	});
	}

Not a big deal, but at some point in the near/far future, we should consolidate all of our walks into one. It's cheap but it's not free.

[levkk]

Very nice!