Allow shard setting with comments #293

jmeagher · 2023-01-24T21:36:17Z

What

Allows shard selection by the client to come in via comments like /* shard_id: 1 */ select * from foo;

Why

We're using a setup in Ruby that makes it tough or impossible to inject commands on the connection to set the shard before it gets to the "real" SQL being run. Instead we have an updated PG adapter that allows injection of comments before each executed SQL statement. We need this support in pgcat in order to keep some complex shard picking logic in Ruby code while using pgcat for connection management.

Local Testing

Run postgres and pgcat with the default options. Run psql < tests/sharding/query_routing_setup.sql to setup the database for the tests and run ./tests/pgbench/external_shard_test.sh as often as needed to exercise the shard setting comment test.

This was developed with @drdrsh

levkk · 2023-01-25T00:54:20Z

Hey there John, nice to be working together again!

I think you're on the right track with this PR. We've been thinking about this functionality for a while now and this is definitely a welcome change.

A few pointers to make this work:

`QueryRouter`

Take a look at QueryRouter, specifically try_execute_command. It could be nice to incorporate this there. You'll find a lot of similar regexes living there as well.

Test suite

The existing test suite is pretty comprehensive, and you'll also note the presence of Ruby and Python (Mostafa likes Ruby while Zain likes Python, personally I'm more of a bash guy, but everyone is welcome here); pick your favorite.

query_routing_setup.sql does all of the setup so you just need to write your own SQL tests if desired. You'll note the presence of partitions that should validate that your query routing addition works as intended (otherwise there will be insertion errors due to partition bounds). They should not have issues working locally, not sure why Circle is broken (isn't it always in every project ever...).

Backwards compatibility

PgCat is being used in production by a few organizations now, so any functionality that changes core behavior should be optional; the easiest way to make that work is to make it config-driven.

In your case, I'd suggest a config to tell the query parser to either use the existing SQL AST parser or your Regex approach, maybe something like query_parser_mode = "regex" or query_parser_mode = "ast_parser". Then the regex should also be configurable, as you may suspect different people write their SQL comments differently. A bit of historical context: #168

Let me know if I can help further. Welcome to the project!

…rable, and include sharding_key setting

jmeagher · 2023-01-25T23:51:03Z

I did some cleanup on this to move the logic somewhere it made more sense. I also made it configurable and optional by default. It now also support setting a sharding_key to help support #168. If other special commands are needed via comments it should be easy to expand this to include other options like role.

Let me know how this looks.

src/pool.rs

src/query_router.rs

src/config.rs

…f sharding commands are found

jmeagher · 2023-01-26T18:13:52Z

pgcat.toml

@@ -85,6 +85,12 @@ query_parser_enabled = true
 # queries. The primary can always be explicitly selected with our custom protocol.
 primary_reads_enabled = true

+# Allow sharding commands to be passed as statement comments instead of
+# separate commands. If these are unset this functionality is disabled.
+# sharding_key_regex = '/\* sharding_key: (\d+) \*/'


These are commented out because when they were uncommented it caused the Ruby TOML parser to fail when loading the file. These aren't used in tests right now, so it should be fine.

We will need to figure this out because otherwise this feature won't be tested in end-to-end tests we have in Ruby. I think it's fine to test it with Bash to start, but the Ruby test suit protects us against bad things and we want it to run there too.

src/query_router.rs

levkk

The LOC LB test is flakey, rerunning it usually helps. We have an action item to fix this. The PR looks good to me to merge. Eventually we would want some tests in bash/Ruby/Python so we don't get hit with a regression. I'll let @drdrsh take a look before merging.

drdrsh · 2023-01-26T21:50:22Z

src/query_router.rs

+        // Check for any sharding regex matches in any queries
+        match code as char {
+            // For Parse and Query messages peek to see if they specify a shard_id as a comment early in the statement


I think we can avoid this entire branch when the number of shards is 1, we can get that from self.pool_settings.shards. This will make this check even cheaper for the base, unsharded use case

Scratch that, the lack of the regex in the confs should be enough of a gate

drdrsh · 2023-01-26T21:59:12Z

We'll test this change internally. I will merge it once we are done @levkk

jmeagher · 2023-02-15T20:49:06Z

src/config.rs

    pub shards: BTreeMap<String, Shard>,
    pub users: BTreeMap<String, User>,
+    // Note, don't put simple fields below these configs. There's a compatability issue with TOML that makes it


Having the new fields below user and shard broke our internal tool that generates the TOML config. Since this is a pretty odd error that we hit I'm leaving behind the breadcrumb that got us to fix it.

jmeagher added 6 commits January 24, 2023 15:25

Allow shard setting with comments

c104ca8

try to trigger a build

fc461e7

try to trigger a build again

91b2176

try to trigger a build yet again

48125a0

last try before switching to other PR

6f7d269

Try fixing circleci tests by adding the right connection info

eda1740

levkk mentioned this pull request Jan 25, 2023

Almost automatic shard detection #165

Closed

jmeagher mentioned this pull request Jan 25, 2023

Special behavior invocation via comment instead of a custom SQL #168

Open

Cleanup, move logic to query_router, and make search patterns configu…

552d124

…rable, and include sharding_key setting

levkk reviewed Jan 25, 2023

View reviewed changes

src/pool.rs Show resolved Hide resolved