From b970668e8f56e60eb1f11533e43b7872dd37a5eb Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 29 Apr 2022 12:52:33 -0500 Subject: [PATCH 01/88] piped "./receive.sh | cargo run -- | jq" --- worker/walrus/.gitignore | 13 + worker/walrus/Cargo.toml | 15 + worker/walrus/diesel.toml | 5 + worker/walrus/migrations/.gitkeep | 0 .../2022-04-29-132611_initial/down.sql | 1 + .../2022-04-29-132611_initial/up.sql | 524 ++++++++++++++++++ worker/walrus/receive.sh | 14 + worker/walrus/src/main.rs | 94 ++++ 8 files changed, 666 insertions(+) create mode 100644 worker/walrus/.gitignore create mode 100644 worker/walrus/Cargo.toml create mode 100644 worker/walrus/diesel.toml create mode 100644 worker/walrus/migrations/.gitkeep create mode 100644 worker/walrus/migrations/2022-04-29-132611_initial/down.sql create mode 100644 worker/walrus/migrations/2022-04-29-132611_initial/up.sql create mode 100755 worker/walrus/receive.sh create mode 100644 worker/walrus/src/main.rs diff --git a/worker/walrus/.gitignore b/worker/walrus/.gitignore new file mode 100644 index 0000000..3797cf3 --- /dev/null +++ b/worker/walrus/.gitignore @@ -0,0 +1,13 @@ +.DS_Store +.idea/ +/target +*.iml +**/*.rs.bk +Cargo.lock +*.swp +*.rs.swp +**/*.swp +*.diff +/results +regression.diffs +regression.out diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml new file mode 100644 index 0000000..da4cef8 --- /dev/null +++ b/worker/walrus/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "walrus" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +clap = { version = "3.1.12", features = ["derive"] } +diesel = { version = "2.0.0-rc.0", features = ["postgres", "serde_json", "uuid"] } +diesel_migrations = { version = "2.0.0-rc.0", features = ["postgres"] } +dotenv = "0.15.0" +serde_json = "1.0" +serde = "1.0" +uuid = { version = "1.0", features = ["serde"] } diff --git a/worker/walrus/diesel.toml b/worker/walrus/diesel.toml new file mode 100644 index 0000000..92267c8 --- /dev/null +++ b/worker/walrus/diesel.toml @@ -0,0 +1,5 @@ +# For documentation on how to configure this file, +# see diesel.rs/guides/configuring-diesel-cli + +[print_schema] +file = "src/schema.rs" diff --git a/worker/walrus/migrations/.gitkeep b/worker/walrus/migrations/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/worker/walrus/migrations/2022-04-29-132611_initial/down.sql b/worker/walrus/migrations/2022-04-29-132611_initial/down.sql new file mode 100644 index 0000000..9b23cd4 --- /dev/null +++ b/worker/walrus/migrations/2022-04-29-132611_initial/down.sql @@ -0,0 +1 @@ +drop schema realtime cascade; diff --git a/worker/walrus/migrations/2022-04-29-132611_initial/up.sql b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql new file mode 100644 index 0000000..e36afbc --- /dev/null +++ b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql @@ -0,0 +1,524 @@ +create schema realtime; + +create type realtime.equality_op as enum( + 'eq', 'neq', 'lt', 'lte', 'gt', 'gte' +); + + +create type realtime.action as enum ('INSERT', 'UPDATE', 'DELETE', 'ERROR'); + + +create function realtime.cast(val text, type_ regtype) + returns jsonb + immutable + language plpgsql +as $$ +declare + res jsonb; +begin + execute format('select to_jsonb(%L::'|| type_::text || ')', val) into res; + return res; +end +$$; + + +create type realtime.user_defined_filter as ( + column_name text, + op realtime.equality_op, + value text +); + + +create function realtime.to_regrole(role_name text) + returns regrole + immutable + language sql + -- required to allow use in generated clause +as $$ select role_name::regrole $$; + + +create table realtime.subscription ( + -- Tracks which subscriptions are active + id bigint generated always as identity primary key, + subscription_id uuid not null, + entity regclass not null, + filters realtime.user_defined_filter[] not null default '{}', + claims jsonb not null, + claims_role regrole not null generated always as (realtime.to_regrole(claims ->> 'role')) stored, + created_at timestamp not null default timezone('utc', now()), + + unique (subscription_id, entity, filters) +); +create index ix_realtime_subscription_entity on realtime.subscription using hash (entity); + + +create or replace function realtime.subscription_check_filters() + returns trigger + language plpgsql +as $$ +/* +Validates that the user defined filters for a subscription: +- refer to valid columns that the claimed role may access +- values are coercable to the correct column type +*/ +declare + col_names text[] = coalesce( + array_agg(c.column_name order by c.ordinal_position), + '{}'::text[] + ) + from + information_schema.columns c + where + format('%I.%I', c.table_schema, c.table_name)::regclass = new.entity + and pg_catalog.has_column_privilege( + (new.claims ->> 'role'), + format('%I.%I', c.table_schema, c.table_name)::regclass, + c.column_name, + 'SELECT' + ); + filter realtime.user_defined_filter; + col_type regtype; +begin + for filter in select * from unnest(new.filters) loop + -- Filtered column is valid + if not filter.column_name = any(col_names) then + raise exception 'invalid column for filter %', filter.column_name; + end if; + + -- Type is sanitized and safe for string interpolation + col_type = ( + select atttypid::regtype + from pg_catalog.pg_attribute + where attrelid = new.entity + and attname = filter.column_name + ); + if col_type is null then + raise exception 'failed to lookup type for column %', filter.column_name; + end if; + -- raises an exception if value is not coercable to type + perform realtime.cast(filter.value, col_type); + end loop; + + -- Apply consistent order to filters so the unique constraint on + -- (subscription_id, entity, filters) can't be tricked by a different filter order + new.filters = coalesce( + array_agg(f order by f.column_name, f.op, f.value), + '{}' + ) from unnest(new.filters) f; + + return new; +end; +$$; + +create trigger tr_check_filters + before insert or update on realtime.subscription + for each row + execute function realtime.subscription_check_filters(); + + +create or replace function realtime.quote_wal2json(entity regclass) + returns text + language sql + immutable + strict +as $$ + select + ( + select string_agg('\' || ch,'') + from unnest(string_to_array(nsp.nspname::text, null)) with ordinality x(ch, idx) + where + not (x.idx = 1 and x.ch = '"') + and not ( + x.idx = array_length(string_to_array(nsp.nspname::text, null), 1) + and x.ch = '"' + ) + ) + || '.' + || ( + select string_agg('\' || ch,'') + from unnest(string_to_array(pc.relname::text, null)) with ordinality x(ch, idx) + where + not (x.idx = 1 and x.ch = '"') + and not ( + x.idx = array_length(string_to_array(nsp.nspname::text, null), 1) + and x.ch = '"' + ) + ) + from + pg_class pc + join pg_namespace nsp + on pc.relnamespace = nsp.oid + where + pc.oid = entity +$$; + + +create or replace function realtime.check_equality_op( + op realtime.equality_op, + type_ regtype, + val_1 text, + val_2 text +) + returns bool + immutable + language plpgsql +as $$ +/* +Casts *val_1* and *val_2* as type *type_* and check the *op* condition for truthiness +*/ +declare + op_symbol text = ( + case + when op = 'eq' then '=' + when op = 'neq' then '!=' + when op = 'lt' then '<' + when op = 'lte' then '<=' + when op = 'gt' then '>' + when op = 'gte' then '>=' + else 'UNKNOWN OP' + end + ); + res boolean; +begin + execute format('select %L::'|| type_::text || ' ' || op_symbol || ' %L::'|| type_::text, val_1, val_2) into res; + return res; +end; +$$; + + +create type realtime.wal_column as ( + name text, + type_name text, + type_oid oid, + value jsonb, + is_pkey boolean, + is_selectable boolean +); + +create or replace function realtime.build_prepared_statement_sql( + prepared_statement_name text, + entity regclass, + columns realtime.wal_column[] +) + returns text + language sql +as $$ +/* +Builds a sql string that, if executed, creates a prepared statement to +tests retrive a row from *entity* by its primary key columns. + +Example + select realtime.build_prepared_statment_sql('public.notes', '{"id"}'::text[], '{"bigint"}'::text[]) +*/ + select +'prepare ' || prepared_statement_name || ' as + select + exists( + select + 1 + from + ' || entity || ' + where + ' || string_agg(quote_ident(pkc.name) || '=' || quote_nullable(pkc.value #>> '{}') , ' and ') || ' + )' + from + unnest(columns) pkc + where + pkc.is_pkey + group by + entity +$$; + + +create type realtime.wal_rls as ( + wal jsonb, + is_rls_enabled boolean, + subscription_ids uuid[], + errors text[] +); + + + +create or replace function realtime.is_visible_through_filters(columns realtime.wal_column[], filters realtime.user_defined_filter[]) + returns bool + language sql + immutable +as $$ +/* +Should the record be visible (true) or filtered out (false) after *filters* are applied +*/ + select + -- Default to allowed when no filters present + coalesce( + sum( + realtime.check_equality_op( + op:=f.op, + type_:=col.type_oid::regtype, + -- cast jsonb to text + val_1:=col.value #>> '{}', + val_2:=f.value + )::int + ) = count(1), + true + ) + from + unnest(filters) f + join unnest(columns) col + on f.column_name = col.name; +$$; + + +create or replace function realtime.apply_rls(wal jsonb, max_record_bytes int = 1024 * 1024) + returns setof realtime.wal_rls + language plpgsql + volatile +as $$ +declare + -- Regclass of the table e.g. public.notes + entity_ regclass = (quote_ident(wal ->> 'schema') || '.' || quote_ident(wal ->> 'table'))::regclass; + + -- I, U, D, T: insert, update ... + action realtime.action = ( + case wal ->> 'action' + when 'I' then 'INSERT' + when 'U' then 'UPDATE' + when 'D' then 'DELETE' + else 'ERROR' + end + ); + + -- Is row level security enabled for the table + is_rls_enabled bool = relrowsecurity from pg_class where oid = entity_; + + subscriptions realtime.subscription[] = array_agg(subs) + from + realtime.subscription subs + where + subs.entity = entity_; + + -- Subscription vars + roles regrole[] = array_agg(distinct us.claims_role) + from + unnest(subscriptions) us; + + working_role regrole; + claimed_role regrole; + claims jsonb; + + subscription_id uuid; + subscription_has_access bool; + visible_to_subscription_ids uuid[] = '{}'; + + -- structured info for wal's columns + columns realtime.wal_column[]; + -- previous identity values for update/delete + old_columns realtime.wal_column[]; + + error_record_exceeds_max_size boolean = octet_length(wal::text) > max_record_bytes; + + -- Primary jsonb output for record + output jsonb; + +begin + perform set_config('role', null, true); + + columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'columns') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + old_columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'identity') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + for working_role in select * from unnest(roles) loop + + -- Update `is_selectable` for columns and old_columns + columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(columns) c; + + old_columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(old_columns) c; + + if action <> 'DELETE' and count(1) = 0 from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + -- subscriptions is already filtered by entity + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 400: Bad Request, no primary key'] + )::realtime.wal_rls; + + -- The claims role does not have SELECT permission to the primary key of entity + elsif action <> 'DELETE' and sum(c.is_selectable::int) <> count(1) from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 401: Unauthorized'] + )::realtime.wal_rls; + + else + output = jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action, + 'commit_timestamp', to_char( + (wal ->> 'timestamp')::timestamptz, + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ), + 'columns', ( + select + jsonb_agg( + jsonb_build_object( + 'name', pa.attname, + 'type', pt.typname + ) + order by pa.attnum asc + ) + from + pg_attribute pa + join pg_type pt + on pa.atttypid = pt.oid + where + attrelid = entity_ + and attnum > 0 + and pg_catalog.has_column_privilege(working_role, entity_, pa.attname, 'SELECT') + ) + ) + -- Add "record" key for insert and update + || case + when error_record_exceeds_max_size then jsonb_build_object('record', '{}'::jsonb) + when action in ('INSERT', 'UPDATE') then + jsonb_build_object( + 'record', + (select jsonb_object_agg((c).name, (c).value) from unnest(columns) c where (c).is_selectable) + ) + else '{}'::jsonb + end + -- Add "old_record" key for update and delete + || case + when error_record_exceeds_max_size then jsonb_build_object('old_record', '{}'::jsonb) + when action in ('UPDATE', 'DELETE') then + jsonb_build_object( + 'old_record', + (select jsonb_object_agg((c).name, (c).value) from unnest(old_columns) c where (c).is_selectable) + ) + else '{}'::jsonb + end; + + -- Create the prepared statement + if is_rls_enabled and action <> 'DELETE' then + if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then + deallocate walrus_rls_stmt; + end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, columns); + end if; + + visible_to_subscription_ids = '{}'; + + for subscription_id, claims in ( + select + subs.subscription_id, + subs.claims + from + unnest(subscriptions) subs + where + subs.entity = entity_ + and subs.claims_role = working_role + and realtime.is_visible_through_filters(columns, subs.filters) + ) loop + + if not is_rls_enabled or action = 'DELETE' then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + else + -- Check if RLS allows the role to see the record + perform + set_config('role', working_role::text, true), + set_config('request.jwt.claims', claims::text, true); + + execute 'execute walrus_rls_stmt' into subscription_has_access; + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end if; + end loop; + + perform set_config('role', null, true); + + return next ( + output, + is_rls_enabled, + visible_to_subscription_ids, + case + when error_record_exceeds_max_size then array['Error 413: Payload Too Large'] + else '{}' + end + )::realtime.wal_rls; + + end if; + end loop; + + perform set_config('role', null, true); +end; +$$; diff --git a/worker/walrus/receive.sh b/worker/walrus/receive.sh new file mode 100755 index 0000000..fbce01d --- /dev/null +++ b/worker/walrus/receive.sh @@ -0,0 +1,14 @@ +pg_recvlogical \ + --file=- \ + --dbname=postgresql://oliverrice@localhost:28814/walrus \ + --plugin=wal2json \ + --option=include-pk=1 \ + --option=include-transaction=false \ + --option=include-timestamp=true \ + --option=include-type-oids=true \ + --option=format-version=2 \ + --option=actions=insert,update,delete \ + --slot=realtime \ + --create-slot \ + --if-not-exists \ + --start diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs new file mode 100644 index 0000000..56d26e8 --- /dev/null +++ b/worker/walrus/src/main.rs @@ -0,0 +1,94 @@ +use clap::Parser; +use diesel::dsl::sql; +use diesel::sql_types::*; +use diesel::*; +use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; +use serde::Serialize; +use serde_json; +use std::io; +use std::io::BufRead; +use uuid; + +pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); + +#[derive(Serialize)] +pub struct WalrusRecord { + wal: serde_json::Value, + is_rls_enabled: bool, + subscription_ids: Vec, + errors: Vec, +} + +#[derive(Parser, Debug)] +#[clap(author, version, about, long_about = None)] +struct Args { + /// Number of times to greet + #[clap(short, long, default_value_t = 1)] + count: u8, +} + +fn main() { + // Parse command line arguments + let _args = Args::parse(); + + // Connect to Postgres + let database_url = "postgresql://oliverrice:@localhost:28814/walrus"; + let conn = &mut PgConnection::establish(&database_url).unwrap(); + + // Run any pending migrations + conn.run_pending_migrations(MIGRATIONS) + .expect("Pending migrations failed to execute"); + + // Reading from stdin + let stdin = io::stdin(); + let stdin_reader = io::BufReader::new(stdin); + let stdin_lines = stdin_reader.lines(); + + // Iterate input data + for input_line in stdin_lines { + match input_line { + Ok(line) => { + let result_json_line = serde_json::from_str::(&line); + match result_json_line { + Ok(json_line) => { + let result_wal_rls_rows = + sql::< + Record<( + sql_types::Jsonb, + sql_types::Bool, + sql_types::Array, + sql_types::Array, + )>, + >("SELECT x from realtime.apply_rls(") + .bind::(json_line) + .sql(") x") + .get_results::<(serde_json::Value, bool, Vec, Vec)>( + conn, + ); + match result_wal_rls_rows { + Ok(rows) => { + for row in rows { + let walrus_rec = WalrusRecord { + wal: row.0, + is_rls_enabled: row.1, + subscription_ids: row.2, + errors: row.3, + }; + match serde_json::to_string(&walrus_rec) { + Ok(walrus_json) => println!("{}", walrus_json), + Err(err) => { + println!("Failed to serialize walrus result: {}", err) + } + } + } + } + Err(err) => println!("WALRUS Error: {}", err), + } + } + Err(err) => println!("Failed to parse: {}", err), + } + } + Err(err) => println!("Error: {}", err), + } + } +} From e49e45757c869e356d06663b1a3fcc14b868b753 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 29 Apr 2022 13:13:17 -0500 Subject: [PATCH 02/88] single command --- worker/walrus/src/main.rs | 121 ++++++++++++++++++++++++-------------- 1 file changed, 78 insertions(+), 43 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 56d26e8..725f21e 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -7,6 +7,7 @@ use serde::Serialize; use serde_json; use std::io; use std::io::BufRead; +use std::process::{Command, Stdio}; use uuid; pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); @@ -22,73 +23,107 @@ pub struct WalrusRecord { #[derive(Parser, Debug)] #[clap(author, version, about, long_about = None)] struct Args { - /// Number of times to greet - #[clap(short, long, default_value_t = 1)] - count: u8, + #[clap(long, default_value = "realtime")] + slot: String, + + #[clap(long, default_value = "postgresql://postgres@localhost:5432/postgres")] + connection: String, } fn main() { // Parse command line arguments - let _args = Args::parse(); + let args = Args::parse(); + + let mut cmd = Command::new("pg_recvlogical") + //&args + .args(vec![ + "--file=-", + "--plugin=wal2json", + &format!("--dbname={}", args.connection), + "--option=include-pk=1", + "--option=include-transaction=false", + "--option=include-timestamp=true", + "--option=include-type-oids=true", + "--option=format-version=2", + "--option=actions=insert,update,delete", + &format!("--slot={}", args.slot), + "--create-slot", + "--if-not-exists", + "--start", + ]) + .stdout(Stdio::piped()) + .spawn() + .unwrap(); - // Connect to Postgres - let database_url = "postgresql://oliverrice:@localhost:28814/walrus"; - let conn = &mut PgConnection::establish(&database_url).unwrap(); + { + // Connect to Postgres + // "postgresql://oliverrice:@localhost:28814/walrus"; + let conn = &mut PgConnection::establish(&args.connection).unwrap(); - // Run any pending migrations - conn.run_pending_migrations(MIGRATIONS) - .expect("Pending migrations failed to execute"); + // Run any pending migrations + conn.run_pending_migrations(MIGRATIONS) + .expect("Pending migrations failed to execute"); - // Reading from stdin - let stdin = io::stdin(); - let stdin_reader = io::BufReader::new(stdin); - let stdin_lines = stdin_reader.lines(); + // Reading from stdin + let stdin = cmd.stdout.as_mut().unwrap(); + let stdin_reader = io::BufReader::new(stdin); + let stdin_lines = stdin_reader.lines(); - // Iterate input data - for input_line in stdin_lines { - match input_line { - Ok(line) => { - let result_json_line = serde_json::from_str::(&line); - match result_json_line { - Ok(json_line) => { - let result_wal_rls_rows = - sql::< + // Iterate input data + for input_line in stdin_lines { + match input_line { + Ok(line) => { + let result_json_line = serde_json::from_str::(&line); + match result_json_line { + Ok(json_line) => { + let result_wal_rls_rows = sql::< Record<( sql_types::Jsonb, sql_types::Bool, sql_types::Array, sql_types::Array, )>, - >("SELECT x from realtime.apply_rls(") + >( + "SELECT x from realtime.apply_rls(" + ) .bind::(json_line) .sql(") x") - .get_results::<(serde_json::Value, bool, Vec, Vec)>( - conn, - ); - match result_wal_rls_rows { - Ok(rows) => { - for row in rows { - let walrus_rec = WalrusRecord { - wal: row.0, - is_rls_enabled: row.1, - subscription_ids: row.2, - errors: row.3, - }; - match serde_json::to_string(&walrus_rec) { - Ok(walrus_json) => println!("{}", walrus_json), - Err(err) => { - println!("Failed to serialize walrus result: {}", err) + .get_results::<( + serde_json::Value, + bool, + Vec, + Vec, + )>(conn); + match result_wal_rls_rows { + Ok(rows) => { + for row in rows { + let walrus_rec = WalrusRecord { + wal: row.0, + is_rls_enabled: row.1, + subscription_ids: row.2, + errors: row.3, + }; + match serde_json::to_string(&walrus_rec) { + Ok(walrus_json) => println!("{}", walrus_json), + Err(err) => { + println!( + "Failed to serialize walrus result: {}", + err + ) + } } } } + Err(err) => println!("WALRUS Error: {}", err), } - Err(err) => println!("WALRUS Error: {}", err), } + Err(err) => println!("Failed to parse: {}", err), } - Err(err) => println!("Failed to parse: {}", err), } + Err(err) => println!("Error: {}", err), } - Err(err) => println!("Error: {}", err), } } + + cmd.wait().unwrap(); } From 15fd585f91d54148e50fe100265000f16244ab75 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 29 Apr 2022 13:13:40 -0500 Subject: [PATCH 03/88] rm shell script --- worker/walrus/receive.sh | 14 -------------- 1 file changed, 14 deletions(-) delete mode 100755 worker/walrus/receive.sh diff --git a/worker/walrus/receive.sh b/worker/walrus/receive.sh deleted file mode 100755 index fbce01d..0000000 --- a/worker/walrus/receive.sh +++ /dev/null @@ -1,14 +0,0 @@ -pg_recvlogical \ - --file=- \ - --dbname=postgresql://oliverrice@localhost:28814/walrus \ - --plugin=wal2json \ - --option=include-pk=1 \ - --option=include-transaction=false \ - --option=include-timestamp=true \ - --option=include-type-oids=true \ - --option=format-version=2 \ - --option=actions=insert,update,delete \ - --slot=realtime \ - --create-slot \ - --if-not-exists \ - --start From 0804b64a58836ef8b274396145bbe3bc92f0ef50 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 29 Apr 2022 14:26:41 -0500 Subject: [PATCH 04/88] minimal readme --- worker/walrus/README.md | 130 +++++++++++++++++++++++++++ worker/walrus/docker-compose.yml | 23 +++++ worker/walrus/dockerfiles/Dockerfile | 8 ++ worker/walrus/src/main.rs | 3 + 4 files changed, 164 insertions(+) create mode 100644 worker/walrus/README.md create mode 100644 worker/walrus/docker-compose.yml create mode 100644 worker/walrus/dockerfiles/Dockerfile diff --git a/worker/walrus/README.md b/worker/walrus/README.md new file mode 100644 index 0000000..a95520e --- /dev/null +++ b/worker/walrus/README.md @@ -0,0 +1,130 @@ +# WALRUS Worker + +WALRUS background worker runs next to a PostgreSQL instance. It applies row level security to the write ahead log (WAL) and optionally forwards those records to external services. + + +``` +walrus 0.1.0 + +USAGE: + walrus [OPTIONS] + +OPTIONS: + --connection [default: postgresql://postgres@localhost:5432/postgres] + -h, --help Print help information + --slot [default: realtime] + -V, --version Print version information +``` + + +## Features + + +### Realtime (no polling) + +The worker wraps [pg_recvlogical](https://www.postgresql.org/docs/current/app-pgrecvlogical.html), a lightweight tool that ships with PostgreSQL, making use of the [streaming replication protocol](https://www.postgresql.org/docs/current/protocol-logical-replication.html#:~:text=The%20logical%20replication%20protocol%20sends,Start%20and%20Stream%20Stop%20messages.) to recieve new WAL messages. + +`walrus` processes each message as soon as they are forwarded from `pg_recvlogical`. That change reduces message latency (up-to) 100ms when compared with the current implementation of polling for changes through [pg_logical_get_changes](https://www.postgresql.org/docs/14/logicaldecoding-example.html) + +### Self-Configuring + +#### Replication Slot +On startup, the requested replication slot if created with the correct wal2json settings if it doesn't exist. + +#### Migrations (Embedded) + +The `walrus` executable contains embedded migrations from `./migrations` and applies any unapplied migrations at process startup. + + +## Try it Out + +Clone and Navigate +```sh +git clone https://github.com/supabase/walrus.git +cd walrus +git checkout worker +cd worker/walrus +``` + +Start the DB +```sh +docker-compose up +``` + +Run the walrus worker +```sh +cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres + +# Note: if you have jq installed, the output is more readable with +# cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres | jq +``` + +Connect to the database at `postgresql://postgres:password@localhost:5501/postgres` + +and execute the following SQL to create a subscription and a WAL record. + +```sql +-- Create a table we can subscribe to +create table book( + id int primary key, + title text +); + +-- Create a dummy subscription to our new table +insert into realtime.subscription(subscription_id, entity, claims) +select + gen_random_uuid(), + 'public.book', + jsonb_build_object( + 'role', 'postgres', + 'email', 'o@r.com', + 'sub', gen_random_uuid() + ); + +-- Create a record +insert into book(id, title) +values (1, 'Foo'); +``` + + +Now, looking back out the output from the `cargo run` command, you see the following printed to stdout +```json +> cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres +... +{ + "wal":{ + "columns":[ + { + "name":"id", + "type":"int4" + }, + { + "name":"title", + "type":"text" + } + ], + "commit_timestamp":"2022-04-29T19:04:23Z", + "record":{ + "id":1, + "title":"Foo" + }, + "schema":"public", + "table":"book", + "type":"INSERT" + }, + "is_rls_enabled":false, + "subscription_ids":[ + "af68a1b5-fbb3-4154-84cc-a2ee1a7048f9", + "93ba49e4-b8e9-4ab9-bedd-9332a6397806" + ], + "errors":[] +} +``` + + +## Possible Enhancements + +- Rate limiting +- Some work that is currently being performed in SQL could be shuffled out rust for a performance improvement + - Reshaping the WAL records + - Any state we want to track between calls to `realtime.apply_rls` diff --git a/worker/walrus/docker-compose.yml b/worker/walrus/docker-compose.yml new file mode 100644 index 0000000..4f9296e --- /dev/null +++ b/worker/walrus/docker-compose.yml @@ -0,0 +1,23 @@ +version: '3' +services: + db: + container_name: walrus_streaming + build: + context: . + dockerfile: ./dockerfiles/Dockerfile + ports: + - "5501:5432" + command: + - postgres + - -c + - wal_level=logical + - -c + - fsync=off + healthcheck: + test: ["CMD-SHELL", "PGUSER=postgres", "pg_isready"] + interval: 1s + timeout: 10s + retries: 5 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: password diff --git a/worker/walrus/dockerfiles/Dockerfile b/worker/walrus/dockerfiles/Dockerfile new file mode 100644 index 0000000..cfea2fb --- /dev/null +++ b/worker/walrus/dockerfiles/Dockerfile @@ -0,0 +1,8 @@ +FROM supabase/postgres:latest + +RUN apt-get update +RUN apt install build-essential postgresql-server-dev-14 -y + +RUN git clone https://github.com/eulerto/wal2json.git --depth 1 53b548a29ebd6119323b6eb2f6013d7c5fe807ec +RUN cd wal2json && make && make install + diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 725f21e..06111dc 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -20,6 +20,9 @@ pub struct WalrusRecord { errors: Vec, } +/// Write-Ahead-Log Realtime Unified Security (WALRUS) background worker +/// runs next to a PostgreSQL instance and forwards its Write-Ahead-Log +/// to external services #[derive(Parser, Debug)] #[clap(author, version, about, long_about = None)] struct Args { From 761068b9545d50156a4a2f6afe8a568951dd26f5 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 29 Apr 2022 14:28:04 -0500 Subject: [PATCH 05/88] add requirements to readme --- worker/walrus/README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/worker/walrus/README.md b/worker/walrus/README.md index a95520e..eb3bcb6 100644 --- a/worker/walrus/README.md +++ b/worker/walrus/README.md @@ -38,6 +38,12 @@ The `walrus` executable contains embedded migrations from `./migrations` and app ## Try it Out + +Requires: +- rust/cargo +- docker-compose +- postgres installed locally (for `pg_recvlogical`) + Clone and Navigate ```sh git clone https://github.com/supabase/walrus.git From 0aa5962a941e24e25801dcb1859c81c456b91adc Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 17 May 2022 11:43:24 -0500 Subject: [PATCH 06/88] realtime transport --- .gitignore | 3 + worker/Cargo.toml | 5 + worker/README.md | 64 ++++++++++ worker/{walrus => }/docker-compose.yml | 0 worker/dockerfiles/Dockerfile | 7 ++ worker/realtime/.gitignore | 13 ++ worker/realtime/Cargo.toml | 25 ++++ worker/realtime/README.md | 22 ++++ worker/realtime/src/main.rs | 160 +++++++++++++++++++++++++ worker/walrus/Cargo.toml | 1 + worker/walrus/README.md | 4 +- worker/walrus/dockerfiles/Dockerfile | 8 -- 12 files changed, 302 insertions(+), 10 deletions(-) create mode 100644 worker/Cargo.toml create mode 100644 worker/README.md rename worker/{walrus => }/docker-compose.yml (100%) create mode 100644 worker/dockerfiles/Dockerfile create mode 100644 worker/realtime/.gitignore create mode 100644 worker/realtime/Cargo.toml create mode 100644 worker/realtime/README.md create mode 100644 worker/realtime/src/main.rs delete mode 100644 worker/walrus/dockerfiles/Dockerfile diff --git a/.gitignore b/.gitignore index c1a10e0..709a198 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,6 @@ regression.* __pycache__/ *.egg-info/ *.swp +target/ +Cargo.lock +.DS_Store diff --git a/worker/Cargo.toml b/worker/Cargo.toml new file mode 100644 index 0000000..cd4533e --- /dev/null +++ b/worker/Cargo.toml @@ -0,0 +1,5 @@ +[workspace] +members = [ + "walrus", + "realtime", +] diff --git a/worker/README.md b/worker/README.md new file mode 100644 index 0000000..000b030 --- /dev/null +++ b/worker/README.md @@ -0,0 +1,64 @@ +# WALRUS Worker + Realtime Transport + +Example showing how to stream WAL from postgres, apply row level security, and push changes to supabase realtime. + +- `walrus/` is responsible for receiving WAL, formatting messages, and applying row level security +- `realtime/` is the websocket trasport layer for supabase realtime. + +More info about each component can be found in their directories' README.md. + +## Example: + +Requires: +- rust/cargo +- docker-compose +- postgres installed locally (for `pg_recvlogical`) + +Clone and Navigate +```sh +git clone https://github.com/supabase/walrus.git +cd walrus +git checkout worker +cd worker +``` + +Start the DB +```sh +docker-compose up +``` + +Run the `walrus` worker, piping its output to `realtime` transport +```sh +cargo run --bin walrus -- \ + --connection=postgresql://postgres:password@localhost:5501/postgres | +cargo run --bin realtime -- \ + --url=wss://sendwal.fly.dev/socket \ + --apikey= +``` + +Connect to the database at `postgresql://postgres:password@localhost:5501/postgres` + +and execute the following SQL to create a subscription and a WAL record. + +```sql +-- Create a table we can subscribe to +create table book( + id int primary key, + title text +); + +-- Create a dummy subscription to our new table +insert into realtime.subscription(subscription_id, entity, claims) +select + gen_random_uuid(), + 'public.book', + jsonb_build_object( + 'role', 'postgres', + 'email', 'o@r.com', + 'sub', gen_random_uuid() + ); + +-- Create a record +insert into book(id, title) +values (1, 'Foo'); +``` \ No newline at end of file diff --git a/worker/walrus/docker-compose.yml b/worker/docker-compose.yml similarity index 100% rename from worker/walrus/docker-compose.yml rename to worker/docker-compose.yml diff --git a/worker/dockerfiles/Dockerfile b/worker/dockerfiles/Dockerfile new file mode 100644 index 0000000..1798011 --- /dev/null +++ b/worker/dockerfiles/Dockerfile @@ -0,0 +1,7 @@ +FROM supabase/postgres:latest +RUN apt-get update +RUN apt-get install build-essential postgresql-server-dev-14 -y + +RUN git clone https://github.com/eulerto/wal2json.git +RUN cd wal2json && git checkout 53b548a29ebd6119323b6eb2f6013d7c5fe807ec && make && make install + diff --git a/worker/realtime/.gitignore b/worker/realtime/.gitignore new file mode 100644 index 0000000..3797cf3 --- /dev/null +++ b/worker/realtime/.gitignore @@ -0,0 +1,13 @@ +.DS_Store +.idea/ +/target +*.iml +**/*.rs.bk +Cargo.lock +*.swp +*.rs.swp +**/*.swp +*.diff +/results +regression.diffs +regression.out diff --git a/worker/realtime/Cargo.toml b/worker/realtime/Cargo.toml new file mode 100644 index 0000000..b032a2c --- /dev/null +++ b/worker/realtime/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "realtime" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +clap = { version = "3.1.12", features = ["derive"] } +dotenv = "0.15.0" +serde_json = "1.0" +serde = { version = "1.0", features=["derive"] } +uuid = { version = "1.0", features = ["serde"] } +tokio = { version = "1", features = ["full"] } +tokio-util = { version="0.7.2", features=["codec"] } +tokio-tungstenite = { verison = "0.17.1", features=["native-tls"] } +futures-util = "0.3.21" +url = "2.2.2" +futures-channel = "0.3.21" +futures = "0.3.21" + + + +#[dependencies.phoenix-channels] +#path = "../phoenix-channels-rs" diff --git a/worker/realtime/README.md b/worker/realtime/README.md new file mode 100644 index 0000000..00b6a02 --- /dev/null +++ b/worker/realtime/README.md @@ -0,0 +1,22 @@ +# Realtime Transport + +Realtime Transport reads JSON form stdin and forwards it to supabase realtime + +See parent directory for example usage + +## CLI + +``` +realtime 0.1.0 +reads JSON from stdin and forwards it to supabase realtime + +USAGE: + realtime [OPTIONS] --apikey + +OPTIONS: + --apikey + -h, --help Print help information + --topic [default: room:test] + --url [default: wss://sendwal.fly.dev/socket] + -V, --version Print version information +``` \ No newline at end of file diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs new file mode 100644 index 0000000..8e5bf74 --- /dev/null +++ b/worker/realtime/src/main.rs @@ -0,0 +1,160 @@ +use clap::Parser; +use futures::stream::SplitSink; +use futures_util::SinkExt; +use futures_util::{future, pin_mut, StreamExt}; +use serde::Serialize; +use serde_json; +use std::str; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt}; +use tokio::time::{sleep, Duration}; +use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream}; +use url; + +/// reads JSON from stdin and forwards it to supabase realtime +#[derive(Parser, Debug)] +#[clap(author, version, about, long_about = None)] +struct Args { + #[clap(long, default_value = "wss://sendwal.fly.dev/socket")] + url: String, + + #[clap(long)] + apikey: String, + + #[clap(long, default_value = "room:test")] + topic: String, +} + +#[derive(Serialize)] +enum PhoenixMessageEvent { + #[serde(rename(serialize = "phx_join"))] + Join, + #[serde(rename(serialize = "new_msg"))] + Message, + #[serde(rename(serialize = "heartbeat"))] + Heartbeat, +} + +#[derive(Serialize)] +struct PhoenixMessage { + event: PhoenixMessageEvent, + payload: serde_json::Value, + #[serde(rename(serialize = "ref"))] + reference: Option, + topic: String, +} + +#[tokio::main] +async fn main() { + // url + let args = Args::parse(); + let addr = build_url(&args.url, &args.apikey); + let url = url::Url::parse(&addr).unwrap(); + + // websocket + let (ws_stream, _) = connect_async(url).await.expect("Failed to connect"); + println!("WebSocket handshake successful"); + let (mut write, read) = ws_stream.split(); + + write = join_topic(write, args.topic.to_string()).await; + + // Futures channel + let (tx, rx) = futures_channel::mpsc::unbounded(); + let heartbeat_tx = tx.clone(); + + tokio::spawn(read_stdin(tx, args.topic.to_string())); + tokio::spawn(heartbeat(heartbeat_tx, args.topic.to_string())); + + // Map + let tx_to_ws = rx.map(Ok).forward(write); + let ws_to_stdout = { + read.for_each(|message| async { + let data = message.unwrap().into_data(); + tokio::io::stdout().write_all(&data).await.unwrap(); + }) + }; + + pin_mut!(tx_to_ws, ws_to_stdout); + future::select(tx_to_ws, ws_to_stdout).await; +} + +pub fn build_url(url: &str, apikey: &str) -> String { + let addr = format!("{}/websocket?vsn={}&apikey={}", url, "1.0.0", apikey); + addr +} + +async fn join_topic( + mut writer: SplitSink< + WebSocketStream>, + Message, + >, + topic_name: String, +) -> SplitSink>, Message> { + // join channel + let join_message = PhoenixMessage { + event: PhoenixMessageEvent::Join, + payload: serde_json::json!({}), + reference: None, + topic: topic_name, + }; + + let join_message = serde_json::to_string(&join_message).unwrap(); + let msg = Message::Text(join_message); + writer.send(msg).await.unwrap(); + writer +} + +// Our helper method which will read data from stdin and send it along the +// sender provided. +async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: String) { + let mut stdin = tokio::io::stdin(); + loop { + let mut buf = vec![0; 61024]; + + // Read stdin + let n = match stdin.read(&mut buf).await { + Err(_) | Ok(0) => break, + Ok(n) => n, + }; + buf.truncate(n); + + let mut lines = buf.lines(); + + while let Some(line) = lines.next_line().await.expect("invalid line") { + // Read stdin as string + + // Parse stdin string as json + let msg_json = + serde_json::from_str(&line).expect(&format!("failed to parse message '{}'", line)); + + // Repack json contents into a phoenix message + let phoenix_msg = PhoenixMessage { + event: PhoenixMessageEvent::Message, + payload: msg_json, + reference: None, + topic: topic.to_string(), + }; + + // Wrap phoenix message in a websocket message + let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); + + // push to futures stream + tx.unbounded_send(msg).unwrap(); + } + } +} + +async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender, topic: String) { + loop { + sleep(Duration::from_secs(3)).await; + let phoenix_msg = PhoenixMessage { + event: PhoenixMessageEvent::Heartbeat, + payload: serde_json::json!({"msg": "ping"}), + reference: None, + topic: topic.to_string(), + }; + // Wrap phoenix message in a websocket message + let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); + // push to futures stream + tx.unbounded_send(msg).unwrap(); + } +} diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index da4cef8..4c2bfd2 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -13,3 +13,4 @@ dotenv = "0.15.0" serde_json = "1.0" serde = "1.0" uuid = { version = "1.0", features = ["serde"] } + diff --git a/worker/walrus/README.md b/worker/walrus/README.md index eb3bcb6..5571ef4 100644 --- a/worker/walrus/README.md +++ b/worker/walrus/README.md @@ -19,7 +19,6 @@ OPTIONS: ## Features - ### Realtime (no polling) The worker wraps [pg_recvlogical](https://www.postgresql.org/docs/current/app-pgrecvlogical.html), a lightweight tool that ships with PostgreSQL, making use of the [streaming replication protocol](https://www.postgresql.org/docs/current/protocol-logical-replication.html#:~:text=The%20logical%20replication%20protocol%20sends,Start%20and%20Stream%20Stop%20messages.) to recieve new WAL messages. @@ -49,7 +48,7 @@ Clone and Navigate git clone https://github.com/supabase/walrus.git cd walrus git checkout worker -cd worker/walrus +cd worker ``` Start the DB @@ -59,6 +58,7 @@ docker-compose up Run the walrus worker ```sh +cd walrus cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres # Note: if you have jq installed, the output is more readable with diff --git a/worker/walrus/dockerfiles/Dockerfile b/worker/walrus/dockerfiles/Dockerfile deleted file mode 100644 index cfea2fb..0000000 --- a/worker/walrus/dockerfiles/Dockerfile +++ /dev/null @@ -1,8 +0,0 @@ -FROM supabase/postgres:latest - -RUN apt-get update -RUN apt install build-essential postgresql-server-dev-14 -y - -RUN git clone https://github.com/eulerto/wal2json.git --depth 1 53b548a29ebd6119323b6eb2f6013d7c5fe807ec -RUN cd wal2json && make && make install - From e94ea548f6d572e80ea21db3bbb4f0201074deb8 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 06:43:18 -0500 Subject: [PATCH 07/88] improve stdio buffer; handle unwraps explicitly --- worker/realtime/Cargo.toml | 1 + worker/realtime/src/main.rs | 138 +++++++++++++++++++++--------------- 2 files changed, 80 insertions(+), 59 deletions(-) diff --git a/worker/realtime/Cargo.toml b/worker/realtime/Cargo.toml index b032a2c..de14453 100644 --- a/worker/realtime/Cargo.toml +++ b/worker/realtime/Cargo.toml @@ -18,6 +18,7 @@ futures-util = "0.3.21" url = "2.2.2" futures-channel = "0.3.21" futures = "0.3.21" +log = "0.4.17" diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 8e5bf74..9fcdf9e 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -2,10 +2,11 @@ use clap::Parser; use futures::stream::SplitSink; use futures_util::SinkExt; use futures_util::{future, pin_mut, StreamExt}; +use log::{error, info, warn}; use serde::Serialize; use serde_json; use std::str; -use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; use tokio::time::{sleep, Duration}; use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream}; use url; @@ -48,33 +49,48 @@ async fn main() { // url let args = Args::parse(); let addr = build_url(&args.url, &args.apikey); - let url = url::Url::parse(&addr).unwrap(); + let url = url::Url::parse(&addr).expect("invalid URL"); - // websocket - let (ws_stream, _) = connect_async(url).await.expect("Failed to connect"); - println!("WebSocket handshake successful"); - let (mut write, read) = ws_stream.split(); - - write = join_topic(write, args.topic.to_string()).await; - - // Futures channel - let (tx, rx) = futures_channel::mpsc::unbounded(); - let heartbeat_tx = tx.clone(); - - tokio::spawn(read_stdin(tx, args.topic.to_string())); - tokio::spawn(heartbeat(heartbeat_tx, args.topic.to_string())); - - // Map - let tx_to_ws = rx.map(Ok).forward(write); - let ws_to_stdout = { - read.for_each(|message| async { - let data = message.unwrap().into_data(); - tokio::io::stdout().write_all(&data).await.unwrap(); - }) - }; - - pin_mut!(tx_to_ws, ws_to_stdout); - future::select(tx_to_ws, ws_to_stdout).await; + loop { + // websocket + info!("Connecting to websocket"); + let ws_connection = connect_async(&url).await; + + match ws_connection { + Err(msg) => { + let n_seconds = 3; + error!("Failed to connect to websocket. Error: {}", msg); + info!("Attempting websocket reconnect in {} seconds", n_seconds); + sleep(Duration::from_secs(n_seconds)).await; + } + Ok((ws_stream, _)) => { + println!("WebSocket handshake successful"); + + let (mut write, read) = ws_stream.split(); + + write = join_topic(write, args.topic.to_string()).await; + + // Futures channel + let (tx, rx) = futures_channel::mpsc::unbounded(); + let heartbeat_tx = tx.clone(); + + tokio::spawn(read_stdin(tx, args.topic.to_string())); + tokio::spawn(heartbeat(heartbeat_tx, args.topic.to_string())); + + // Map + let tx_to_ws = rx.map(Ok).forward(write); + let ws_to_stdout = { + read.for_each(|message| async { + let data = message.unwrap().into_data(); + tokio::io::stdout().write_all(&data).await.unwrap(); + }) + }; + + pin_mut!(tx_to_ws, ws_to_stdout); + future::select(tx_to_ws, ws_to_stdout).await; + } + } + } } pub fn build_url(url: &str, apikey: &str) -> String { @@ -106,39 +122,43 @@ async fn join_topic( // Our helper method which will read data from stdin and send it along the // sender provided. async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: String) { - let mut stdin = tokio::io::stdin(); - loop { - let mut buf = vec![0; 61024]; - - // Read stdin - let n = match stdin.read(&mut buf).await { - Err(_) | Ok(0) => break, - Ok(n) => n, - }; - buf.truncate(n); - - let mut lines = buf.lines(); - - while let Some(line) = lines.next_line().await.expect("invalid line") { - // Read stdin as string + let stdin = tokio::io::stdin(); + let buf = BufReader::new(stdin); + let mut lines = buf.lines(); - // Parse stdin string as json - let msg_json = - serde_json::from_str(&line).expect(&format!("failed to parse message '{}'", line)); - - // Repack json contents into a phoenix message - let phoenix_msg = PhoenixMessage { - event: PhoenixMessageEvent::Message, - payload: msg_json, - reference: None, - topic: topic.to_string(), - }; - - // Wrap phoenix message in a websocket message - let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); - - // push to futures stream - tx.unbounded_send(msg).unwrap(); + loop { + let line_res_opt: Result, std::io::Error> = lines.next_line().await; + + match line_res_opt { + Ok(line_opt) => { + match line_opt { + Some(line) => { + // Parse stdin string as json + let msg_json = serde_json::from_str(&line) + .expect(&format!("failed to parse message '{}'", line)); + + // Repack json contents into a phoenix message + let phoenix_msg = PhoenixMessage { + event: PhoenixMessageEvent::Message, + payload: msg_json, + reference: None, + topic: topic.to_string(), + }; + + // Wrap phoenix message in a websocket message + let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); + + // push to futures stream + tx.unbounded_send(msg).unwrap(); + } + None => { + warn!("Received empty line from stdin"); + } + } + } + Err(err) => { + error!("Error reading line from stdin: {}", err); + } } } } From 1c80b2c51e4cf49c082beb458e449899f84f24c4 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 06:53:01 -0500 Subject: [PATCH 08/88] handle non-json stdin --- worker/realtime/src/main.rs | 40 +++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 9fcdf9e..6e862eb 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -6,7 +6,7 @@ use log::{error, info, warn}; use serde::Serialize; use serde_json; use std::str; -use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::time::{sleep, Duration}; use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream}; use url; @@ -134,22 +134,28 @@ async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: match line_opt { Some(line) => { // Parse stdin string as json - let msg_json = serde_json::from_str(&line) - .expect(&format!("failed to parse message '{}'", line)); - - // Repack json contents into a phoenix message - let phoenix_msg = PhoenixMessage { - event: PhoenixMessageEvent::Message, - payload: msg_json, - reference: None, - topic: topic.to_string(), - }; - - // Wrap phoenix message in a websocket message - let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); - - // push to futures stream - tx.unbounded_send(msg).unwrap(); + match serde_json::from_str(&line) { + Ok(msg_json) => { + // Repack json contents into a phoenix message + let phoenix_msg = PhoenixMessage { + event: PhoenixMessageEvent::Message, + payload: msg_json, + reference: None, + topic: topic.to_string(), + }; + // Wrap phoenix message in a websocket message + let msg = + Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); + // push to output stream + tx.unbounded_send(msg).unwrap(); + } + Err(err) => { + warn!( + "Error parsing stdin line to json: error={}, line={}", + err, line + ) + } + } } None => { warn!("Received empty line from stdin"); From 8e1069ef86ce6a46561b3680203a798fd5a5b022 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 08:28:06 -0500 Subject: [PATCH 09/88] enable logger --- worker/realtime/Cargo.toml | 1 + worker/realtime/src/main.rs | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/worker/realtime/Cargo.toml b/worker/realtime/Cargo.toml index de14453..6c3b637 100644 --- a/worker/realtime/Cargo.toml +++ b/worker/realtime/Cargo.toml @@ -19,6 +19,7 @@ url = "2.2.2" futures-channel = "0.3.21" futures = "0.3.21" log = "0.4.17" +env_logger = "0.9.0" diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 6e862eb..7a141ed 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -1,4 +1,5 @@ use clap::Parser; +use env_logger; use futures::stream::SplitSink; use futures_util::SinkExt; use futures_util::{future, pin_mut, StreamExt}; @@ -51,6 +52,9 @@ async fn main() { let addr = build_url(&args.url, &args.apikey); let url = url::Url::parse(&addr).expect("invalid URL"); + // enable logger + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); + loop { // websocket info!("Connecting to websocket"); @@ -64,8 +68,7 @@ async fn main() { sleep(Duration::from_secs(n_seconds)).await; } Ok((ws_stream, _)) => { - println!("WebSocket handshake successful"); - + info!("WebSocket handshake successful"); let (mut write, read) = ws_stream.split(); write = join_topic(write, args.topic.to_string()).await; From 001a37a292f0b2666fff44b39c65a07633596bcf Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 09:33:50 -0500 Subject: [PATCH 10/88] idempotent migrations --- .../2022-04-29-132611_initial/up.sql | 172 +++++++++++++----- worker/walrus/src/main.rs | 19 +- 2 files changed, 146 insertions(+), 45 deletions(-) diff --git a/worker/walrus/migrations/2022-04-29-132611_initial/up.sql b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql index e36afbc..693e3a7 100644 --- a/worker/walrus/migrations/2022-04-29-132611_initial/up.sql +++ b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql @@ -1,14 +1,61 @@ -create schema realtime; +create schema if not exists realtime; -create type realtime.equality_op as enum( - 'eq', 'neq', 'lt', 'lte', 'gt', 'gte' -); +-- Temporary functions to assist with idempotency +create or replace function realtime.type_exists(schema_name text, type_name text) + returns bool + language plpgsql +as $$ +begin + begin + perform format('%I.%I', schema_name, type_name)::regtype; + return true; + exception when others then + return false; + end; +end +$$; -create type realtime.action as enum ('INSERT', 'UPDATE', 'DELETE', 'ERROR'); +create or replace function realtime.table_exists(schema_name text, table_name text) + returns bool + language plpgsql +as $$ +begin + begin + perform format('%I.%I', schema_name, table_name)::regclass; + return true; + exception when others then + return false; + end; +end +$$; -create function realtime.cast(val text, type_ regtype) +-- realtime.equality_op +do $$ + begin + if not realtime.type_exists('realtime', 'equality_op') then + + create type realtime.equality_op as enum( + 'eq', 'neq', 'lt', 'lte', 'gt', 'gte' + ); + + end if; + end +$$; + +-- realtime.action +do $$ + begin + if not realtime.type_exists('realtime', 'action') then + + create type realtime.action as enum ('INSERT', 'UPDATE', 'DELETE', 'ERROR'); + + end if; + end +$$; + +create or replace function realtime.cast(val text, type_ regtype) returns jsonb immutable language plpgsql @@ -22,14 +69,8 @@ end $$; -create type realtime.user_defined_filter as ( - column_name text, - op realtime.equality_op, - value text -); - -create function realtime.to_regrole(role_name text) +create or replace function realtime.to_regrole(role_name text) returns regrole immutable language sql @@ -37,19 +78,28 @@ create function realtime.to_regrole(role_name text) as $$ select role_name::regrole $$; -create table realtime.subscription ( - -- Tracks which subscriptions are active - id bigint generated always as identity primary key, - subscription_id uuid not null, - entity regclass not null, - filters realtime.user_defined_filter[] not null default '{}', - claims jsonb not null, - claims_role regrole not null generated always as (realtime.to_regrole(claims ->> 'role')) stored, - created_at timestamp not null default timezone('utc', now()), +-- realtime.subscription +do $$ + begin + if not realtime.table_exists('realtime', 'subscription') then + + create table realtime.subscription ( + -- Tracks which subscriptions are active + id bigint generated always as identity primary key, + subscription_id uuid not null, + entity regclass not null, + filters realtime.user_defined_filter[] not null default '{}', + claims jsonb not null, + claims_role regrole not null generated always as (realtime.to_regrole(claims ->> 'role')) stored, + created_at timestamp not null default timezone('utc', now()), - unique (subscription_id, entity, filters) -); -create index ix_realtime_subscription_entity on realtime.subscription using hash (entity); + unique (subscription_id, entity, filters) + ); + create index ix_realtime_subscription_entity on realtime.subscription using hash (entity); + + end if; + end +$$; create or replace function realtime.subscription_check_filters() @@ -110,11 +160,25 @@ begin end; $$; -create trigger tr_check_filters - before insert or update on realtime.subscription - for each row - execute function realtime.subscription_check_filters(); +-- tr_check_filters on realtime.subscription +do $$ + begin + if not exists( + select 1 + from pg_trigger + where + tgrelid = 'realtime.subscription'::regclass::oid + and tgname = 'tr_check_filters' + ) then + + create trigger tr_check_filters + before insert or update on realtime.subscription + for each row + execute function realtime.subscription_check_filters(); + end if; + end +$$; create or replace function realtime.quote_wal2json(entity regclass) returns text @@ -186,14 +250,24 @@ end; $$; -create type realtime.wal_column as ( - name text, - type_name text, - type_oid oid, - value jsonb, - is_pkey boolean, - is_selectable boolean -); +-- realtime.wal_column +do $$ + begin + if not realtime.type_exists('realtime', 'wal_column') then + + create type realtime.wal_column as ( + name text, + type_name text, + type_oid oid, + value jsonb, + is_pkey boolean, + is_selectable boolean + ); + + end if; + end +$$; + create or replace function realtime.build_prepared_statement_sql( prepared_statement_name text, @@ -229,13 +303,21 @@ Example entity $$; +-- realtime.wal_rls +do $$ + begin + if not realtime.type_exists('realtime', 'wal_rls') then -create type realtime.wal_rls as ( - wal jsonb, - is_rls_enabled boolean, - subscription_ids uuid[], - errors text[] -); + create type realtime.wal_rls as ( + wal jsonb, + is_rls_enabled boolean, + subscription_ids uuid[], + errors text[] + ); + + end if; + end +$$; @@ -522,3 +604,7 @@ begin perform set_config('role', null, true); end; $$; + + +drop function realtime.type_exists; +drop function realtime.table_exists; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 06111dc..fdd42ae 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -5,6 +5,7 @@ use diesel::*; use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use serde::Serialize; use serde_json; +use std::error::Error; use std::io; use std::io::BufRead; use std::process::{Command, Stdio}; @@ -12,6 +13,21 @@ use uuid; pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); +fn run_migrations( + connection: &mut PgConnection, +) -> Result<(), Box> { + sql_query("create schema if not exists realtime") + .execute(connection) + .expect("failed to create 'realtime' schema"); + + sql_query("set search_path='realtime'") + .execute(connection) + .expect("failed to set search path"); + + connection.run_pending_migrations(MIGRATIONS)?; + Ok(()) +} + #[derive(Serialize)] pub struct WalrusRecord { wal: serde_json::Value, @@ -64,8 +80,7 @@ fn main() { let conn = &mut PgConnection::establish(&args.connection).unwrap(); // Run any pending migrations - conn.run_pending_migrations(MIGRATIONS) - .expect("Pending migrations failed to execute"); + run_migrations(conn).expect("Pending migrations failed to execute"); // Reading from stdin let stdin = cmd.stdout.as_mut().unwrap(); From 2d63ec5425f3db2c2418c0b16234749293a50692 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 11:13:41 -0500 Subject: [PATCH 11/88] sync from master --- .../down.sql | 254 ++++++++++++++++ .../up.sql | 276 ++++++++++++++++++ 2 files changed, 530 insertions(+) create mode 100644 worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/down.sql create mode 100644 worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/up.sql diff --git a/worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/down.sql b/worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/down.sql new file mode 100644 index 0000000..24d4e07 --- /dev/null +++ b/worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/down.sql @@ -0,0 +1,254 @@ +create or replace function realtime.apply_rls(wal jsonb, max_record_bytes int = 1024 * 1024) + returns setof realtime.wal_rls + language plpgsql + volatile +as $$ +declare + -- Regclass of the table e.g. public.notes + entity_ regclass = (quote_ident(wal ->> 'schema') || '.' || quote_ident(wal ->> 'table'))::regclass; + + -- I, U, D, T: insert, update ... + action realtime.action = ( + case wal ->> 'action' + when 'I' then 'INSERT' + when 'U' then 'UPDATE' + when 'D' then 'DELETE' + else 'ERROR' + end + ); + + -- Is row level security enabled for the table + is_rls_enabled bool = relrowsecurity from pg_class where oid = entity_; + + subscriptions realtime.subscription[] = array_agg(subs) + from + realtime.subscription subs + where + subs.entity = entity_; + + -- Subscription vars + roles regrole[] = array_agg(distinct us.claims_role) + from + unnest(subscriptions) us; + + working_role regrole; + claimed_role regrole; + claims jsonb; + + subscription_id uuid; + subscription_has_access bool; + visible_to_subscription_ids uuid[] = '{}'; + + -- structured info for wal's columns + columns realtime.wal_column[]; + -- previous identity values for update/delete + old_columns realtime.wal_column[]; + + error_record_exceeds_max_size boolean = octet_length(wal::text) > max_record_bytes; + + -- Primary jsonb output for record + output jsonb; + +begin + perform set_config('role', null, true); + + columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'columns') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + old_columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'identity') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + for working_role in select * from unnest(roles) loop + + -- Update `is_selectable` for columns and old_columns + columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(columns) c; + + old_columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(old_columns) c; + + if action <> 'DELETE' and count(1) = 0 from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + -- subscriptions is already filtered by entity + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 400: Bad Request, no primary key'] + )::realtime.wal_rls; + + -- The claims role does not have SELECT permission to the primary key of entity + elsif action <> 'DELETE' and sum(c.is_selectable::int) <> count(1) from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 401: Unauthorized'] + )::realtime.wal_rls; + + else + output = jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action, + 'commit_timestamp', to_char( + (wal ->> 'timestamp')::timestamptz, + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ), + 'columns', ( + select + jsonb_agg( + jsonb_build_object( + 'name', pa.attname, + 'type', pt.typname + ) + order by pa.attnum asc + ) + from + pg_attribute pa + join pg_type pt + on pa.atttypid = pt.oid + where + attrelid = entity_ + and attnum > 0 + and pg_catalog.has_column_privilege(working_role, entity_, pa.attname, 'SELECT') + ) + ) + -- Add "record" key for insert and update + || case + when error_record_exceeds_max_size then jsonb_build_object('record', '{}'::jsonb) + when action in ('INSERT', 'UPDATE') then + jsonb_build_object( + 'record', + (select jsonb_object_agg((c).name, (c).value) from unnest(columns) c where (c).is_selectable) + ) + else '{}'::jsonb + end + -- Add "old_record" key for update and delete + || case + when error_record_exceeds_max_size then jsonb_build_object('old_record', '{}'::jsonb) + when action in ('UPDATE', 'DELETE') then + jsonb_build_object( + 'old_record', + (select jsonb_object_agg((c).name, (c).value) from unnest(old_columns) c where (c).is_selectable) + ) + else '{}'::jsonb + end; + + -- Create the prepared statement + if is_rls_enabled and action <> 'DELETE' then + if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then + deallocate walrus_rls_stmt; + end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, columns); + end if; + + visible_to_subscription_ids = '{}'; + + for subscription_id, claims in ( + select + subs.subscription_id, + subs.claims + from + unnest(subscriptions) subs + where + subs.entity = entity_ + and subs.claims_role = working_role + and realtime.is_visible_through_filters(columns, subs.filters) + ) loop + + if not is_rls_enabled or action = 'DELETE' then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + else + -- Check if RLS allows the role to see the record + perform + set_config('role', working_role::text, true), + set_config('request.jwt.claims', claims::text, true); + + execute 'execute walrus_rls_stmt' into subscription_has_access; + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end if; + end loop; + + perform set_config('role', null, true); + + return next ( + output, + is_rls_enabled, + visible_to_subscription_ids, + case + when error_record_exceeds_max_size then array['Error 413: Payload Too Large'] + else '{}' + end + )::realtime.wal_rls; + + end if; + end loop; + + perform set_config('role', null, true); +end; +$$; diff --git a/worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/up.sql b/worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/up.sql new file mode 100644 index 0000000..289b2b6 --- /dev/null +++ b/worker/walrus/migrations/2022-05-18-143641_incl_small_columns_on_larg_payload/up.sql @@ -0,0 +1,276 @@ +create or replace function realtime.apply_rls(wal jsonb, max_record_bytes int = 1024 * 1024) + returns setof realtime.wal_rls + language plpgsql + volatile +as $$ +declare + -- Regclass of the table e.g. public.notes + entity_ regclass = (quote_ident(wal ->> 'schema') || '.' || quote_ident(wal ->> 'table'))::regclass; + + -- I, U, D, T: insert, update ... + action realtime.action = ( + case wal ->> 'action' + when 'I' then 'INSERT' + when 'U' then 'UPDATE' + when 'D' then 'DELETE' + else 'ERROR' + end + ); + + -- Is row level security enabled for the table + is_rls_enabled bool = relrowsecurity from pg_class where oid = entity_; + + subscriptions realtime.subscription[] = array_agg(subs) + from + realtime.subscription subs + where + subs.entity = entity_; + + -- Subscription vars + roles regrole[] = array_agg(distinct us.claims_role) + from + unnest(subscriptions) us; + + working_role regrole; + claimed_role regrole; + claims jsonb; + + subscription_id uuid; + subscription_has_access bool; + visible_to_subscription_ids uuid[] = '{}'; + + -- structured info for wal's columns + columns realtime.wal_column[]; + -- previous identity values for update/delete + old_columns realtime.wal_column[]; + + error_record_exceeds_max_size boolean = octet_length(wal::text) > max_record_bytes; + + -- Primary jsonb output for record + output jsonb; + +begin + perform set_config('role', null, true); + + columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'columns') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + old_columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'identity') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + for working_role in select * from unnest(roles) loop + + -- Update `is_selectable` for columns and old_columns + columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(columns) c; + + old_columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(old_columns) c; + + if action <> 'DELETE' and count(1) = 0 from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + -- subscriptions is already filtered by entity + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 400: Bad Request, no primary key'] + )::realtime.wal_rls; + + -- The claims role does not have SELECT permission to the primary key of entity + elsif action <> 'DELETE' and sum(c.is_selectable::int) <> count(1) from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 401: Unauthorized'] + )::realtime.wal_rls; + + else + output = jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action, + 'commit_timestamp', to_char( + (wal ->> 'timestamp')::timestamptz, + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ), + 'columns', ( + select + jsonb_agg( + jsonb_build_object( + 'name', pa.attname, + 'type', pt.typname + ) + order by pa.attnum asc + ) + from + pg_attribute pa + join pg_type pt + on pa.atttypid = pt.oid + where + attrelid = entity_ + and attnum > 0 + and pg_catalog.has_column_privilege(working_role, entity_, pa.attname, 'SELECT') + ) + ) + -- Add "record" key for insert and update + || case + when action in ('INSERT', 'UPDATE') then + case + when error_record_exceeds_max_size then + jsonb_build_object( + 'record', + ( + select jsonb_object_agg((c).name, (c).value) + from unnest(columns) c + where (c).is_selectable and (octet_length((c).value::text) <= 64) + ) + ) + else + jsonb_build_object( + 'record', + (select jsonb_object_agg((c).name, (c).value) from unnest(columns) c where (c).is_selectable) + ) + end + else '{}'::jsonb + end + -- Add "old_record" key for update and delete + || case + when action in ('UPDATE', 'DELETE') then + case + when error_record_exceeds_max_size then + jsonb_build_object( + 'old_record', + ( + select jsonb_object_agg((c).name, (c).value) + from unnest(old_columns) c + where (c).is_selectable and (octet_length((c).value::text) <= 64) + ) + ) + else + jsonb_build_object( + 'old_record', + (select jsonb_object_agg((c).name, (c).value) from unnest(old_columns) c where (c).is_selectable) + ) + end + else '{}'::jsonb + end; + + -- Create the prepared statement + if is_rls_enabled and action <> 'DELETE' then + if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then + deallocate walrus_rls_stmt; + end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, columns); + end if; + + visible_to_subscription_ids = '{}'; + + for subscription_id, claims in ( + select + subs.subscription_id, + subs.claims + from + unnest(subscriptions) subs + where + subs.entity = entity_ + and subs.claims_role = working_role + and realtime.is_visible_through_filters(columns, subs.filters) + ) loop + + if not is_rls_enabled or action = 'DELETE' then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + else + -- Check if RLS allows the role to see the record + perform + set_config('role', working_role::text, true), + set_config('request.jwt.claims', claims::text, true); + + execute 'execute walrus_rls_stmt' into subscription_has_access; + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end if; + end loop; + + perform set_config('role', null, true); + + return next ( + output, + is_rls_enabled, + visible_to_subscription_ids, + case + when error_record_exceeds_max_size then array['Error 413: Payload Too Large'] + else '{}' + end + )::realtime.wal_rls; + + end if; + end loop; + + perform set_config('role', null, true); +end; +$$; From e836d938653d33ce3d7b717fcdcaa93790529ddd Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 11:33:39 -0500 Subject: [PATCH 12/88] handle postgres restarts/shutdowns gracefully --- worker/walrus/src/main.rs | 156 +++++++++++++++++++++++--------------- 1 file changed, 95 insertions(+), 61 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index fdd42ae..700cce8 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -9,6 +9,8 @@ use std::error::Error; use std::io; use std::io::BufRead; use std::process::{Command, Stdio}; +use std::thread::sleep; +use std::time; use uuid; pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); @@ -53,7 +55,32 @@ fn main() { // Parse command line arguments let args = Args::parse(); - let mut cmd = Command::new("pg_recvlogical") + loop { + match run(&args) { + Err(err) => { + println!("Error: {}", err); + } + _ => continue, + }; + println!("Stream interrupted. Restarting pg_recvlogical in 5 seconds"); + sleep(time::Duration::from_secs(5)); + } +} + +fn run(args: &Args) -> Result<(), String> { + // Connect to Postgres + let conn_result = &mut PgConnection::establish(&args.connection); + + let conn = match conn_result { + Ok(c) => c, + Err(_) => { + return Err("failed to make postgres connection".to_string()); + } + }; + // Run pending migrations + run_migrations(conn).expect("Pending migrations failed to execute"); + + let cmd = Command::new("pg_recvlogical") //&args .args(vec![ "--file=-", @@ -69,79 +96,86 @@ fn main() { "--create-slot", "--if-not-exists", "--start", + "--no-loop", ]) .stdout(Stdio::piped()) - .spawn() - .unwrap(); - - { - // Connect to Postgres - // "postgresql://oliverrice:@localhost:28814/walrus"; - let conn = &mut PgConnection::establish(&args.connection).unwrap(); - - // Run any pending migrations - run_migrations(conn).expect("Pending migrations failed to execute"); + .stderr(Stdio::piped()) + .spawn(); - // Reading from stdin - let stdin = cmd.stdout.as_mut().unwrap(); - let stdin_reader = io::BufReader::new(stdin); - let stdin_lines = stdin_reader.lines(); + match cmd { + Err(err) => Err(format!("{}", err)), + Ok(mut cmd) => { + println!("Connection established"); + // Reading from stdin + let stdin = cmd.stdout.as_mut().unwrap(); + let stdin_reader = io::BufReader::new(stdin); + let stdin_lines = stdin_reader.lines(); - // Iterate input data - for input_line in stdin_lines { - match input_line { - Ok(line) => { - let result_json_line = serde_json::from_str::(&line); - match result_json_line { - Ok(json_line) => { - let result_wal_rls_rows = sql::< - Record<( - sql_types::Jsonb, - sql_types::Bool, - sql_types::Array, - sql_types::Array, - )>, - >( - "SELECT x from realtime.apply_rls(" - ) - .bind::(json_line) - .sql(") x") - .get_results::<( - serde_json::Value, - bool, - Vec, - Vec, - )>(conn); - match result_wal_rls_rows { - Ok(rows) => { - for row in rows { - let walrus_rec = WalrusRecord { - wal: row.0, - is_rls_enabled: row.1, - subscription_ids: row.2, - errors: row.3, - }; - match serde_json::to_string(&walrus_rec) { - Ok(walrus_json) => println!("{}", walrus_json), - Err(err) => { - println!( - "Failed to serialize walrus result: {}", - err - ) + // Iterate input data + for input_line in stdin_lines { + match input_line { + Ok(line) => { + let result_json_line = serde_json::from_str::(&line); + match result_json_line { + Ok(json_line) => { + let result_wal_rls_rows = + sql::< + Record<( + sql_types::Jsonb, + sql_types::Bool, + sql_types::Array, + sql_types::Array, + )>, + >( + "SELECT x from realtime.apply_rls(" + ) + .bind::(json_line) + .sql(") x") + .get_results::<( + serde_json::Value, + bool, + Vec, + Vec, + )>( + conn + ); + match result_wal_rls_rows { + Ok(rows) => { + for row in rows { + let walrus_rec = WalrusRecord { + wal: row.0, + is_rls_enabled: row.1, + subscription_ids: row.2, + errors: row.3, + }; + match serde_json::to_string(&walrus_rec) { + Ok(walrus_json) => println!("{}", walrus_json), + Err(err) => { + println!( + "Failed to serialize walrus result: {}", + err + ) + } } } } + Err(err) => { + cmd.kill().unwrap(); + println!("WALRUS Error: {}", err); + return Err("walrus error".to_string()); + } } - Err(err) => println!("WALRUS Error: {}", err), } + Err(err) => println!("Failed to parse: {}", err), } - Err(err) => println!("Failed to parse: {}", err), } + Err(err) => println!("Error: {}", err), } - Err(err) => println!("Error: {}", err), + } + match cmd.wait() { + Ok(_) => Ok(()), + Err(err) => Err(format!("{}", err)), } } } - - cmd.wait().unwrap(); } From 5415a21ee5b62d5a5349be7ae23ace1ce76b94ea Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 12:15:56 -0500 Subject: [PATCH 13/88] walrus uses logging --- worker/realtime/Cargo.toml | 7 ------- worker/walrus/Cargo.toml | 5 ++--- worker/walrus/src/main.rs | 19 ++++++++++++------- 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/worker/realtime/Cargo.toml b/worker/realtime/Cargo.toml index 6c3b637..b63d157 100644 --- a/worker/realtime/Cargo.toml +++ b/worker/realtime/Cargo.toml @@ -3,8 +3,6 @@ name = "realtime" version = "0.1.0" edition = "2021" -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - [dependencies] clap = { version = "3.1.12", features = ["derive"] } dotenv = "0.15.0" @@ -20,8 +18,3 @@ futures-channel = "0.3.21" futures = "0.3.21" log = "0.4.17" env_logger = "0.9.0" - - - -#[dependencies.phoenix-channels] -#path = "../phoenix-channels-rs" diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index 4c2bfd2..f60177f 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -3,8 +3,6 @@ name = "walrus" version = "0.1.0" edition = "2021" -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - [dependencies] clap = { version = "3.1.12", features = ["derive"] } diesel = { version = "2.0.0-rc.0", features = ["postgres", "serde_json", "uuid"] } @@ -13,4 +11,5 @@ dotenv = "0.15.0" serde_json = "1.0" serde = "1.0" uuid = { version = "1.0", features = ["serde"] } - +log = "0.4.17" +env_logger = "0.9.0" diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 700cce8..53a2e14 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -3,6 +3,8 @@ use diesel::dsl::sql; use diesel::sql_types::*; use diesel::*; use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; +use env_logger; +use log::{error, info, warn}; use serde::Serialize; use serde_json; use std::error::Error; @@ -55,14 +57,17 @@ fn main() { // Parse command line arguments let args = Args::parse(); + // enable logger + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); + loop { match run(&args) { Err(err) => { - println!("Error: {}", err); + warn!("Error: {}", err); } _ => continue, }; - println!("Stream interrupted. Restarting pg_recvlogical in 5 seconds"); + info!("Stream interrupted. Restarting pg_recvlogical in 5 seconds"); sleep(time::Duration::from_secs(5)); } } @@ -105,7 +110,7 @@ fn run(args: &Args) -> Result<(), String> { match cmd { Err(err) => Err(format!("{}", err)), Ok(mut cmd) => { - println!("Connection established"); + info!("Connection established"); // Reading from stdin let stdin = cmd.stdout.as_mut().unwrap(); let stdin_reader = io::BufReader::new(stdin); @@ -151,7 +156,7 @@ fn run(args: &Args) -> Result<(), String> { match serde_json::to_string(&walrus_rec) { Ok(walrus_json) => println!("{}", walrus_json), Err(err) => { - println!( + error!( "Failed to serialize walrus result: {}", err ) @@ -161,15 +166,15 @@ fn run(args: &Args) -> Result<(), String> { } Err(err) => { cmd.kill().unwrap(); - println!("WALRUS Error: {}", err); + error!("WALRUS Error: {}", err); return Err("walrus error".to_string()); } } } - Err(err) => println!("Failed to parse: {}", err), + Err(err) => error!("Failed to parse: {}", err), } } - Err(err) => println!("Error: {}", err), + Err(err) => error!("Error: {}", err), } } match cmd.wait() { From cd0e06c987be677f6950d0649182d97b53978cd9 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 18 May 2022 12:19:53 -0500 Subject: [PATCH 14/88] log walrus startup --- worker/walrus/src/main.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 53a2e14..f73aed9 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -84,6 +84,7 @@ fn run(args: &Args) -> Result<(), String> { }; // Run pending migrations run_migrations(conn).expect("Pending migrations failed to execute"); + info!("Postgres connection established"); let cmd = Command::new("pg_recvlogical") //&args @@ -110,7 +111,7 @@ fn run(args: &Args) -> Result<(), String> { match cmd { Err(err) => Err(format!("{}", err)), Ok(mut cmd) => { - info!("Connection established"); + info!("pg_recvlogical started"); // Reading from stdin let stdin = cmd.stdout.as_mut().unwrap(); let stdin_reader = io::BufReader::new(stdin); From 128c0cf441cd962c76f1abe27e02027c522e0c03 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 19 May 2022 05:54:23 -0500 Subject: [PATCH 15/88] generic headers --- worker/realtime/src/main.rs | 36 +++++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 7a141ed..5f2584b 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -19,13 +19,33 @@ struct Args { #[clap(long, default_value = "wss://sendwal.fly.dev/socket")] url: String, - #[clap(long)] - apikey: String, + #[clap( + long, + value_name = "HEADER>=, #[clap(long, default_value = "room:test")] topic: String, } +fn parse_header(user_input: &str) -> Result<(String, String), String> { + let mut splitter = user_input.splitn(2, "="); + + let key = splitter.next(); + let val = splitter.next(); + + match (key, val) { + (Some(k), Some(v)) => Ok((k.to_string(), v.to_string())), + _ => Err(format!( + "Could not parse header key-value pair: {}", + user_input + )), + } +} + #[derive(Serialize)] enum PhoenixMessageEvent { #[serde(rename(serialize = "phx_join"))] @@ -49,9 +69,11 @@ struct PhoenixMessage { async fn main() { // url let args = Args::parse(); - let addr = build_url(&args.url, &args.apikey); + let addr = build_url(&args.url, &args.header); let url = url::Url::parse(&addr).expect("invalid URL"); + println!("{:?}", args); + // enable logger env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); @@ -96,8 +118,12 @@ async fn main() { } } -pub fn build_url(url: &str, apikey: &str) -> String { - let addr = format!("{}/websocket?vsn={}&apikey={}", url, "1.0.0", apikey); +pub fn build_url(url: &str, params: &Vec<(String, String)>) -> String { + let mut params_uri: String = "".to_owned(); + for (k, v) in params { + params_uri.push_str(&format!("&{}={}", k, v)); + } + let addr = format!("{}/websocket?vsn={}{}", url, "1.0.0", params_uri); addr } From d0af2697fcc61822946520c7518a84ccbf83d44e Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 19 May 2022 06:00:32 -0500 Subject: [PATCH 16/88] generic header in readme --- worker/README.md | 2 +- worker/realtime/README.md | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/worker/README.md b/worker/README.md index 000b030..f1d8283 100644 --- a/worker/README.md +++ b/worker/README.md @@ -33,7 +33,7 @@ cargo run --bin walrus -- \ --connection=postgresql://postgres:password@localhost:5501/postgres | cargo run --bin realtime -- \ --url=wss://sendwal.fly.dev/socket \ - --apikey= + --header=apikey= ``` Connect to the database at `postgresql://postgres:password@localhost:5501/postgres` diff --git a/worker/realtime/README.md b/worker/realtime/README.md index 00b6a02..a07eccb 100644 --- a/worker/realtime/README.md +++ b/worker/realtime/README.md @@ -11,12 +11,13 @@ realtime 0.1.0 reads JSON from stdin and forwards it to supabase realtime USAGE: - realtime [OPTIONS] --apikey + realtime [OPTIONS] OPTIONS: - --apikey - -h, --help Print help information - --topic [default: room:test] - --url [default: wss://sendwal.fly.dev/socket] - -V, --version Print version information + -h, --help Print help information + --header
= + --topic [default: room:test] + --url [default: wss://sendwal.fly.dev/socket] + -V, --version Print version information + ``` \ No newline at end of file From 36fea60acb852cfef320baa15ddb22a56df0cd08 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 25 May 2022 11:27:27 -0500 Subject: [PATCH 17/88] topic tweaks --- worker/realtime/src/main.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 5f2584b..65557bf 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -50,7 +50,7 @@ fn parse_header(user_input: &str) -> Result<(String, String), String> { enum PhoenixMessageEvent { #[serde(rename(serialize = "phx_join"))] Join, - #[serde(rename(serialize = "new_msg"))] + #[serde(rename(serialize = "changes"))] Message, #[serde(rename(serialize = "heartbeat"))] Heartbeat, @@ -205,7 +205,7 @@ async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender, topic: S event: PhoenixMessageEvent::Heartbeat, payload: serde_json::json!({"msg": "ping"}), reference: None, - topic: topic.to_string(), + topic: "phoenix".to_string(), }; // Wrap phoenix message in a websocket message let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); From 8f2cb85948e5e311e8f8ab8944ae61d41c1f021e Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 25 May 2022 12:27:42 -0500 Subject: [PATCH 18/88] remove topic from heartbeat --- worker/realtime/src/main.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 65557bf..f8cbe8f 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -100,7 +100,7 @@ async fn main() { let heartbeat_tx = tx.clone(); tokio::spawn(read_stdin(tx, args.topic.to_string())); - tokio::spawn(heartbeat(heartbeat_tx, args.topic.to_string())); + tokio::spawn(heartbeat(heartbeat_tx)); // Map let tx_to_ws = rx.map(Ok).forward(write); @@ -198,7 +198,7 @@ async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: } } -async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender, topic: String) { +async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender) { loop { sleep(Duration::from_secs(3)).await; let phoenix_msg = PhoenixMessage { From 9dd24038e3fe37d95c98ca0f4fcebd0d214f47ea Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 1 Jun 2022 07:13:24 -0500 Subject: [PATCH 19/88] handle protocol errors from realtime --- worker/realtime/src/main.rs | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index f8cbe8f..9bdef7f 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -7,7 +7,7 @@ use log::{error, info, warn}; use serde::Serialize; use serde_json; use std::str; -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::time::{sleep, Duration}; use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream}; use url; @@ -106,8 +106,19 @@ async fn main() { let tx_to_ws = rx.map(Ok).forward(write); let ws_to_stdout = { read.for_each(|message| async { - let data = message.unwrap().into_data(); - tokio::io::stdout().write_all(&data).await.unwrap(); + match message { + Ok(msg) => match msg.into_text() { + Ok(msg_text) => info!("{}", msg_text), + Err(err) => error!( + "Failed to parse message from realtime service: Error: {}", + err + ), + }, + Err(err) => error!( + "Failed to read message from realtime service: Error: {}", + err + ), + }; }) }; From 9a48ae9d62898c0ebc239c0f31aeb408679e5d40 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 1 Jun 2022 09:13:17 -0500 Subject: [PATCH 20/88] structured wal2json parser --- worker/walrus/Cargo.toml | 2 ++ .../2022-04-29-132611_initial/up.sql | 17 +++++++++ worker/walrus/src/main.rs | 10 ++++-- worker/walrus/src/wal2json.rs | 36 +++++++++++++++++++ 4 files changed, 62 insertions(+), 3 deletions(-) create mode 100644 worker/walrus/src/wal2json.rs diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index f60177f..ee49f11 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -13,3 +13,5 @@ serde = "1.0" uuid = { version = "1.0", features = ["serde"] } log = "0.4.17" env_logger = "0.9.0" +itertools = "0.10.3" +cached = "0.34.0" diff --git a/worker/walrus/migrations/2022-04-29-132611_initial/up.sql b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql index 693e3a7..667d54b 100644 --- a/worker/walrus/migrations/2022-04-29-132611_initial/up.sql +++ b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql @@ -44,6 +44,22 @@ do $$ end $$; +-- realtime.user_defined_filter +do $$ + begin + if not realtime.type_exists('realtime', 'user_defined_filter') then + + create type realtime.user_defined_filter as ( + column_name text, + op realtime.equality_op, + value text + ); + + end if; + end +$$; + + -- realtime.action do $$ begin @@ -249,6 +265,7 @@ begin end; $$; +drop type if exists realtime.wal_column cascade; -- realtime.wal_column do $$ diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index f73aed9..c70356e 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -15,6 +15,8 @@ use std::thread::sleep; use std::time; use uuid; +mod wal2json; + pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); fn run_migrations( @@ -121,9 +123,11 @@ fn run(args: &Args) -> Result<(), String> { for input_line in stdin_lines { match input_line { Ok(line) => { - let result_json_line = serde_json::from_str::(&line); - match result_json_line { - Ok(json_line) => { + let result_record = serde_json::from_str::(&line); + match result_record { + Ok(wal2json_record) => { + let json_line = serde_json::json!(wal2json_record); + let result_wal_rls_rows = sql::< Record<( diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs new file mode 100644 index 0000000..7fe546f --- /dev/null +++ b/worker/walrus/src/wal2json.rs @@ -0,0 +1,36 @@ +use serde::{Deserialize, Serialize}; +use std::*; + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Column { + pub name: String, + pub r#type: String, + pub typeoid: i32, + pub value: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct PrimaryKeyRef { + pub name: String, + pub r#type: String, + pub typeoid: i32, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub enum Action { + I, + U, + D, + T, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Record { + pub action: String, + pub schema: String, + pub table: String, + pub pk: Vec, + pub columns: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub identity: Option>, +} From a004790c23f94a6c8d8bf50f0fb136647d8a70d8 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 1 Jun 2022 12:56:18 -0500 Subject: [PATCH 21/88] towards cached rust impl --- worker/walrus/Cargo.toml | 1 + .../down.sql | 1 + .../2022-06-01-154923_helper_functions/up.sql | 120 ++++++++++ worker/walrus/src/main.rs | 205 +++++++++++++++++- worker/walrus/src/realtime_fmt.rs | 40 ++++ worker/walrus/src/wal2json.rs | 4 +- 6 files changed, 369 insertions(+), 2 deletions(-) create mode 100644 worker/walrus/migrations/2022-06-01-154923_helper_functions/down.sql create mode 100644 worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql create mode 100644 worker/walrus/src/realtime_fmt.rs diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index ee49f11..2e96ee6 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -15,3 +15,4 @@ log = "0.4.17" env_logger = "0.9.0" itertools = "0.10.3" cached = "0.34.0" +chrono = { version = "0.4", features = ["serde"] } diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/down.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/down.sql new file mode 100644 index 0000000..291a97c --- /dev/null +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/down.sql @@ -0,0 +1 @@ +-- This file should undo anything in `up.sql` \ No newline at end of file diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql new file mode 100644 index 0000000..f336be1 --- /dev/null +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -0,0 +1,120 @@ +create function realtime.is_in_publication( + schema_name text, + table_name text, + publication_name text +) + returns bool + language sql +as $$ + select + exists( + select + 1 + from + pg_publication pp + left join pg_publication_tables ppt + on pp.pubname = ppt.pubname + where + pp.pubname = publication_name + and ppt.schemaname = schema_name + and ppt.tablename = table_name + limit 1 + ) +$$; + +create function realtime.is_subscribed_to( + schema_name text, + table_name text +) + returns bool + language sql +as $$ + select + exists( + select + 1 + from + realtime.subscription s + where + s.entity = format('%I.%I', schema_name, table_name)::regclass + limit 1 + ) +$$; + +create function realtime.is_rls_enabled(schema_name text, table_name text) + returns bool + language sql +as $$ + select + relrowsecurity + from + pg_class + where + oid = format('%I.%I', schema_name, table_name)::regclass + limit 1; +$$; + +create function realtime.subscribed_roles( + schema_name text, + table_name text +) + returns text[] + language sql +as $$ + select + coalesce(array_agg(distinct claims_role), '{}') + from + realtime.subscription s + where + s.entity = format('%I.%I', schema_name, table_name)::regclass + limit 1 +$$; + +create function realtime.selectable_columns( + schema_name text, + table_name text, + role_name text +) + returns text[] + language sql +as $$ + select + coalesce( + array_agg( + pa.attname::text + ) filter ( + where pg_catalog.has_column_privilege( + $3, + format('%I.%I', $1, $2)::regclass, + pa.attname, + 'SELECT' + ) + ), + '{}' + ) + from + pg_class e + join pg_attribute pa + on e.oid = pa.attrelid + where + e.oid = format('%I.%I', $1, $2)::regclass + and pa.attnum > 0 + and not pa.attisdropped; +$$; + + +create function realtime.get_subscription_ids( + schema_name text, + table_name text, +) + returns uuid[] + language sql +as $$ + select + coalesce(array_agg(distinct claims_role), '{}') + from + realtime.subscription s + where + s.entity = format('%I.%I', schema_name, table_name)::regclass + limit 1 +$$; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index c70356e..ba1fe6e 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -1,3 +1,5 @@ +use cached::proc_macro::cached; +use cached::TimedSizedCache; use clap::Parser; use diesel::dsl::sql; use diesel::sql_types::*; @@ -7,14 +9,15 @@ use env_logger; use log::{error, info, warn}; use serde::Serialize; use serde_json; +use std::collections::HashMap; use std::error::Error; use std::io; use std::io::BufRead; use std::process::{Command, Stdio}; use std::thread::sleep; use std::time; -use uuid; +mod realtime_fmt; mod wal2json; pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); @@ -84,6 +87,7 @@ fn run(args: &Args) -> Result<(), String> { return Err("failed to make postgres connection".to_string()); } }; + // Run pending migrations run_migrations(conn).expect("Pending migrations failed to execute"); info!("Postgres connection established"); @@ -123,9 +127,14 @@ fn run(args: &Args) -> Result<(), String> { for input_line in stdin_lines { match input_line { Ok(line) => { + println!("{}", line); let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { + // New + let walrus = process_record(&wal2json_record, conn); + + // Old let json_line = serde_json::json!(wal2json_record); let result_wal_rls_rows = @@ -189,3 +198,197 @@ fn run(args: &Args) -> Result<(), String> { } } } + +fn process_record( + rec: &wal2json::Record, + conn: &mut PgConnection, +) -> Result { + let is_in_publication = is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; + let is_subscribed_to = is_subscribed_to(&rec.schema, &rec.table, conn)?; + let is_rls_enabled = is_rls_enabled(&rec.schema, &rec.table, conn)?; + let subscribed_roles = subscribed_roles(&rec.schema, &rec.table, conn)?; + let pkey_cols: Vec<&String> = (&rec).pk.iter().map(|x| &x.name).collect(); + + println!("Published {}", is_in_publication); + println!("Subscribed {}", is_subscribed_to); + println!("Secured {}", is_rls_enabled); + println!("Subscribed Roles {}", subscribed_roles.join(", ")); + + let action = match rec.action { + wal2json::Action::I => realtime_fmt::Action::INSERT, + wal2json::Action::U => realtime_fmt::Action::UPDATE, + wal2json::Action::D => realtime_fmt::Action::DELETE, + wal2json::Action::T => realtime_fmt::Action::TRUNCATE, + }; + + if action != realtime_fmt::Action::DELETE && pkey_cols.len() == 0 { + return Ok(realtime_fmt::WALRLS { + wal: realtime_fmt::Data { + schema: rec.schema.to_string(), + table: rec.table.to_string(), + r#type: action, + commit_timestamp: rec.timestamp.to_string(), + columns: vec![], + record: HashMap::new(), + old_record: None, + }, + is_rls_enabled, + subscription_ids: get_subscription_ids(&rec.schema, &rec.table, conn)?, + errors: vec!["Error 400: Bad Request, no primary key".to_string()], + }); + } + + for role in &subscribed_roles { + let selectable_columns = selectable_columns(&rec.schema, &rec.table, role, conn)?; + println!("Selectable Columns {}", selectable_columns.join(", ")); + + //let column_data = rec.columns.into_iter().filter(|x| selectable_columns.contains(&x.name)).collect() + + // For user in subscribed users + // // check column permissions + } + + Ok(realtime_fmt::WALRLS { + wal: realtime_fmt::Data { + schema: rec.schema.to_string(), + table: rec.table.to_string(), + r#type: action, + commit_timestamp: rec.timestamp.to_string(), + columns: vec![], + record: HashMap::new(), + old_record: None, + }, + is_rls_enabled, + subscription_ids: vec![], + errors: vec![], + }) +} + +pub mod sql_functions { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; + } + + sql_function! { + fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; + } + + sql_function! { + fn is_subscribed_to(schema_name: Text, table_name: Text) -> Bool; + } + + sql_function! { + fn subscribed_roles(schema_name: Text, table_name: Text) -> Array; + } + + sql_function! { + fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; + } + + sql_function! { + fn get_subscription_ids(schema_name: Text, table_name: Text) -> Array; + } +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +fn is_rls_enabled( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql_functions::is_rls_enabled(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}-{}", schema_name, table_name, publication_name) }"#, + sync_writes = true +)] +fn is_in_publication( + schema_name: &str, + table_name: &str, + publication_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql_functions::is_in_publication( + schema_name, + table_name, + publication_name, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +fn is_subscribed_to( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql_functions::is_subscribed_to(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "TimedSizedCache, String>>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +fn subscribed_roles( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::subscribed_roles(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "TimedSizedCache, String>>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, + sync_writes = true +)] +fn selectable_columns( + schema_name: &str, + table_name: &str, + role_name: &str, + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::selectable_columns( + schema_name, + table_name, + role_name, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +fn get_subscription_ids( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::get_subscription_ids(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs new file mode 100644 index 0000000..f84ce3a --- /dev/null +++ b/worker/walrus/src/realtime_fmt.rs @@ -0,0 +1,40 @@ +use chrono; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::*; +use uuid; + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +pub enum Action { + INSERT, + UPDATE, + DELETE, + TRUNCATE, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Column { + pub name: String, + pub r#type: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Data { + pub schema: String, + pub table: String, + pub r#type: Action, + // TODO + pub commit_timestamp: String, //chrono::DateTime, + pub columns: Vec, + pub record: HashMap, + #[serde(skip_serializing_if = "Option::is_none")] + pub old_record: Option>, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct WALRLS { + pub wal: Data, + pub is_rls_enabled: bool, + pub subscription_ids: Vec, + pub errors: Vec, +} diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index 7fe546f..d73c7a2 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -1,3 +1,4 @@ +use chrono; use serde::{Deserialize, Serialize}; use std::*; @@ -26,11 +27,12 @@ pub enum Action { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Record { - pub action: String, + pub action: Action, pub schema: String, pub table: String, pub pk: Vec, pub columns: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub identity: Option>, + pub timestamp: String, //chrono::DateTime, } From c180a33ccb29b7f5ee18461c5d9545efa88da7b1 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 1 Jun 2022 17:09:44 -0500 Subject: [PATCH 22/88] partial move. all except rls and filters --- .../2022-06-01-154923_helper_functions/up.sql | 24 ++- worker/walrus/src/main.rs | 166 ++++++++++++++---- 2 files changed, 151 insertions(+), 39 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index f336be1..fdafb7d 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -82,6 +82,7 @@ as $$ coalesce( array_agg( pa.attname::text + order by pa.attnum asc ) filter ( where pg_catalog.has_column_privilege( $3, @@ -99,22 +100,41 @@ as $$ where e.oid = format('%I.%I', $1, $2)::regclass and pa.attnum > 0 - and not pa.attisdropped; + and not pa.attisdropped + $$; create function realtime.get_subscription_ids( + schema_name text, + table_name text +) + returns uuid[] + language sql +as $$ + select + coalesce(array_agg(subscription_id), '{}') + from + realtime.subscription s + where + s.entity = format('%I.%I', schema_name, table_name)::regclass + limit 1 +$$; + +create function realtime.get_subscription_ids_by_role( schema_name text, table_name text, + role_name text ) returns uuid[] language sql as $$ select - coalesce(array_agg(distinct claims_role), '{}') + coalesce(array_agg(subscription_id), '{}') from realtime.subscription s where s.entity = format('%I.%I', schema_name, table_name)::regclass + and claims_role = role_name::regrole limit 1 $$; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index ba1fe6e..7f4777a 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -127,13 +127,13 @@ fn run(args: &Args) -> Result<(), String> { for input_line in stdin_lines { match input_line { Ok(line) => { - println!("{}", line); let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { // New - let walrus = process_record(&wal2json_record, conn); + let walrus = process_record(&wal2json_record, 1024 * 1024, conn); + /* // Old let json_line = serde_json::json!(wal2json_record); @@ -158,16 +158,12 @@ fn run(args: &Args) -> Result<(), String> { )>( conn ); - match result_wal_rls_rows { + */ + //match result_wal_rls_rows { + match walrus { Ok(rows) => { for row in rows { - let walrus_rec = WalrusRecord { - wal: row.0, - is_rls_enabled: row.1, - subscription_ids: row.2, - errors: row.3, - }; - match serde_json::to_string(&walrus_rec) { + match serde_json::to_string(&row) { Ok(walrus_json) => println!("{}", walrus_json), Err(err) => { error!( @@ -201,19 +197,29 @@ fn run(args: &Args) -> Result<(), String> { fn process_record( rec: &wal2json::Record, + max_record_bytes: usize, conn: &mut PgConnection, -) -> Result { +) -> Result, String> { let is_in_publication = is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; let is_subscribed_to = is_subscribed_to(&rec.schema, &rec.table, conn)?; let is_rls_enabled = is_rls_enabled(&rec.schema, &rec.table, conn)?; let subscribed_roles = subscribed_roles(&rec.schema, &rec.table, conn)?; - let pkey_cols: Vec<&String> = (&rec).pk.iter().map(|x| &x.name).collect(); - println!("Published {}", is_in_publication); - println!("Subscribed {}", is_subscribed_to); - println!("Secured {}", is_rls_enabled); - println!("Subscribed Roles {}", subscribed_roles.join(", ")); + let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; + + //println!("Published {}", is_in_publication); + //println!("Subscribed {}", is_subscribed_to); + //println!("Secured {}", is_rls_enabled); + //println!("Subscribed Roles {}", subscribed_roles.join(", ")); + + let mut result = vec![]; + // If the table isn't in the publication or no one is listening, return + if !(is_in_publication & is_subscribed_to) { + return Ok(vec![]); + } + + let pkey_cols: Vec<&String> = (&rec).pk.iter().map(|x| &x.name).collect(); let action = match rec.action { wal2json::Action::I => realtime_fmt::Action::INSERT, wal2json::Action::U => realtime_fmt::Action::UPDATE, @@ -221,12 +227,13 @@ fn process_record( wal2json::Action::T => realtime_fmt::Action::TRUNCATE, }; + // If the table has no primary key, return if action != realtime_fmt::Action::DELETE && pkey_cols.len() == 0 { - return Ok(realtime_fmt::WALRLS { + let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { schema: rec.schema.to_string(), table: rec.table.to_string(), - r#type: action, + r#type: action.clone(), commit_timestamp: rec.timestamp.to_string(), columns: vec![], record: HashMap::new(), @@ -235,33 +242,98 @@ fn process_record( is_rls_enabled, subscription_ids: get_subscription_ids(&rec.schema, &rec.table, conn)?, errors: vec!["Error 400: Bad Request, no primary key".to_string()], - }); + }; + result.push(r); + return Ok(result); } for role in &subscribed_roles { let selectable_columns = selectable_columns(&rec.schema, &rec.table, role, conn)?; - println!("Selectable Columns {}", selectable_columns.join(", ")); - //let column_data = rec.columns.into_iter().filter(|x| selectable_columns.contains(&x.name)).collect() + //println!("Selectable Columns {}", selectable_columns.join(", ")); + + let mut record_elem = HashMap::new(); + let mut old_record_elem = None; + let mut old_record_elem_content = HashMap::new(); + + // If the role select any columns in the table, return + if action != realtime_fmt::Action::DELETE && selectable_columns.len() == 0 { + let r = realtime_fmt::WALRLS { + wal: realtime_fmt::Data { + schema: rec.schema.to_string(), + table: rec.table.to_string(), + r#type: action.clone(), + commit_timestamp: rec.timestamp.to_string(), + columns: vec![], + record: HashMap::new(), + old_record: None, + }, + is_rls_enabled, + subscription_ids: get_subscription_ids_by_role( + &rec.schema, + &rec.table, + &role, + conn, + )?, + errors: vec!["Error 401: Unauthorized".to_string()], + }; + result.push(r); + } else { + if vec![realtime_fmt::Action::INSERT, realtime_fmt::Action::UPDATE].contains(&action) { + for col_name in &selectable_columns { + 'record: for col in &rec.columns { + if col_name == &col.name { + if !exceeds_max_size || col.value.to_string().len() < 64 { + record_elem.insert(col_name.to_string(), col.value.clone()); + break 'record; + } + } + } + } + } - // For user in subscribed users - // // check column permissions + if vec![realtime_fmt::Action::UPDATE, realtime_fmt::Action::DELETE].contains(&action) { + for col_name in &selectable_columns { + match &rec.identity { + Some(identity) => { + 'old_record: for col in identity { + if col_name == &col.name { + if !exceeds_max_size || col.value.to_string().len() < 64 { + old_record_elem_content + .insert(col_name.to_string(), col.value.clone()); + break 'old_record; + } + } + } + } + None => (), + } + } + old_record_elem = Some(old_record_elem_content); + } + } + + // TODO CHECK RLS + // TODO FILTERS + + let r = realtime_fmt::WALRLS { + wal: realtime_fmt::Data { + schema: rec.schema.to_string(), + table: rec.table.to_string(), + r#type: action.clone(), + commit_timestamp: rec.timestamp.to_string(), + columns: vec![], + record: record_elem, + old_record: old_record_elem, + }, + is_rls_enabled, + subscription_ids: get_subscription_ids_by_role(&rec.schema, &rec.table, &role, conn)?, + errors: vec!["Error 401: Unauthorized".to_string()], + }; + result.push(r); } - Ok(realtime_fmt::WALRLS { - wal: realtime_fmt::Data { - schema: rec.schema.to_string(), - table: rec.table.to_string(), - r#type: action, - commit_timestamp: rec.timestamp.to_string(), - columns: vec![], - record: HashMap::new(), - old_record: None, - }, - is_rls_enabled, - subscription_ids: vec![], - errors: vec![], - }) + Ok(result) } pub mod sql_functions { @@ -291,6 +363,10 @@ pub mod sql_functions { sql_function! { fn get_subscription_ids(schema_name: Text, table_name: Text) -> Array; } + + sql_function! { + fn get_subscription_ids_by_role(schema_name: Text, table_name: Text, role_name: Text) -> Array; + } } #[cached( @@ -304,6 +380,7 @@ fn is_rls_enabled( table_name: &str, conn: &mut PgConnection, ) -> Result { + println!("IN IS_RLS_ENABLED"); select(sql_functions::is_rls_enabled(schema_name, table_name)) .first(conn) .map_err(|x| format!("{}", x)) @@ -392,3 +469,18 @@ fn get_subscription_ids( .first(conn) .map_err(|x| format!("{}", x)) } + +fn get_subscription_ids_by_role( + schema_name: &str, + table_name: &str, + role_name: &str, + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::get_subscription_ids_by_role( + schema_name, + table_name, + role_name, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} From 317e83bc2e0bf7c49f2c021fa36e1664b2975344 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 2 Jun 2022 13:46:09 -0500 Subject: [PATCH 23/88] load subscriptions --- .../2022-06-01-154923_helper_functions/up.sql | 35 ++++- worker/walrus/src/main.rs | 137 +++++++++++------- worker/walrus/src/realtime_fmt.rs | 35 ++++- worker/walrus/src/wal2json.rs | 6 +- 4 files changed, 150 insertions(+), 63 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index fdafb7d..f205f97 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -83,15 +83,8 @@ as $$ array_agg( pa.attname::text order by pa.attnum asc - ) filter ( - where pg_catalog.has_column_privilege( - $3, - format('%I.%I', $1, $2)::regclass, - pa.attname, - 'SELECT' - ) ), - '{}' + array['abc'] ) from pg_class e @@ -138,3 +131,29 @@ as $$ and claims_role = role_name::regrole limit 1 $$; + + +create function realtime.get_subscriptions( + schema_name text, + table_name text +) + returns jsonb[] + language sql +as $$ + select + coalesce( + array_agg( + jsonb_build_object( + 'subscription_id', subscription_id, + 'filters', filters, + 'claims_role', claims_role + ) + ), + '{}' + ) + from + realtime.subscription s + where + s.entity = format('%I.%I', schema_name, table_name)::regclass + limit 1 +$$; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 7f4777a..36f9341 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -11,8 +11,7 @@ use serde::Serialize; use serde_json; use std::collections::HashMap; use std::error::Error; -use std::io; -use std::io::BufRead; +use std::io::{self, BufRead, Write}; use std::process::{Command, Stdio}; use std::thread::sleep; use std::time; @@ -133,33 +132,6 @@ fn run(args: &Args) -> Result<(), String> { // New let walrus = process_record(&wal2json_record, 1024 * 1024, conn); - /* - // Old - let json_line = serde_json::json!(wal2json_record); - - let result_wal_rls_rows = - sql::< - Record<( - sql_types::Jsonb, - sql_types::Bool, - sql_types::Array, - sql_types::Array, - )>, - >( - "SELECT x from realtime.apply_rls(" - ) - .bind::(json_line) - .sql(") x") - .get_results::<( - serde_json::Value, - bool, - Vec, - Vec, - )>( - conn - ); - */ - //match result_wal_rls_rows { match walrus { Ok(rows) => { for row in rows { @@ -203,6 +175,9 @@ fn process_record( let is_in_publication = is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; let is_subscribed_to = is_subscribed_to(&rec.schema, &rec.table, conn)?; let is_rls_enabled = is_rls_enabled(&rec.schema, &rec.table, conn)?; + + let subscriptions = get_subscriptions(&rec.schema, &rec.table, conn)?; + let subscribed_roles = subscribed_roles(&rec.schema, &rec.table, conn)?; let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; @@ -212,7 +187,7 @@ fn process_record( //println!("Secured {}", is_rls_enabled); //println!("Subscribed Roles {}", subscribed_roles.join(", ")); - let mut result = vec![]; + let mut result: Vec = vec![]; // If the table isn't in the publication or no one is listening, return if !(is_in_publication & is_subscribed_to) { @@ -250,7 +225,16 @@ fn process_record( for role in &subscribed_roles { let selectable_columns = selectable_columns(&rec.schema, &rec.table, role, conn)?; - //println!("Selectable Columns {}", selectable_columns.join(", ")); + let mut columns = vec![]; + + for col in &rec.columns { + if selectable_columns.contains(&col.name) { + columns.push(realtime_fmt::Column { + name: col.name.to_string(), + type_: col.type_.to_string(), + }) + } + } let mut record_elem = HashMap::new(); let mut old_record_elem = None; @@ -264,7 +248,7 @@ fn process_record( table: rec.table.to_string(), r#type: action.clone(), commit_timestamp: rec.timestamp.to_string(), - columns: vec![], + columns, record: HashMap::new(), old_record: None, }, @@ -311,26 +295,35 @@ fn process_record( } old_record_elem = Some(old_record_elem_content); } - } - // TODO CHECK RLS - // TODO FILTERS + // TODO FILTERS - let r = realtime_fmt::WALRLS { - wal: realtime_fmt::Data { - schema: rec.schema.to_string(), - table: rec.table.to_string(), - r#type: action.clone(), - commit_timestamp: rec.timestamp.to_string(), - columns: vec![], - record: record_elem, - old_record: old_record_elem, - }, - is_rls_enabled, - subscription_ids: get_subscription_ids_by_role(&rec.schema, &rec.table, &role, conn)?, - errors: vec!["Error 401: Unauthorized".to_string()], - }; - result.push(r); + // TODO CHECK RLS + + let r = realtime_fmt::WALRLS { + wal: realtime_fmt::Data { + schema: rec.schema.to_string(), + table: rec.table.to_string(), + r#type: action.clone(), + commit_timestamp: rec.timestamp.to_string(), + columns, + record: record_elem, + old_record: old_record_elem, + }, + is_rls_enabled, + subscription_ids: get_subscription_ids_by_role( + &rec.schema, + &rec.table, + &role, + conn, + )?, + errors: match exceeds_max_size { + true => vec!["Error 413: Payload Too Large".to_string()], + false => vec![], + }, + }; + result.push(r); + } } Ok(result) @@ -367,6 +360,10 @@ pub mod sql_functions { sql_function! { fn get_subscription_ids_by_role(schema_name: Text, table_name: Text, role_name: Text) -> Array; } + + sql_function! { + fn get_subscriptions(schema_name: Text, table_name: Text) -> Array; + } } #[cached( @@ -380,7 +377,6 @@ fn is_rls_enabled( table_name: &str, conn: &mut PgConnection, ) -> Result { - println!("IN IS_RLS_ENABLED"); select(sql_functions::is_rls_enabled(schema_name, table_name)) .first(conn) .map_err(|x| format!("{}", x)) @@ -460,6 +456,12 @@ fn selectable_columns( .map_err(|x| format!("{}", x)) } +#[cached( + type = "TimedSizedCache, String>>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] fn get_subscription_ids( schema_name: &str, table_name: &str, @@ -470,6 +472,12 @@ fn get_subscription_ids( .map_err(|x| format!("{}", x)) } +#[cached( + type = "TimedSizedCache, String>>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, + sync_writes = true +)] fn get_subscription_ids_by_role( schema_name: &str, table_name: &str, @@ -484,3 +492,30 @@ fn get_subscription_ids_by_role( .first(conn) .map_err(|x| format!("{}", x)) } + +#[cached( + type = "TimedSizedCache, String>>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +fn get_subscriptions( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result, String> { + let subs: Vec = + select(sql_functions::get_subscriptions(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x))?; + + let mut res = vec![]; + + for sub_json in subs { + let sub: realtime_fmt::Subscription = + serde_json::from_value(sub_json).map_err(|x| format!("{}", x))?; + res.push(sub); + } + + Ok(res) +} diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index f84ce3a..c571924 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -15,7 +15,8 @@ pub enum Action { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Column { pub name: String, - pub r#type: String, + #[serde(alias = "type")] + pub type_: String, } #[derive(Serialize, Deserialize, Clone, Debug)] @@ -27,7 +28,6 @@ pub struct Data { pub commit_timestamp: String, //chrono::DateTime, pub columns: Vec, pub record: HashMap, - #[serde(skip_serializing_if = "Option::is_none")] pub old_record: Option>, } @@ -38,3 +38,34 @@ pub struct WALRLS { pub subscription_ids: Vec, pub errors: Vec, } + +// Subscriptions +#[derive(Serialize, Deserialize, Clone, Debug)] +pub enum Op { + #[serde(alias = "eq")] + Equal, + #[serde(alias = "neq")] + NotEqual, + #[serde(alias = "lt")] + LessThan, + #[serde(alias = "lte")] + LessThanOrEqual, + #[serde(alias = "gt")] + GreaterThan, + #[serde(alias = "gte")] + GreaterThanOrEqual, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct UserDefinedFiltern { + pub column_name: String, + pub op: Op, + pub value: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Subscription { + pub subscription_id: uuid::Uuid, + pub filters: Vec, + pub claims_role: String, +} diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index d73c7a2..525816f 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -5,7 +5,8 @@ use std::*; #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Column { pub name: String, - pub r#type: String, + #[serde(alias = "type")] + pub type_: String, pub typeoid: i32, pub value: serde_json::Value, } @@ -13,7 +14,8 @@ pub struct Column { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct PrimaryKeyRef { pub name: String, - pub r#type: String, + #[serde(alias = "type")] + pub type_: String, pub typeoid: i32, } From cb39f2a5b5a0218c463457e05fd8fcad7f3423a4 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 2 Jun 2022 15:59:58 -0500 Subject: [PATCH 24/88] manage subscriptions using the wal stream --- .../2022-06-01-154923_helper_functions/up.sql | 93 ++++--- worker/walrus/src/main.rs | 244 +++++++++--------- worker/walrus/src/realtime_fmt.rs | 8 +- 3 files changed, 173 insertions(+), 172 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index f205f97..72342ca 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -22,25 +22,6 @@ as $$ ) $$; -create function realtime.is_subscribed_to( - schema_name text, - table_name text -) - returns bool - language sql -as $$ - select - exists( - select - 1 - from - realtime.subscription s - where - s.entity = format('%I.%I', schema_name, table_name)::regclass - limit 1 - ) -$$; - create function realtime.is_rls_enabled(schema_name text, table_name text) returns bool language sql @@ -98,44 +79,62 @@ as $$ $$; -create function realtime.get_subscription_ids( - schema_name text, - table_name text -) - returns uuid[] +create function realtime.to_table_name(regclass) + returns text language sql -as $$ + immutable +as +$$ + with x(maybe_quoted_name) as ( + select + coalesce(nullif(split_part($1::text, '.', 2), ''), $1::text) + ) select - coalesce(array_agg(subscription_id), '{}') + case + when x.maybe_quoted_name like '"%"' then substring( + x.maybe_quoted_name, + 2, + character_length(x.maybe_quoted_name)-2 + ) + else x.maybe_quoted_name + end from - realtime.subscription s - where - s.entity = format('%I.%I', schema_name, table_name)::regclass - limit 1 + x $$; -create function realtime.get_subscription_ids_by_role( - schema_name text, - table_name text, - role_name text -) - returns uuid[] +create function realtime.to_schema_name(regclass) + returns text language sql -as $$ + immutable +as +$$ + with x(maybe_quoted_name) as ( + select + coalesce( + nullif(split_part($1::text, '.', 1), ''), + ( + select relnamespace::regnamespace::text + from pg_class + where oid = $1 + limit 1 + ) + ) + ) select - coalesce(array_agg(subscription_id), '{}') + case + when maybe_quoted_name like '"%"' then substring( + maybe_quoted_name, + 1, + character_length(maybe_quoted_name)-2 + ) + else maybe_quoted_name + end from - realtime.subscription s - where - s.entity = format('%I.%I', schema_name, table_name)::regclass - and claims_role = role_name::regrole - limit 1 + x $$; create function realtime.get_subscriptions( - schema_name text, - table_name text ) returns jsonb[] language sql @@ -144,6 +143,8 @@ as $$ coalesce( array_agg( jsonb_build_object( + 'schema_name', realtime.to_schema_name(entity), + 'table_name', realtime.to_table_name(entity), 'subscription_id', subscription_id, 'filters', filters, 'claims_role', claims_role @@ -153,7 +154,5 @@ as $$ ) from realtime.subscription s - where - s.entity = format('%I.%I', schema_name, table_name)::regclass limit 1 $$; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 36f9341..cfea52c 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -1,11 +1,12 @@ use cached::proc_macro::cached; -use cached::TimedSizedCache; +use cached::{SizedCache, TimedSizedCache}; use clap::Parser; use diesel::dsl::sql; use diesel::sql_types::*; use diesel::*; use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use env_logger; +use itertools::Itertools; use log::{error, info, warn}; use serde::Serialize; use serde_json; @@ -122,6 +123,9 @@ fn run(args: &Args) -> Result<(), String> { let stdin_reader = io::BufReader::new(stdin); let stdin_lines = stdin_reader.lines(); + // Load initial snapshot of subscriptions + let mut subscriptions = get_subscriptions(conn)?; + // Iterate input data for input_line in stdin_lines { match input_line { @@ -130,7 +134,12 @@ fn run(args: &Args) -> Result<(), String> { match result_record { Ok(wal2json_record) => { // New - let walrus = process_record(&wal2json_record, 1024 * 1024, conn); + let walrus = process_record( + &wal2json_record, + &mut subscriptions, + 1024 * 1024, + conn, + ); match walrus { Ok(rows) => { @@ -169,18 +178,48 @@ fn run(args: &Args) -> Result<(), String> { fn process_record( rec: &wal2json::Record, + mut subscriptions: &mut Vec, max_record_bytes: usize, conn: &mut PgConnection, ) -> Result, String> { let is_in_publication = is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; - let is_subscribed_to = is_subscribed_to(&rec.schema, &rec.table, conn)?; + let is_subscribed_to = subscriptions.len() > 0; let is_rls_enabled = is_rls_enabled(&rec.schema, &rec.table, conn)?; + let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; - let subscriptions = get_subscriptions(&rec.schema, &rec.table, conn)?; + // If the record is a new subscription. Handle it and return + if rec.schema == "realtime" && rec.table == "subscription" { + //TODO manage the subscriptions vector from the WAL stream + match rec.action { + wal2json::Action::I => { + /* + + realtime_fmt::Subscription{ + schema_name: rec.schema.to_string(), + table_name: rec.table_name.to_string(), + subscription_id: + filters: + claims_role: + } + */ + } + wal2json::Action::U => { + panic!("subscriptions should not be updated"); + } + wal2json::Action::D => {} + wal2json::Action::T => { + subscriptions.clear(); + } + } - let subscribed_roles = subscribed_roles(&rec.schema, &rec.table, conn)?; + return Ok(vec![]); + } - let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; + let subscribed_roles: Vec<&String> = subscriptions + .iter() + .map(|x| &x.claims_role) + .unique() + .collect(); //println!("Published {}", is_in_publication); //println!("Subscribed {}", is_subscribed_to); @@ -215,16 +254,25 @@ fn process_record( old_record: None, }, is_rls_enabled, - subscription_ids: get_subscription_ids(&rec.schema, &rec.table, conn)?, + subscription_ids: subscriptions + .iter() + .map(|x| x.subscription_id.clone()) + .collect(), errors: vec!["Error 400: Bad Request, no primary key".to_string()], }; result.push(r); return Ok(result); } - for role in &subscribed_roles { + for role in subscribed_roles { let selectable_columns = selectable_columns(&rec.schema, &rec.table, role, conn)?; + let role_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions + .iter() + .filter(|x| &x.claims_role == role) + .map(|x| x) + .collect(); + let mut columns = vec![]; for col in &rec.columns { @@ -253,12 +301,10 @@ fn process_record( old_record: None, }, is_rls_enabled, - subscription_ids: get_subscription_ids_by_role( - &rec.schema, - &rec.table, - &role, - conn, - )?, + subscription_ids: role_subscriptions + .iter() + .map(|x| x.subscription_id.clone()) + .collect(), errors: vec!["Error 401: Unauthorized".to_string()], }; result.push(r); @@ -296,9 +342,27 @@ fn process_record( old_record_elem = Some(old_record_elem_content); } - // TODO FILTERS + // FILTERS + let mut delegate_to_sql = vec![]; + + let mut subscription_id_is_visible_through_filters = vec![]; + for sub in role_subscriptions { + match visible_through_filters(&sub.filters, &rec.columns) { + Ok(true) => { + subscription_id_is_visible_through_filters.push(sub.subscription_id) + } + Ok(false) => (), + // TODO: delegate to SQL when we can't handle the comparison in rust + Err(_) => { + delegate_to_sql.push(sub.subscription_id); + } + } + } // TODO CHECK RLS + if is_rls_enabled { + panic!("RLS tables not yet implemented"); + } let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { @@ -311,12 +375,9 @@ fn process_record( old_record: old_record_elem, }, is_rls_enabled, - subscription_ids: get_subscription_ids_by_role( - &rec.schema, - &rec.table, - &role, - conn, - )?, + // TODO should be the intersection of visible through filters and RLS (if + // applicable) + subscription_ids: subscription_id_is_visible_through_filters, errors: match exceeds_max_size { true => vec!["Error 413: Payload Too Large".to_string()], false => vec![], @@ -329,6 +390,41 @@ fn process_record( Ok(result) } +fn visible_through_filters( + filters: &Vec, + columns: &Vec, +) -> Result { + for filter in filters { + match columns + .iter() + .filter(|x| x.name == filter.column_name) + .next() + { + Some(column) => match column.type_.as_ref() { + "integer" | "bigint" | "varchar" | "uuid" => match filter.op { + realtime_fmt::Op::Equal => { + match filter.value.to_string() != column.value.to_string() { + true => { + return Ok(false); + } + false => (), + } + } + _ => return Err("Could not handle op. Delegate comparison to SQL".to_string()), + }, + _ => { + return Err(format!( + "Could not handle type {}. Delegate comparison to SQL", + column.type_ + )) + } + }, + None => return Err("Filtered on non-existent column".to_string()), + } + } + Ok(true) +} + pub mod sql_functions { use diesel::sql_types::*; use diesel::*; @@ -341,28 +437,12 @@ pub mod sql_functions { fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; } - sql_function! { - fn is_subscribed_to(schema_name: Text, table_name: Text) -> Bool; - } - - sql_function! { - fn subscribed_roles(schema_name: Text, table_name: Text) -> Array; - } - sql_function! { fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; } sql_function! { - fn get_subscription_ids(schema_name: Text, table_name: Text) -> Array; - } - - sql_function! { - fn get_subscription_ids_by_role(schema_name: Text, table_name: Text, role_name: Text) -> Array; - } - - sql_function! { - fn get_subscriptions(schema_name: Text, table_name: Text) -> Array; + fn get_subscriptions() -> Array; } } @@ -403,41 +483,9 @@ fn is_in_publication( .map_err(|x| format!("{}", x)) } -#[cached( - type = "TimedSizedCache>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, - sync_writes = true -)] -fn is_subscribed_to( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result { - select(sql_functions::is_subscribed_to(schema_name, table_name)) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -#[cached( - type = "TimedSizedCache, String>>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, - sync_writes = true -)] -fn subscribed_roles( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::subscribed_roles(schema_name, table_name)) - .first(conn) - .map_err(|x| format!("{}", x)) -} - #[cached( type = "TimedSizedCache, String>>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + create = "{ TimedSizedCache::with_size_and_lifespan(500, 1)}", convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, sync_writes = true )] @@ -456,58 +504,10 @@ fn selectable_columns( .map_err(|x| format!("{}", x)) } -#[cached( - type = "TimedSizedCache, String>>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, - sync_writes = true -)] -fn get_subscription_ids( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::get_subscription_ids(schema_name, table_name)) +fn get_subscriptions(conn: &mut PgConnection) -> Result, String> { + let subs: Vec = select(sql_functions::get_subscriptions()) .first(conn) - .map_err(|x| format!("{}", x)) -} - -#[cached( - type = "TimedSizedCache, String>>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, - sync_writes = true -)] -fn get_subscription_ids_by_role( - schema_name: &str, - table_name: &str, - role_name: &str, - conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::get_subscription_ids_by_role( - schema_name, - table_name, - role_name, - )) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -#[cached( - type = "TimedSizedCache, String>>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, - sync_writes = true -)] -fn get_subscriptions( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result, String> { - let subs: Vec = - select(sql_functions::get_subscriptions(schema_name, table_name)) - .first(conn) - .map_err(|x| format!("{}", x))?; + .map_err(|x| format!("{}", x))?; let mut res = vec![]; diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index c571924..73ead7b 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -57,15 +57,17 @@ pub enum Op { } #[derive(Serialize, Deserialize, Clone, Debug)] -pub struct UserDefinedFiltern { +pub struct UserDefinedFilter { pub column_name: String, pub op: Op, - pub value: serde_json::Value, + pub value: String, // Why did I make this a text field?, } #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Subscription { + pub schema_name: String, + pub table_name: String, pub subscription_id: uuid::Uuid, - pub filters: Vec, + pub filters: Vec, pub claims_role: String, } From a8bbbf816ce7b25001e39a79001184b42553bd5d Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 14 Jun 2022 08:38:41 -0500 Subject: [PATCH 25/88] hooks for replacement of pg_recvlogical --- worker/Cargo.toml | 1 + worker/stream_logical/Cargo.toml | 12 +++ worker/stream_logical/src/main.rs | 172 ++++++++++++++++++++++++++++++ worker/walrus/Cargo.toml | 2 +- 4 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 worker/stream_logical/Cargo.toml create mode 100644 worker/stream_logical/src/main.rs diff --git a/worker/Cargo.toml b/worker/Cargo.toml index cd4533e..d5db734 100644 --- a/worker/Cargo.toml +++ b/worker/Cargo.toml @@ -2,4 +2,5 @@ members = [ "walrus", "realtime", + "stream_logical", ] diff --git a/worker/stream_logical/Cargo.toml b/worker/stream_logical/Cargo.toml new file mode 100644 index 0000000..a05fa29 --- /dev/null +++ b/worker/stream_logical/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "stream_logical" +version = "0.1.0" +edition = "2021" + +[dependencies] +tokio = { version = "1", features = ["full"] } +serde = { version = "1" } +futures = "0.3" +bytes = "1.0" +tokio-postgres = { git = "https://github.com/MaterializeInc/rust-postgres" } +postgres-protocol = { git = "https://github.com/MaterializeInc/rust-postgres" } diff --git a/worker/stream_logical/src/main.rs b/worker/stream_logical/src/main.rs new file mode 100644 index 0000000..64137a3 --- /dev/null +++ b/worker/stream_logical/src/main.rs @@ -0,0 +1,172 @@ +use futures::StreamExt; +use postgres_protocol::message::backend::LogicalReplicationMessage::*; +use postgres_protocol::message::backend::ReplicationMessage::*; +use postgres_protocol::message::backend::{LogicalReplicationMessage, ReplicationMessage}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::time::Instant; +use tokio_postgres::replication::LogicalReplicationStream; +use tokio_postgres::types::PgLsn; +use tokio_postgres::NoTls; +use tokio_postgres::SimpleQueryMessage::Row; + +/// Describes a table in a PostgreSQL database. +#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] +pub struct PostgresTableDesc { + /// The OID of the table. + pub oid: u32, + /// The name of the schema that the table belongs to. + pub namespace: String, + /// The name of the table. + pub name: String, + /// The description of each column, in order. + pub columns: Vec, +} + +/// Describes a column in a [`PostgresTableDesc`]. +#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] +pub struct PostgresColumnDesc { + /// The name of the column. + pub name: String, + /// The OID of the column's type. + pub type_oid: u32, + /// The modifier for the column's type. + pub type_mod: i32, + /// True if the column lacks a `NOT NULL` constraint. + pub nullable: bool, + /// Whether the column is part of the table's primary key. + /// TODO The order of the columns in the primary key matters too. + pub primary_key: bool, +} + +struct PostgresTaskInfo { + /// Our cursor into the WAL + lsn: PgLsn, + source_tables: HashMap, +} + +#[tokio::main] +async fn main() { + let conninfo = "host=127.0.0.1 port=5501 user=postgres password=password replication=database"; + let (client, connection) = tokio_postgres::connect(conninfo, NoTls).await.unwrap(); + + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {}", e); + } + }); + + client + .simple_query("DROP TABLE IF EXISTS test_logical_replication") + .await + .unwrap(); + client + .simple_query("CREATE TABLE test_logical_replication(i int)") + .await + .unwrap(); + let res = client + .simple_query("SELECT 'test_logical_replication'::regclass::oid") + .await + .unwrap(); + + let rel_id: u32 = if let Row(row) = &res[0] { + row.get("oid").unwrap().parse().unwrap() + } else { + panic!("unexpeced query message"); + }; + + client + .simple_query("DROP PUBLICATION IF EXISTS test_pub") + .await + .unwrap(); + client + .simple_query("CREATE PUBLICATION test_pub FOR ALL TABLES") + .await + .unwrap(); + + let slot = "test_logical_slot"; + + let query = format!( + r#"CREATE_REPLICATION_SLOT {:?} TEMPORARY LOGICAL "pgoutput""#, + slot + ); + let slot_query = client.simple_query(&query).await.unwrap(); + let lsn = if let Row(row) = &slot_query[0] { + row.get("consistent_point").unwrap() + } else { + panic!("unexpeced query message"); + }; + + println!("lsn {}", lsn); + + // issue a query that will appear in the slot's stream since it happened after its creation + client + .simple_query("INSERT INTO test_logical_replication VALUES (42)") + .await + .unwrap(); + + let options = r#"("proto_version" '1', "publication_names" 'test_pub')"#; + let query = format!( + r#"START_REPLICATION SLOT {:?} LOGICAL {} {}"#, + slot, lsn, options + ); + let copy_stream = client + .copy_both_simple::(&query) + .await + .unwrap(); + + let stream = LogicalReplicationStream::new(copy_stream); + tokio::pin!(stream); + + let task = PostgresTaskInfo { + lsn: 0.into(), + source_tables: HashMap::new(), + }; + let mut last_keepalive = Instant::now(); + let mut inserts: Vec = vec![]; + let mut deletes: Vec = vec![]; + + // TODO(olirice) track table defs mapping + // TODO(olirice) send keepalive every 20 seconds + + loop { + let msg: Option< + Result, tokio_postgres::Error>, + > = stream.next().await; + + let msg_res = match msg { + Some(Ok(XLogData(xlog_data))) => match xlog_data.data() { + Begin(_) => { + println!("begin"); + } + Insert(insert) => { + println!("insert"); + } + Update(update) => { + println!("update"); + } + Delete(delete) => { + println!("delete"); + } + Commit(commit) => { + println!("commit"); + } + Relation(relation) => { + println!("relation"); + } + Origin(_) | Type(_) => {} + Truncate(truncate) => { + println!("truncate"); + } + _ => println!("unknown logical replication message type"), + }, + Some(Err(_)) => panic!("unexpected replication stream error"), + None => panic!("unexpected replication stream end"), + Some(Ok(PrimaryKeepAlive(_))) => { + println!("keep alive") + } + Some(Ok(_)) => (), + _ => println!("Unexpected replication message"), + }; + } +} diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index 2e96ee6..b92082b 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -14,5 +14,5 @@ uuid = { version = "1.0", features = ["serde"] } log = "0.4.17" env_logger = "0.9.0" itertools = "0.10.3" -cached = "0.34.0" +cached = "0.34.1" chrono = { version = "0.4", features = ["serde"] } From 6531cab26441d9c95ef7b1a3c36e4f7dce358728 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 14 Jun 2022 08:50:30 -0500 Subject: [PATCH 26/88] readme --- worker/stream_logical/README.md | 44 +++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 worker/stream_logical/README.md diff --git a/worker/stream_logical/README.md b/worker/stream_logical/README.md new file mode 100644 index 0000000..e27b3e3 --- /dev/null +++ b/worker/stream_logical/README.md @@ -0,0 +1,44 @@ +# stream_logical + +stream_logical is an option for listening to the logical replication stream from postgres without the existing dependencies on: +- pg_recvlogical (local dependency) +- wal2json (server dependency) + +The dependency on pg_recvlogical requires that postgres is installed on the same machine as the walrus server. The wal2json dependency is not installed by default on many postgres setups which reduces compatibility. + +This solution would enable walrus to run against unmodified postgres. + +### Current State + +- Connects to postgres (from parent dir docker-compose) +- Creates a publication and temporary replication slot +- Creates a small amount of WAL +- Receives the WAL messages and prints the message type + +### Known TODOS + +- A keepalive message needs to be sent on a schedule +- Tracking table state +- Convert logical replication stream to a useable format +- Serialize messages + + +### Key Commands / Notes + +```shell +cargo run --bin stream_logical +``` + +Output +```text +lsn 0/173C578 +keep alive +begin +relation +insert +commit +keep alive +keep alive +``` + + From 6e269499585cb1a346b80823284f72d328a884fc Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 17 Jun 2022 04:58:22 -0500 Subject: [PATCH 27/88] port regtype failover --- .../down.sql | 307 +++++++++++++++++ .../up.sql | 314 ++++++++++++++++++ 2 files changed, 621 insertions(+) create mode 100644 worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/down.sql create mode 100644 worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/up.sql diff --git a/worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/down.sql b/worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/down.sql new file mode 100644 index 0000000..d1b1586 --- /dev/null +++ b/worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/down.sql @@ -0,0 +1,307 @@ +create or replace function realtime.apply_rls(wal jsonb, max_record_bytes int = 1024 * 1024) + returns setof realtime.wal_rls + language plpgsql + volatile +as $$ +declare + -- Regclass of the table e.g. public.notes + entity_ regclass = (quote_ident(wal ->> 'schema') || '.' || quote_ident(wal ->> 'table'))::regclass; + + -- I, U, D, T: insert, update ... + action realtime.action = ( + case wal ->> 'action' + when 'I' then 'INSERT' + when 'U' then 'UPDATE' + when 'D' then 'DELETE' + else 'ERROR' + end + ); + + -- Is row level security enabled for the table + is_rls_enabled bool = relrowsecurity from pg_class where oid = entity_; + + subscriptions realtime.subscription[] = array_agg(subs) + from + realtime.subscription subs + where + subs.entity = entity_; + + -- Subscription vars + roles regrole[] = array_agg(distinct us.claims_role) + from + unnest(subscriptions) us; + + working_role regrole; + claimed_role regrole; + claims jsonb; + + subscription_id uuid; + subscription_has_access bool; + visible_to_subscription_ids uuid[] = '{}'; + + -- structured info for wal's columns + columns realtime.wal_column[]; + -- previous identity values for update/delete + old_columns realtime.wal_column[]; + + error_record_exceeds_max_size boolean = octet_length(wal::text) > max_record_bytes; + + -- Primary jsonb output for record + output jsonb; + +begin + perform set_config('role', null, true); + + columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'columns') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + old_columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + (x->>'typeoid')::regtype + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'identity') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + for working_role in select * from unnest(roles) loop + + -- Update `is_selectable` for columns and old_columns + columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(columns) c; + + old_columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(old_columns) c; + + if action <> 'DELETE' and count(1) = 0 from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + -- subscriptions is already filtered by entity + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 400: Bad Request, no primary key'] + )::realtime.wal_rls; + + -- The claims role does not have SELECT permission to the primary key of entity + elsif action <> 'DELETE' and sum(c.is_selectable::int) <> count(1) from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 401: Unauthorized'] + )::realtime.wal_rls; + + else + output = jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action, + 'commit_timestamp', to_char( + (wal ->> 'timestamp')::timestamptz, + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ), + 'columns', ( + select + jsonb_agg( + jsonb_build_object( + 'name', pa.attname, + 'type', pt.typname + ) + order by pa.attnum asc + ) + from + pg_attribute pa + join pg_type pt + on pa.atttypid = pt.oid + where + attrelid = entity_ + and attnum > 0 + and pg_catalog.has_column_privilege(working_role, entity_, pa.attname, 'SELECT') + ) + ) + -- Add "record" key for insert and update + || case + when action in ('INSERT', 'UPDATE') then + case + when error_record_exceeds_max_size then + jsonb_build_object( + 'record', + ( + select jsonb_object_agg((c).name, (c).value) + from unnest(columns) c + where (c).is_selectable and (octet_length((c).value::text) <= 64) + ) + ) + else + jsonb_build_object( + 'record', + (select jsonb_object_agg((c).name, (c).value) from unnest(columns) c where (c).is_selectable) + ) + end + else '{}'::jsonb + end + -- Add "old_record" key for update and delete + || case + when action in ('UPDATE', 'DELETE') then + case + when error_record_exceeds_max_size then + jsonb_build_object( + 'old_record', + ( + select jsonb_object_agg((c).name, (c).value) + from unnest(old_columns) c + where (c).is_selectable and (octet_length((c).value::text) <= 64) + ) + ) + else + jsonb_build_object( + 'old_record', + (select jsonb_object_agg((c).name, (c).value) from unnest(old_columns) c where (c).is_selectable) + ) + end + else '{}'::jsonb + end; + + -- Create the prepared statement + if is_rls_enabled and action <> 'DELETE' then + if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then + deallocate walrus_rls_stmt; + end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, columns); + end if; + + visible_to_subscription_ids = '{}'; + + for subscription_id, claims in ( + select + subs.subscription_id, + subs.claims + from + unnest(subscriptions) subs + where + subs.entity = entity_ + and subs.claims_role = working_role + and realtime.is_visible_through_filters(columns, subs.filters) + ) loop + + if not is_rls_enabled or action = 'DELETE' then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + else + -- Check if RLS allows the role to see the record + perform + set_config('role', working_role::text, true), + set_config('request.jwt.claims', claims::text, true); + + execute 'execute walrus_rls_stmt' into subscription_has_access; + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end if; + end loop; + + perform set_config('role', null, true); + + return next ( + output, + is_rls_enabled, + visible_to_subscription_ids, + case + when error_record_exceeds_max_size then array['Error 413: Payload Too Large'] + else '{}' + end + )::realtime.wal_rls; + + end if; + end loop; + + perform set_config('role', null, true); +end; +$$; + + +create or replace function realtime.is_visible_through_filters(columns realtime.wal_column[], filters realtime.user_defined_filter[]) + returns bool + language sql + immutable +as $$ +/* +Should the record be visible (true) or filtered out (false) after *filters* are applied +*/ + select + -- Default to allowed when no filters present + coalesce( + sum( + realtime.check_equality_op( + op:=f.op, + type_:=col.type_oid::regtype, + -- cast jsonb to text + val_1:=col.value #>> '{}', + val_2:=f.value + )::int + ) = count(1), + true + ) + from + unnest(filters) f + join unnest(columns) col + on f.column_name = col.name; +$$; + + diff --git a/worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/up.sql b/worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/up.sql new file mode 100644 index 0000000..7afab03 --- /dev/null +++ b/worker/walrus/migrations/2022-06-17-094511_regtype casting failover to name/up.sql @@ -0,0 +1,314 @@ +create or replace function realtime.is_visible_through_filters(columns realtime.wal_column[], filters realtime.user_defined_filter[]) + returns bool + language sql + immutable +as $$ +/* +Should the record be visible (true) or filtered out (false) after *filters* are applied +*/ + select + -- Default to allowed when no filters present + coalesce( + sum( + realtime.check_equality_op( + op:=f.op, + type_:=coalesce( + col.type_oid::regtype, -- null when wal2json version <= 2.4 + col.type_name::regtype + ), + -- cast jsonb to text + val_1:=col.value #>> '{}', + val_2:=f.value + )::int + ) = count(1), + true + ) + from + unnest(filters) f + join unnest(columns) col + on f.column_name = col.name; +$$; + + +create or replace function realtime.apply_rls(wal jsonb, max_record_bytes int = 1024 * 1024) + returns setof realtime.wal_rls + language plpgsql + volatile +as $$ +declare + -- Regclass of the table e.g. public.notes + entity_ regclass = (quote_ident(wal ->> 'schema') || '.' || quote_ident(wal ->> 'table'))::regclass; + + -- I, U, D, T: insert, update ... + action realtime.action = ( + case wal ->> 'action' + when 'I' then 'INSERT' + when 'U' then 'UPDATE' + when 'D' then 'DELETE' + else 'ERROR' + end + ); + + -- Is row level security enabled for the table + is_rls_enabled bool = relrowsecurity from pg_class where oid = entity_; + + subscriptions realtime.subscription[] = array_agg(subs) + from + realtime.subscription subs + where + subs.entity = entity_; + + -- Subscription vars + roles regrole[] = array_agg(distinct us.claims_role) + from + unnest(subscriptions) us; + + working_role regrole; + claimed_role regrole; + claims jsonb; + + subscription_id uuid; + subscription_has_access bool; + visible_to_subscription_ids uuid[] = '{}'; + + -- structured info for wal's columns + columns realtime.wal_column[]; + -- previous identity values for update/delete + old_columns realtime.wal_column[]; + + error_record_exceeds_max_size boolean = octet_length(wal::text) > max_record_bytes; + + -- Primary jsonb output for record + output jsonb; + +begin + perform set_config('role', null, true); + + columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + coalesce( + (x->>'typeoid')::regtype, -- null when wal2json version <= 2.4 + (x->>'type')::regtype + ) + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'columns') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + old_columns = + array_agg( + ( + x->>'name', + x->>'type', + x->>'typeoid', + realtime.cast( + (x->'value') #>> '{}', + coalesce( + (x->>'typeoid')::regtype, -- null when wal2json version <= 2.4 + (x->>'type')::regtype + ) + ), + (pks ->> 'name') is not null, + true + )::realtime.wal_column + ) + from + jsonb_array_elements(wal -> 'identity') x + left join jsonb_array_elements(wal -> 'pk') pks + on (x ->> 'name') = (pks ->> 'name'); + + for working_role in select * from unnest(roles) loop + + -- Update `is_selectable` for columns and old_columns + columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(columns) c; + + old_columns = + array_agg( + ( + c.name, + c.type_name, + c.type_oid, + c.value, + c.is_pkey, + pg_catalog.has_column_privilege(working_role, entity_, c.name, 'SELECT') + )::realtime.wal_column + ) + from + unnest(old_columns) c; + + if action <> 'DELETE' and count(1) = 0 from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + -- subscriptions is already filtered by entity + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 400: Bad Request, no primary key'] + )::realtime.wal_rls; + + -- The claims role does not have SELECT permission to the primary key of entity + elsif action <> 'DELETE' and sum(c.is_selectable::int) <> count(1) from unnest(columns) c where c.is_pkey then + return next ( + jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action + ), + is_rls_enabled, + (select array_agg(s.subscription_id) from unnest(subscriptions) as s where claims_role = working_role), + array['Error 401: Unauthorized'] + )::realtime.wal_rls; + + else + output = jsonb_build_object( + 'schema', wal ->> 'schema', + 'table', wal ->> 'table', + 'type', action, + 'commit_timestamp', to_char( + (wal ->> 'timestamp')::timestamptz, + 'YYYY-MM-DD"T"HH24:MI:SS"Z"' + ), + 'columns', ( + select + jsonb_agg( + jsonb_build_object( + 'name', pa.attname, + 'type', pt.typname + ) + order by pa.attnum asc + ) + from + pg_attribute pa + join pg_type pt + on pa.atttypid = pt.oid + where + attrelid = entity_ + and attnum > 0 + and pg_catalog.has_column_privilege(working_role, entity_, pa.attname, 'SELECT') + ) + ) + -- Add "record" key for insert and update + || case + when action in ('INSERT', 'UPDATE') then + case + when error_record_exceeds_max_size then + jsonb_build_object( + 'record', + ( + select jsonb_object_agg((c).name, (c).value) + from unnest(columns) c + where (c).is_selectable and (octet_length((c).value::text) <= 64) + ) + ) + else + jsonb_build_object( + 'record', + (select jsonb_object_agg((c).name, (c).value) from unnest(columns) c where (c).is_selectable) + ) + end + else '{}'::jsonb + end + -- Add "old_record" key for update and delete + || case + when action in ('UPDATE', 'DELETE') then + case + when error_record_exceeds_max_size then + jsonb_build_object( + 'old_record', + ( + select jsonb_object_agg((c).name, (c).value) + from unnest(old_columns) c + where (c).is_selectable and (octet_length((c).value::text) <= 64) + ) + ) + else + jsonb_build_object( + 'old_record', + (select jsonb_object_agg((c).name, (c).value) from unnest(old_columns) c where (c).is_selectable) + ) + end + else '{}'::jsonb + end; + + -- Create the prepared statement + if is_rls_enabled and action <> 'DELETE' then + if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then + deallocate walrus_rls_stmt; + end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, columns); + end if; + + visible_to_subscription_ids = '{}'; + + for subscription_id, claims in ( + select + subs.subscription_id, + subs.claims + from + unnest(subscriptions) subs + where + subs.entity = entity_ + and subs.claims_role = working_role + and realtime.is_visible_through_filters(columns, subs.filters) + ) loop + + if not is_rls_enabled or action = 'DELETE' then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + else + -- Check if RLS allows the role to see the record + perform + set_config('role', working_role::text, true), + set_config('request.jwt.claims', claims::text, true); + + execute 'execute walrus_rls_stmt' into subscription_has_access; + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end if; + end loop; + + perform set_config('role', null, true); + + return next ( + output, + is_rls_enabled, + visible_to_subscription_ids, + case + when error_record_exceeds_max_size then array['Error 413: Payload Too Large'] + else '{}' + end + )::realtime.wal_rls; + + end if; + end loop; + + perform set_config('role', null, true); +end; +$$; From 835cd51903f67d784232ee24dc273308a9c1ff18 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 20 Jun 2022 17:36:45 -0500 Subject: [PATCH 28/88] rls delegate works. filter delegate sql untested --- .../2022-06-01-154923_helper_functions/up.sql | 109 ++++++ worker/walrus/src/main.rs | 310 ++++++++---------- worker/walrus/src/migrations.rs | 20 ++ worker/walrus/src/sql_functions.rs | 144 ++++++++ worker/walrus/src/wal2json.rs | 2 +- worker/walrus/src/walrus_fmt.rs | 20 ++ 6 files changed, 425 insertions(+), 180 deletions(-) create mode 100644 worker/walrus/src/migrations.rs create mode 100644 worker/walrus/src/sql_functions.rs create mode 100644 worker/walrus/src/walrus_fmt.rs diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 72342ca..76d2e6d 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -156,3 +156,112 @@ as $$ realtime.subscription s limit 1 $$; + + +create function realtime.is_visible_through_filters( + columns jsonb, + filters jsonb +) + returns bool + language sql +as $$ + select + realtime.is_visible_through_filters( + columns := ( + select + array_agg( + ( + c ->> 'name', + c ->> 'type_name', + c ->> 'type_oid', + c ->> 'value', + c ->> 'is_pkey', + c ->> 'is_selectable' + )::realtime.wal_column + ) + from + jsonb_array_elements(columns) c + ), + filters := ( + select + array_agg( + ( + f ->> 'column_name', + f ->> 'op', + f ->> 'value' + )::realtime.user_defined_filter + ) + from + jsonb_array_elements(filters) f + ) + ) +$$; + + +create function realtime.is_visible_through_rls( + schema_name text, + table_name text, + columns jsonb, + subscription_ids uuid[] +) + returns uuid[] + language plpgsql +as $$ +declare + entity_ regclass = format('%I.%I', schema_name, table_name)::regclass; + cols realtime.wal_column[]; + visible_to_subscription_ids uuid[] = '{}'; + subscription_id uuid; + subscription_has_access bool; + claims jsonb; +begin + raise exception '%', columns; + cols = ( + select + array_agg( + ( + c ->> 'name', + c ->> 'type_name', + c ->> 'type_oid', + c -> 'value', + c ->> 'is_pkey', + c ->> 'is_selectable' + )::realtime.wal_column + ) + from + jsonb_array_elements(columns) c + ); + + -- Create the prepared statement + if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then + deallocate walrus_rls_stmt; + end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, cols); + + for subscription_id, claims in ( + select + subs.subscription_id, + subs.claims + from + realtime.subscription subs + where + subs.subscription_id = any(subscription_ids) + ) + loop + -- Check if RLS allows the role to see the record + perform + set_config('role', (claims ->> 'role')::text, true), + set_config('request.jwt.claims', claims::text, true); + + execute 'execute walrus_rls_stmt' into subscription_has_access; + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end loop; + + perform set_config('role', null, true); + + return visible_to_subscription_ids; +end; +$$ diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index cfea52c..630b476 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -1,49 +1,20 @@ -use cached::proc_macro::cached; -use cached::{SizedCache, TimedSizedCache}; use clap::Parser; -use diesel::dsl::sql; -use diesel::sql_types::*; use diesel::*; -use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use env_logger; use itertools::Itertools; use log::{error, info, warn}; -use serde::Serialize; use serde_json; use std::collections::HashMap; -use std::error::Error; use std::io::{self, BufRead, Write}; use std::process::{Command, Stdio}; use std::thread::sleep; use std::time; +mod migrations; mod realtime_fmt; +mod sql_functions; mod wal2json; - -pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); - -fn run_migrations( - connection: &mut PgConnection, -) -> Result<(), Box> { - sql_query("create schema if not exists realtime") - .execute(connection) - .expect("failed to create 'realtime' schema"); - - sql_query("set search_path='realtime'") - .execute(connection) - .expect("failed to set search path"); - - connection.run_pending_migrations(MIGRATIONS)?; - Ok(()) -} - -#[derive(Serialize)] -pub struct WalrusRecord { - wal: serde_json::Value, - is_rls_enabled: bool, - subscription_ids: Vec, - errors: Vec, -} +mod walrus_fmt; /// Write-Ahead-Log Realtime Unified Security (WALRUS) background worker /// runs next to a PostgreSQL instance and forwards its Write-Ahead-Log @@ -89,7 +60,7 @@ fn run(args: &Args) -> Result<(), String> { }; // Run pending migrations - run_migrations(conn).expect("Pending migrations failed to execute"); + migrations::run_migrations(conn).expect("Pending migrations failed to execute"); info!("Postgres connection established"); let cmd = Command::new("pg_recvlogical") @@ -124,7 +95,7 @@ fn run(args: &Args) -> Result<(), String> { let stdin_lines = stdin_reader.lines(); // Load initial snapshot of subscriptions - let mut subscriptions = get_subscriptions(conn)?; + let mut subscriptions = sql_functions::get_subscriptions(conn)?; // Iterate input data for input_line in stdin_lines { @@ -133,10 +104,13 @@ fn run(args: &Args) -> Result<(), String> { let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { + // Update subscriptions if needed + update_subscriptions(&wal2json_record, &mut subscriptions); + // New let walrus = process_record( &wal2json_record, - &mut subscriptions, + &subscriptions, 1024 * 1024, conn, ); @@ -176,17 +150,12 @@ fn run(args: &Args) -> Result<(), String> { } } -fn process_record( +/// Checks to see if the new record is a change to realtime.subscriptions +/// and updates the subscriptions variable if a change is detected +fn update_subscriptions( rec: &wal2json::Record, mut subscriptions: &mut Vec, - max_record_bytes: usize, - conn: &mut PgConnection, -) -> Result, String> { - let is_in_publication = is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; - let is_subscribed_to = subscriptions.len() > 0; - let is_rls_enabled = is_rls_enabled(&rec.schema, &rec.table, conn)?; - let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; - +) -> () { // If the record is a new subscription. Handle it and return if rec.schema == "realtime" && rec.table == "subscription" { //TODO manage the subscriptions vector from the WAL stream @@ -211,9 +180,28 @@ fn process_record( subscriptions.clear(); } } - - return Ok(vec![]); } +} + +fn pkey_cols(rec: &wal2json::Record) -> Vec<&String> { + rec.pk.iter().map(|x| &x.name).collect() +} + +fn has_primary_key(rec: &wal2json::Record) -> bool { + pkey_cols(rec).len() != 0 +} + +fn process_record( + rec: &wal2json::Record, + subscriptions: &Vec, + max_record_bytes: usize, + conn: &mut PgConnection, +) -> Result, String> { + let is_in_publication = + sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; + let is_subscribed_to = subscriptions.len() > 0; + let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; + let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; let subscribed_roles: Vec<&String> = subscriptions .iter() @@ -221,19 +209,13 @@ fn process_record( .unique() .collect(); - //println!("Published {}", is_in_publication); - //println!("Subscribed {}", is_subscribed_to); - //println!("Secured {}", is_rls_enabled); - //println!("Subscribed Roles {}", subscribed_roles.join(", ")); - let mut result: Vec = vec![]; - // If the table isn't in the publication or no one is listening, return + // If the table isn't in the publication or no one is subscribed, do no work if !(is_in_publication & is_subscribed_to) { return Ok(vec![]); } - let pkey_cols: Vec<&String> = (&rec).pk.iter().map(|x| &x.name).collect(); let action = match rec.action { wal2json::Action::I => realtime_fmt::Action::INSERT, wal2json::Action::U => realtime_fmt::Action::UPDATE, @@ -242,7 +224,7 @@ fn process_record( }; // If the table has no primary key, return - if action != realtime_fmt::Action::DELETE && pkey_cols.len() == 0 { + if action != realtime_fmt::Action::DELETE && !has_primary_key(rec) { let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { schema: rec.schema.to_string(), @@ -265,7 +247,8 @@ fn process_record( } for role in subscribed_roles { - let selectable_columns = selectable_columns(&rec.schema, &rec.table, role, conn)?; + let selectable_columns = + sql_functions::selectable_columns(&rec.schema, &rec.table, role, conn)?; let role_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions .iter() @@ -273,22 +256,21 @@ fn process_record( .map(|x| x) .collect(); - let mut columns = vec![]; - - for col in &rec.columns { - if selectable_columns.contains(&col.name) { - columns.push(realtime_fmt::Column { - name: col.name.to_string(), - type_: col.type_.to_string(), - }) - } - } + let columns = rec + .columns + .iter() + .filter(|col| selectable_columns.contains(&col.name)) + .map(|w2j_col| realtime_fmt::Column { + name: w2j_col.name.to_string(), + type_: w2j_col.type_.to_string(), + }) + .collect(); let mut record_elem = HashMap::new(); let mut old_record_elem = None; let mut old_record_elem_content = HashMap::new(); - // If the role select any columns in the table, return + // If the role can not select any columns in the table, return if action != realtime_fmt::Action::DELETE && selectable_columns.len() == 0 { let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { @@ -342,27 +324,74 @@ fn process_record( old_record_elem = Some(old_record_elem_content); } - // FILTERS - let mut delegate_to_sql = vec![]; + let walcols: Vec = rec + .columns + .iter() + .map(|col| { + walrus_fmt::WALColumn { + name: col.name.to_string(), + type_name: col.type_.to_string(), + type_oid: col.typeoid.clone(), + value: col.value.clone(), + is_pkey: pkey_cols(rec).contains(&&col.name), + is_selectable: false, // stub: unused, + } + }) + .collect(); + // User Defined Filters let mut subscription_id_is_visible_through_filters = vec![]; + for sub in role_subscriptions { match visible_through_filters(&sub.filters, &rec.columns) { Ok(true) => { subscription_id_is_visible_through_filters.push(sub.subscription_id) } Ok(false) => (), - // TODO: delegate to SQL when we can't handle the comparison in rust + // delegate to SQL when we can't handle the comparison in rust Err(_) => { - delegate_to_sql.push(sub.subscription_id); + match sql_functions::is_visible_through_filters( + &walcols, + &sub.filters, + conn, + ) { + Ok(true) => { + subscription_id_is_visible_through_filters.push(sub.subscription_id) + } + Ok(false) => (), + Err(_) => { + panic!("error from sql during filter"); + } + }; } } } - // TODO CHECK RLS - if is_rls_enabled { - panic!("RLS tables not yet implemented"); - } + println!("past filters"); + + // Row Level Security + let subscription_ids_to_notify = match is_rls_enabled + && subscription_id_is_visible_through_filters.len() > 0 + && !vec![realtime_fmt::Action::DELETE, realtime_fmt::Action::TRUNCATE] + .contains(&action) + { + false => subscription_id_is_visible_through_filters, + true => { + match sql_functions::is_visible_through_rls( + &rec.schema, + &rec.table, + &walcols, + &subscription_id_is_visible_through_filters, + conn, + ) { + Ok(sub_ids) => sub_ids, + Err(err) => { + println!("error from sql during RLS {}", err); + panic!("rls"); + } + } + } + }; let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { @@ -375,9 +404,7 @@ fn process_record( old_record: old_record_elem, }, is_rls_enabled, - // TODO should be the intersection of visible through filters and RLS (if - // applicable) - subscription_ids: subscription_id_is_visible_through_filters, + subscription_ids: subscription_ids_to_notify, errors: match exceeds_max_size { true => vec!["Error 413: Payload Too Large".to_string()], false => vec![], @@ -390,26 +417,46 @@ fn process_record( Ok(result) } +fn is_null(v: &serde_json::Value) -> bool { + v == &serde_json::Value::Null +} + fn visible_through_filters( filters: &Vec, columns: &Vec, ) -> Result { + use crate::realtime_fmt::Op; + for filter in filters { + let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { + Ok(v) => v, + Err(err) => return Err(format!("{}", err)), + }; + match columns .iter() .filter(|x| x.name == filter.column_name) .next() { Some(column) => match column.type_.as_ref() { - "integer" | "bigint" | "varchar" | "uuid" => match filter.op { - realtime_fmt::Op::Equal => { - match filter.value.to_string() != column.value.to_string() { - true => { - return Ok(false); - } - false => (), + "integer" | "bigint" | "character varying" | "text" | "uuid" => match filter.op { + Op::Equal => { + if column.value != filter_value + || !is_null(&filter_value) + || !is_null(&column.value) + { + return Ok(false); + } + } + Op::NotEqual => { + if column.value == filter_value + || !is_null(&filter_value) + || !is_null(&column.value) + { + return Ok(false); } } + // TODO LT, LTE, GT, GTE _ => return Err("Could not handle op. Delegate comparison to SQL".to_string()), }, _ => { @@ -424,98 +471,3 @@ fn visible_through_filters( } Ok(true) } - -pub mod sql_functions { - use diesel::sql_types::*; - use diesel::*; - - sql_function! { - fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; - } - - sql_function! { - fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; - } - - sql_function! { - fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; - } - - sql_function! { - fn get_subscriptions() -> Array; - } -} - -#[cached( - type = "TimedSizedCache>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, - sync_writes = true -)] -fn is_rls_enabled( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result { - select(sql_functions::is_rls_enabled(schema_name, table_name)) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -#[cached( - type = "TimedSizedCache>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}-{}", schema_name, table_name, publication_name) }"#, - sync_writes = true -)] -fn is_in_publication( - schema_name: &str, - table_name: &str, - publication_name: &str, - conn: &mut PgConnection, -) -> Result { - select(sql_functions::is_in_publication( - schema_name, - table_name, - publication_name, - )) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -#[cached( - type = "TimedSizedCache, String>>", - create = "{ TimedSizedCache::with_size_and_lifespan(500, 1)}", - convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, - sync_writes = true -)] -fn selectable_columns( - schema_name: &str, - table_name: &str, - role_name: &str, - conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::selectable_columns( - schema_name, - table_name, - role_name, - )) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -fn get_subscriptions(conn: &mut PgConnection) -> Result, String> { - let subs: Vec = select(sql_functions::get_subscriptions()) - .first(conn) - .map_err(|x| format!("{}", x))?; - - let mut res = vec![]; - - for sub_json in subs { - let sub: realtime_fmt::Subscription = - serde_json::from_value(sub_json).map_err(|x| format!("{}", x))?; - res.push(sub); - } - - Ok(res) -} diff --git a/worker/walrus/src/migrations.rs b/worker/walrus/src/migrations.rs new file mode 100644 index 0000000..dbb7385 --- /dev/null +++ b/worker/walrus/src/migrations.rs @@ -0,0 +1,20 @@ +use diesel::*; +use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; +use std::error::Error; + +pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations"); + +pub fn run_migrations( + connection: &mut PgConnection, +) -> Result<(), Box> { + sql_query("create schema if not exists realtime") + .execute(connection) + .expect("failed to create 'realtime' schema"); + + sql_query("set search_path='realtime'") + .execute(connection) + .expect("failed to set search path"); + + connection.run_pending_migrations(MIGRATIONS)?; + Ok(()) +} diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs new file mode 100644 index 0000000..2735b37 --- /dev/null +++ b/worker/walrus/src/sql_functions.rs @@ -0,0 +1,144 @@ +use cached::proc_macro::cached; +use cached::{SizedCache, TimedSizedCache}; +use diesel::*; + +pub mod sql_functions { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; + } + + sql_function! { + fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; + } + + sql_function! { + fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; + } + + sql_function! { + fn get_subscriptions() -> Array; + } + + sql_function! { + fn is_visible_through_filters(columns: Jsonb, filters: Jsonb ) -> Bool + } + + sql_function! { + fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, subscription_ids: Array) -> Array + } +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +pub fn is_rls_enabled( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql_functions::is_rls_enabled(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}-{}", schema_name, table_name, publication_name) }"#, + sync_writes = true +)] +pub fn is_in_publication( + schema_name: &str, + table_name: &str, + publication_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql_functions::is_in_publication( + schema_name, + table_name, + publication_name, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "TimedSizedCache, String>>", + create = "{ TimedSizedCache::with_size_and_lifespan(500, 1)}", + convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, + sync_writes = true +)] +pub fn selectable_columns( + schema_name: &str, + table_name: &str, + role_name: &str, + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::selectable_columns( + schema_name, + table_name, + role_name, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +#[cached( + type = "SizedCache>", + create = "{ SizedCache::with_size(1500)}", + convert = r#"{ format!("{:?}-{:?}", columns, filters) }"#, + sync_writes = true +)] +pub fn is_visible_through_filters( + columns: &Vec, + filters: &Vec, + conn: &mut PgConnection, +) -> Result { + select(sql_functions::is_visible_through_filters( + serde_json::json!(columns), + serde_json::json!(filters), + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +pub fn is_visible_through_rls( + schema_name: &str, + table_name: &str, + columns: &Vec, + subscription_ids: &Vec, + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::is_visible_through_rls( + schema_name, + table_name, + serde_json::to_value(columns).unwrap(), + subscription_ids, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +pub fn get_subscriptions( + conn: &mut PgConnection, +) -> Result, String> { + let subs: Vec = select(sql_functions::get_subscriptions()) + .first(conn) + .map_err(|x| format!("{}", x))?; + + let mut res = vec![]; + + for sub_json in subs { + let sub: crate::realtime_fmt::Subscription = + serde_json::from_value(sub_json).map_err(|x| format!("{}", x))?; + res.push(sub); + } + + Ok(res) +} diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index 525816f..1314011 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -7,7 +7,7 @@ pub struct Column { pub name: String, #[serde(alias = "type")] pub type_: String, - pub typeoid: i32, + pub typeoid: Option, pub value: serde_json::Value, } diff --git a/worker/walrus/src/walrus_fmt.rs b/worker/walrus/src/walrus_fmt.rs new file mode 100644 index 0000000..0de3b20 --- /dev/null +++ b/worker/walrus/src/walrus_fmt.rs @@ -0,0 +1,20 @@ +use chrono; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize)] +pub struct WalrusRecord { + wal: serde_json::Value, + is_rls_enabled: bool, + subscription_ids: Vec, + errors: Vec, +} + +#[derive(Serialize, Debug)] +pub struct WALColumn { + pub name: String, + pub type_name: String, + pub type_oid: Option, + pub value: serde_json::Value, + pub is_pkey: bool, + pub is_selectable: bool, +} From f27c56bc9b2cdd34e0b60429fcc4e83889f3560a Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 21 Jun 2022 09:25:44 -0500 Subject: [PATCH 29/88] reasonably complete local filtering --- worker/walrus/src/filters.rs | 120 ++++++++++++++++++++++++++++++ worker/walrus/src/main.rs | 89 ++++++---------------- worker/walrus/src/realtime_fmt.rs | 2 +- worker/walrus/src/walrus_fmt.rs | 2 +- 4 files changed, 144 insertions(+), 69 deletions(-) create mode 100644 worker/walrus/src/filters.rs diff --git a/worker/walrus/src/filters.rs b/worker/walrus/src/filters.rs new file mode 100644 index 0000000..fb1434c --- /dev/null +++ b/worker/walrus/src/filters.rs @@ -0,0 +1,120 @@ +use crate::realtime_fmt::UserDefinedFilter; +use crate::wal2json::Column; +use log::warn; + +fn is_null(v: &serde_json::Value) -> bool { + v == &serde_json::Value::Null +} + +pub fn visible_through_filters( + filters: &Vec, + columns: &Vec, +) -> Result { + use crate::realtime_fmt::Op; + + for filter in filters { + let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { + Ok(v) => v, + Err(err) => return Err(format!("{}", err)), + }; + + let column = match columns + .iter() + .filter(|x| x.name == filter.column_name) + .next() + { + Some(col) => col, + // The filter references a column that does not exist + None => { + warn!( + "Attempted to filter non-existing column {}", + filter.column_name + ); + return Ok(false); + } + }; + + // Null column or filter values always result in a false because return is alway null + // until `IS` op is implemented + if is_null(&filter_value) || is_null(&column.value) { + return Ok(false); + }; + + match column.type_.as_ref() { + // All operations supported + "boolean" | "smallint" | "integer" | "bigint" | "serial" | "bigserial" | "numeric" + | "double precision" | "character" | "character varying" | "text" => match &filter.op { + // List is currently exhastive but that may be possible if types/ops expand + Op::Equal + | Op::NotEqual + | Op::LessThan + | Op::LessThanOrEqual + | Op::GreaterThan + | Op::GreaterThanOrEqual => { + let valid_ops = get_valid_ops(&column.value, &filter_value)?; + if valid_ops.contains(&filter.op) { + return Ok(false); + } + } + }, + // Only Equality ops supported + "uuid" => match &filter.op { + Op::Equal | Op::NotEqual => { + let valid_ops = get_valid_ops(&column.value, &filter_value)?; + if valid_ops.contains(&filter.op) { + return Ok(false); + }; + } + _ => return Err("could not handle filter op for allowed types".to_string()), + }, + _ => return Err("Could not handle type. Delegate comparison to SQL".to_string()), + }; + } + Ok(true) +} + +/// Returns a vector of realtime_fmt::Op that match for a OP b +fn get_valid_ops( + a: &serde_json::Value, + b: &serde_json::Value, +) -> Result, String> { + use serde_json::Value; + + match (a, b) { + (Value::Null, Value::Null) => Ok(vec![]), + (Value::Bool(a_), Value::Bool(b_)) => Ok(get_matching_ops(a_, b_)), + (Value::Number(a_), Value::Number(b_)) => Ok(get_matching_ops(&a_.as_f64(), &b_.as_f64())), + (Value::String(a_), Value::String(b_)) => Ok(get_matching_ops(a_, b_)), + // Array possible + // Object possible + _ => return Err("non-scalar or mismatched json value types".to_string()), + } +} + +fn get_matching_ops(a: &T, b: &T) -> Vec +where + T: PartialEq + PartialOrd, +{ + use crate::realtime_fmt::Op; + let mut ops = vec![]; + + if a == b { + ops.push(Op::Equal); + }; + if a != b { + ops.push(Op::NotEqual); + }; + if a < b { + ops.push(Op::LessThan); + }; + if a <= b { + ops.push(Op::LessThanOrEqual); + }; + if a > b { + ops.push(Op::GreaterThan); + }; + if a >= b { + ops.push(Op::GreaterThanOrEqual); + }; + ops +} diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 630b476..18b4b49 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -5,11 +5,12 @@ use itertools::Itertools; use log::{error, info, warn}; use serde_json; use std::collections::HashMap; -use std::io::{self, BufRead, Write}; +use std::io::{self, BufRead}; use std::process::{Command, Stdio}; use std::thread::sleep; use std::time; +mod filters; mod migrations; mod realtime_fmt; mod sql_functions; @@ -203,8 +204,18 @@ fn process_record( let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; - let subscribed_roles: Vec<&String> = subscriptions + // Subscriptions to the current entity + let entity_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions .iter() + .filter(|x| &x.schema_name == &rec.schema) + .filter(|x| &x.table_name == &rec.table) + .map(|x| x) + .collect(); + + let subscribed_roles: Vec<&String> = entity_subscriptions + .iter() + .filter(|x| &x.schema_name == &rec.schema) + .filter(|x| &x.table_name == &rec.table) .map(|x| &x.claims_role) .unique() .collect(); @@ -247,15 +258,16 @@ fn process_record( } for role in subscribed_roles { - let selectable_columns = - sql_functions::selectable_columns(&rec.schema, &rec.table, role, conn)?; - - let role_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions + // Subscriptions to current entity + role + let entity_role_subscriptions: Vec<&realtime_fmt::Subscription> = entity_subscriptions .iter() .filter(|x| &x.claims_role == role) - .map(|x| x) + .map(|x| *x) .collect(); + let selectable_columns = + sql_functions::selectable_columns(&rec.schema, &rec.table, role, conn)?; + let columns = rec .columns .iter() @@ -283,7 +295,7 @@ fn process_record( old_record: None, }, is_rls_enabled, - subscription_ids: role_subscriptions + subscription_ids: entity_role_subscriptions .iter() .map(|x| x.subscription_id.clone()) .collect(), @@ -342,8 +354,8 @@ fn process_record( // User Defined Filters let mut subscription_id_is_visible_through_filters = vec![]; - for sub in role_subscriptions { - match visible_through_filters(&sub.filters, &rec.columns) { + for sub in entity_role_subscriptions { + match filters::visible_through_filters(&sub.filters, &rec.columns) { Ok(true) => { subscription_id_is_visible_through_filters.push(sub.subscription_id) } @@ -367,8 +379,6 @@ fn process_record( } } - println!("past filters"); - // Row Level Security let subscription_ids_to_notify = match is_rls_enabled && subscription_id_is_visible_through_filters.len() > 0 @@ -416,58 +426,3 @@ fn process_record( Ok(result) } - -fn is_null(v: &serde_json::Value) -> bool { - v == &serde_json::Value::Null -} - -fn visible_through_filters( - filters: &Vec, - columns: &Vec, -) -> Result { - use crate::realtime_fmt::Op; - - for filter in filters { - let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { - Ok(v) => v, - Err(err) => return Err(format!("{}", err)), - }; - - match columns - .iter() - .filter(|x| x.name == filter.column_name) - .next() - { - Some(column) => match column.type_.as_ref() { - "integer" | "bigint" | "character varying" | "text" | "uuid" => match filter.op { - Op::Equal => { - if column.value != filter_value - || !is_null(&filter_value) - || !is_null(&column.value) - { - return Ok(false); - } - } - Op::NotEqual => { - if column.value == filter_value - || !is_null(&filter_value) - || !is_null(&column.value) - { - return Ok(false); - } - } - // TODO LT, LTE, GT, GTE - _ => return Err("Could not handle op. Delegate comparison to SQL".to_string()), - }, - _ => { - return Err(format!( - "Could not handle type {}. Delegate comparison to SQL", - column.type_ - )) - } - }, - None => return Err("Filtered on non-existent column".to_string()), - } - } - Ok(true) -} diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 73ead7b..1867706 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -40,7 +40,7 @@ pub struct WALRLS { } // Subscriptions -#[derive(Serialize, Deserialize, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub enum Op { #[serde(alias = "eq")] Equal, diff --git a/worker/walrus/src/walrus_fmt.rs b/worker/walrus/src/walrus_fmt.rs index 0de3b20..891b9c2 100644 --- a/worker/walrus/src/walrus_fmt.rs +++ b/worker/walrus/src/walrus_fmt.rs @@ -1,5 +1,5 @@ use chrono; -use serde::{Deserialize, Serialize}; +use serde::Serialize; #[derive(Serialize)] pub struct WalrusRecord { From 921096cb7e16e222afa0ff9ae0296bc83593a12d Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 21 Jun 2022 09:35:51 -0500 Subject: [PATCH 30/88] bugfix invert filters --- worker/walrus/src/filters.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/worker/walrus/src/filters.rs b/worker/walrus/src/filters.rs index fb1434c..ac8b1e3 100644 --- a/worker/walrus/src/filters.rs +++ b/worker/walrus/src/filters.rs @@ -52,7 +52,7 @@ pub fn visible_through_filters( | Op::GreaterThan | Op::GreaterThanOrEqual => { let valid_ops = get_valid_ops(&column.value, &filter_value)?; - if valid_ops.contains(&filter.op) { + if !valid_ops.contains(&filter.op) { return Ok(false); } } @@ -61,7 +61,7 @@ pub fn visible_through_filters( "uuid" => match &filter.op { Op::Equal | Op::NotEqual => { let valid_ops = get_valid_ops(&column.value, &filter_value)?; - if valid_ops.contains(&filter.op) { + if !valid_ops.contains(&filter.op) { return Ok(false); }; } From f117706b315befb09d7e51628ea5bde979459768 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 21 Jun 2022 16:03:18 -0500 Subject: [PATCH 31/88] parse subscriptions from wal stream --- .../2022-06-01-154923_helper_functions/up.sql | 5 +- worker/walrus/src/main.rs | 67 +++---- worker/walrus/src/realtime_fmt.rs | 179 ++++++++++++++++++ worker/walrus/src/wal2json.rs | 8 +- 4 files changed, 210 insertions(+), 49 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 76d2e6d..c72c61e 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -143,10 +143,11 @@ as $$ coalesce( array_agg( jsonb_build_object( - 'schema_name', realtime.to_schema_name(entity), + 'id', id, + 'schema_name', realtime.to_schema_name(entity), 'table_name', realtime.to_table_name(entity), 'subscription_id', subscription_id, - 'filters', filters, + 'filters', to_jsonb(filters), 'claims_role', claims_role ) ), diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 18b4b49..fe1f87f 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -75,7 +75,7 @@ fn run(args: &Args) -> Result<(), String> { "--option=include-timestamp=true", "--option=include-type-oids=true", "--option=format-version=2", - "--option=actions=insert,update,delete", + "--option=actions=insert,update,delete,truncate", &format!("--slot={}", args.slot), "--create-slot", "--if-not-exists", @@ -106,7 +106,10 @@ fn run(args: &Args) -> Result<(), String> { match result_record { Ok(wal2json_record) => { // Update subscriptions if needed - update_subscriptions(&wal2json_record, &mut subscriptions); + realtime_fmt::update_subscriptions( + &wal2json_record, + &mut subscriptions, + ); // New let walrus = process_record( @@ -151,41 +154,11 @@ fn run(args: &Args) -> Result<(), String> { } } -/// Checks to see if the new record is a change to realtime.subscriptions -/// and updates the subscriptions variable if a change is detected -fn update_subscriptions( - rec: &wal2json::Record, - mut subscriptions: &mut Vec, -) -> () { - // If the record is a new subscription. Handle it and return - if rec.schema == "realtime" && rec.table == "subscription" { - //TODO manage the subscriptions vector from the WAL stream - match rec.action { - wal2json::Action::I => { - /* - - realtime_fmt::Subscription{ - schema_name: rec.schema.to_string(), - table_name: rec.table_name.to_string(), - subscription_id: - filters: - claims_role: - } - */ - } - wal2json::Action::U => { - panic!("subscriptions should not be updated"); - } - wal2json::Action::D => {} - wal2json::Action::T => { - subscriptions.clear(); - } - } - } -} - fn pkey_cols(rec: &wal2json::Record) -> Vec<&String> { - rec.pk.iter().map(|x| &x.name).collect() + match &rec.pk { + Some(pkey_refs) => pkey_refs.iter().map(|x| &x.name).collect(), + None => vec![], + } } fn has_primary_key(rec: &wal2json::Record) -> bool { @@ -202,6 +175,7 @@ fn process_record( sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; let is_subscribed_to = subscriptions.len() > 0; let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; + let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; // Subscriptions to the current entity @@ -222,11 +196,6 @@ fn process_record( let mut result: Vec = vec![]; - // If the table isn't in the publication or no one is subscribed, do no work - if !(is_in_publication & is_subscribed_to) { - return Ok(vec![]); - } - let action = match rec.action { wal2json::Action::I => realtime_fmt::Action::INSERT, wal2json::Action::U => realtime_fmt::Action::UPDATE, @@ -234,6 +203,11 @@ fn process_record( wal2json::Action::T => realtime_fmt::Action::TRUNCATE, }; + // If the table isn't in the publication or no one is subscribed, do no work + if !(is_in_publication && is_subscribed_to && action != realtime_fmt::Action::TRUNCATE) { + return Ok(vec![]); + } + // If the table has no primary key, return if action != realtime_fmt::Action::DELETE && !has_primary_key(rec) { let r = realtime_fmt::WALRLS { @@ -270,6 +244,8 @@ fn process_record( let columns = rec .columns + .as_ref() + .unwrap_or(&vec![]) .iter() .filter(|col| selectable_columns.contains(&col.name)) .map(|w2j_col| realtime_fmt::Column { @@ -305,7 +281,7 @@ fn process_record( } else { if vec![realtime_fmt::Action::INSERT, realtime_fmt::Action::UPDATE].contains(&action) { for col_name in &selectable_columns { - 'record: for col in &rec.columns { + 'record: for col in rec.columns.as_ref().unwrap_or(&vec![]) { if col_name == &col.name { if !exceeds_max_size || col.value.to_string().len() < 64 { record_elem.insert(col_name.to_string(), col.value.clone()); @@ -338,6 +314,8 @@ fn process_record( let walcols: Vec = rec .columns + .as_ref() + .unwrap_or(&vec![]) .iter() .map(|col| { walrus_fmt::WALColumn { @@ -355,7 +333,10 @@ fn process_record( let mut subscription_id_is_visible_through_filters = vec![]; for sub in entity_role_subscriptions { - match filters::visible_through_filters(&sub.filters, &rec.columns) { + match filters::visible_through_filters( + &sub.filters, + rec.columns.as_ref().unwrap_or(&vec![]), + ) { Ok(true) => { subscription_id_is_visible_through_filters.push(sub.subscription_id) } diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 1867706..9ce6ac4 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -1,4 +1,7 @@ +use crate::wal2json; use chrono; +use log::error; +use log::info; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::*; @@ -65,9 +68,185 @@ pub struct UserDefinedFilter { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Subscription { + pub id: u64, pub schema_name: String, pub table_name: String, pub subscription_id: uuid::Uuid, pub filters: Vec, pub claims_role: String, } + +/// Checks to see if the new record is a change to realtime.subscriptions +/// and updates the subscriptions variable if a change is detected +pub fn update_subscriptions(rec: &wal2json::Record, subscriptions: &mut Vec) -> () { + fn get_column_value_by_name( + columns: &Vec, + column_name: &str, + ) -> Option { + columns + .iter() + .filter(|x| x.name == column_name) + .map(|x| x.value.clone()) + .next() + } + + enum NewOrOld { + New, + Old, + } + + fn wal2json_to_subscription( + rec: &wal2json::Record, + new_or_old: NewOrOld, + ) -> Result { + let columns = match new_or_old { + NewOrOld::New => match &rec.columns { + Some(cols) => cols.clone(), + None => { + return Err("Failed to parse wal2json record. No columns contents".to_string()); + } + }, + NewOrOld::Old => match &rec.identity { + Some(ident) => ident.clone(), + None => { + return Err("Failed to parse wal2json record. No identity contents".to_string()); + } + }, + }; + + let id: u64 = match get_column_value_by_name(&columns, "id") { + Some(id_json) => match id_json { + serde_json::Value::Number(id_num) => match id_num.as_u64() { + Some(id) => id, + None => { + return Err(format!( + "Invalid id in realtime.subscription. Expected u64: {}", + id_num, + )); + } + }, + _ => { + return Err(format!( + "Invalid id in realtime.subscription. Expected number, got: {}", + id_json + )); + } + }, + None => { + return Err("No id column found on realtime.subscription".to_string()); + } + }; + + let subscription_id = match get_column_value_by_name(&columns, "subscription_id") { + Some(uuid_json) => match uuid::Uuid::parse_str(&uuid_json.as_str().unwrap()) { + Ok(uuid) => uuid, + Err(err) => { + return Err(format!( + "Invalid subscription_id realtime.subscription: {}", + err + )); + } + }, + None => { + return Err("No subscription_id found on realtime.subscription".to_string()); + } + }; + + let filters: Vec = match get_column_value_by_name(&columns, "filters") { + Some(fs_json) => match serde_json::from_value(fs_json.clone()) { + Ok(udfs) => udfs, + Err(err) => { + return Err(format!("Invalid filters in realtime.subscription: {}", err)); + } + }, + None => { + return Err("No filters column found on realtime.subscription".to_string()); + } + }; + + let claims_role = match get_column_value_by_name(&columns, "claims_role") { + Some(role_json) => match role_json { + serde_json::Value::String(role_name) => role_name, + _ => { + return Err("Invalid claims_role in realtime.subscription".to_string()); + } + }, + None => { + return Err("No claims_role column found on realtime.subscription".to_string()); + } + } + .to_string(); + + Ok(Subscription { + id, + schema_name: rec.schema.to_string(), + table_name: rec.table.to_string(), + subscription_id, + filters, + claims_role, + }) + } + + // If the record is not a subscription, return + if rec.schema != "realtime" || rec.table != "subscription" { + return (); + } + //TODO manage the subscriptions vector from the WAL stream + match rec.action { + wal2json::Action::I => { + info!("SUBSCRIPTION: GOT INSERT"); + match wal2json_to_subscription(rec, NewOrOld::New) { + Ok(new_sub) => subscriptions.push(new_sub), + Err(err) => error!( + "Failed to parse wal2json record to realtime.subscription: {} ", + err + ), + }; + } + wal2json::Action::U => { + info!("SUBSCRIPTION: GOT UPDATE"); + // Delete old sub + let old_sub = match wal2json_to_subscription(rec, NewOrOld::New) { + Ok(old_sub) => old_sub, + Err(err) => { + error!( + "Failed to parse wal2json record to realtime.subscription update new: {} ", + err + ); + return (); + } + }; + subscriptions.retain_mut(|x| x.id != old_sub.id); + + // Add new sub + match wal2json_to_subscription(rec, NewOrOld::New) { + Ok(new_sub) => subscriptions.push(new_sub), + Err(err) => { + error!( + "Failed to parse wal2json record to realtime.subscription update new: {} ", + err + ); + return (); + } + }; + } + wal2json::Action::D => { + info!("SUBSCRIPTION: GOT DELETE"); + let old_sub = match wal2json_to_subscription(rec, NewOrOld::New) { + Ok(old_sub) => old_sub, + Err(err) => { + error!( + "Failed to parse wal2json record to realtime.subscription update new: {} ", + err + ); + return (); + } + }; + subscriptions.retain(|x| x.id != old_sub.id); + } + wal2json::Action::T => { + info!("SUBSCRIPTION: GOT TRUNCATE"); + subscriptions.clear(); + } + } +} diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index 1314011..7e9fdb0 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -32,9 +32,9 @@ pub struct Record { pub action: Action, pub schema: String, pub table: String, - pub pk: Vec, - pub columns: Vec, + pub pk: Option>, + pub columns: Option>, // option is for truncate #[serde(skip_serializing_if = "Option::is_none")] - pub identity: Option>, - pub timestamp: String, //chrono::DateTime, + pub identity: Option>, // option is for insert/update + pub timestamp: String, //chrono::DateTime, } From f1ec1a4aa88efab6a1b3aa9254a63890fa0979f5 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 21 Jun 2022 16:39:34 -0500 Subject: [PATCH 32/88] subscription manager --- .../2022-06-01-154923_helper_functions/up.sql | 19 ++ worker/walrus/src/main.rs | 1 + worker/walrus/src/realtime_fmt.rs | 199 +++++------------- worker/walrus/src/sql_functions.rs | 18 ++ worker/walrus/src/wal2json.rs | 2 +- 5 files changed, 93 insertions(+), 146 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index c72c61e..0467f5b 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -158,6 +158,25 @@ as $$ limit 1 $$; +create function realtime.get_subscription_by_id(id bigint) + returns jsonb + language sql +as $$ + select + jsonb_build_object( + 'id', id, + 'schema_name', realtime.to_schema_name(entity), + 'table_name', realtime.to_table_name(entity), + 'subscription_id', subscription_id, + 'filters', to_jsonb(filters), + 'claims_role', claims_role + ) + from + realtime.subscription s + where + s.id = $1 + limit 1 +$$; create function realtime.is_visible_through_filters( columns jsonb, diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index fe1f87f..d8bb6b9 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -109,6 +109,7 @@ fn run(args: &Args) -> Result<(), String> { realtime_fmt::update_subscriptions( &wal2json_record, &mut subscriptions, + conn, ); // New diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 9ce6ac4..ea55712 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -1,7 +1,7 @@ use crate::wal2json; use chrono; +use diesel::*; use log::error; -use log::info; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::*; @@ -68,7 +68,7 @@ pub struct UserDefinedFilter { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Subscription { - pub id: u64, + pub id: i64, pub schema_name: String, pub table_name: String, pub subscription_id: uuid::Uuid, @@ -78,124 +78,59 @@ pub struct Subscription { /// Checks to see if the new record is a change to realtime.subscriptions /// and updates the subscriptions variable if a change is detected -pub fn update_subscriptions(rec: &wal2json::Record, subscriptions: &mut Vec) -> () { - fn get_column_value_by_name( - columns: &Vec, - column_name: &str, - ) -> Option { - columns - .iter() - .filter(|x| x.name == column_name) - .map(|x| x.value.clone()) - .next() +pub fn update_subscriptions( + rec: &wal2json::Record, + subscriptions: &mut Vec, + conn: &mut PgConnection, +) -> () { + // If the record is not a subscription, return + if rec.schema != "realtime" || rec.table != "subscription" { + return (); } - enum NewOrOld { - New, - Old, + if rec.action == wal2json::Action::T { + subscriptions.clear(); + return (); } - fn wal2json_to_subscription( - rec: &wal2json::Record, - new_or_old: NewOrOld, - ) -> Result { - let columns = match new_or_old { - NewOrOld::New => match &rec.columns { - Some(cols) => cols.clone(), - None => { - return Err("Failed to parse wal2json record. No columns contents".to_string()); - } - }, - NewOrOld::Old => match &rec.identity { - Some(ident) => ident.clone(), + let id: i64 = match rec + .columns + .as_ref() + // Deletes have the id value in the identity field + .unwrap_or(rec.identity.as_ref().unwrap_or(&vec![])) + .iter() + .filter(|x| x.name == "id") + .map(|x| x.value.clone()) + .next() + { + Some(id_json) => match id_json { + serde_json::Value::Number(id_num) => match id_num.as_i64() { + Some(id) => id, None => { - return Err("Failed to parse wal2json record. No identity contents".to_string()); - } - }, - }; - - let id: u64 = match get_column_value_by_name(&columns, "id") { - Some(id_json) => match id_json { - serde_json::Value::Number(id_num) => match id_num.as_u64() { - Some(id) => id, - None => { - return Err(format!( - "Invalid id in realtime.subscription. Expected u64: {}", - id_num, - )); - } - }, - _ => { - return Err(format!( - "Invalid id in realtime.subscription. Expected number, got: {}", - id_json - )); - } - }, - None => { - return Err("No id column found on realtime.subscription".to_string()); - } - }; - - let subscription_id = match get_column_value_by_name(&columns, "subscription_id") { - Some(uuid_json) => match uuid::Uuid::parse_str(&uuid_json.as_str().unwrap()) { - Ok(uuid) => uuid, - Err(err) => { - return Err(format!( - "Invalid subscription_id realtime.subscription: {}", - err - )); - } - }, - None => { - return Err("No subscription_id found on realtime.subscription".to_string()); - } - }; - - let filters: Vec = match get_column_value_by_name(&columns, "filters") { - Some(fs_json) => match serde_json::from_value(fs_json.clone()) { - Ok(udfs) => udfs, - Err(err) => { - return Err(format!("Invalid filters in realtime.subscription: {}", err)); - } - }, - None => { - return Err("No filters column found on realtime.subscription".to_string()); - } - }; - - let claims_role = match get_column_value_by_name(&columns, "claims_role") { - Some(role_json) => match role_json { - serde_json::Value::String(role_name) => role_name, - _ => { - return Err("Invalid claims_role in realtime.subscription".to_string()); + error!( + "Invalid id in realtime.subscription. Expected i64, got: {}", + id_num + ); + return (); } }, - None => { - return Err("No claims_role column found on realtime.subscription".to_string()); + _ => { + error!( + "Invalid id in realtime.subscription. Expected number, got: {}", + id_json + ); + return (); } + }, + None => { + error!("No id column found on realtime.subscription"); + return (); } - .to_string(); + }; - Ok(Subscription { - id, - schema_name: rec.schema.to_string(), - table_name: rec.table.to_string(), - subscription_id, - filters, - claims_role, - }) - } - - // If the record is not a subscription, return - if rec.schema != "realtime" || rec.table != "subscription" { - return (); - } - //TODO manage the subscriptions vector from the WAL stream match rec.action { wal2json::Action::I => { - info!("SUBSCRIPTION: GOT INSERT"); - match wal2json_to_subscription(rec, NewOrOld::New) { + match crate::sql_functions::get_subscription_by_id(id, conn) { Ok(new_sub) => subscriptions.push(new_sub), Err(err) => error!( "Failed to parse wal2json record to realtime.subscription: {} ", @@ -204,49 +139,23 @@ pub fn update_subscriptions(rec: &wal2json::Record, subscriptions: &mut Vec { - info!("SUBSCRIPTION: GOT UPDATE"); - // Delete old sub - let old_sub = match wal2json_to_subscription(rec, NewOrOld::New) { - Ok(old_sub) => old_sub, - Err(err) => { - error!( - "Failed to parse wal2json record to realtime.subscription update new: {} ", - err - ); - return (); - } - }; - subscriptions.retain_mut(|x| x.id != old_sub.id); + // Delete existing sub + subscriptions.retain_mut(|x| x.id != id); - // Add new sub - match wal2json_to_subscription(rec, NewOrOld::New) { + // Add updated sub + match crate::sql_functions::get_subscription_by_id(id, conn) { Ok(new_sub) => subscriptions.push(new_sub), - Err(err) => { - error!( - "Failed to parse wal2json record to realtime.subscription update new: {} ", - err - ); - return (); - } + Err(err) => error!( + "Failed to parse wal2json record to realtime.subscription: {} ", + err + ), }; } wal2json::Action::D => { - info!("SUBSCRIPTION: GOT DELETE"); - let old_sub = match wal2json_to_subscription(rec, NewOrOld::New) { - Ok(old_sub) => old_sub, - Err(err) => { - error!( - "Failed to parse wal2json record to realtime.subscription update new: {} ", - err - ); - return (); - } - }; - subscriptions.retain(|x| x.id != old_sub.id); + subscriptions.retain(|x| x.id != id); } wal2json::Action::T => { - info!("SUBSCRIPTION: GOT TRUNCATE"); - subscriptions.clear(); + // Handled above } - } + }; } diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index 2735b37..c04ffdf 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -22,6 +22,10 @@ pub mod sql_functions { fn get_subscriptions() -> Array; } + sql_function! { + fn get_subscription_by_id(id: BigInt) -> Jsonb; + } + sql_function! { fn is_visible_through_filters(columns: Jsonb, filters: Jsonb ) -> Bool } @@ -142,3 +146,17 @@ pub fn get_subscriptions( Ok(res) } + +pub fn get_subscription_by_id( + id: i64, + conn: &mut PgConnection, +) -> Result { + let sub_value: serde_json::Value = select(sql_functions::get_subscription_by_id(id)) + .first(conn) + .map_err(|x| format!("{}", x))?; + + let sub: crate::realtime_fmt::Subscription = + serde_json::from_value(sub_value).map_err(|x| format!("{}", x))?; + + Ok(sub) +} diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index 7e9fdb0..90a4297 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -19,7 +19,7 @@ pub struct PrimaryKeyRef { pub typeoid: i32, } -#[derive(Serialize, Deserialize, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] pub enum Action { I, U, From c13d483f4f2cb0fe82d956afbbfc7ba9a1a9305b Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 10:05:28 -0500 Subject: [PATCH 33/88] batch filters delegated to sql --- .../2022-06-01-154923_helper_functions/up.sql | 83 +++++++++++-------- worker/walrus/src/main.rs | 46 +++++----- worker/walrus/src/sql_functions.rs | 15 ++-- 3 files changed, 81 insertions(+), 63 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 0467f5b..875db43 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -147,7 +147,7 @@ as $$ 'schema_name', realtime.to_schema_name(entity), 'table_name', realtime.to_table_name(entity), 'subscription_id', subscription_id, - 'filters', to_jsonb(filters), + 'filters', filters, 'claims_role', claims_role ) ), @@ -168,7 +168,7 @@ as $$ 'schema_name', realtime.to_schema_name(entity), 'table_name', realtime.to_table_name(entity), 'subscription_id', subscription_id, - 'filters', to_jsonb(filters), + 'filters', filters, 'claims_role', claims_role ) from @@ -180,41 +180,57 @@ $$; create function realtime.is_visible_through_filters( columns jsonb, - filters jsonb + subscription_ids uuid[] ) - returns bool - language sql + returns uuid[] + language plpgsql as $$ - select - realtime.is_visible_through_filters( - columns := ( - select - array_agg( - ( - c ->> 'name', - c ->> 'type_name', - c ->> 'type_oid', - c ->> 'value', - c ->> 'is_pkey', - c ->> 'is_selectable' - )::realtime.wal_column - ) - from - jsonb_array_elements(columns) c - ), - filters := ( - select - array_agg( - ( - f ->> 'column_name', - f ->> 'op', - f ->> 'value' - )::realtime.user_defined_filter - ) - from - jsonb_array_elements(filters) f +declare + cols realtime.wal_column[]; + visible_to_subscription_ids uuid[] = '{}'; + subscription_id uuid; + filters realtime.user_defined_filter[]; + subscription_has_access bool; +begin + cols = ( + select + array_agg( + ( + c ->> 'name', + c ->> 'type_name', + c ->> 'type_oid', + c -> 'value', + c ->> 'is_pkey', + c ->> 'is_selectable' + )::realtime.wal_column ) + from + jsonb_array_elements(columns) c + ); + + for subscription_id, filters in ( + select + subs.subscription_id, + subs.filters + from + realtime.subscription subs + where + subs.subscription_id = any(subscription_ids) ) + loop + + subscription_has_access = realtime.is_visible_through_filters( + columns := cols, + filters := filters + ); + + if subscription_has_access then + visible_to_subscription_ids = visible_to_subscription_ids || subscription_id; + end if; + end loop; + + return visible_to_subscription_ids; +end; $$; @@ -235,7 +251,6 @@ declare subscription_has_access bool; claims jsonb; begin - raise exception '%', columns; cols = ( select array_agg( diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index d8bb6b9..eb5df89 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -2,7 +2,7 @@ use clap::Parser; use diesel::*; use env_logger; use itertools::Itertools; -use log::{error, info, warn}; +use log::{debug, error, info, warn}; use serde_json; use std::collections::HashMap; use std::io::{self, BufRead}; @@ -35,7 +35,8 @@ fn main() { let args = Args::parse(); // enable logger - env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); + //env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("debug")).init(); loop { match run(&args) { @@ -96,7 +97,9 @@ fn run(args: &Args) -> Result<(), String> { let stdin_lines = stdin_reader.lines(); // Load initial snapshot of subscriptions + info!("Snapshot of subscriptions loading"); let mut subscriptions = sql_functions::get_subscriptions(conn)?; + info!("Snapshot of subscriptions loaded"); // Iterate input data for input_line in stdin_lines { @@ -332,6 +335,7 @@ fn process_record( // User Defined Filters let mut subscription_id_is_visible_through_filters = vec![]; + let mut subscription_id_delegate_to_sql = vec![]; for sub in entity_role_subscriptions { match filters::visible_through_filters( @@ -339,28 +343,32 @@ fn process_record( rec.columns.as_ref().unwrap_or(&vec![]), ) { Ok(true) => { - subscription_id_is_visible_through_filters.push(sub.subscription_id) + debug!("Filters handled in rust: {:?}", &sub.filters); + subscription_id_is_visible_through_filters.push(sub.subscription_id); } Ok(false) => (), // delegate to SQL when we can't handle the comparison in rust - Err(_) => { - match sql_functions::is_visible_through_filters( - &walcols, - &sub.filters, - conn, - ) { - Ok(true) => { - subscription_id_is_visible_through_filters.push(sub.subscription_id) - } - Ok(false) => (), - Err(_) => { - panic!("error from sql during filter"); - } - }; + Err(err) => { + debug!( + "Filters delegated to SQL: {:?}. Error: {}", + &sub.filters, err + ); + subscription_id_delegate_to_sql.push(sub.subscription_id); } } } + match sql_functions::is_visible_through_filters( + &walcols, + &subscription_id_delegate_to_sql, + conn, + ) { + Ok(sub_ids) => subscription_id_is_visible_through_filters.extend(&sub_ids), + Err(err) => { + error!("Failed to deletegate some filters to SQL: {}", err) + } + } + // Row Level Security let subscription_ids_to_notify = match is_rls_enabled && subscription_id_is_visible_through_filters.len() > 0 @@ -378,8 +386,8 @@ fn process_record( ) { Ok(sub_ids) => sub_ids, Err(err) => { - println!("error from sql during RLS {}", err); - panic!("rls"); + error!("Failed to delegate RLS to SQL: {}", err); + vec![] } } } diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index c04ffdf..9fe623e 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -27,7 +27,7 @@ pub mod sql_functions { } sql_function! { - fn is_visible_through_filters(columns: Jsonb, filters: Jsonb ) -> Bool + fn is_visible_through_filters(columns: Jsonb, subscription_ids: Array ) -> Array } sql_function! { @@ -93,20 +93,15 @@ pub fn selectable_columns( .map_err(|x| format!("{}", x)) } -#[cached( - type = "SizedCache>", - create = "{ SizedCache::with_size(1500)}", - convert = r#"{ format!("{:?}-{:?}", columns, filters) }"#, - sync_writes = true -)] pub fn is_visible_through_filters( columns: &Vec, - filters: &Vec, + subscription_ids: &Vec, + // TODO: convert this to use subscription_ids to reduce n calls conn: &mut PgConnection, -) -> Result { +) -> Result, String> { select(sql_functions::is_visible_through_filters( serde_json::json!(columns), - serde_json::json!(filters), + subscription_ids, )) .first(conn) .map_err(|x| format!("{}", x)) From 6feb6522c7feb7e9f2c413f1e12f6dbcfae41d3a Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 11:20:50 -0500 Subject: [PATCH 34/88] chrono for timestamps --- worker/walrus/src/main.rs | 20 ++++++++------- worker/walrus/src/realtime_fmt.rs | 6 ++--- worker/walrus/src/sql_functions.rs | 2 +- worker/walrus/src/timestamp_fmt.rs | 40 ++++++++++++++++++++++++++++++ worker/walrus/src/wal2json.rs | 6 +++-- worker/walrus/src/walrus_fmt.rs | 1 - 6 files changed, 59 insertions(+), 16 deletions(-) create mode 100644 worker/walrus/src/timestamp_fmt.rs diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index eb5df89..c77a6a3 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -14,6 +14,7 @@ mod filters; mod migrations; mod realtime_fmt; mod sql_functions; +mod timestamp_fmt; mod wal2json; mod walrus_fmt; @@ -36,7 +37,7 @@ fn main() { // enable logger //env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); - env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("debug")).init(); + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); loop { match run(&args) { @@ -105,6 +106,7 @@ fn run(args: &Args) -> Result<(), String> { for input_line in stdin_lines { match input_line { Ok(line) => { + println!("{}", line); let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { @@ -219,7 +221,7 @@ fn process_record( schema: rec.schema.to_string(), table: rec.table.to_string(), r#type: action.clone(), - commit_timestamp: rec.timestamp.to_string(), + commit_timestamp: rec.timestamp, columns: vec![], record: HashMap::new(), old_record: None, @@ -269,7 +271,7 @@ fn process_record( schema: rec.schema.to_string(), table: rec.table.to_string(), r#type: action.clone(), - commit_timestamp: rec.timestamp.to_string(), + commit_timestamp: rec.timestamp, columns, record: HashMap::new(), old_record: None, @@ -343,16 +345,16 @@ fn process_record( rec.columns.as_ref().unwrap_or(&vec![]), ) { Ok(true) => { - debug!("Filters handled in rust: {:?}", &sub.filters); + //debug!("Filters handled in rust: {:?}", &sub.filters); subscription_id_is_visible_through_filters.push(sub.subscription_id); } Ok(false) => (), // delegate to SQL when we can't handle the comparison in rust Err(err) => { - debug!( - "Filters delegated to SQL: {:?}. Error: {}", - &sub.filters, err - ); + //debug!( + // "Filters delegated to SQL: {:?}. Error: {}", + // &sub.filters, err + //); subscription_id_delegate_to_sql.push(sub.subscription_id); } } @@ -398,7 +400,7 @@ fn process_record( schema: rec.schema.to_string(), table: rec.table.to_string(), r#type: action.clone(), - commit_timestamp: rec.timestamp.to_string(), + commit_timestamp: rec.timestamp, columns, record: record_elem, old_record: old_record_elem, diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index ea55712..b7f56bd 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -1,5 +1,5 @@ use crate::wal2json; -use chrono; +use chrono::{DateTime, Utc}; use diesel::*; use log::error; use serde::{Deserialize, Serialize}; @@ -27,8 +27,8 @@ pub struct Data { pub schema: String, pub table: String, pub r#type: Action, - // TODO - pub commit_timestamp: String, //chrono::DateTime, + #[serde(with = "crate::timestamp_fmt")] + pub commit_timestamp: DateTime, pub columns: Vec, pub record: HashMap, pub old_record: Option>, diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index 9fe623e..49b62a6 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -1,5 +1,5 @@ use cached::proc_macro::cached; -use cached::{SizedCache, TimedSizedCache}; +use cached::TimedSizedCache; use diesel::*; pub mod sql_functions { diff --git a/worker/walrus/src/timestamp_fmt.rs b/worker/walrus/src/timestamp_fmt.rs new file mode 100644 index 0000000..b7b37f2 --- /dev/null +++ b/worker/walrus/src/timestamp_fmt.rs @@ -0,0 +1,40 @@ +use chrono::{DateTime, TimeZone, Utc}; +use serde::{self, Deserialize, Deserializer, Serializer}; + +// Example: 2022-06-22 15:38:19.695275+00 +// https://docs.rs/chrono/latest/chrono/format/strftime/index.html +const DESER_FORMAT: &'static str = "%Y-%m-%d %H:%M:%S%.f%#z"; + +// Example: 2000-01-01T00:01:01Z +const SER_FORMAT: &'static str = "%+"; + +// The signature of a serialize_with function must follow the pattern: +// +// fn serialize(&T, S) -> Result +// where +// S: Serializer +// +// although it may also be generic over the input types T. +pub fn serialize(date: &DateTime, serializer: S) -> Result +where + S: Serializer, +{ + let s = format!("{}", date.format(SER_FORMAT)); + serializer.serialize_str(&s) +} + +// The signature of a deserialize_with function must follow the pattern: +// +// fn deserialize<'de, D>(D) -> Result +// where +// D: Deserializer<'de> +// +// although it may also be generic over the output types T. +pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let s = String::deserialize(deserializer)?; + Utc.datetime_from_str(&s, DESER_FORMAT) + .map_err(serde::de::Error::custom) +} diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index 90a4297..9eafb98 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -1,4 +1,4 @@ -use chrono; +use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use std::*; @@ -36,5 +36,7 @@ pub struct Record { pub columns: Option>, // option is for truncate #[serde(skip_serializing_if = "Option::is_none")] pub identity: Option>, // option is for insert/update - pub timestamp: String, //chrono::DateTime, + // Example: 2022-06-22 15:38:19.695275+00 + #[serde(with = "crate::timestamp_fmt")] + pub timestamp: DateTime, } diff --git a/worker/walrus/src/walrus_fmt.rs b/worker/walrus/src/walrus_fmt.rs index 891b9c2..2a9feb0 100644 --- a/worker/walrus/src/walrus_fmt.rs +++ b/worker/walrus/src/walrus_fmt.rs @@ -1,4 +1,3 @@ -use chrono; use serde::Serialize; #[derive(Serialize)] From 8691e89d79136f5edba1cce3c787a6636f88864d Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 11:22:42 -0500 Subject: [PATCH 35/88] reduce heartbeat freq --- worker/realtime/src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 9bdef7f..42aca72 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -211,7 +211,7 @@ async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender) { loop { - sleep(Duration::from_secs(3)).await; + sleep(Duration::from_secs(20)).await; let phoenix_msg = PhoenixMessage { event: PhoenixMessageEvent::Heartbeat, payload: serde_json::json!({"msg": "ping"}), From ccfd3d3e0e560bc16b25ab170673d26f850546ca Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 11:38:45 -0500 Subject: [PATCH 36/88] rm wal2json output; --- worker/walrus/src/main.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index c77a6a3..078d428 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -106,7 +106,6 @@ fn run(args: &Args) -> Result<(), String> { for input_line in stdin_lines { match input_line { Ok(line) => { - println!("{}", line); let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { From 865016067d66c8a920eea37cd2f2b8d66531a5ff Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 11:45:21 -0500 Subject: [PATCH 37/88] sub name and realtime_fmt alias --- worker/walrus/src/main.rs | 2 +- worker/walrus/src/realtime_fmt.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 078d428..c7134ce 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -177,7 +177,7 @@ fn process_record( conn: &mut PgConnection, ) -> Result, String> { let is_in_publication = - sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_realtime", conn)?; + sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_multiplayer", conn)?; let is_subscribed_to = subscriptions.len() > 0; let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index b7f56bd..7abc173 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -18,7 +18,7 @@ pub enum Action { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Column { pub name: String, - #[serde(alias = "type")] + #[serde(rename(serialize = "type", deserialize = "type"))] pub type_: String, } From c520d897921e30a8e22b59ca6743d2f0e009d719 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 11:56:15 -0500 Subject: [PATCH 38/88] bugfix timestamp format --- worker/walrus/src/main.rs | 1 + worker/walrus/src/timestamp_fmt.rs | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index c7134ce..064a6bb 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -176,6 +176,7 @@ fn process_record( max_record_bytes: usize, conn: &mut PgConnection, ) -> Result, String> { + // TODO subscription name as argument let is_in_publication = sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_multiplayer", conn)?; let is_subscribed_to = subscriptions.len() > 0; diff --git a/worker/walrus/src/timestamp_fmt.rs b/worker/walrus/src/timestamp_fmt.rs index b7b37f2..47c3208 100644 --- a/worker/walrus/src/timestamp_fmt.rs +++ b/worker/walrus/src/timestamp_fmt.rs @@ -6,7 +6,7 @@ use serde::{self, Deserialize, Deserializer, Serializer}; const DESER_FORMAT: &'static str = "%Y-%m-%d %H:%M:%S%.f%#z"; // Example: 2000-01-01T00:01:01Z -const SER_FORMAT: &'static str = "%+"; +const SER_FORMAT: &'static str = "%Y-%m-%dT%H:%M:%S%.fZ"; // The signature of a serialize_with function must follow the pattern: // From 1cbb7f7e70e859326a57bcf4da321053c46f2b00 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 22 Jun 2022 16:33:33 -0500 Subject: [PATCH 39/88] skip roundtrip when no sql filter delegates --- worker/walrus/src/main.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 064a6bb..32f39e5 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -360,14 +360,16 @@ fn process_record( } } - match sql_functions::is_visible_through_filters( - &walcols, - &subscription_id_delegate_to_sql, - conn, - ) { - Ok(sub_ids) => subscription_id_is_visible_through_filters.extend(&sub_ids), - Err(err) => { - error!("Failed to deletegate some filters to SQL: {}", err) + if subscription_id_delegate_to_sql.len() > 0 { + match sql_functions::is_visible_through_filters( + &walcols, + &subscription_id_delegate_to_sql, + conn, + ) { + Ok(sub_ids) => subscription_id_is_visible_through_filters.extend(&sub_ids), + Err(err) => { + error!("Failed to deletegate some filters to SQL: {}", err) + } } } From cf16611a21c6f9fa5bf971419052631d5a3ed9ec Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 24 Jun 2022 09:50:49 -0500 Subject: [PATCH 40/88] diesel table defs --- worker/walrus/Cargo.toml | 2 +- .../2022-06-01-154923_helper_functions/up.sql | 24 +-- worker/walrus/src/main.rs | 26 ++- worker/walrus/src/realtime_fmt.rs | 150 +++++++++++++----- worker/walrus/src/schema.rs | 17 ++ 5 files changed, 161 insertions(+), 58 deletions(-) create mode 100644 worker/walrus/src/schema.rs diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index b92082b..e743baf 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] clap = { version = "3.1.12", features = ["derive"] } -diesel = { version = "2.0.0-rc.0", features = ["postgres", "serde_json", "uuid"] } +diesel = { version = "2.0.0-rc.0", features = ["postgres", "serde_json", "uuid", "chrono"] } diesel_migrations = { version = "2.0.0-rc.0", features = ["postgres"] } dotenv = "0.15.0" serde_json = "1.0" diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 875db43..3266c9f 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -110,15 +110,10 @@ as $$ with x(maybe_quoted_name) as ( select - coalesce( - nullif(split_part($1::text, '.', 1), ''), - ( - select relnamespace::regnamespace::text - from pg_class - where oid = $1 - limit 1 - ) - ) + relnamespace::regnamespace::text + from pg_class + where oid = $1 + limit 1 ) select case @@ -300,3 +295,14 @@ begin return visible_to_subscription_ids; end; $$ + +alter table realtime.subscription add column schema_name text generated always as (realtime.to_schema_name(entity)) stored; +alter table realtime.subscription add column table_name text generated always as (realtime.to_table_name(entity)) stored; +alter table realtime.subscription add column claims_role_name text generated always as (realtime.to_regrole((claims ->> 'role'::text))) stored; + +alter table realtime.subscription alter schema_name set not null; +alter table realtime.subscription alter table_name set not null; +alter table realtime.subscription alter claims_role_name set not null; +-- alter table realtime.subscription alter filters drop default; +-- alter table realtime.subscription alter filters type jsonb using to_jsonb(filters); +-- alter table realtime.subscription alter filters set default '[]'; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 32f39e5..1b208ed 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -1,8 +1,11 @@ +#[macro_use] +extern crate diesel; use clap::Parser; +use diesel::prelude::*; use diesel::*; use env_logger; use itertools::Itertools; -use log::{debug, error, info, warn}; +use log::{error, info, warn}; use serde_json; use std::collections::HashMap; use std::io::{self, BufRead}; @@ -13,6 +16,7 @@ use std::time; mod filters; mod migrations; mod realtime_fmt; +mod schema; mod sql_functions; mod timestamp_fmt; mod wal2json; @@ -36,7 +40,6 @@ fn main() { let args = Args::parse(); // enable logger - //env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); loop { @@ -99,9 +102,19 @@ fn run(args: &Args) -> Result<(), String> { // Load initial snapshot of subscriptions info!("Snapshot of subscriptions loading"); - let mut subscriptions = sql_functions::get_subscriptions(conn)?; + use schema::realtime::subscription::dsl::*; + let mut subscriptions = match subscription.load::(conn) { + Ok(subscriptions) => subscriptions, + Err(err) => { + cmd.kill().unwrap(); + error!("Error loading subscriptions: {}", err); + return Err("Error loading subscriptions".to_string()); + } + }; info!("Snapshot of subscriptions loaded"); + //println!("subs {:?}", subscriptions); + // Iterate input data for input_line in stdin_lines { match input_line { @@ -109,6 +122,7 @@ fn run(args: &Args) -> Result<(), String> { let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { + //println!("rec {:?}", wal2json_record); // Update subscriptions if needed realtime_fmt::update_subscriptions( &wal2json_record, @@ -176,7 +190,7 @@ fn process_record( max_record_bytes: usize, conn: &mut PgConnection, ) -> Result, String> { - // TODO subscription name as argument + // TODO publication name as argument let is_in_publication = sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_multiplayer", conn)?; let is_subscribed_to = subscriptions.len() > 0; @@ -196,7 +210,7 @@ fn process_record( .iter() .filter(|x| &x.schema_name == &rec.schema) .filter(|x| &x.table_name == &rec.table) - .map(|x| &x.claims_role) + .map(|x| &x.claims_role_name) .unique() .collect(); @@ -241,7 +255,7 @@ fn process_record( // Subscriptions to current entity + role let entity_role_subscriptions: Vec<&realtime_fmt::Subscription> = entity_subscriptions .iter() - .filter(|x| &x.claims_role == role) + .filter(|x| &x.claims_role_name == role) .map(|x| *x) .collect(); diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 7abc173..090af6d 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -1,9 +1,14 @@ use crate::wal2json; -use chrono::{DateTime, Utc}; +use chrono::{DateTime, NaiveDateTime, Utc}; +use diesel::deserialize::{self, FromSql}; +use diesel::pg::{Pg, PgValue}; +use diesel::serialize::{self, IsNull, Output, ToSql, WriteTuple}; +use diesel::sql_types::{Record, Text}; use diesel::*; use log::error; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::io::Write; use std::*; use uuid; @@ -43,37 +48,20 @@ pub struct WALRLS { } // Subscriptions -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] -pub enum Op { - #[serde(alias = "eq")] - Equal, - #[serde(alias = "neq")] - NotEqual, - #[serde(alias = "lt")] - LessThan, - #[serde(alias = "lte")] - LessThanOrEqual, - #[serde(alias = "gt")] - GreaterThan, - #[serde(alias = "gte")] - GreaterThanOrEqual, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct UserDefinedFilter { - pub column_name: String, - pub op: Op, - pub value: String, // Why did I make this a text field?, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug, Queryable)] pub struct Subscription { pub id: i64, - pub schema_name: String, - pub table_name: String, pub subscription_id: uuid::Uuid, + pub entity: i32, + // This also works for anonymous deser of filters (schema.rs also must change) + //pub filters: Vec<(String, EqualityOp, String)>, pub filters: Vec, - pub claims_role: String, + pub claims: serde_json::Value, + pub claims_role: i32, + pub created_at: NaiveDateTime, + pub schema_name: String, + pub table_name: String, + pub claims_role_name: String, } /// Checks to see if the new record is a change to realtime.subscriptions @@ -93,7 +81,7 @@ pub fn update_subscriptions( return (); } - let id: i64 = match rec + let id_val: i64 = match rec .columns .as_ref() // Deletes have the id value in the identity field @@ -105,7 +93,7 @@ pub fn update_subscriptions( { Some(id_json) => match id_json { serde_json::Value::Number(id_num) => match id_num.as_i64() { - Some(id) => id, + Some(id_val) => id_val, None => { error!( "Invalid id in realtime.subscription. Expected i64, got: {}", @@ -127,35 +115,113 @@ pub fn update_subscriptions( return (); } }; + use crate::schema::realtime::subscription::dsl::*; match rec.action { wal2json::Action::I => { - match crate::sql_functions::get_subscription_by_id(id, conn) { + use crate::schema::realtime::subscription::dsl::*; + match subscription.filter(id.eq(id)).first::(conn) { Ok(new_sub) => subscriptions.push(new_sub), - Err(err) => error!( - "Failed to parse wal2json record to realtime.subscription: {} ", - err - ), + Err(err) => error!("No subscription found: {} ", err), }; } wal2json::Action::U => { // Delete existing sub - subscriptions.retain_mut(|x| x.id != id); + subscriptions.retain_mut(|x| x.id != id_val); // Add updated sub - match crate::sql_functions::get_subscription_by_id(id, conn) { + match subscription.filter(id.eq(id)).first::(conn) { Ok(new_sub) => subscriptions.push(new_sub), - Err(err) => error!( - "Failed to parse wal2json record to realtime.subscription: {} ", - err - ), + Err(err) => error!("No subscription found: {} ", err), }; } wal2json::Action::D => { - subscriptions.retain(|x| x.id != id); + subscriptions.retain(|x| x.id != id_val); } wal2json::Action::T => { // Handled above } }; } + +#[derive(SqlType, PartialEq)] +#[diesel(postgres_type(schema = "realtime", name = "equality_op"))] +pub struct OpType; + +#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize)] +#[diesel(sql_type = OpType)] +pub enum Op { + #[serde(alias = "eq")] + Equal, + #[serde(alias = "neq")] + NotEqual, + #[serde(alias = "lt")] + LessThan, + #[serde(alias = "lte")] + LessThanOrEqual, + #[serde(alias = "gt")] + GreaterThan, + #[serde(alias = "gte")] + GreaterThanOrEqual, +} + +impl ToSql for Op { + fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { + match *self { + Op::Equal => out.write_all(b"eq")?, + Op::NotEqual => out.write_all(b"neq")?, + Op::LessThan => out.write_all(b"lt")?, + Op::LessThanOrEqual => out.write_all(b"lte")?, + Op::GreaterThan => out.write_all(b"gt")?, + Op::GreaterThanOrEqual => out.write_all(b"gt")?, + } + Ok(IsNull::No) + } +} + +impl FromSql for Op { + fn from_sql(bytes: PgValue<'_>) -> deserialize::Result { + match bytes.as_bytes() { + b"eq" => Ok(Op::Equal), + b"neq" => Ok(Op::NotEqual), + b"lt" => Ok(Op::LessThan), + b"lte" => Ok(Op::LessThanOrEqual), + b"gt" => Ok(Op::GreaterThan), + b"gte" => Ok(Op::GreaterThanOrEqual), + _ => Err("Unrecognized enum variant".into()), + } + } +} + +#[derive(SqlType, PartialEq)] +#[diesel(postgres_type(schema = "realtime", name = "user_defined_filter"))] +pub struct UserDefinedFilterType; + +#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize)] +#[diesel(sql_type = UserDefinedFilterType)] +pub struct UserDefinedFilter { + pub column_name: String, + pub op: Op, + pub value: String, // Why did I make this a text field?, +} + +impl ToSql for UserDefinedFilter { + fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { + WriteTuple::<(Text, OpType, Text)>::write_tuple( + &(self.column_name.as_str(), &self.op, self.value.as_str()), + out, + ) + } +} + +impl FromSql for UserDefinedFilter { + fn from_sql(bytes: PgValue<'_>) -> deserialize::Result { + let (column_name, op, value) = + FromSql::, Pg>::from_sql(bytes)?; + Ok(UserDefinedFilter { + column_name, + op, + value, + }) + } +} diff --git a/worker/walrus/src/schema.rs b/worker/walrus/src/schema.rs new file mode 100644 index 0000000..42879d3 --- /dev/null +++ b/worker/walrus/src/schema.rs @@ -0,0 +1,17 @@ +pub mod realtime { + table! { + realtime.subscription (id) { + id -> Int8, + subscription_id -> Uuid, + entity -> Int4, + //filters -> Array>, + filters -> Array, + claims -> Jsonb, + claims_role -> Int4, + created_at -> Timestamp, + schema_name -> Text, + table_name -> Text, + claims_role_name -> Text, + } + } +} From 64216828eae0d8ab5e82c39acbab9e925d18ab48 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 24 Jun 2022 09:53:56 -0500 Subject: [PATCH 41/88] rm deprecated sql functions --- .../2022-06-01-154923_helper_functions/up.sql | 44 ------------------- worker/walrus/src/main.rs | 3 +- worker/walrus/src/sql_functions.rs | 40 ----------------- 3 files changed, 1 insertion(+), 86 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 3266c9f..88cc25c 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -129,50 +129,6 @@ $$ $$; -create function realtime.get_subscriptions( -) - returns jsonb[] - language sql -as $$ - select - coalesce( - array_agg( - jsonb_build_object( - 'id', id, - 'schema_name', realtime.to_schema_name(entity), - 'table_name', realtime.to_table_name(entity), - 'subscription_id', subscription_id, - 'filters', filters, - 'claims_role', claims_role - ) - ), - '{}' - ) - from - realtime.subscription s - limit 1 -$$; - -create function realtime.get_subscription_by_id(id bigint) - returns jsonb - language sql -as $$ - select - jsonb_build_object( - 'id', id, - 'schema_name', realtime.to_schema_name(entity), - 'table_name', realtime.to_table_name(entity), - 'subscription_id', subscription_id, - 'filters', filters, - 'claims_role', claims_role - ) - from - realtime.subscription s - where - s.id = $1 - limit 1 -$$; - create function realtime.is_visible_through_filters( columns jsonb, subscription_ids uuid[] diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 1b208ed..ca05242 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -2,7 +2,6 @@ extern crate diesel; use clap::Parser; use diesel::prelude::*; -use diesel::*; use env_logger; use itertools::Itertools; use log::{error, info, warn}; @@ -364,7 +363,7 @@ fn process_record( } Ok(false) => (), // delegate to SQL when we can't handle the comparison in rust - Err(err) => { + Err(_) => { //debug!( // "Filters delegated to SQL: {:?}. Error: {}", // &sub.filters, err diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index 49b62a6..d07372d 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -18,14 +18,6 @@ pub mod sql_functions { fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; } - sql_function! { - fn get_subscriptions() -> Array; - } - - sql_function! { - fn get_subscription_by_id(id: BigInt) -> Jsonb; - } - sql_function! { fn is_visible_through_filters(columns: Jsonb, subscription_ids: Array ) -> Array } @@ -123,35 +115,3 @@ pub fn is_visible_through_rls( .first(conn) .map_err(|x| format!("{}", x)) } - -pub fn get_subscriptions( - conn: &mut PgConnection, -) -> Result, String> { - let subs: Vec = select(sql_functions::get_subscriptions()) - .first(conn) - .map_err(|x| format!("{}", x))?; - - let mut res = vec![]; - - for sub_json in subs { - let sub: crate::realtime_fmt::Subscription = - serde_json::from_value(sub_json).map_err(|x| format!("{}", x))?; - res.push(sub); - } - - Ok(res) -} - -pub fn get_subscription_by_id( - id: i64, - conn: &mut PgConnection, -) -> Result { - let sub_value: serde_json::Value = select(sql_functions::get_subscription_by_id(id)) - .first(conn) - .map_err(|x| format!("{}", x))?; - - let sub: crate::realtime_fmt::Subscription = - serde_json::from_value(sub_value).map_err(|x| format!("{}", x))?; - - Ok(sub) -} From 1f85e237dc8c9deaf6c87deca0cf1371c61f1022 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 24 Jun 2022 10:03:02 -0500 Subject: [PATCH 42/88] publication is argument --- worker/walrus/src/main.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index ca05242..c94915c 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -32,6 +32,9 @@ struct Args { #[clap(long, default_value = "postgresql://postgres@localhost:5432/postgres")] connection: String, + + #[clap(long, default_value = "supabase_multiplayer")] + publication: String, } fn main() { @@ -56,6 +59,7 @@ fn main() { fn run(args: &Args) -> Result<(), String> { // Connect to Postgres let conn_result = &mut PgConnection::establish(&args.connection); + let publication = &args.publication; let conn = match conn_result { Ok(c) => c, @@ -133,6 +137,7 @@ fn run(args: &Args) -> Result<(), String> { let walrus = process_record( &wal2json_record, &subscriptions, + publication, 1024 * 1024, conn, ); @@ -186,12 +191,12 @@ fn has_primary_key(rec: &wal2json::Record) -> bool { fn process_record( rec: &wal2json::Record, subscriptions: &Vec, + publication: &str, max_record_bytes: usize, conn: &mut PgConnection, ) -> Result, String> { - // TODO publication name as argument let is_in_publication = - sql_functions::is_in_publication(&rec.schema, &rec.table, "supabase_multiplayer", conn)?; + sql_functions::is_in_publication(&rec.schema, &rec.table, publication, conn)?; let is_subscribed_to = subscriptions.len() > 0; let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; From 0a2b8d14f2cde368c774e3bcb06620e706630fe9 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 24 Jun 2022 14:19:51 -0500 Subject: [PATCH 43/88] reduce workload on early exit --- .../2022-06-01-154923_helper_functions/up.sql | 5 ++-- worker/walrus/src/main.rs | 24 +++++++++---------- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 88cc25c..db02463 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -259,6 +259,5 @@ alter table realtime.subscription add column claims_role_name text generated alw alter table realtime.subscription alter schema_name set not null; alter table realtime.subscription alter table_name set not null; alter table realtime.subscription alter claims_role_name set not null; --- alter table realtime.subscription alter filters drop default; --- alter table realtime.subscription alter filters type jsonb using to_jsonb(filters); --- alter table realtime.subscription alter filters set default '[]'; + +create index ix_realtime_subscription_subscription_id on realtime.subscription (subscription_id); diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index c94915c..9461bb4 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -202,6 +202,18 @@ fn process_record( let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; + let action = match rec.action { + wal2json::Action::I => realtime_fmt::Action::INSERT, + wal2json::Action::U => realtime_fmt::Action::UPDATE, + wal2json::Action::D => realtime_fmt::Action::DELETE, + wal2json::Action::T => realtime_fmt::Action::TRUNCATE, + }; + + // If the table isn't in the publication or no one is subscribed, do no work + if !(is_in_publication && is_subscribed_to && action != realtime_fmt::Action::TRUNCATE) { + return Ok(vec![]); + } + // Subscriptions to the current entity let entity_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions .iter() @@ -220,18 +232,6 @@ fn process_record( let mut result: Vec = vec![]; - let action = match rec.action { - wal2json::Action::I => realtime_fmt::Action::INSERT, - wal2json::Action::U => realtime_fmt::Action::UPDATE, - wal2json::Action::D => realtime_fmt::Action::DELETE, - wal2json::Action::T => realtime_fmt::Action::TRUNCATE, - }; - - // If the table isn't in the publication or no one is subscribed, do no work - if !(is_in_publication && is_subscribed_to && action != realtime_fmt::Action::TRUNCATE) { - return Ok(vec![]); - } - // If the table has no primary key, return if action != realtime_fmt::Action::DELETE && !has_primary_key(rec) { let r = realtime_fmt::WALRLS { From d181e23a835daa9d0181075fede19bb46f263fe2 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 24 Jun 2022 14:41:05 -0500 Subject: [PATCH 44/88] missing semicolon --- .../walrus/migrations/2022-06-01-154923_helper_functions/up.sql | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index db02463..4ffa556 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -250,7 +250,7 @@ begin return visible_to_subscription_ids; end; -$$ +$$; alter table realtime.subscription add column schema_name text generated always as (realtime.to_schema_name(entity)) stored; alter table realtime.subscription add column table_name text generated always as (realtime.to_table_name(entity)) stored; From 294e49fa7e92d061ad617be0e90601cef746c0f5 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 24 Jun 2022 16:47:03 -0500 Subject: [PATCH 45/88] subscription query bugfix --- worker/walrus/src/main.rs | 56 ++++++++++++++++++++++++------ worker/walrus/src/migrations.rs | 1 + worker/walrus/src/realtime_fmt.rs | 36 ++++++++++++++++--- worker/walrus/src/sql_functions.rs | 5 +++ 4 files changed, 83 insertions(+), 15 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 9461bb4..8a0d8c6 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -2,9 +2,10 @@ extern crate diesel; use clap::Parser; use diesel::prelude::*; +use diesel::*; use env_logger; use itertools::Itertools; -use log::{error, info, warn}; +use log::{debug, error, info, warn}; use serde_json; use std::collections::HashMap; use std::io::{self, BufRead}; @@ -72,6 +73,11 @@ fn run(args: &Args) -> Result<(), String> { migrations::run_migrations(conn).expect("Pending migrations failed to execute"); info!("Postgres connection established"); + // Empty search path + sql_query("set search_path=''") + .execute(conn) + .expect("failed to set search path"); + let cmd = Command::new("pg_recvlogical") //&args .args(vec![ @@ -146,7 +152,10 @@ fn run(args: &Args) -> Result<(), String> { Ok(rows) => { for row in rows { match serde_json::to_string(&row) { - Ok(walrus_json) => println!("{}", walrus_json), + Ok(walrus_json) => { + println!("{}", walrus_json); + debug!("Pushed record for {}.{} with {} subscribers", row.wal.schema, row.wal.table, row.subscription_ids.len()); + } Err(err) => { error!( "Failed to serialize walrus result: {}", @@ -197,9 +206,41 @@ fn process_record( ) -> Result, String> { let is_in_publication = sql_functions::is_in_publication(&rec.schema, &rec.table, publication, conn)?; - let is_subscribed_to = subscriptions.len() > 0; + + let load_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions + .iter() + .filter(|x| &x.schema_name == "load_messages") + .map(|x| x) + .collect(); + + debug!("N load subs {}", &load_subscriptions.len(),); + + // Subscriptions to the current entity + let entity_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions + .iter() + .filter(|x| &x.schema_name == &rec.schema) + .filter(|x| &x.table_name == &rec.table) + .map(|x| x) + .collect(); + + if subscriptions.len() > 0 { + debug!( + "Rec {} {} table {} {}", + &rec.schema, + &rec.table, + &subscriptions.first().unwrap().schema_name, + &subscriptions.first().unwrap().table_name, + ); + } + + let is_subscribed_to = entity_subscriptions.len() > 0; let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; + debug!( + "Processing record: {}.{} inpub: {}, entity_subs {}, rls_on {}", + &rec.schema, &rec.table, is_in_publication, is_subscribed_to, is_rls_enabled + ); + let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; let action = match rec.action { @@ -211,17 +252,10 @@ fn process_record( // If the table isn't in the publication or no one is subscribed, do no work if !(is_in_publication && is_subscribed_to && action != realtime_fmt::Action::TRUNCATE) { + debug!("Early exit. Not in pub or no one listening"); return Ok(vec![]); } - // Subscriptions to the current entity - let entity_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions - .iter() - .filter(|x| &x.schema_name == &rec.schema) - .filter(|x| &x.table_name == &rec.table) - .map(|x| x) - .collect(); - let subscribed_roles: Vec<&String> = entity_subscriptions .iter() .filter(|x| &x.schema_name == &rec.schema) diff --git a/worker/walrus/src/migrations.rs b/worker/walrus/src/migrations.rs index dbb7385..4ed4657 100644 --- a/worker/walrus/src/migrations.rs +++ b/worker/walrus/src/migrations.rs @@ -16,5 +16,6 @@ pub fn run_migrations( .expect("failed to set search path"); connection.run_pending_migrations(MIGRATIONS)?; + Ok(()) } diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 090af6d..7268b8a 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -5,7 +5,7 @@ use diesel::pg::{Pg, PgValue}; use diesel::serialize::{self, IsNull, Output, ToSql, WriteTuple}; use diesel::sql_types::{Record, Text}; use diesel::*; -use log::error; +use log::{debug, error}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::io::Write; @@ -76,8 +76,11 @@ pub fn update_subscriptions( return (); } + debug!("Subscription record detected"); + if rec.action == wal2json::Action::T { subscriptions.clear(); + debug!("Subscription truncate. Total {}", subscriptions.len()); return (); } @@ -120,23 +123,48 @@ pub fn update_subscriptions( match rec.action { wal2json::Action::I => { use crate::schema::realtime::subscription::dsl::*; - match subscription.filter(id.eq(id)).first::(conn) { - Ok(new_sub) => subscriptions.push(new_sub), + match subscription + .filter(id.eq(id_val)) + .first::(conn) + { + Ok(new_sub) => { + subscriptions.push(new_sub); + debug!("Subscription inserted. Total {}", subscriptions.len()); + } Err(err) => error!("No subscription found: {} ", err), }; } wal2json::Action::U => { // Delete existing sub + let before_update_count = subscriptions.len(); + subscriptions.retain_mut(|x| x.id != id_val); // Add updated sub - match subscription.filter(id.eq(id)).first::(conn) { + match subscription + .filter(id.eq(id_val)) + .first::(conn) + { Ok(new_sub) => subscriptions.push(new_sub), Err(err) => error!("No subscription found: {} ", err), }; + + debug!( + "Subscription update. Total before {}, after {}. id_val {}", + before_update_count, + subscriptions.len(), + id_val, + ); } wal2json::Action::D => { + let before_delete_count = subscriptions.len(); subscriptions.retain(|x| x.id != id_val); + + debug!( + "Subscription delete. Total before {}, after {}", + before_delete_count, + subscriptions.len() + ); } wal2json::Action::T => { // Handled above diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index d07372d..0f90a72 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -7,22 +7,27 @@ pub mod sql_functions { use diesel::*; sql_function! { + #[sql_name = "realtime.is_rls_enabled"] fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; } sql_function! { + #[sql_name = "realtime.is_in_publication"] fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; } sql_function! { + #[sql_name = "realtime.selectable_columns"] fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; } sql_function! { + #[sql_name = "realtime.is_visible_through_filters"] fn is_visible_through_filters(columns: Jsonb, subscription_ids: Array ) -> Array } sql_function! { + #[sql_name = "realtime.is_visible_through_rls"] fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, subscription_ids: Array) -> Array } } From 9aa7048c0fcd4b5e6ec66171723c835f8a39387d Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 27 Jun 2022 11:50:38 -0500 Subject: [PATCH 46/88] filter subs, not ids --- .../2022-06-01-154923_helper_functions/up.sql | 24 +++++------ worker/walrus/src/main.rs | 43 +++++++++++++------ worker/walrus/src/realtime_fmt.rs | 6 +-- worker/walrus/src/sql_functions.rs | 16 +++---- 4 files changed, 52 insertions(+), 37 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 4ffa556..86033dc 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -131,15 +131,15 @@ $$; create function realtime.is_visible_through_filters( columns jsonb, - subscription_ids uuid[] + ids int8[] -- realtime.subscription.id ) - returns uuid[] + returns int8[] language plpgsql as $$ declare cols realtime.wal_column[]; - visible_to_subscription_ids uuid[] = '{}'; - subscription_id uuid; + visible_to_subscription_ids int8[] = '{}'; + subscription_id int8; filters realtime.user_defined_filter[]; subscription_has_access bool; begin @@ -161,12 +161,12 @@ begin for subscription_id, filters in ( select - subs.subscription_id, + subs.id, subs.filters from realtime.subscription subs where - subs.subscription_id = any(subscription_ids) + subs.id = any(ids) ) loop @@ -189,16 +189,16 @@ create function realtime.is_visible_through_rls( schema_name text, table_name text, columns jsonb, - subscription_ids uuid[] + ids int8[] ) - returns uuid[] + returns int8[] language plpgsql as $$ declare entity_ regclass = format('%I.%I', schema_name, table_name)::regclass; cols realtime.wal_column[]; - visible_to_subscription_ids uuid[] = '{}'; - subscription_id uuid; + visible_to_subscription_ids int8[] = '{}'; + subscription_id int8; subscription_has_access bool; claims jsonb; begin @@ -226,12 +226,12 @@ begin for subscription_id, claims in ( select - subs.subscription_id, + subs.id, subs.claims from realtime.subscription subs where - subs.subscription_id = any(subscription_ids) + subs.id = any(ids) ) loop -- Check if RLS allows the role to see the record diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 8a0d8c6..008b917 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -122,7 +122,7 @@ fn run(args: &Args) -> Result<(), String> { }; info!("Snapshot of subscriptions loaded"); - //println!("subs {:?}", subscriptions); + // println!("subs {:?}", subscriptions); // Iterate input data for input_line in stdin_lines { @@ -388,8 +388,8 @@ fn process_record( .collect(); // User Defined Filters - let mut subscription_id_is_visible_through_filters = vec![]; - let mut subscription_id_delegate_to_sql = vec![]; + let mut visible_through_filters = vec![]; + let mut delegate_to_sql_filters = vec![]; for sub in entity_role_subscriptions { match filters::visible_through_filters( @@ -398,7 +398,7 @@ fn process_record( ) { Ok(true) => { //debug!("Filters handled in rust: {:?}", &sub.filters); - subscription_id_is_visible_through_filters.push(sub.subscription_id); + visible_through_filters.push(sub); } Ok(false) => (), // delegate to SQL when we can't handle the comparison in rust @@ -407,18 +407,25 @@ fn process_record( // "Filters delegated to SQL: {:?}. Error: {}", // &sub.filters, err //); - subscription_id_delegate_to_sql.push(sub.subscription_id); + delegate_to_sql_filters.push(sub); } } } - if subscription_id_delegate_to_sql.len() > 0 { + if delegate_to_sql_filters.len() > 0 { match sql_functions::is_visible_through_filters( &walcols, - &subscription_id_delegate_to_sql, + &delegate_to_sql_filters.iter().map(|x| x.id).collect(), conn, ) { - Ok(sub_ids) => subscription_id_is_visible_through_filters.extend(&sub_ids), + Ok(sub_ids) => { + for sub in delegate_to_sql_filters + .iter() + .filter(|x| sub_ids.contains(&x.id)) + { + visible_through_filters.push(sub) + } + } Err(err) => { error!("Failed to deletegate some filters to SQL: {}", err) } @@ -426,21 +433,25 @@ fn process_record( } // Row Level Security - let subscription_ids_to_notify = match is_rls_enabled - && subscription_id_is_visible_through_filters.len() > 0 + let subscriptions_to_notify: Vec<&realtime_fmt::Subscription> = match is_rls_enabled + && visible_through_filters.len() > 0 && !vec![realtime_fmt::Action::DELETE, realtime_fmt::Action::TRUNCATE] .contains(&action) { - false => subscription_id_is_visible_through_filters, + false => visible_through_filters, true => { match sql_functions::is_visible_through_rls( &rec.schema, &rec.table, &walcols, - &subscription_id_is_visible_through_filters, + &visible_through_filters.iter().map(|x| x.id).collect(), conn, ) { - Ok(sub_ids) => sub_ids, + Ok(sub_ids) => visible_through_filters + .iter() + .filter(|x| sub_ids.contains(&x.id)) + .map(|x| *x) + .collect(), Err(err) => { error!("Failed to delegate RLS to SQL: {}", err); vec![] @@ -460,7 +471,11 @@ fn process_record( old_record: old_record_elem, }, is_rls_enabled, - subscription_ids: subscription_ids_to_notify, + subscription_ids: subscriptions_to_notify + .iter() + .map(|x| x.subscription_id) + .unique() + .collect(), errors: match exceeds_max_size { true => vec!["Error 413: Payload Too Large".to_string()], false => vec![], diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 7268b8a..d37fa50 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -48,7 +48,7 @@ pub struct WALRLS { } // Subscriptions -#[derive(Serialize, Deserialize, Clone, Debug, Queryable)] +#[derive(Serialize, Deserialize, Clone, Debug, Queryable, Eq, PartialEq)] pub struct Subscription { pub id: i64, pub subscription_id: uuid::Uuid, @@ -176,7 +176,7 @@ pub fn update_subscriptions( #[diesel(postgres_type(schema = "realtime", name = "equality_op"))] pub struct OpType; -#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize)] +#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] #[diesel(sql_type = OpType)] pub enum Op { #[serde(alias = "eq")] @@ -225,7 +225,7 @@ impl FromSql for Op { #[diesel(postgres_type(schema = "realtime", name = "user_defined_filter"))] pub struct UserDefinedFilterType; -#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize)] +#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] #[diesel(sql_type = UserDefinedFilterType)] pub struct UserDefinedFilter { pub column_name: String, diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index 0f90a72..0f42b33 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -23,12 +23,12 @@ pub mod sql_functions { sql_function! { #[sql_name = "realtime.is_visible_through_filters"] - fn is_visible_through_filters(columns: Jsonb, subscription_ids: Array ) -> Array + fn is_visible_through_filters(columns: Jsonb, ids: Array ) -> Array } sql_function! { #[sql_name = "realtime.is_visible_through_rls"] - fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, subscription_ids: Array) -> Array + fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, ids: Array) -> Array } } @@ -92,13 +92,13 @@ pub fn selectable_columns( pub fn is_visible_through_filters( columns: &Vec, - subscription_ids: &Vec, + ids: &Vec, // TODO: convert this to use subscription_ids to reduce n calls conn: &mut PgConnection, -) -> Result, String> { +) -> Result, String> { select(sql_functions::is_visible_through_filters( serde_json::json!(columns), - subscription_ids, + ids, )) .first(conn) .map_err(|x| format!("{}", x)) @@ -108,14 +108,14 @@ pub fn is_visible_through_rls( schema_name: &str, table_name: &str, columns: &Vec, - subscription_ids: &Vec, + ids: &Vec, conn: &mut PgConnection, -) -> Result, String> { +) -> Result, String> { select(sql_functions::is_visible_through_rls( schema_name, table_name, serde_json::to_value(columns).unwrap(), - subscription_ids, + ids, )) .first(conn) .map_err(|x| format!("{}", x)) From 0476873c98d91117bbbb1ba87389ee735d80ffe9 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 28 Jun 2022 10:51:11 -0500 Subject: [PATCH 47/88] extract topic --- worker/realtime/src/main.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 42aca72..171e8a7 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -71,8 +71,9 @@ async fn main() { let args = Args::parse(); let addr = build_url(&args.url, &args.header); let url = url::Url::parse(&addr).expect("invalid URL"); + let topic = args.topic.to_string(); - println!("{:?}", args); + info!("{:?}", args); // enable logger env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); @@ -93,13 +94,13 @@ async fn main() { info!("WebSocket handshake successful"); let (mut write, read) = ws_stream.split(); - write = join_topic(write, args.topic.to_string()).await; + write = join_topic(write, topic.to_string()).await; // Futures channel let (tx, rx) = futures_channel::mpsc::unbounded(); let heartbeat_tx = tx.clone(); - tokio::spawn(read_stdin(tx, args.topic.to_string())); + tokio::spawn(read_stdin(tx, topic.to_string())); tokio::spawn(heartbeat(heartbeat_tx)); // Map @@ -221,6 +222,9 @@ async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender) { // Wrap phoenix message in a websocket message let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); // push to futures stream - tx.unbounded_send(msg).unwrap(); + match tx.unbounded_send(msg) { + Ok(()) => (), + Err(err) => error!("Error sending heatbeat: {}", err), + }; } } From 720c53259d8c2e59be74889415e04c145fe8de01 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 28 Jun 2022 15:25:54 -0500 Subject: [PATCH 48/88] show id on failed sub loc --- worker/walrus/src/realtime_fmt.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index d37fa50..60d8c5b 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -131,7 +131,7 @@ pub fn update_subscriptions( subscriptions.push(new_sub); debug!("Subscription inserted. Total {}", subscriptions.len()); } - Err(err) => error!("No subscription found: {} ", err), + Err(err) => error!("No subscription found: id={}, Error: {} ", id_val, err), }; } wal2json::Action::U => { From 9e1b2c033f5229092f9fdae285d2cd7eff012057 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Tue, 5 Jul 2022 17:49:21 -0500 Subject: [PATCH 49/88] beginning of test port --- worker/walrus/Cargo.toml | 2 +- worker/walrus/src/main.rs | 95 +++++++++++++++++++++++++++---- worker/walrus/src/realtime_fmt.rs | 8 +-- 3 files changed, 90 insertions(+), 15 deletions(-) diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index e743baf..66ac441 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -10,7 +10,7 @@ diesel_migrations = { version = "2.0.0-rc.0", features = ["postgres"] } dotenv = "0.15.0" serde_json = "1.0" serde = "1.0" -uuid = { version = "1.0", features = ["serde"] } +uuid = { version = "1.0", features = ["serde", "v4"] } log = "0.4.17" env_logger = "0.9.0" itertools = "0.10.3" diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 008b917..f7656b5 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -36,6 +36,10 @@ struct Args { #[clap(long, default_value = "supabase_multiplayer")] publication: String, + + // Exit when no work remains + #[clap(long)] + exit_on_no_work: bool, } fn main() { @@ -49,6 +53,10 @@ fn main() { match run(&args) { Err(err) => { warn!("Error: {}", err); + + if args.exit_on_no_work { + return (); + } } _ => continue, }; @@ -122,8 +130,6 @@ fn run(args: &Args) -> Result<(), String> { }; info!("Snapshot of subscriptions loaded"); - // println!("subs {:?}", subscriptions); - // Iterate input data for input_line in stdin_lines { match input_line { @@ -207,14 +213,6 @@ fn process_record( let is_in_publication = sql_functions::is_in_publication(&rec.schema, &rec.table, publication, conn)?; - let load_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions - .iter() - .filter(|x| &x.schema_name == "load_messages") - .map(|x| x) - .collect(); - - debug!("N load subs {}", &load_subscriptions.len(),); - // Subscriptions to the current entity let entity_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions .iter() @@ -487,3 +485,80 @@ fn process_record( Ok(result) } + +#[cfg(test)] +mod tests { + extern crate diesel; + use crate::realtime_fmt::Subscription; + use crate::schema::realtime::subscription::dsl::*; + use crate::wal2json; + use chrono::Utc; + use diesel::prelude::*; + use diesel::*; + use serde_json::json; + use uuid; + + fn establish_connection() -> PgConnection { + let database_url = "postgresql://postgres:password@localhost:5501/postgres"; + PgConnection::establish(&database_url).unwrap() + } + + fn clean(conn: &mut PgConnection) { + delete(subscription).execute(conn).unwrap(); + } + + #[test] + fn test_basic() { + let mut conn = establish_connection(); + + let claim_sub = uuid::Uuid::new_v4(); + + insert_into(subscription) + .values(( + subscription_id.eq(uuid::Uuid::new_v4()), + entity.eq(16487), + claims.eq(json!({ + "role": "postgres", + "email": "example@example.com", + "sub": claim_sub + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscriptions = subscription.load::(&mut conn).unwrap(); + clean(&mut conn); + + assert_eq!(subscriptions.len(), 1); + } + + #[test] + fn test_no_one_listening() { + let mut conn = establish_connection(); + + let rec = wal2json::Record { + action: wal2json::Action::I, + schema: "public".to_string(), + table: "notes".to_string(), + pk: Some(vec![wal2json::PrimaryKeyRef { + name: "id".to_string(), + type_: "int4".to_string(), // todo + typeoid: 4, // todo + }]), + columns: Some(vec![]), + identity: None, + timestamp: Utc::now(), + }; + + let res = crate::process_record( + &rec, + &vec![], + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + assert_eq!(res, vec![]); + } +} diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 60d8c5b..600418e 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -12,7 +12,7 @@ use std::io::Write; use std::*; use uuid; -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] pub enum Action { INSERT, UPDATE, @@ -20,14 +20,14 @@ pub enum Action { TRUNCATE, } -#[derive(Serialize, Deserialize, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] pub struct Column { pub name: String, #[serde(rename(serialize = "type", deserialize = "type"))] pub type_: String, } -#[derive(Serialize, Deserialize, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] pub struct Data { pub schema: String, pub table: String, @@ -39,7 +39,7 @@ pub struct Data { pub old_record: Option>, } -#[derive(Serialize, Deserialize, Clone, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] pub struct WALRLS { pub wal: Data, pub is_rls_enabled: bool, From b9c9d6e7ffb94846949050d26324cb57dc0e18f2 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 08:46:57 -0500 Subject: [PATCH 50/88] value in tests --- worker/walrus/src/main.rs | 27 +++++++++++++++++++++++---- worker/walrus/src/wal2json.rs | 4 ++-- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index f7656b5..03d05c8 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -192,12 +192,21 @@ fn run(args: &Args) -> Result<(), String> { } } -fn pkey_cols(rec: &wal2json::Record) -> Vec<&String> { +/* +fn pkey_cols<'a>(rec: &'a wal2json::Record) -> Vec<&'a str> { match &rec.pk { Some(pkey_refs) => pkey_refs.iter().map(|x| &x.name).collect(), None => vec![], } } +*/ + +fn pkey_cols(rec: &wal2json::Record) -> Vec { + match &rec.pk { + Some(pkey_refs) => pkey_refs.iter().map(|x| x.name.clone()).collect(), + None => vec![], + } +} fn has_primary_key(rec: &wal2json::Record) -> bool { pkey_cols(rec).len() != 0 @@ -498,6 +507,11 @@ mod tests { use serde_json::json; use uuid; + const BOOLOID: i32 = 16; + const INT4OID: u32 = 23; + const INT8OID: i32 = 20; + const TEXTOID: i32 = 25; + fn establish_connection() -> PgConnection { let database_url = "postgresql://postgres:password@localhost:5501/postgres"; PgConnection::establish(&database_url).unwrap() @@ -542,10 +556,15 @@ mod tests { table: "notes".to_string(), pk: Some(vec![wal2json::PrimaryKeyRef { name: "id".to_string(), - type_: "int4".to_string(), // todo - typeoid: 4, // todo + type_: "int4".to_string(), + typeoid: INT4OID, + }]), + columns: Some(vec![wal2json::Column { + name: "id".to_string(), + type_: "int4".to_string(), + typeoid: Some(INT4OID), + value: json!(1), }]), - columns: Some(vec![]), identity: None, timestamp: Utc::now(), }; diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index 9eafb98..c0ba722 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -16,7 +16,7 @@ pub struct PrimaryKeyRef { pub name: String, #[serde(alias = "type")] pub type_: String, - pub typeoid: i32, + pub typeoid: u32, } #[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] @@ -30,7 +30,7 @@ pub enum Action { #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Record { pub action: Action, - pub schema: String, + pub schema: String, //&str, pub table: String, pub pk: Option>, pub columns: Option>, // option is for truncate From 8312c546b319a65779c6cc66520b1b95bfa248bd Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 09:01:54 -0500 Subject: [PATCH 51/88] references throughout structs --- worker/walrus/src/main.rs | 35 ++++++++++++------------------- worker/walrus/src/realtime_fmt.rs | 22 +++++++++---------- worker/walrus/src/wal2json.rs | 24 ++++++++++----------- 3 files changed, 36 insertions(+), 45 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 03d05c8..60866be 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -192,18 +192,9 @@ fn run(args: &Args) -> Result<(), String> { } } -/* fn pkey_cols<'a>(rec: &'a wal2json::Record) -> Vec<&'a str> { match &rec.pk { - Some(pkey_refs) => pkey_refs.iter().map(|x| &x.name).collect(), - None => vec![], - } -} -*/ - -fn pkey_cols(rec: &wal2json::Record) -> Vec { - match &rec.pk { - Some(pkey_refs) => pkey_refs.iter().map(|x| x.name.clone()).collect(), + Some(pkey_refs) => pkey_refs.iter().map(|x| x.name).collect(), None => vec![], } } @@ -212,13 +203,13 @@ fn has_primary_key(rec: &wal2json::Record) -> bool { pkey_cols(rec).len() != 0 } -fn process_record( - rec: &wal2json::Record, +fn process_record<'a>( + rec: &'a wal2json::Record, subscriptions: &Vec, publication: &str, max_record_bytes: usize, conn: &mut PgConnection, -) -> Result, String> { +) -> Result>, String> { let is_in_publication = sql_functions::is_in_publication(&rec.schema, &rec.table, publication, conn)?; @@ -312,10 +303,10 @@ fn process_record( .as_ref() .unwrap_or(&vec![]) .iter() - .filter(|col| selectable_columns.contains(&col.name)) + .filter(|col| selectable_columns.contains(&col.name.to_string())) .map(|w2j_col| realtime_fmt::Column { - name: w2j_col.name.to_string(), - type_: w2j_col.type_.to_string(), + name: w2j_col.name, + type_: w2j_col.type_, }) .collect(); @@ -552,16 +543,16 @@ mod tests { let rec = wal2json::Record { action: wal2json::Action::I, - schema: "public".to_string(), - table: "notes".to_string(), + schema: "public", + table: "notes", pk: Some(vec![wal2json::PrimaryKeyRef { - name: "id".to_string(), - type_: "int4".to_string(), + name: "id", + type_: "int4", typeoid: INT4OID, }]), columns: Some(vec![wal2json::Column { - name: "id".to_string(), - type_: "int4".to_string(), + name: "id", + type_: "int4", typeoid: Some(INT4OID), value: json!(1), }]), diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 600418e..9b53c23 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -12,7 +12,7 @@ use std::io::Write; use std::*; use uuid; -#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] pub enum Action { INSERT, UPDATE, @@ -20,28 +20,28 @@ pub enum Action { TRUNCATE, } -#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] -pub struct Column { - pub name: String, +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub struct Column<'a> { + pub name: &'a str, #[serde(rename(serialize = "type", deserialize = "type"))] - pub type_: String, + pub type_: &'a str, } -#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] -pub struct Data { +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub struct Data<'a> { pub schema: String, pub table: String, pub r#type: Action, #[serde(with = "crate::timestamp_fmt")] pub commit_timestamp: DateTime, - pub columns: Vec, + pub columns: Vec>, pub record: HashMap, pub old_record: Option>, } -#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] -pub struct WALRLS { - pub wal: Data, +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub struct WALRLS<'a> { + pub wal: Data<'a>, pub is_rls_enabled: bool, pub subscription_ids: Vec, pub errors: Vec, diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs index c0ba722..60cb8d6 100644 --- a/worker/walrus/src/wal2json.rs +++ b/worker/walrus/src/wal2json.rs @@ -3,19 +3,19 @@ use serde::{Deserialize, Serialize}; use std::*; #[derive(Serialize, Deserialize, Clone, Debug)] -pub struct Column { - pub name: String, +pub struct Column<'a> { + pub name: &'a str, #[serde(alias = "type")] - pub type_: String, + pub type_: &'a str, pub typeoid: Option, pub value: serde_json::Value, } #[derive(Serialize, Deserialize, Clone, Debug)] -pub struct PrimaryKeyRef { - pub name: String, +pub struct PrimaryKeyRef<'a> { + pub name: &'a str, #[serde(alias = "type")] - pub type_: String, + pub type_: &'a str, pub typeoid: u32, } @@ -28,14 +28,14 @@ pub enum Action { } #[derive(Serialize, Deserialize, Clone, Debug)] -pub struct Record { +pub struct Record<'a> { pub action: Action, - pub schema: String, //&str, - pub table: String, - pub pk: Option>, - pub columns: Option>, // option is for truncate + pub schema: &'a str, + pub table: &'a str, + pub pk: Option>>, + pub columns: Option>>, // option is for truncate #[serde(skip_serializing_if = "Option::is_none")] - pub identity: Option>, // option is for insert/update + pub identity: Option>>, // option is for insert/update // Example: 2022-06-22 15:38:19.695275+00 #[serde(with = "crate::timestamp_fmt")] pub timestamp: DateTime, From 3791c8c45f46a2038b2c54f5ae40a9c97156df5b Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 09:23:05 -0500 Subject: [PATCH 52/88] reduce allocs --- worker/walrus/src/main.rs | 39 +++++++++++++++---------------- worker/walrus/src/realtime_fmt.rs | 12 +++++----- worker/walrus/src/walrus_fmt.rs | 12 +++++----- 3 files changed, 31 insertions(+), 32 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 60866be..ddd0273 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -268,10 +268,10 @@ fn process_record<'a>( if action != realtime_fmt::Action::DELETE && !has_primary_key(rec) { let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { - schema: rec.schema.to_string(), - table: rec.table.to_string(), + schema: rec.schema, + table: rec.table, r#type: action.clone(), - commit_timestamp: rec.timestamp, + commit_timestamp: &rec.timestamp, columns: vec![], record: HashMap::new(), old_record: None, @@ -281,7 +281,7 @@ fn process_record<'a>( .iter() .map(|x| x.subscription_id.clone()) .collect(), - errors: vec!["Error 400: Bad Request, no primary key".to_string()], + errors: vec!["Error 400: Bad Request, no primary key"], }; result.push(r); return Ok(result); @@ -318,10 +318,10 @@ fn process_record<'a>( if action != realtime_fmt::Action::DELETE && selectable_columns.len() == 0 { let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { - schema: rec.schema.to_string(), - table: rec.table.to_string(), + schema: rec.schema, + table: rec.table, r#type: action.clone(), - commit_timestamp: rec.timestamp, + commit_timestamp: &rec.timestamp, columns, record: HashMap::new(), old_record: None, @@ -331,16 +331,16 @@ fn process_record<'a>( .iter() .map(|x| x.subscription_id.clone()) .collect(), - errors: vec!["Error 401: Unauthorized".to_string()], + errors: vec!["Error 401: Unauthorized"], }; result.push(r); } else { if vec![realtime_fmt::Action::INSERT, realtime_fmt::Action::UPDATE].contains(&action) { for col_name in &selectable_columns { 'record: for col in rec.columns.as_ref().unwrap_or(&vec![]) { - if col_name == &col.name { + if col_name == col.name { if !exceeds_max_size || col.value.to_string().len() < 64 { - record_elem.insert(col_name.to_string(), col.value.clone()); + record_elem.insert(col.name, col.value.clone()); break 'record; } } @@ -353,10 +353,9 @@ fn process_record<'a>( match &rec.identity { Some(identity) => { 'old_record: for col in identity { - if col_name == &col.name { + if col_name == col.name { if !exceeds_max_size || col.value.to_string().len() < 64 { - old_record_elem_content - .insert(col_name.to_string(), col.value.clone()); + old_record_elem_content.insert(col.name, col.value.clone()); break 'old_record; } } @@ -375,9 +374,9 @@ fn process_record<'a>( .iter() .map(|col| { walrus_fmt::WALColumn { - name: col.name.to_string(), - type_name: col.type_.to_string(), - type_oid: col.typeoid.clone(), + name: col.name, + type_name: col.type_, + type_oid: col.typeoid, value: col.value.clone(), is_pkey: pkey_cols(rec).contains(&&col.name), is_selectable: false, // stub: unused, @@ -460,10 +459,10 @@ fn process_record<'a>( let r = realtime_fmt::WALRLS { wal: realtime_fmt::Data { - schema: rec.schema.to_string(), - table: rec.table.to_string(), + schema: rec.schema, + table: rec.table, r#type: action.clone(), - commit_timestamp: rec.timestamp, + commit_timestamp: &rec.timestamp, columns, record: record_elem, old_record: old_record_elem, @@ -475,7 +474,7 @@ fn process_record<'a>( .unique() .collect(), errors: match exceeds_max_size { - true => vec!["Error 413: Payload Too Large".to_string()], + true => vec!["Error 413: Payload Too Large"], false => vec![], }, }; diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index 9b53c23..fd4145c 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -29,14 +29,14 @@ pub struct Column<'a> { #[derive(Serialize, Clone, Debug, Eq, PartialEq)] pub struct Data<'a> { - pub schema: String, - pub table: String, + pub schema: &'a str, + pub table: &'a str, pub r#type: Action, #[serde(with = "crate::timestamp_fmt")] - pub commit_timestamp: DateTime, + pub commit_timestamp: &'a DateTime, pub columns: Vec>, - pub record: HashMap, - pub old_record: Option>, + pub record: HashMap<&'a str, serde_json::Value>, + pub old_record: Option>, } #[derive(Serialize, Clone, Debug, Eq, PartialEq)] @@ -44,7 +44,7 @@ pub struct WALRLS<'a> { pub wal: Data<'a>, pub is_rls_enabled: bool, pub subscription_ids: Vec, - pub errors: Vec, + pub errors: Vec<&'a str>, } // Subscriptions diff --git a/worker/walrus/src/walrus_fmt.rs b/worker/walrus/src/walrus_fmt.rs index 2a9feb0..5437d2a 100644 --- a/worker/walrus/src/walrus_fmt.rs +++ b/worker/walrus/src/walrus_fmt.rs @@ -1,17 +1,17 @@ use serde::Serialize; #[derive(Serialize)] -pub struct WalrusRecord { +pub struct WalrusRecord<'a> { wal: serde_json::Value, is_rls_enabled: bool, - subscription_ids: Vec, - errors: Vec, + subscription_ids: Vec<&'a uuid::Uuid>, + errors: Vec<&'a str>, } #[derive(Serialize, Debug)] -pub struct WALColumn { - pub name: String, - pub type_name: String, +pub struct WALColumn<'a> { + pub name: &'a str, + pub type_name: &'a str, pub type_oid: Option, pub value: serde_json::Value, pub is_pkey: bool, From 8d6ccdd103bd44e7c0eabb487e60d949acc0f4ce Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 09:55:20 -0500 Subject: [PATCH 53/88] relocate sql migrations and schema --- worker/walrus/src/filters/record.rs | 0 worker/walrus/src/filters/table.rs | 0 worker/walrus/src/main.rs | 9 ++++----- worker/walrus/src/models/realtime.rs | 0 worker/walrus/src/models/wal2json.rs | 0 worker/walrus/src/models/walrus.rs | 0 worker/walrus/src/realtime_fmt.rs | 3 +-- worker/walrus/src/sql.rs | 2 ++ worker/walrus/src/{ => sql}/migrations.rs | 0 worker/walrus/src/{ => sql}/schema.rs | 0 10 files changed, 7 insertions(+), 7 deletions(-) create mode 100644 worker/walrus/src/filters/record.rs create mode 100644 worker/walrus/src/filters/table.rs create mode 100644 worker/walrus/src/models/realtime.rs create mode 100644 worker/walrus/src/models/wal2json.rs create mode 100644 worker/walrus/src/models/walrus.rs create mode 100644 worker/walrus/src/sql.rs rename worker/walrus/src/{ => sql}/migrations.rs (100%) rename worker/walrus/src/{ => sql}/schema.rs (100%) diff --git a/worker/walrus/src/filters/record.rs b/worker/walrus/src/filters/record.rs new file mode 100644 index 0000000..e69de29 diff --git a/worker/walrus/src/filters/table.rs b/worker/walrus/src/filters/table.rs new file mode 100644 index 0000000..e69de29 diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index ddd0273..b48407e 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -7,6 +7,7 @@ use env_logger; use itertools::Itertools; use log::{debug, error, info, warn}; use serde_json; +use sql::schema::realtime::subscription::dsl::*; use std::collections::HashMap; use std::io::{self, BufRead}; use std::process::{Command, Stdio}; @@ -14,9 +15,8 @@ use std::thread::sleep; use std::time; mod filters; -mod migrations; mod realtime_fmt; -mod schema; +mod sql; mod sql_functions; mod timestamp_fmt; mod wal2json; @@ -78,7 +78,7 @@ fn run(args: &Args) -> Result<(), String> { }; // Run pending migrations - migrations::run_migrations(conn).expect("Pending migrations failed to execute"); + sql::migrations::run_migrations(conn).expect("Pending migrations failed to execute"); info!("Postgres connection established"); // Empty search path @@ -119,7 +119,6 @@ fn run(args: &Args) -> Result<(), String> { // Load initial snapshot of subscriptions info!("Snapshot of subscriptions loading"); - use schema::realtime::subscription::dsl::*; let mut subscriptions = match subscription.load::(conn) { Ok(subscriptions) => subscriptions, Err(err) => { @@ -489,7 +488,7 @@ fn process_record<'a>( mod tests { extern crate diesel; use crate::realtime_fmt::Subscription; - use crate::schema::realtime::subscription::dsl::*; + use crate::sql::schema::realtime::subscription::dsl::*; use crate::wal2json; use chrono::Utc; use diesel::prelude::*; diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs new file mode 100644 index 0000000..e69de29 diff --git a/worker/walrus/src/models/wal2json.rs b/worker/walrus/src/models/wal2json.rs new file mode 100644 index 0000000..e69de29 diff --git a/worker/walrus/src/models/walrus.rs b/worker/walrus/src/models/walrus.rs new file mode 100644 index 0000000..e69de29 diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs index fd4145c..c77595b 100644 --- a/worker/walrus/src/realtime_fmt.rs +++ b/worker/walrus/src/realtime_fmt.rs @@ -1,3 +1,4 @@ +use crate::sql::schema::realtime::subscription::dsl::*; use crate::wal2json; use chrono::{DateTime, NaiveDateTime, Utc}; use diesel::deserialize::{self, FromSql}; @@ -118,11 +119,9 @@ pub fn update_subscriptions( return (); } }; - use crate::schema::realtime::subscription::dsl::*; match rec.action { wal2json::Action::I => { - use crate::schema::realtime::subscription::dsl::*; match subscription .filter(id.eq(id_val)) .first::(conn) diff --git a/worker/walrus/src/sql.rs b/worker/walrus/src/sql.rs new file mode 100644 index 0000000..c80cbca --- /dev/null +++ b/worker/walrus/src/sql.rs @@ -0,0 +1,2 @@ +pub mod migrations; +pub mod schema; diff --git a/worker/walrus/src/migrations.rs b/worker/walrus/src/sql/migrations.rs similarity index 100% rename from worker/walrus/src/migrations.rs rename to worker/walrus/src/sql/migrations.rs diff --git a/worker/walrus/src/schema.rs b/worker/walrus/src/sql/schema.rs similarity index 100% rename from worker/walrus/src/schema.rs rename to worker/walrus/src/sql/schema.rs From 587d6d7c4300c4124e6252a65ef66f630732a1c1 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 10:04:22 -0500 Subject: [PATCH 54/88] relocate models --- worker/walrus/src/filters.rs | 17 +- worker/walrus/src/main.rs | 62 ++++--- worker/walrus/src/models.rs | 3 + worker/walrus/src/models/realtime.rs | 254 +++++++++++++++++++++++++++ worker/walrus/src/models/wal2json.rs | 42 +++++ worker/walrus/src/models/walrus.rs | 19 ++ worker/walrus/src/realtime_fmt.rs | 254 --------------------------- worker/walrus/src/sql/schema.rs | 2 +- worker/walrus/src/sql_functions.rs | 5 +- worker/walrus/src/wal2json.rs | 42 ----- worker/walrus/src/walrus_fmt.rs | 19 -- 11 files changed, 360 insertions(+), 359 deletions(-) create mode 100644 worker/walrus/src/models.rs delete mode 100644 worker/walrus/src/realtime_fmt.rs delete mode 100644 worker/walrus/src/wal2json.rs delete mode 100644 worker/walrus/src/walrus_fmt.rs diff --git a/worker/walrus/src/filters.rs b/worker/walrus/src/filters.rs index ac8b1e3..86d33eb 100644 --- a/worker/walrus/src/filters.rs +++ b/worker/walrus/src/filters.rs @@ -1,5 +1,4 @@ -use crate::realtime_fmt::UserDefinedFilter; -use crate::wal2json::Column; +use crate::models::{realtime, wal2json}; use log::warn; fn is_null(v: &serde_json::Value) -> bool { @@ -7,10 +6,10 @@ fn is_null(v: &serde_json::Value) -> bool { } pub fn visible_through_filters( - filters: &Vec, - columns: &Vec, + filters: &Vec, + columns: &Vec, ) -> Result { - use crate::realtime_fmt::Op; + use realtime::Op; for filter in filters { let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { @@ -73,11 +72,11 @@ pub fn visible_through_filters( Ok(true) } -/// Returns a vector of realtime_fmt::Op that match for a OP b +/// Returns a vector of realtime::Op that match for a OP b fn get_valid_ops( a: &serde_json::Value, b: &serde_json::Value, -) -> Result, String> { +) -> Result, String> { use serde_json::Value; match (a, b) { @@ -91,11 +90,11 @@ fn get_valid_ops( } } -fn get_matching_ops(a: &T, b: &T) -> Vec +fn get_matching_ops(a: &T, b: &T) -> Vec where T: PartialEq + PartialOrd, { - use crate::realtime_fmt::Op; + use realtime::Op; let mut ops = vec![]; if a == b { diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index b48407e..0212f1b 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -6,6 +6,7 @@ use diesel::*; use env_logger; use itertools::Itertools; use log::{debug, error, info, warn}; +use models::{realtime, wal2json, walrus}; use serde_json; use sql::schema::realtime::subscription::dsl::*; use std::collections::HashMap; @@ -15,12 +16,10 @@ use std::thread::sleep; use std::time; mod filters; -mod realtime_fmt; +mod models; mod sql; mod sql_functions; mod timestamp_fmt; -mod wal2json; -mod walrus_fmt; /// Write-Ahead-Log Realtime Unified Security (WALRUS) background worker /// runs next to a PostgreSQL instance and forwards its Write-Ahead-Log @@ -119,7 +118,7 @@ fn run(args: &Args) -> Result<(), String> { // Load initial snapshot of subscriptions info!("Snapshot of subscriptions loading"); - let mut subscriptions = match subscription.load::(conn) { + let mut subscriptions = match subscription.load::(conn) { Ok(subscriptions) => subscriptions, Err(err) => { cmd.kill().unwrap(); @@ -138,7 +137,7 @@ fn run(args: &Args) -> Result<(), String> { Ok(wal2json_record) => { //println!("rec {:?}", wal2json_record); // Update subscriptions if needed - realtime_fmt::update_subscriptions( + realtime::update_subscriptions( &wal2json_record, &mut subscriptions, conn, @@ -204,16 +203,16 @@ fn has_primary_key(rec: &wal2json::Record) -> bool { fn process_record<'a>( rec: &'a wal2json::Record, - subscriptions: &Vec, + subscriptions: &Vec, publication: &str, max_record_bytes: usize, conn: &mut PgConnection, -) -> Result>, String> { +) -> Result>, String> { let is_in_publication = sql_functions::is_in_publication(&rec.schema, &rec.table, publication, conn)?; // Subscriptions to the current entity - let entity_subscriptions: Vec<&realtime_fmt::Subscription> = subscriptions + let entity_subscriptions: Vec<&realtime::Subscription> = subscriptions .iter() .filter(|x| &x.schema_name == &rec.schema) .filter(|x| &x.table_name == &rec.table) @@ -241,14 +240,14 @@ fn process_record<'a>( let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; let action = match rec.action { - wal2json::Action::I => realtime_fmt::Action::INSERT, - wal2json::Action::U => realtime_fmt::Action::UPDATE, - wal2json::Action::D => realtime_fmt::Action::DELETE, - wal2json::Action::T => realtime_fmt::Action::TRUNCATE, + wal2json::Action::I => realtime::Action::INSERT, + wal2json::Action::U => realtime::Action::UPDATE, + wal2json::Action::D => realtime::Action::DELETE, + wal2json::Action::T => realtime::Action::TRUNCATE, }; // If the table isn't in the publication or no one is subscribed, do no work - if !(is_in_publication && is_subscribed_to && action != realtime_fmt::Action::TRUNCATE) { + if !(is_in_publication && is_subscribed_to && action != realtime::Action::TRUNCATE) { debug!("Early exit. Not in pub or no one listening"); return Ok(vec![]); } @@ -261,12 +260,12 @@ fn process_record<'a>( .unique() .collect(); - let mut result: Vec = vec![]; + let mut result: Vec = vec![]; // If the table has no primary key, return - if action != realtime_fmt::Action::DELETE && !has_primary_key(rec) { - let r = realtime_fmt::WALRLS { - wal: realtime_fmt::Data { + if action != realtime::Action::DELETE && !has_primary_key(rec) { + let r = realtime::WALRLS { + wal: realtime::Data { schema: rec.schema, table: rec.table, r#type: action.clone(), @@ -288,7 +287,7 @@ fn process_record<'a>( for role in subscribed_roles { // Subscriptions to current entity + role - let entity_role_subscriptions: Vec<&realtime_fmt::Subscription> = entity_subscriptions + let entity_role_subscriptions: Vec<&realtime::Subscription> = entity_subscriptions .iter() .filter(|x| &x.claims_role_name == role) .map(|x| *x) @@ -303,7 +302,7 @@ fn process_record<'a>( .unwrap_or(&vec![]) .iter() .filter(|col| selectable_columns.contains(&col.name.to_string())) - .map(|w2j_col| realtime_fmt::Column { + .map(|w2j_col| realtime::Column { name: w2j_col.name, type_: w2j_col.type_, }) @@ -314,9 +313,9 @@ fn process_record<'a>( let mut old_record_elem_content = HashMap::new(); // If the role can not select any columns in the table, return - if action != realtime_fmt::Action::DELETE && selectable_columns.len() == 0 { - let r = realtime_fmt::WALRLS { - wal: realtime_fmt::Data { + if action != realtime::Action::DELETE && selectable_columns.len() == 0 { + let r = realtime::WALRLS { + wal: realtime::Data { schema: rec.schema, table: rec.table, r#type: action.clone(), @@ -334,7 +333,7 @@ fn process_record<'a>( }; result.push(r); } else { - if vec![realtime_fmt::Action::INSERT, realtime_fmt::Action::UPDATE].contains(&action) { + if vec![realtime::Action::INSERT, realtime::Action::UPDATE].contains(&action) { for col_name in &selectable_columns { 'record: for col in rec.columns.as_ref().unwrap_or(&vec![]) { if col_name == col.name { @@ -347,7 +346,7 @@ fn process_record<'a>( } } - if vec![realtime_fmt::Action::UPDATE, realtime_fmt::Action::DELETE].contains(&action) { + if vec![realtime::Action::UPDATE, realtime::Action::DELETE].contains(&action) { for col_name in &selectable_columns { match &rec.identity { Some(identity) => { @@ -366,13 +365,13 @@ fn process_record<'a>( old_record_elem = Some(old_record_elem_content); } - let walcols: Vec = rec + let walcols: Vec = rec .columns .as_ref() .unwrap_or(&vec![]) .iter() .map(|col| { - walrus_fmt::WALColumn { + walrus::WALColumn { name: col.name, type_name: col.type_, type_oid: col.typeoid, @@ -429,10 +428,9 @@ fn process_record<'a>( } // Row Level Security - let subscriptions_to_notify: Vec<&realtime_fmt::Subscription> = match is_rls_enabled + let subscriptions_to_notify: Vec<&realtime::Subscription> = match is_rls_enabled && visible_through_filters.len() > 0 - && !vec![realtime_fmt::Action::DELETE, realtime_fmt::Action::TRUNCATE] - .contains(&action) + && !vec![realtime::Action::DELETE, realtime::Action::TRUNCATE].contains(&action) { false => visible_through_filters, true => { @@ -456,8 +454,8 @@ fn process_record<'a>( } }; - let r = realtime_fmt::WALRLS { - wal: realtime_fmt::Data { + let r = realtime::WALRLS { + wal: realtime::Data { schema: rec.schema, table: rec.table, r#type: action.clone(), @@ -487,7 +485,7 @@ fn process_record<'a>( #[cfg(test)] mod tests { extern crate diesel; - use crate::realtime_fmt::Subscription; + use crate::realtime::Subscription; use crate::sql::schema::realtime::subscription::dsl::*; use crate::wal2json; use chrono::Utc; diff --git a/worker/walrus/src/models.rs b/worker/walrus/src/models.rs new file mode 100644 index 0000000..6c4d81f --- /dev/null +++ b/worker/walrus/src/models.rs @@ -0,0 +1,3 @@ +pub mod realtime; +pub mod wal2json; +pub mod walrus; diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index e69de29..c77595b 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -0,0 +1,254 @@ +use crate::sql::schema::realtime::subscription::dsl::*; +use crate::wal2json; +use chrono::{DateTime, NaiveDateTime, Utc}; +use diesel::deserialize::{self, FromSql}; +use diesel::pg::{Pg, PgValue}; +use diesel::serialize::{self, IsNull, Output, ToSql, WriteTuple}; +use diesel::sql_types::{Record, Text}; +use diesel::*; +use log::{debug, error}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::io::Write; +use std::*; +use uuid; + +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub enum Action { + INSERT, + UPDATE, + DELETE, + TRUNCATE, +} + +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub struct Column<'a> { + pub name: &'a str, + #[serde(rename(serialize = "type", deserialize = "type"))] + pub type_: &'a str, +} + +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub struct Data<'a> { + pub schema: &'a str, + pub table: &'a str, + pub r#type: Action, + #[serde(with = "crate::timestamp_fmt")] + pub commit_timestamp: &'a DateTime, + pub columns: Vec>, + pub record: HashMap<&'a str, serde_json::Value>, + pub old_record: Option>, +} + +#[derive(Serialize, Clone, Debug, Eq, PartialEq)] +pub struct WALRLS<'a> { + pub wal: Data<'a>, + pub is_rls_enabled: bool, + pub subscription_ids: Vec, + pub errors: Vec<&'a str>, +} + +// Subscriptions +#[derive(Serialize, Deserialize, Clone, Debug, Queryable, Eq, PartialEq)] +pub struct Subscription { + pub id: i64, + pub subscription_id: uuid::Uuid, + pub entity: i32, + // This also works for anonymous deser of filters (schema.rs also must change) + //pub filters: Vec<(String, EqualityOp, String)>, + pub filters: Vec, + pub claims: serde_json::Value, + pub claims_role: i32, + pub created_at: NaiveDateTime, + pub schema_name: String, + pub table_name: String, + pub claims_role_name: String, +} + +/// Checks to see if the new record is a change to realtime.subscriptions +/// and updates the subscriptions variable if a change is detected +pub fn update_subscriptions( + rec: &wal2json::Record, + subscriptions: &mut Vec, + conn: &mut PgConnection, +) -> () { + // If the record is not a subscription, return + if rec.schema != "realtime" || rec.table != "subscription" { + return (); + } + + debug!("Subscription record detected"); + + if rec.action == wal2json::Action::T { + subscriptions.clear(); + debug!("Subscription truncate. Total {}", subscriptions.len()); + return (); + } + + let id_val: i64 = match rec + .columns + .as_ref() + // Deletes have the id value in the identity field + .unwrap_or(rec.identity.as_ref().unwrap_or(&vec![])) + .iter() + .filter(|x| x.name == "id") + .map(|x| x.value.clone()) + .next() + { + Some(id_json) => match id_json { + serde_json::Value::Number(id_num) => match id_num.as_i64() { + Some(id_val) => id_val, + None => { + error!( + "Invalid id in realtime.subscription. Expected i64, got: {}", + id_num + ); + return (); + } + }, + _ => { + error!( + "Invalid id in realtime.subscription. Expected number, got: {}", + id_json + ); + return (); + } + }, + None => { + error!("No id column found on realtime.subscription"); + return (); + } + }; + + match rec.action { + wal2json::Action::I => { + match subscription + .filter(id.eq(id_val)) + .first::(conn) + { + Ok(new_sub) => { + subscriptions.push(new_sub); + debug!("Subscription inserted. Total {}", subscriptions.len()); + } + Err(err) => error!("No subscription found: id={}, Error: {} ", id_val, err), + }; + } + wal2json::Action::U => { + // Delete existing sub + let before_update_count = subscriptions.len(); + + subscriptions.retain_mut(|x| x.id != id_val); + + // Add updated sub + match subscription + .filter(id.eq(id_val)) + .first::(conn) + { + Ok(new_sub) => subscriptions.push(new_sub), + Err(err) => error!("No subscription found: {} ", err), + }; + + debug!( + "Subscription update. Total before {}, after {}. id_val {}", + before_update_count, + subscriptions.len(), + id_val, + ); + } + wal2json::Action::D => { + let before_delete_count = subscriptions.len(); + subscriptions.retain(|x| x.id != id_val); + + debug!( + "Subscription delete. Total before {}, after {}", + before_delete_count, + subscriptions.len() + ); + } + wal2json::Action::T => { + // Handled above + } + }; +} + +#[derive(SqlType, PartialEq)] +#[diesel(postgres_type(schema = "realtime", name = "equality_op"))] +pub struct OpType; + +#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] +#[diesel(sql_type = OpType)] +pub enum Op { + #[serde(alias = "eq")] + Equal, + #[serde(alias = "neq")] + NotEqual, + #[serde(alias = "lt")] + LessThan, + #[serde(alias = "lte")] + LessThanOrEqual, + #[serde(alias = "gt")] + GreaterThan, + #[serde(alias = "gte")] + GreaterThanOrEqual, +} + +impl ToSql for Op { + fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { + match *self { + Op::Equal => out.write_all(b"eq")?, + Op::NotEqual => out.write_all(b"neq")?, + Op::LessThan => out.write_all(b"lt")?, + Op::LessThanOrEqual => out.write_all(b"lte")?, + Op::GreaterThan => out.write_all(b"gt")?, + Op::GreaterThanOrEqual => out.write_all(b"gt")?, + } + Ok(IsNull::No) + } +} + +impl FromSql for Op { + fn from_sql(bytes: PgValue<'_>) -> deserialize::Result { + match bytes.as_bytes() { + b"eq" => Ok(Op::Equal), + b"neq" => Ok(Op::NotEqual), + b"lt" => Ok(Op::LessThan), + b"lte" => Ok(Op::LessThanOrEqual), + b"gt" => Ok(Op::GreaterThan), + b"gte" => Ok(Op::GreaterThanOrEqual), + _ => Err("Unrecognized enum variant".into()), + } + } +} + +#[derive(SqlType, PartialEq)] +#[diesel(postgres_type(schema = "realtime", name = "user_defined_filter"))] +pub struct UserDefinedFilterType; + +#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] +#[diesel(sql_type = UserDefinedFilterType)] +pub struct UserDefinedFilter { + pub column_name: String, + pub op: Op, + pub value: String, // Why did I make this a text field?, +} + +impl ToSql for UserDefinedFilter { + fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { + WriteTuple::<(Text, OpType, Text)>::write_tuple( + &(self.column_name.as_str(), &self.op, self.value.as_str()), + out, + ) + } +} + +impl FromSql for UserDefinedFilter { + fn from_sql(bytes: PgValue<'_>) -> deserialize::Result { + let (column_name, op, value) = + FromSql::, Pg>::from_sql(bytes)?; + Ok(UserDefinedFilter { + column_name, + op, + value, + }) + } +} diff --git a/worker/walrus/src/models/wal2json.rs b/worker/walrus/src/models/wal2json.rs index e69de29..60cb8d6 100644 --- a/worker/walrus/src/models/wal2json.rs +++ b/worker/walrus/src/models/wal2json.rs @@ -0,0 +1,42 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use std::*; + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Column<'a> { + pub name: &'a str, + #[serde(alias = "type")] + pub type_: &'a str, + pub typeoid: Option, + pub value: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct PrimaryKeyRef<'a> { + pub name: &'a str, + #[serde(alias = "type")] + pub type_: &'a str, + pub typeoid: u32, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +pub enum Action { + I, + U, + D, + T, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct Record<'a> { + pub action: Action, + pub schema: &'a str, + pub table: &'a str, + pub pk: Option>>, + pub columns: Option>>, // option is for truncate + #[serde(skip_serializing_if = "Option::is_none")] + pub identity: Option>>, // option is for insert/update + // Example: 2022-06-22 15:38:19.695275+00 + #[serde(with = "crate::timestamp_fmt")] + pub timestamp: DateTime, +} diff --git a/worker/walrus/src/models/walrus.rs b/worker/walrus/src/models/walrus.rs index e69de29..5437d2a 100644 --- a/worker/walrus/src/models/walrus.rs +++ b/worker/walrus/src/models/walrus.rs @@ -0,0 +1,19 @@ +use serde::Serialize; + +#[derive(Serialize)] +pub struct WalrusRecord<'a> { + wal: serde_json::Value, + is_rls_enabled: bool, + subscription_ids: Vec<&'a uuid::Uuid>, + errors: Vec<&'a str>, +} + +#[derive(Serialize, Debug)] +pub struct WALColumn<'a> { + pub name: &'a str, + pub type_name: &'a str, + pub type_oid: Option, + pub value: serde_json::Value, + pub is_pkey: bool, + pub is_selectable: bool, +} diff --git a/worker/walrus/src/realtime_fmt.rs b/worker/walrus/src/realtime_fmt.rs deleted file mode 100644 index c77595b..0000000 --- a/worker/walrus/src/realtime_fmt.rs +++ /dev/null @@ -1,254 +0,0 @@ -use crate::sql::schema::realtime::subscription::dsl::*; -use crate::wal2json; -use chrono::{DateTime, NaiveDateTime, Utc}; -use diesel::deserialize::{self, FromSql}; -use diesel::pg::{Pg, PgValue}; -use diesel::serialize::{self, IsNull, Output, ToSql, WriteTuple}; -use diesel::sql_types::{Record, Text}; -use diesel::*; -use log::{debug, error}; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::io::Write; -use std::*; -use uuid; - -#[derive(Serialize, Clone, Debug, Eq, PartialEq)] -pub enum Action { - INSERT, - UPDATE, - DELETE, - TRUNCATE, -} - -#[derive(Serialize, Clone, Debug, Eq, PartialEq)] -pub struct Column<'a> { - pub name: &'a str, - #[serde(rename(serialize = "type", deserialize = "type"))] - pub type_: &'a str, -} - -#[derive(Serialize, Clone, Debug, Eq, PartialEq)] -pub struct Data<'a> { - pub schema: &'a str, - pub table: &'a str, - pub r#type: Action, - #[serde(with = "crate::timestamp_fmt")] - pub commit_timestamp: &'a DateTime, - pub columns: Vec>, - pub record: HashMap<&'a str, serde_json::Value>, - pub old_record: Option>, -} - -#[derive(Serialize, Clone, Debug, Eq, PartialEq)] -pub struct WALRLS<'a> { - pub wal: Data<'a>, - pub is_rls_enabled: bool, - pub subscription_ids: Vec, - pub errors: Vec<&'a str>, -} - -// Subscriptions -#[derive(Serialize, Deserialize, Clone, Debug, Queryable, Eq, PartialEq)] -pub struct Subscription { - pub id: i64, - pub subscription_id: uuid::Uuid, - pub entity: i32, - // This also works for anonymous deser of filters (schema.rs also must change) - //pub filters: Vec<(String, EqualityOp, String)>, - pub filters: Vec, - pub claims: serde_json::Value, - pub claims_role: i32, - pub created_at: NaiveDateTime, - pub schema_name: String, - pub table_name: String, - pub claims_role_name: String, -} - -/// Checks to see if the new record is a change to realtime.subscriptions -/// and updates the subscriptions variable if a change is detected -pub fn update_subscriptions( - rec: &wal2json::Record, - subscriptions: &mut Vec, - conn: &mut PgConnection, -) -> () { - // If the record is not a subscription, return - if rec.schema != "realtime" || rec.table != "subscription" { - return (); - } - - debug!("Subscription record detected"); - - if rec.action == wal2json::Action::T { - subscriptions.clear(); - debug!("Subscription truncate. Total {}", subscriptions.len()); - return (); - } - - let id_val: i64 = match rec - .columns - .as_ref() - // Deletes have the id value in the identity field - .unwrap_or(rec.identity.as_ref().unwrap_or(&vec![])) - .iter() - .filter(|x| x.name == "id") - .map(|x| x.value.clone()) - .next() - { - Some(id_json) => match id_json { - serde_json::Value::Number(id_num) => match id_num.as_i64() { - Some(id_val) => id_val, - None => { - error!( - "Invalid id in realtime.subscription. Expected i64, got: {}", - id_num - ); - return (); - } - }, - _ => { - error!( - "Invalid id in realtime.subscription. Expected number, got: {}", - id_json - ); - return (); - } - }, - None => { - error!("No id column found on realtime.subscription"); - return (); - } - }; - - match rec.action { - wal2json::Action::I => { - match subscription - .filter(id.eq(id_val)) - .first::(conn) - { - Ok(new_sub) => { - subscriptions.push(new_sub); - debug!("Subscription inserted. Total {}", subscriptions.len()); - } - Err(err) => error!("No subscription found: id={}, Error: {} ", id_val, err), - }; - } - wal2json::Action::U => { - // Delete existing sub - let before_update_count = subscriptions.len(); - - subscriptions.retain_mut(|x| x.id != id_val); - - // Add updated sub - match subscription - .filter(id.eq(id_val)) - .first::(conn) - { - Ok(new_sub) => subscriptions.push(new_sub), - Err(err) => error!("No subscription found: {} ", err), - }; - - debug!( - "Subscription update. Total before {}, after {}. id_val {}", - before_update_count, - subscriptions.len(), - id_val, - ); - } - wal2json::Action::D => { - let before_delete_count = subscriptions.len(); - subscriptions.retain(|x| x.id != id_val); - - debug!( - "Subscription delete. Total before {}, after {}", - before_delete_count, - subscriptions.len() - ); - } - wal2json::Action::T => { - // Handled above - } - }; -} - -#[derive(SqlType, PartialEq)] -#[diesel(postgres_type(schema = "realtime", name = "equality_op"))] -pub struct OpType; - -#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] -#[diesel(sql_type = OpType)] -pub enum Op { - #[serde(alias = "eq")] - Equal, - #[serde(alias = "neq")] - NotEqual, - #[serde(alias = "lt")] - LessThan, - #[serde(alias = "lte")] - LessThanOrEqual, - #[serde(alias = "gt")] - GreaterThan, - #[serde(alias = "gte")] - GreaterThanOrEqual, -} - -impl ToSql for Op { - fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { - match *self { - Op::Equal => out.write_all(b"eq")?, - Op::NotEqual => out.write_all(b"neq")?, - Op::LessThan => out.write_all(b"lt")?, - Op::LessThanOrEqual => out.write_all(b"lte")?, - Op::GreaterThan => out.write_all(b"gt")?, - Op::GreaterThanOrEqual => out.write_all(b"gt")?, - } - Ok(IsNull::No) - } -} - -impl FromSql for Op { - fn from_sql(bytes: PgValue<'_>) -> deserialize::Result { - match bytes.as_bytes() { - b"eq" => Ok(Op::Equal), - b"neq" => Ok(Op::NotEqual), - b"lt" => Ok(Op::LessThan), - b"lte" => Ok(Op::LessThanOrEqual), - b"gt" => Ok(Op::GreaterThan), - b"gte" => Ok(Op::GreaterThanOrEqual), - _ => Err("Unrecognized enum variant".into()), - } - } -} - -#[derive(SqlType, PartialEq)] -#[diesel(postgres_type(schema = "realtime", name = "user_defined_filter"))] -pub struct UserDefinedFilterType; - -#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] -#[diesel(sql_type = UserDefinedFilterType)] -pub struct UserDefinedFilter { - pub column_name: String, - pub op: Op, - pub value: String, // Why did I make this a text field?, -} - -impl ToSql for UserDefinedFilter { - fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { - WriteTuple::<(Text, OpType, Text)>::write_tuple( - &(self.column_name.as_str(), &self.op, self.value.as_str()), - out, - ) - } -} - -impl FromSql for UserDefinedFilter { - fn from_sql(bytes: PgValue<'_>) -> deserialize::Result { - let (column_name, op, value) = - FromSql::, Pg>::from_sql(bytes)?; - Ok(UserDefinedFilter { - column_name, - op, - value, - }) - } -} diff --git a/worker/walrus/src/sql/schema.rs b/worker/walrus/src/sql/schema.rs index 42879d3..c2dba16 100644 --- a/worker/walrus/src/sql/schema.rs +++ b/worker/walrus/src/sql/schema.rs @@ -5,7 +5,7 @@ pub mod realtime { subscription_id -> Uuid, entity -> Int4, //filters -> Array>, - filters -> Array, + filters -> Array, claims -> Jsonb, claims_role -> Int4, created_at -> Timestamp, diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index 0f42b33..8379395 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -1,3 +1,4 @@ +use crate::models::walrus; use cached::proc_macro::cached; use cached::TimedSizedCache; use diesel::*; @@ -91,7 +92,7 @@ pub fn selectable_columns( } pub fn is_visible_through_filters( - columns: &Vec, + columns: &Vec, ids: &Vec, // TODO: convert this to use subscription_ids to reduce n calls conn: &mut PgConnection, @@ -107,7 +108,7 @@ pub fn is_visible_through_filters( pub fn is_visible_through_rls( schema_name: &str, table_name: &str, - columns: &Vec, + columns: &Vec, ids: &Vec, conn: &mut PgConnection, ) -> Result, String> { diff --git a/worker/walrus/src/wal2json.rs b/worker/walrus/src/wal2json.rs deleted file mode 100644 index 60cb8d6..0000000 --- a/worker/walrus/src/wal2json.rs +++ /dev/null @@ -1,42 +0,0 @@ -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use std::*; - -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct Column<'a> { - pub name: &'a str, - #[serde(alias = "type")] - pub type_: &'a str, - pub typeoid: Option, - pub value: serde_json::Value, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct PrimaryKeyRef<'a> { - pub name: &'a str, - #[serde(alias = "type")] - pub type_: &'a str, - pub typeoid: u32, -} - -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] -pub enum Action { - I, - U, - D, - T, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct Record<'a> { - pub action: Action, - pub schema: &'a str, - pub table: &'a str, - pub pk: Option>>, - pub columns: Option>>, // option is for truncate - #[serde(skip_serializing_if = "Option::is_none")] - pub identity: Option>>, // option is for insert/update - // Example: 2022-06-22 15:38:19.695275+00 - #[serde(with = "crate::timestamp_fmt")] - pub timestamp: DateTime, -} diff --git a/worker/walrus/src/walrus_fmt.rs b/worker/walrus/src/walrus_fmt.rs deleted file mode 100644 index 5437d2a..0000000 --- a/worker/walrus/src/walrus_fmt.rs +++ /dev/null @@ -1,19 +0,0 @@ -use serde::Serialize; - -#[derive(Serialize)] -pub struct WalrusRecord<'a> { - wal: serde_json::Value, - is_rls_enabled: bool, - subscription_ids: Vec<&'a uuid::Uuid>, - errors: Vec<&'a str>, -} - -#[derive(Serialize, Debug)] -pub struct WALColumn<'a> { - pub name: &'a str, - pub type_name: &'a str, - pub type_oid: Option, - pub value: serde_json::Value, - pub is_pkey: bool, - pub is_selectable: bool, -} From ec2587189b282fe22555c4c4ad8bec59589fef59 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 10:27:34 -0500 Subject: [PATCH 55/88] towards filter relocation --- worker/walrus/src/filters.rs | 121 +-------------- worker/walrus/src/filters/record.rs | 2 + .../src/filters/record/row_level_security.rs | 29 ++++ .../walrus/src/filters/record/user_defined.rs | 145 ++++++++++++++++++ worker/walrus/src/filters/table.rs | 2 + .../walrus/src/filters/table/publication.rs | 34 ++++ .../src/filters/table/row_level_security.rs | 29 ++++ worker/walrus/src/main.rs | 27 ++-- worker/walrus/src/sql_functions.rs | 89 ----------- 9 files changed, 252 insertions(+), 226 deletions(-) create mode 100644 worker/walrus/src/filters/record/row_level_security.rs create mode 100644 worker/walrus/src/filters/record/user_defined.rs create mode 100644 worker/walrus/src/filters/table/publication.rs create mode 100644 worker/walrus/src/filters/table/row_level_security.rs diff --git a/worker/walrus/src/filters.rs b/worker/walrus/src/filters.rs index 86d33eb..659f88f 100644 --- a/worker/walrus/src/filters.rs +++ b/worker/walrus/src/filters.rs @@ -1,119 +1,2 @@ -use crate::models::{realtime, wal2json}; -use log::warn; - -fn is_null(v: &serde_json::Value) -> bool { - v == &serde_json::Value::Null -} - -pub fn visible_through_filters( - filters: &Vec, - columns: &Vec, -) -> Result { - use realtime::Op; - - for filter in filters { - let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { - Ok(v) => v, - Err(err) => return Err(format!("{}", err)), - }; - - let column = match columns - .iter() - .filter(|x| x.name == filter.column_name) - .next() - { - Some(col) => col, - // The filter references a column that does not exist - None => { - warn!( - "Attempted to filter non-existing column {}", - filter.column_name - ); - return Ok(false); - } - }; - - // Null column or filter values always result in a false because return is alway null - // until `IS` op is implemented - if is_null(&filter_value) || is_null(&column.value) { - return Ok(false); - }; - - match column.type_.as_ref() { - // All operations supported - "boolean" | "smallint" | "integer" | "bigint" | "serial" | "bigserial" | "numeric" - | "double precision" | "character" | "character varying" | "text" => match &filter.op { - // List is currently exhastive but that may be possible if types/ops expand - Op::Equal - | Op::NotEqual - | Op::LessThan - | Op::LessThanOrEqual - | Op::GreaterThan - | Op::GreaterThanOrEqual => { - let valid_ops = get_valid_ops(&column.value, &filter_value)?; - if !valid_ops.contains(&filter.op) { - return Ok(false); - } - } - }, - // Only Equality ops supported - "uuid" => match &filter.op { - Op::Equal | Op::NotEqual => { - let valid_ops = get_valid_ops(&column.value, &filter_value)?; - if !valid_ops.contains(&filter.op) { - return Ok(false); - }; - } - _ => return Err("could not handle filter op for allowed types".to_string()), - }, - _ => return Err("Could not handle type. Delegate comparison to SQL".to_string()), - }; - } - Ok(true) -} - -/// Returns a vector of realtime::Op that match for a OP b -fn get_valid_ops( - a: &serde_json::Value, - b: &serde_json::Value, -) -> Result, String> { - use serde_json::Value; - - match (a, b) { - (Value::Null, Value::Null) => Ok(vec![]), - (Value::Bool(a_), Value::Bool(b_)) => Ok(get_matching_ops(a_, b_)), - (Value::Number(a_), Value::Number(b_)) => Ok(get_matching_ops(&a_.as_f64(), &b_.as_f64())), - (Value::String(a_), Value::String(b_)) => Ok(get_matching_ops(a_, b_)), - // Array possible - // Object possible - _ => return Err("non-scalar or mismatched json value types".to_string()), - } -} - -fn get_matching_ops(a: &T, b: &T) -> Vec -where - T: PartialEq + PartialOrd, -{ - use realtime::Op; - let mut ops = vec![]; - - if a == b { - ops.push(Op::Equal); - }; - if a != b { - ops.push(Op::NotEqual); - }; - if a < b { - ops.push(Op::LessThan); - }; - if a <= b { - ops.push(Op::LessThanOrEqual); - }; - if a > b { - ops.push(Op::GreaterThan); - }; - if a >= b { - ops.push(Op::GreaterThanOrEqual); - }; - ops -} +pub mod record; +pub mod table; diff --git a/worker/walrus/src/filters/record.rs b/worker/walrus/src/filters/record.rs index e69de29..b8bbaa5 100644 --- a/worker/walrus/src/filters/record.rs +++ b/worker/walrus/src/filters/record.rs @@ -0,0 +1,2 @@ +pub mod row_level_security; +pub mod user_defined; diff --git a/worker/walrus/src/filters/record/row_level_security.rs b/worker/walrus/src/filters/record/row_level_security.rs new file mode 100644 index 0000000..a43c741 --- /dev/null +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -0,0 +1,29 @@ +use crate::models::walrus; +use diesel::*; + +pub mod sql { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + #[sql_name = "realtime.is_visible_through_rls"] + fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, ids: Array) -> Array + } +} + +pub fn is_visible_through_rls( + schema_name: &str, + table_name: &str, + columns: &Vec, + ids: &Vec, + conn: &mut PgConnection, +) -> Result, String> { + select(sql::is_visible_through_rls( + schema_name, + table_name, + serde_json::to_value(columns).unwrap(), + ids, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} diff --git a/worker/walrus/src/filters/record/user_defined.rs b/worker/walrus/src/filters/record/user_defined.rs new file mode 100644 index 0000000..8969b32 --- /dev/null +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -0,0 +1,145 @@ +use crate::models::walrus; +use crate::models::{realtime, wal2json}; +use diesel::*; +use log::warn; + +pub mod sql_functions { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + #[sql_name = "realtime.is_visible_through_filters"] + fn is_visible_through_filters(columns: Jsonb, ids: Array ) -> Array + } +} + +pub fn is_visible_through_filters_sql( + columns: &Vec, + ids: &Vec, + // TODO: convert this to use subscription_ids to reduce n calls + conn: &mut PgConnection, +) -> Result, String> { + select(sql_functions::is_visible_through_filters( + serde_json::json!(columns), + ids, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} + +fn is_null(v: &serde_json::Value) -> bool { + v == &serde_json::Value::Null +} + +pub fn visible_through_filters( + filters: &Vec, + columns: &Vec, +) -> Result { + use realtime::Op; + + for filter in filters { + let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { + Ok(v) => v, + Err(err) => return Err(format!("{}", err)), + }; + + let column = match columns + .iter() + .filter(|x| x.name == filter.column_name) + .next() + { + Some(col) => col, + // The filter references a column that does not exist + None => { + warn!( + "Attempted to filter non-existing column {}", + filter.column_name + ); + return Ok(false); + } + }; + + // Null column or filter values always result in a false because return is alway null + // until `IS` op is implemented + if is_null(&filter_value) || is_null(&column.value) { + return Ok(false); + }; + + match column.type_.as_ref() { + // All operations supported + "boolean" | "smallint" | "integer" | "bigint" | "serial" | "bigserial" | "numeric" + | "double precision" | "character" | "character varying" | "text" => match &filter.op { + // List is currently exhastive but that may be possible if types/ops expand + Op::Equal + | Op::NotEqual + | Op::LessThan + | Op::LessThanOrEqual + | Op::GreaterThan + | Op::GreaterThanOrEqual => { + let valid_ops = get_valid_ops(&column.value, &filter_value)?; + if !valid_ops.contains(&filter.op) { + return Ok(false); + } + } + }, + // Only Equality ops supported + "uuid" => match &filter.op { + Op::Equal | Op::NotEqual => { + let valid_ops = get_valid_ops(&column.value, &filter_value)?; + if !valid_ops.contains(&filter.op) { + return Ok(false); + }; + } + _ => return Err("could not handle filter op for allowed types".to_string()), + }, + _ => return Err("Could not handle type. Delegate comparison to SQL".to_string()), + }; + } + Ok(true) +} + +/// Returns a vector of realtime::Op that match for a OP b +fn get_valid_ops( + a: &serde_json::Value, + b: &serde_json::Value, +) -> Result, String> { + use serde_json::Value; + + match (a, b) { + (Value::Null, Value::Null) => Ok(vec![]), + (Value::Bool(a_), Value::Bool(b_)) => Ok(get_matching_ops(a_, b_)), + (Value::Number(a_), Value::Number(b_)) => Ok(get_matching_ops(&a_.as_f64(), &b_.as_f64())), + (Value::String(a_), Value::String(b_)) => Ok(get_matching_ops(a_, b_)), + // Array possible + // Object possible + _ => return Err("non-scalar or mismatched json value types".to_string()), + } +} + +fn get_matching_ops(a: &T, b: &T) -> Vec +where + T: PartialEq + PartialOrd, +{ + use realtime::Op; + let mut ops = vec![]; + + if a == b { + ops.push(Op::Equal); + }; + if a != b { + ops.push(Op::NotEqual); + }; + if a < b { + ops.push(Op::LessThan); + }; + if a <= b { + ops.push(Op::LessThanOrEqual); + }; + if a > b { + ops.push(Op::GreaterThan); + }; + if a >= b { + ops.push(Op::GreaterThanOrEqual); + }; + ops +} diff --git a/worker/walrus/src/filters/table.rs b/worker/walrus/src/filters/table.rs index e69de29..e0df5cc 100644 --- a/worker/walrus/src/filters/table.rs +++ b/worker/walrus/src/filters/table.rs @@ -0,0 +1,2 @@ +pub mod publication; +pub mod row_level_security; diff --git a/worker/walrus/src/filters/table/publication.rs b/worker/walrus/src/filters/table/publication.rs new file mode 100644 index 0000000..190ed4e --- /dev/null +++ b/worker/walrus/src/filters/table/publication.rs @@ -0,0 +1,34 @@ +use cached::proc_macro::cached; +use cached::TimedSizedCache; +use diesel::*; + +pub mod sql { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + #[sql_name = "realtime.is_in_publication"] + fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; + } +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}-{}", schema_name, table_name, publication_name) }"#, + sync_writes = true +)] +pub fn is_in_publication( + schema_name: &str, + table_name: &str, + publication_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql::is_in_publication( + schema_name, + table_name, + publication_name, + )) + .first(conn) + .map_err(|x| format!("{}", x)) +} diff --git a/worker/walrus/src/filters/table/row_level_security.rs b/worker/walrus/src/filters/table/row_level_security.rs new file mode 100644 index 0000000..6ca6e18 --- /dev/null +++ b/worker/walrus/src/filters/table/row_level_security.rs @@ -0,0 +1,29 @@ +use cached::proc_macro::cached; +use cached::TimedSizedCache; +use diesel::*; + +pub mod sql { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + #[sql_name = "realtime.is_rls_enabled"] + fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; + } +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +pub fn is_rls_enabled( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql::is_rls_enabled(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 0212f1b..6177c52 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -209,7 +209,7 @@ fn process_record<'a>( conn: &mut PgConnection, ) -> Result>, String> { let is_in_publication = - sql_functions::is_in_publication(&rec.schema, &rec.table, publication, conn)?; + filters::table::publication::is_in_publication(&rec.schema, &rec.table, publication, conn)?; // Subscriptions to the current entity let entity_subscriptions: Vec<&realtime::Subscription> = subscriptions @@ -219,18 +219,9 @@ fn process_record<'a>( .map(|x| x) .collect(); - if subscriptions.len() > 0 { - debug!( - "Rec {} {} table {} {}", - &rec.schema, - &rec.table, - &subscriptions.first().unwrap().schema_name, - &subscriptions.first().unwrap().table_name, - ); - } - let is_subscribed_to = entity_subscriptions.len() > 0; - let is_rls_enabled = sql_functions::is_rls_enabled(&rec.schema, &rec.table, conn)?; + let is_rls_enabled = + filters::table::row_level_security::is_rls_enabled(&rec.schema, &rec.table, conn)?; debug!( "Processing record: {}.{} inpub: {}, entity_subs {}, rls_on {}", @@ -387,7 +378,7 @@ fn process_record<'a>( let mut delegate_to_sql_filters = vec![]; for sub in entity_role_subscriptions { - match filters::visible_through_filters( + match filters::record::user_defined::visible_through_filters( &sub.filters, rec.columns.as_ref().unwrap_or(&vec![]), ) { @@ -408,7 +399,7 @@ fn process_record<'a>( } if delegate_to_sql_filters.len() > 0 { - match sql_functions::is_visible_through_filters( + match filters::record::user_defined::is_visible_through_filters_sql( &walcols, &delegate_to_sql_filters.iter().map(|x| x.id).collect(), conn, @@ -434,7 +425,7 @@ fn process_record<'a>( { false => visible_through_filters, true => { - match sql_functions::is_visible_through_rls( + match filters::record::row_level_security::is_visible_through_rls( &rec.schema, &rec.table, &walcols, @@ -494,10 +485,10 @@ mod tests { use serde_json::json; use uuid; - const BOOLOID: i32 = 16; + //const BOOLOID: i32 = 16; const INT4OID: u32 = 23; - const INT8OID: i32 = 20; - const TEXTOID: i32 = 25; + //const INT8OID: i32 = 20; + //const TEXTOID: i32 = 25; fn establish_connection() -> PgConnection { let database_url = "postgresql://postgres:password@localhost:5501/postgres"; diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/sql_functions.rs index 8379395..ddec726 100644 --- a/worker/walrus/src/sql_functions.rs +++ b/worker/walrus/src/sql_functions.rs @@ -1,4 +1,3 @@ -use crate::models::walrus; use cached::proc_macro::cached; use cached::TimedSizedCache; use diesel::*; @@ -7,67 +6,10 @@ pub mod sql_functions { use diesel::sql_types::*; use diesel::*; - sql_function! { - #[sql_name = "realtime.is_rls_enabled"] - fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; - } - - sql_function! { - #[sql_name = "realtime.is_in_publication"] - fn is_in_publication(schema_name: Text, table_name: Text, publication_name: Text) -> Bool; - } - sql_function! { #[sql_name = "realtime.selectable_columns"] fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; } - - sql_function! { - #[sql_name = "realtime.is_visible_through_filters"] - fn is_visible_through_filters(columns: Jsonb, ids: Array ) -> Array - } - - sql_function! { - #[sql_name = "realtime.is_visible_through_rls"] - fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, ids: Array) -> Array - } -} - -#[cached( - type = "TimedSizedCache>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, - sync_writes = true -)] -pub fn is_rls_enabled( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result { - select(sql_functions::is_rls_enabled(schema_name, table_name)) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -#[cached( - type = "TimedSizedCache>", - create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}-{}", schema_name, table_name, publication_name) }"#, - sync_writes = true -)] -pub fn is_in_publication( - schema_name: &str, - table_name: &str, - publication_name: &str, - conn: &mut PgConnection, -) -> Result { - select(sql_functions::is_in_publication( - schema_name, - table_name, - publication_name, - )) - .first(conn) - .map_err(|x| format!("{}", x)) } #[cached( @@ -90,34 +32,3 @@ pub fn selectable_columns( .first(conn) .map_err(|x| format!("{}", x)) } - -pub fn is_visible_through_filters( - columns: &Vec, - ids: &Vec, - // TODO: convert this to use subscription_ids to reduce n calls - conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::is_visible_through_filters( - serde_json::json!(columns), - ids, - )) - .first(conn) - .map_err(|x| format!("{}", x)) -} - -pub fn is_visible_through_rls( - schema_name: &str, - table_name: &str, - columns: &Vec, - ids: &Vec, - conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::is_visible_through_rls( - schema_name, - table_name, - serde_json::to_value(columns).unwrap(), - ids, - )) - .first(conn) - .map_err(|x| format!("{}", x)) -} From f693bb6a58df8d1debbde97ef59b7665664e48bf Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 10:44:19 -0500 Subject: [PATCH 56/88] reloc pkey helper functions --- worker/walrus/src/filters/table.rs | 1 + .../table/column_security.rs} | 0 worker/walrus/src/main.rs | 26 +++++++------------ worker/walrus/src/models/wal2json.rs | 13 ++++++++++ 4 files changed, 23 insertions(+), 17 deletions(-) rename worker/walrus/src/{sql_functions.rs => filters/table/column_security.rs} (100%) diff --git a/worker/walrus/src/filters/table.rs b/worker/walrus/src/filters/table.rs index e0df5cc..b1f5f83 100644 --- a/worker/walrus/src/filters/table.rs +++ b/worker/walrus/src/filters/table.rs @@ -1,2 +1,3 @@ +pub mod column_security; pub mod publication; pub mod row_level_security; diff --git a/worker/walrus/src/sql_functions.rs b/worker/walrus/src/filters/table/column_security.rs similarity index 100% rename from worker/walrus/src/sql_functions.rs rename to worker/walrus/src/filters/table/column_security.rs diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 6177c52..0b3f927 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -18,7 +18,6 @@ use std::time; mod filters; mod models; mod sql; -mod sql_functions; mod timestamp_fmt; /// Write-Ahead-Log Realtime Unified Security (WALRUS) background worker @@ -59,7 +58,7 @@ fn main() { } _ => continue, }; - info!("Stream interrupted. Restarting pg_recvlogical in 5 seconds"); + info!("Stream interrupted. Restarting in 5 seconds"); sleep(time::Duration::from_secs(5)); } } @@ -190,17 +189,6 @@ fn run(args: &Args) -> Result<(), String> { } } -fn pkey_cols<'a>(rec: &'a wal2json::Record) -> Vec<&'a str> { - match &rec.pk { - Some(pkey_refs) => pkey_refs.iter().map(|x| x.name).collect(), - None => vec![], - } -} - -fn has_primary_key(rec: &wal2json::Record) -> bool { - pkey_cols(rec).len() != 0 -} - fn process_record<'a>( rec: &'a wal2json::Record, subscriptions: &Vec, @@ -254,7 +242,7 @@ fn process_record<'a>( let mut result: Vec = vec![]; // If the table has no primary key, return - if action != realtime::Action::DELETE && !has_primary_key(rec) { + if action != realtime::Action::DELETE && !rec.has_primary_key() { let r = realtime::WALRLS { wal: realtime::Data { schema: rec.schema, @@ -284,8 +272,12 @@ fn process_record<'a>( .map(|x| *x) .collect(); - let selectable_columns = - sql_functions::selectable_columns(&rec.schema, &rec.table, role, conn)?; + let selectable_columns = filters::table::column_security::selectable_columns( + &rec.schema, + &rec.table, + role, + conn, + )?; let columns = rec .columns @@ -367,7 +359,7 @@ fn process_record<'a>( type_name: col.type_, type_oid: col.typeoid, value: col.value.clone(), - is_pkey: pkey_cols(rec).contains(&&col.name), + is_pkey: rec.pkey_cols().contains(&col.name), is_selectable: false, // stub: unused, } }) diff --git a/worker/walrus/src/models/wal2json.rs b/worker/walrus/src/models/wal2json.rs index 60cb8d6..a2140d6 100644 --- a/worker/walrus/src/models/wal2json.rs +++ b/worker/walrus/src/models/wal2json.rs @@ -40,3 +40,16 @@ pub struct Record<'a> { #[serde(with = "crate::timestamp_fmt")] pub timestamp: DateTime, } + +impl<'a> Record<'a> { + pub fn pkey_cols(&self) -> Vec<&'a str> { + match &self.pk { + Some(pkey_refs) => pkey_refs.iter().map(|x| x.name).collect(), + None => vec![], + } + } + + pub fn has_primary_key(&self) -> bool { + self.pkey_cols().len() != 0 + } +} From 359a4a936223d53e026275fef223a58be580e537 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 10:50:53 -0500 Subject: [PATCH 57/88] action casting --- worker/walrus/src/main.rs | 8 +------- worker/walrus/src/models/realtime.rs | 13 ++++++++++++- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 0b3f927..61e5545 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -217,13 +217,7 @@ fn process_record<'a>( ); let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; - - let action = match rec.action { - wal2json::Action::I => realtime::Action::INSERT, - wal2json::Action::U => realtime::Action::UPDATE, - wal2json::Action::D => realtime::Action::DELETE, - wal2json::Action::T => realtime::Action::TRUNCATE, - }; + let action = realtime::Action::from_wal2json(&rec.action); // If the table isn't in the publication or no one is subscribed, do no work if !(is_in_publication && is_subscribed_to && action != realtime::Action::TRUNCATE) { diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index c77595b..f95d6c6 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -1,5 +1,5 @@ +use crate::models::wal2json; use crate::sql::schema::realtime::subscription::dsl::*; -use crate::wal2json; use chrono::{DateTime, NaiveDateTime, Utc}; use diesel::deserialize::{self, FromSql}; use diesel::pg::{Pg, PgValue}; @@ -21,6 +21,17 @@ pub enum Action { TRUNCATE, } +impl Action { + pub fn from_wal2json(action: &wal2json::Action) -> Self { + match action { + wal2json::Action::I => Self::INSERT, + wal2json::Action::U => Self::UPDATE, + wal2json::Action::D => Self::DELETE, + wal2json::Action::T => Self::TRUNCATE, + } + } +} + #[derive(Serialize, Clone, Debug, Eq, PartialEq)] pub struct Column<'a> { pub name: &'a str, From ea26d3a372ad49b2ce4023da889fbad00e9b054f Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 11:00:25 -0500 Subject: [PATCH 58/88] remove redundant filters --- worker/walrus/src/main.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 61e5545..f4aab9e 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -196,6 +196,10 @@ fn process_record<'a>( max_record_bytes: usize, conn: &mut PgConnection, ) -> Result>, String> { + /* + * Table Level Filters + */ + let is_in_publication = filters::table::publication::is_in_publication(&rec.schema, &rec.table, publication, conn)?; @@ -227,8 +231,6 @@ fn process_record<'a>( let subscribed_roles: Vec<&String> = entity_subscriptions .iter() - .filter(|x| &x.schema_name == &rec.schema) - .filter(|x| &x.table_name == &rec.table) .map(|x| &x.claims_role_name) .unique() .collect(); @@ -259,6 +261,10 @@ fn process_record<'a>( } for role in subscribed_roles { + /* + * Role Level Filters + */ + // Subscriptions to current entity + role let entity_role_subscriptions: Vec<&realtime::Subscription> = entity_subscriptions .iter() @@ -285,6 +291,10 @@ fn process_record<'a>( }) .collect(); + /* + * Record Level Filters + */ + let mut record_elem = HashMap::new(); let mut old_record_elem = None; let mut old_record_elem_content = HashMap::new(); From 6caaaf9755e00c93d0cbc89d69efac5fbb36250e Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 6 Jul 2022 15:43:05 -0500 Subject: [PATCH 59/88] end-to-end test --- worker/walrus/Cargo.toml | 4 + .../2022-06-01-154923_helper_functions/up.sql | 12 ++ worker/walrus/src/filters/table.rs | 1 + worker/walrus/src/filters/table/table_oid.rs | 29 +++ worker/walrus/src/main.rs | 168 +++++++++++++++--- worker/walrus/src/models/realtime.rs | 2 +- worker/walrus/src/sql/schema.rs | 2 +- 7 files changed, 194 insertions(+), 24 deletions(-) create mode 100644 worker/walrus/src/filters/table/table_oid.rs diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index 66ac441..1897253 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -16,3 +16,7 @@ env_logger = "0.9.0" itertools = "0.10.3" cached = "0.34.1" chrono = { version = "0.4", features = ["serde"] } + + +[dev-dependencies] +pretty_assertions = "1.2.1" diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 86033dc..0db286c 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -252,6 +252,18 @@ begin end; $$; + +create function realtime.get_table_oid( + schema_name text, + table_name text +) + returns oid + language sql +as $$ + select format('%I.%I', schema_name, table_name)::regclass::oid; +$$; + + alter table realtime.subscription add column schema_name text generated always as (realtime.to_schema_name(entity)) stored; alter table realtime.subscription add column table_name text generated always as (realtime.to_table_name(entity)) stored; alter table realtime.subscription add column claims_role_name text generated always as (realtime.to_regrole((claims ->> 'role'::text))) stored; diff --git a/worker/walrus/src/filters/table.rs b/worker/walrus/src/filters/table.rs index b1f5f83..f121281 100644 --- a/worker/walrus/src/filters/table.rs +++ b/worker/walrus/src/filters/table.rs @@ -1,3 +1,4 @@ pub mod column_security; pub mod publication; pub mod row_level_security; +pub mod table_oid; diff --git a/worker/walrus/src/filters/table/table_oid.rs b/worker/walrus/src/filters/table/table_oid.rs new file mode 100644 index 0000000..43474a3 --- /dev/null +++ b/worker/walrus/src/filters/table/table_oid.rs @@ -0,0 +1,29 @@ +use cached::proc_macro::cached; +use cached::TimedSizedCache; +use diesel::*; + +pub mod sql { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + #[sql_name = "realtime.get_table_oid"] + fn get_table_oid(schema_name: Text, table_name: Text) -> Oid; + } +} + +#[cached( + type = "TimedSizedCache>", + create = "{ TimedSizedCache::with_size_and_lifespan(10000, 1)}", + convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + sync_writes = true +)] +pub fn get_table_oid( + schema_name: &str, + table_name: &str, + conn: &mut PgConnection, +) -> Result { + select(sql::get_table_oid(schema_name, table_name)) + .first(conn) + .map_err(|x| format!("{}", x)) +} diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index f4aab9e..e322517 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -472,17 +472,19 @@ fn process_record<'a>( #[cfg(test)] mod tests { extern crate diesel; + use crate::models::{realtime, wal2json}; use crate::realtime::Subscription; use crate::sql::schema::realtime::subscription::dsl::*; - use crate::wal2json; - use chrono::Utc; + use chrono::{TimeZone, Utc}; use diesel::prelude::*; use diesel::*; + use pretty_assertions::assert_eq; use serde_json::json; + use std::collections::HashMap; use uuid; //const BOOLOID: i32 = 16; - const INT4OID: u32 = 23; + const INTEGER_OID: u32 = 23; //const INT8OID: i32 = 20; //const TEXTOID: i32 = 25; @@ -491,20 +493,126 @@ mod tests { PgConnection::establish(&database_url).unwrap() } - fn clean(conn: &mut PgConnection) { - delete(subscription).execute(conn).unwrap(); + #[test] + fn test_no_one_listening() { + let mut conn = establish_connection(); + + let rec = wal2json::Record { + action: wal2json::Action::I, + schema: "public", + table: "notes", + pk: Some(vec![wal2json::PrimaryKeyRef { + name: "id", + type_: "int4", + typeoid: INTEGER_OID, + }]), + columns: Some(vec![wal2json::Column { + name: "id", + type_: "int4", + typeoid: Some(INTEGER_OID), + value: json!(1), + }]), + identity: None, + timestamp: Utc::now(), + }; + + let res = crate::process_record( + &rec, + &vec![], + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + assert_eq!(res, vec![]); } #[test] - fn test_basic() { + fn test_simple_subscriber() { let mut conn = establish_connection(); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + + diesel::sql_query( + " + DO + $do$ + BEGIN + IF EXISTS ( + SELECT FROM pg_catalog.pg_roles + WHERE rolname = 'authenticated') THEN + RAISE NOTICE 'Role authenticated already exists. Skipping.'; + ELSE + CREATE ROLE authenticated; + END IF; + END + $do$;", + ) + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("create schema if not exists auth;") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query( + " + create or replace function auth.uid() + returns uuid + language 'sql' + AS $$ + select + coalesce( + nullif(current_setting('request.jwt.claim.sub', true), ''), + (nullif(current_setting('request.jwt.claims', true), '')::jsonb ->> 'sub') + )::uuid + $$; + ", + ) + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("grant all on schema auth to authenticated;") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("truncate table realtime.subscription;") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query( + "create table if not exists public.notes(id int primary key, body text);", + ) + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("truncate table public.notes;") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("drop publication if exists supabase_multiplayer;") + .execute(&mut conn) + .unwrap(); + diesel::sql_query("create publication supabase_multiplayer for all tables;") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("insert into public.notes(id, body) values ( 1, 'hello world') ;") + .execute(&mut conn) + .unwrap(); + + let notes_oid = + crate::filters::table::table_oid::get_table_oid("public", "notes", &mut conn).unwrap(); + let claim_sub = uuid::Uuid::new_v4(); + let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); insert_into(subscription) .values(( - subscription_id.eq(uuid::Uuid::new_v4()), - entity.eq(16487), + subscription_id.eq(sub_id), + entity.eq(notes_oid), claims.eq(json!({ "role": "postgres", "email": "example@example.com", @@ -515,14 +623,9 @@ mod tests { .unwrap(); let subscriptions = subscription.load::(&mut conn).unwrap(); - clean(&mut conn); - - assert_eq!(subscriptions.len(), 1); - } - #[test] - fn test_no_one_listening() { - let mut conn = establish_connection(); + let note_id: i32 = 1; + let ts = Utc.timestamp(61, 0); let rec = wal2json::Record { action: wal2json::Action::I, @@ -531,27 +634,48 @@ mod tests { pk: Some(vec![wal2json::PrimaryKeyRef { name: "id", type_: "int4", - typeoid: INT4OID, + typeoid: INTEGER_OID, }]), columns: Some(vec![wal2json::Column { name: "id", - type_: "int4", - typeoid: Some(INT4OID), - value: json!(1), + type_: "integer", + typeoid: Some(INTEGER_OID), + value: json!(note_id), }]), identity: None, - timestamp: Utc::now(), + timestamp: ts, }; let res = crate::process_record( &rec, - &vec![], + &subscriptions, "supabase_multiplayer", 1024 * 1024, &mut conn, ) .unwrap(); - assert_eq!(res, vec![]); + assert!(res.len() == 1); + assert!(res[0].subscription_ids.len() == 1); + + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes", + r#type: realtime::Action::INSERT, + commit_timestamp: &ts, + columns: vec![realtime::Column { + name: "id", + type_: "integer", + }], + record: HashMap::from([("id", json!(note_id))]), + old_record: None, + }, + is_rls_enabled: false, + subscription_ids: vec![sub_id], + errors: vec![], + }]; + + assert_eq!(res, expected); } } diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index f95d6c6..39d1769 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -64,7 +64,7 @@ pub struct WALRLS<'a> { pub struct Subscription { pub id: i64, pub subscription_id: uuid::Uuid, - pub entity: i32, + pub entity: u32, // This also works for anonymous deser of filters (schema.rs also must change) //pub filters: Vec<(String, EqualityOp, String)>, pub filters: Vec, diff --git a/worker/walrus/src/sql/schema.rs b/worker/walrus/src/sql/schema.rs index c2dba16..f64ae3a 100644 --- a/worker/walrus/src/sql/schema.rs +++ b/worker/walrus/src/sql/schema.rs @@ -3,7 +3,7 @@ pub mod realtime { realtime.subscription (id) { id -> Int8, subscription_id -> Uuid, - entity -> Int4, + entity -> Oid, //filters -> Array>, filters -> Array, claims -> Jsonb, From 8d9bf4c6154183a0ad5967cc53697fe579ae17ef Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 7 Jul 2022 07:04:03 -0500 Subject: [PATCH 60/88] refactor nested loop to compute record and old_record --- .../src/filters/record/row_level_security.rs | 2 +- .../walrus/src/filters/record/user_defined.rs | 2 +- worker/walrus/src/main.rs | 296 +++++++++++------- worker/walrus/src/models/walrus.rs | 12 +- 4 files changed, 193 insertions(+), 119 deletions(-) diff --git a/worker/walrus/src/filters/record/row_level_security.rs b/worker/walrus/src/filters/record/row_level_security.rs index a43c741..5821fa8 100644 --- a/worker/walrus/src/filters/record/row_level_security.rs +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -14,7 +14,7 @@ pub mod sql { pub fn is_visible_through_rls( schema_name: &str, table_name: &str, - columns: &Vec, + columns: &Vec, ids: &Vec, conn: &mut PgConnection, ) -> Result, String> { diff --git a/worker/walrus/src/filters/record/user_defined.rs b/worker/walrus/src/filters/record/user_defined.rs index 8969b32..6001daf 100644 --- a/worker/walrus/src/filters/record/user_defined.rs +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -14,7 +14,7 @@ pub mod sql_functions { } pub fn is_visible_through_filters_sql( - columns: &Vec, + columns: &Vec, ids: &Vec, // TODO: convert this to use subscription_ids to reduce n calls conn: &mut PgConnection, diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index e322517..b49c0e4 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -229,6 +229,7 @@ fn process_record<'a>( return Ok(vec![]); } + // Postgres role names of subscribed users let subscribed_roles: Vec<&String> = entity_subscriptions .iter() .map(|x| &x.claims_role_name) @@ -279,26 +280,10 @@ fn process_record<'a>( conn, )?; - let columns = rec - .columns - .as_ref() - .unwrap_or(&vec![]) - .iter() - .filter(|col| selectable_columns.contains(&col.name.to_string())) - .map(|w2j_col| realtime::Column { - name: w2j_col.name, - type_: w2j_col.type_, - }) - .collect(); - /* * Record Level Filters */ - let mut record_elem = HashMap::new(); - let mut old_record_elem = None; - let mut old_record_elem_content = HashMap::new(); - // If the role can not select any columns in the table, return if action != realtime::Action::DELETE && selectable_columns.len() == 0 { let r = realtime::WALRLS { @@ -307,7 +292,7 @@ fn process_record<'a>( table: rec.table, r#type: action.clone(), commit_timestamp: &rec.timestamp, - columns, + columns: vec![], record: HashMap::new(), old_record: None, }, @@ -320,52 +305,58 @@ fn process_record<'a>( }; result.push(r); } else { - if vec![realtime::Action::INSERT, realtime::Action::UPDATE].contains(&action) { - for col_name in &selectable_columns { - 'record: for col in rec.columns.as_ref().unwrap_or(&vec![]) { - if col_name == col.name { - if !exceeds_max_size || col.value.to_string().len() < 64 { - record_elem.insert(col.name, col.value.clone()); - break 'record; - } - } - } - } - } - - if vec![realtime::Action::UPDATE, realtime::Action::DELETE].contains(&action) { - for col_name in &selectable_columns { - match &rec.identity { - Some(identity) => { - 'old_record: for col in identity { - if col_name == col.name { - if !exceeds_max_size || col.value.to_string().len() < 64 { - old_record_elem_content.insert(col.name, col.value.clone()); - break 'old_record; - } - } - } - } - None => (), - } - } - old_record_elem = Some(old_record_elem_content); - } - - let walcols: Vec = rec + // Repr of columns for internal use + let walcols: Vec = rec .columns .as_ref() .unwrap_or(&vec![]) .iter() - .map(|col| { - walrus::WALColumn { - name: col.name, - type_name: col.type_, - type_oid: col.typeoid, - value: col.value.clone(), - is_pkey: rec.pkey_cols().contains(&col.name), - is_selectable: false, // stub: unused, - } + .map(|col| walrus::Column { + name: col.name, + type_name: col.type_, + type_oid: col.typeoid, + value: col.value.clone(), + is_pkey: rec.pkey_cols().contains(&col.name), + is_selectable: selectable_columns.contains(&col.name.to_string()), + }) + .collect(); + + // Populates for realtime::Action::INSERT | realtime::Action::UPDATE + let record_elem: HashMap<&str, serde_json::Value> = walcols + .iter() + // Column must be selectable by role + .filter(|walcol| walcol.is_selectable) + // Filter out large column values if the record exceeds maximum size + .filter(|walcol| !exceeds_max_size || walcol.value.to_string().len() < 64) + .map(|walcol| (walcol.name, walcol.value.clone())) + .collect(); + + // Populates for realtime::Action::UPDATE, realtime::Action::DELETE + let old_record_elem: Option> = + rec.identity.as_ref().map(|cols| { + cols.iter() + .map(|col| walrus::Column { + name: col.name, + type_name: col.type_, + type_oid: col.typeoid, + value: col.value.clone(), + is_pkey: rec.pkey_cols().contains(&col.name), + is_selectable: selectable_columns.contains(&col.name.to_string()), + }) + // Column must be selectable by role + .filter(|walcol| walcol.is_selectable) + // Filter out large column values if the record exceeds maximum size + .filter(|walcol| !exceeds_max_size || walcol.value.to_string().len() < 64) + .map(|walcol| (walcol.name, walcol.value.clone())) + .collect() + }); + + // Repr of columns for external use + let columns: Vec = walcols + .iter() + .map(|walcol| realtime::Column { + name: walcol.name, + type_: walcol.type_name, }) .collect(); @@ -493,47 +484,8 @@ mod tests { PgConnection::establish(&database_url).unwrap() } - #[test] - fn test_no_one_listening() { - let mut conn = establish_connection(); - - let rec = wal2json::Record { - action: wal2json::Action::I, - schema: "public", - table: "notes", - pk: Some(vec![wal2json::PrimaryKeyRef { - name: "id", - type_: "int4", - typeoid: INTEGER_OID, - }]), - columns: Some(vec![wal2json::Column { - name: "id", - type_: "int4", - typeoid: Some(INTEGER_OID), - value: json!(1), - }]), - identity: None, - timestamp: Utc::now(), - }; - - let res = crate::process_record( - &rec, - &vec![], - "supabase_multiplayer", - 1024 * 1024, - &mut conn, - ) - .unwrap(); - - assert_eq!(res, vec![]); - } - - #[test] - fn test_simple_subscriber() { - let mut conn = establish_connection(); - - crate::sql::migrations::run_migrations(&mut conn) - .expect("Pending migrations failed to execute"); + fn setup(conn: &mut PgConnection) { + crate::sql::migrations::run_migrations(conn).expect("Pending migrations failed to execute"); diesel::sql_query( " @@ -550,11 +502,11 @@ mod tests { END $do$;", ) - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query("create schema if not exists auth;") - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query( @@ -571,34 +523,162 @@ mod tests { $$; ", ) - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query("grant all on schema auth to authenticated;") - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query("truncate table realtime.subscription;") - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query( "create table if not exists public.notes(id int primary key, body text);", ) - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query("truncate table public.notes;") - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query("drop publication if exists supabase_multiplayer;") - .execute(&mut conn) + .execute(conn) .unwrap(); diesel::sql_query("create publication supabase_multiplayer for all tables;") + .execute(conn) + .unwrap(); + } + + #[test] + fn test_no_one_listening() { + let mut conn = establish_connection(); + + let rec = wal2json::Record { + action: wal2json::Action::I, + schema: "public", + table: "notes", + pk: Some(vec![wal2json::PrimaryKeyRef { + name: "id", + type_: "int4", + typeoid: INTEGER_OID, + }]), + columns: Some(vec![wal2json::Column { + name: "id", + type_: "int4", + typeoid: Some(INTEGER_OID), + value: json!(1), + }]), + identity: None, + timestamp: Utc::now(), + }; + + let res = crate::process_record( + &rec, + &vec![], + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + assert_eq!(res, vec![]); + } + + #[test] + fn test_simple_insert() { + let mut conn = establish_connection(); + + setup(&mut conn); + + diesel::sql_query("insert into public.notes(id, body) values ( 1, 'hello world') ;") .execute(&mut conn) .unwrap(); + let notes_oid = + crate::filters::table::table_oid::get_table_oid("public", "notes", &mut conn).unwrap(); + + let claim_sub = uuid::Uuid::new_v4(); + let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); + + insert_into(subscription) + .values(( + subscription_id.eq(sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "postgres", + "email": "example@example.com", + "sub": claim_sub + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscriptions = subscription.load::(&mut conn).unwrap(); + + let note_id: i32 = 1; + let ts = Utc.timestamp(61, 0); + + let rec = wal2json::Record { + action: wal2json::Action::I, + schema: "public", + table: "notes", + pk: Some(vec![wal2json::PrimaryKeyRef { + name: "id", + type_: "int4", + typeoid: INTEGER_OID, + }]), + columns: Some(vec![wal2json::Column { + name: "id", + type_: "integer", + typeoid: Some(INTEGER_OID), + value: json!(note_id), + }]), + identity: None, + timestamp: ts, + }; + + let res = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + assert!(res.len() == 1); + assert!(res[0].subscription_ids.len() == 1); + + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes", + r#type: realtime::Action::INSERT, + commit_timestamp: &ts, + columns: vec![realtime::Column { + name: "id", + type_: "integer", + }], + record: HashMap::from([("id", json!(note_id))]), + old_record: None, + }, + is_rls_enabled: false, + subscription_ids: vec![sub_id], + errors: vec![], + }]; + + assert_eq!(res, expected); + } + + #[test] + fn test_simple_update() { + let mut conn = establish_connection(); + + setup(&mut conn); + diesel::sql_query("insert into public.notes(id, body) values ( 1, 'hello world') ;") .execute(&mut conn) .unwrap(); diff --git a/worker/walrus/src/models/walrus.rs b/worker/walrus/src/models/walrus.rs index 5437d2a..7b0f663 100644 --- a/worker/walrus/src/models/walrus.rs +++ b/worker/walrus/src/models/walrus.rs @@ -1,15 +1,9 @@ use serde::Serialize; -#[derive(Serialize)] -pub struct WalrusRecord<'a> { - wal: serde_json::Value, - is_rls_enabled: bool, - subscription_ids: Vec<&'a uuid::Uuid>, - errors: Vec<&'a str>, -} - +/// An internal representation of columns used when passing column data to SQL is required +/// (user defined filters) #[derive(Serialize, Debug)] -pub struct WALColumn<'a> { +pub struct Column<'a> { pub name: &'a str, pub type_name: &'a str, pub type_oid: Option, From 7b93647ec4a4ebec2aada064db0a81f62fa4904e Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 7 Jul 2022 12:35:22 -0500 Subject: [PATCH 61/88] selectable columns refactor. simple_* tests --- .../2022-06-01-154923_helper_functions/up.sql | 20 +- .../src/filters/table/column_security.rs | 31 +- worker/walrus/src/main.rs | 314 ++++++++++++------ worker/walrus/src/models/realtime.rs | 10 +- 4 files changed, 241 insertions(+), 134 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 0db286c..a6b2d87 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -51,31 +51,36 @@ as $$ limit 1 $$; + create function realtime.selectable_columns( - schema_name text, - table_name text, + table_oid oid, role_name text ) - returns text[] + returns jsonb[] language sql as $$ select coalesce( array_agg( - pa.attname::text + jsonb_build_object( + 'name', pa.attname::text, + 'type', pt.typname::text + ) order by pa.attnum asc ), - array['abc'] + array[]::jsonb[] ) from pg_class e join pg_attribute pa on e.oid = pa.attrelid + join pg_type pt + on pa.atttypid = pt.oid where - e.oid = format('%I.%I', $1, $2)::regclass + e.oid = table_oid --format('%I.%I', $1, $2)::regclass and pa.attnum > 0 + and pg_catalog.has_column_privilege(role_name, table_oid, pa.attname, 'SELECT') and not pa.attisdropped - $$; @@ -263,7 +268,6 @@ as $$ select format('%I.%I', schema_name, table_name)::regclass::oid; $$; - alter table realtime.subscription add column schema_name text generated always as (realtime.to_schema_name(entity)) stored; alter table realtime.subscription add column table_name text generated always as (realtime.to_table_name(entity)) stored; alter table realtime.subscription add column claims_role_name text generated always as (realtime.to_regrole((claims ->> 'role'::text))) stored; diff --git a/worker/walrus/src/filters/table/column_security.rs b/worker/walrus/src/filters/table/column_security.rs index ddec726..6c8d37a 100644 --- a/worker/walrus/src/filters/table/column_security.rs +++ b/worker/walrus/src/filters/table/column_security.rs @@ -1,34 +1,37 @@ +use crate::models::realtime; use cached::proc_macro::cached; use cached::TimedSizedCache; use diesel::*; -pub mod sql_functions { +pub mod sql { use diesel::sql_types::*; use diesel::*; sql_function! { #[sql_name = "realtime.selectable_columns"] - fn selectable_columns(schema_name: Text, table_name: Text, role_name: Text) -> Array; + fn selectable_columns(table_oid: Oid, role_name: Text) -> Array; } } #[cached( - type = "TimedSizedCache, String>>", + type = "TimedSizedCache, String>>", create = "{ TimedSizedCache::with_size_and_lifespan(500, 1)}", - convert = r#"{ format!("{}.{}-{}", schema_name, table_name, role_name) }"#, + convert = r#"{ format!("{}-{}", table_oid, role_name) }"#, sync_writes = true )] pub fn selectable_columns( - schema_name: &str, - table_name: &str, + table_oid: u32, role_name: &str, conn: &mut PgConnection, -) -> Result, String> { - select(sql_functions::selectable_columns( - schema_name, - table_name, - role_name, - )) - .first(conn) - .map_err(|x| format!("{}", x)) +) -> Result, String> { + let cols_as_json: Vec = + select(sql::selectable_columns(table_oid, role_name)) + .first(conn) + .map_err(|x| format!("{}", x))?; + + let r: Result, String> = cols_as_json + .into_iter() + .map(|col_json| serde_json::from_value(col_json).map_err(|x| format!("{}", x))) + .collect(); + r } diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index b49c0e4..a876d46 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -220,6 +220,11 @@ fn process_record<'a>( &rec.schema, &rec.table, is_in_publication, is_subscribed_to, is_rls_enabled ); + println!( + "Processing record: {}.{} inpub: {}, entity_subs {}, rls_on {}", + &rec.schema, &rec.table, is_in_publication, is_subscribed_to, is_rls_enabled + ); + let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; let action = realtime::Action::from_wal2json(&rec.action); @@ -273,19 +278,22 @@ fn process_record<'a>( .map(|x| *x) .collect(); - let selectable_columns = filters::table::column_security::selectable_columns( - &rec.schema, - &rec.table, - role, - conn, - )?; + let table_oid = + crate::filters::table::table_oid::get_table_oid(rec.schema, rec.table, conn)?; + + // realtime columns to output with correct type name conversion (e.g. interger -> int4) + // and filtered to columns selectable by the current user + let columns: Vec = + filters::table::column_security::selectable_columns(table_oid, role, conn)?; + + let selectable_column_names: Vec<&str> = columns.iter().map(|x| x.name.as_ref()).collect(); /* * Record Level Filters */ // If the role can not select any columns in the table, return - if action != realtime::Action::DELETE && selectable_columns.len() == 0 { + if action != realtime::Action::DELETE && columns.len() == 0 { let r = realtime::WALRLS { wal: realtime::Data { schema: rec.schema, @@ -310,6 +318,7 @@ fn process_record<'a>( .columns .as_ref() .unwrap_or(&vec![]) + //.unwrap_or(rec.identity.as_ref().unwrap_or(&vec![])) .iter() .map(|col| walrus::Column { name: col.name, @@ -317,7 +326,7 @@ fn process_record<'a>( type_oid: col.typeoid, value: col.value.clone(), is_pkey: rec.pkey_cols().contains(&col.name), - is_selectable: selectable_columns.contains(&col.name.to_string()), + is_selectable: selectable_column_names.contains(&col.name), }) .collect(); @@ -341,7 +350,7 @@ fn process_record<'a>( type_oid: col.typeoid, value: col.value.clone(), is_pkey: rec.pkey_cols().contains(&col.name), - is_selectable: selectable_columns.contains(&col.name.to_string()), + is_selectable: selectable_column_names.contains(&col.name), }) // Column must be selectable by role .filter(|walcol| walcol.is_selectable) @@ -351,15 +360,6 @@ fn process_record<'a>( .collect() }); - // Repr of columns for external use - let columns: Vec = walcols - .iter() - .map(|walcol| realtime::Column { - name: walcol.name, - type_: walcol.type_name, - }) - .collect(); - // User Defined Filters let mut visible_through_filters = vec![]; let mut delegate_to_sql_filters = vec![]; @@ -484,27 +484,28 @@ mod tests { PgConnection::establish(&database_url).unwrap() } - fn setup(conn: &mut PgConnection) { - crate::sql::migrations::run_migrations(conn).expect("Pending migrations failed to execute"); - - diesel::sql_query( + fn create_role(role: &str, conn: &mut PgConnection) { + diesel::sql_query(format!( " DO $do$ BEGIN IF EXISTS ( SELECT FROM pg_catalog.pg_roles - WHERE rolname = 'authenticated') THEN - RAISE NOTICE 'Role authenticated already exists. Skipping.'; + WHERE rolname = '{}') THEN + RAISE NOTICE 'Role {} already exists. Skipping.'; ELSE - CREATE ROLE authenticated; + CREATE ROLE {}; END IF; END $do$;", - ) + role, role, role + )) .execute(conn) .unwrap(); + } + fn create_auth_schema(conn: &mut PgConnection) { diesel::sql_query("create schema if not exists auth;") .execute(conn) .unwrap(); @@ -525,41 +526,61 @@ mod tests { ) .execute(conn) .unwrap(); + } - diesel::sql_query("grant all on schema auth to authenticated;") + fn grant_all_on_schema(schema: &str, role: &str, conn: &mut PgConnection) { + diesel::sql_query(format!("grant all on schema {} to {};", schema, role)) .execute(conn) .unwrap(); + } - diesel::sql_query("truncate table realtime.subscription;") + fn truncate(schema: &str, table: &str, conn: &mut PgConnection) { + diesel::sql_query(format!("truncate table \"{}\".\"{}\";", schema, table)) .execute(conn) .unwrap(); + } - diesel::sql_query( - "create table if not exists public.notes(id int primary key, body text);", - ) + fn drop_table(schema: &str, table: &str, conn: &mut PgConnection) { + diesel::sql_query(format!( + "drop table if exists \"{}\".\"{}\";", + schema, table + )) .execute(conn) .unwrap(); + } - diesel::sql_query("truncate table public.notes;") - .execute(conn) - .unwrap(); + fn create_publication_for_all_tables(publication_name: &str, conn: &mut PgConnection) { + diesel::sql_query(format!( + "drop publication if exists \"{}\";", + publication_name + )) + .execute(conn) + .unwrap(); - diesel::sql_query("drop publication if exists supabase_multiplayer;") - .execute(conn) - .unwrap(); - diesel::sql_query("create publication supabase_multiplayer for all tables;") - .execute(conn) - .unwrap(); + diesel::sql_query(format!( + "create publication \"{}\" for all tables;", + publication_name + )) + .execute(conn) + .unwrap(); } #[test] fn test_no_one_listening() { let mut conn = establish_connection(); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + + drop_table("public", "notes1", &mut conn); + diesel::sql_query("create table if not exists public.notes1(id int primary key);") + .execute(&mut conn) + .unwrap(); + let rec = wal2json::Record { action: wal2json::Action::I, schema: "public", - table: "notes", + table: "notes1", pk: Some(vec![wal2json::PrimaryKeyRef { name: "id", type_: "int4", @@ -591,14 +612,19 @@ mod tests { fn test_simple_insert() { let mut conn = establish_connection(); - setup(&mut conn); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + create_auth_schema(&mut conn); + truncate("realtime", "subscription", &mut conn); + create_publication_for_all_tables("supabase_multiplayer", &mut conn); - diesel::sql_query("insert into public.notes(id, body) values ( 1, 'hello world') ;") + drop_table("public", "notes2", &mut conn); + diesel::sql_query("create table if not exists public.notes2(id int primary key);") .execute(&mut conn) .unwrap(); let notes_oid = - crate::filters::table::table_oid::get_table_oid("public", "notes", &mut conn).unwrap(); + crate::filters::table::table_oid::get_table_oid("public", "notes2", &mut conn).unwrap(); let claim_sub = uuid::Uuid::new_v4(); let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); @@ -619,26 +645,21 @@ mod tests { let subscriptions = subscription.load::(&mut conn).unwrap(); let note_id: i32 = 1; - let ts = Utc.timestamp(61, 0); - let rec = wal2json::Record { - action: wal2json::Action::I, - schema: "public", - table: "notes", - pk: Some(vec![wal2json::PrimaryKeyRef { - name: "id", - type_: "int4", - typeoid: INTEGER_OID, - }]), - columns: Some(vec![wal2json::Column { - name: "id", - type_: "integer", - typeoid: Some(INTEGER_OID), - value: json!(note_id), - }]), - identity: None, - timestamp: ts, - }; + let wal2json_json = r#"{ + "action":"I", + "timestamp":"2022-07-07 14:52:58.092695+00", + "schema":"public", + "table":"notes2", + "columns":[ + {"name":"id","type":"integer","typeoid":23,"value":1} + ], + "pk":[ + {"name":"id","type":"integer","typeoid":23} + ] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_json).unwrap(); let res = crate::process_record( &rec, @@ -649,24 +670,21 @@ mod tests { ) .unwrap(); - assert!(res.len() == 1); - assert!(res[0].subscription_ids.len() == 1); - let expected = vec![realtime::WALRLS { wal: realtime::Data { schema: "public", - table: "notes", + table: "notes2", r#type: realtime::Action::INSERT, - commit_timestamp: &ts, + commit_timestamp: &rec.timestamp, columns: vec![realtime::Column { - name: "id", - type_: "integer", + name: "id".to_string(), + type_: "int4".to_string(), }], record: HashMap::from([("id", json!(note_id))]), old_record: None, }, is_rls_enabled: false, - subscription_ids: vec![sub_id], + subscription_ids: vec![sub_id], // A SUBSCRIBER EXISTS errors: vec![], }]; @@ -676,15 +694,19 @@ mod tests { #[test] fn test_simple_update() { let mut conn = establish_connection(); - - setup(&mut conn); - - diesel::sql_query("insert into public.notes(id, body) values ( 1, 'hello world') ;") + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + create_auth_schema(&mut conn); + truncate("realtime", "subscription", &mut conn); + create_publication_for_all_tables("supabase_multiplayer", &mut conn); + + drop_table("public", "notes4", &mut conn); + diesel::sql_query("create table if not exists public.notes4(id int primary key);") .execute(&mut conn) .unwrap(); let notes_oid = - crate::filters::table::table_oid::get_table_oid("public", "notes", &mut conn).unwrap(); + crate::filters::table::table_oid::get_table_oid("public", "notes4", &mut conn).unwrap(); let claim_sub = uuid::Uuid::new_v4(); let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); @@ -705,28 +727,109 @@ mod tests { let subscriptions = subscription.load::(&mut conn).unwrap(); let note_id: i32 = 1; - let ts = Utc.timestamp(61, 0); + let old_note_id: i32 = 0; + + let wal2json_json = r#"{ + "action":"U", + "timestamp":"2022-07-07 14:52:58.092695+00", + "schema":"public", + "table":"notes4", + "columns":[ + {"name":"id","type":"integer","typeoid":23,"value":1} + ], + "identity":[ + {"name":"id","type":"integer","typeoid":23,"value":0} + ], + "pk":[ + {"name":"id","type":"integer","typeoid":23} + ] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_json).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); - let rec = wal2json::Record { - action: wal2json::Action::I, - schema: "public", - table: "notes", - pk: Some(vec![wal2json::PrimaryKeyRef { - name: "id", - type_: "int4", - typeoid: INTEGER_OID, - }]), - columns: Some(vec![wal2json::Column { - name: "id", - type_: "integer", - typeoid: Some(INTEGER_OID), - value: json!(note_id), - }]), - identity: None, - timestamp: ts, - }; + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes4", + r#type: realtime::Action::UPDATE, + commit_timestamp: &rec.timestamp, + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + record: HashMap::from([("id", json!(note_id))]), + old_record: Some(HashMap::from([("id", json!(old_note_id))])), + }, + is_rls_enabled: false, + subscription_ids: vec![sub_id], + errors: vec![], + }]; - let res = crate::process_record( + assert_eq!(walrus_output, expected); + } + + #[test] + fn test_simple_delete() { + let mut conn = establish_connection(); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + create_auth_schema(&mut conn); + truncate("realtime", "subscription", &mut conn); + create_publication_for_all_tables("supabase_multiplayer", &mut conn); + + drop_table("public", "notes5", &mut conn); + diesel::sql_query("create table if not exists public.notes5(id int primary key);") + .execute(&mut conn) + .unwrap(); + + let notes_oid: u32 = + crate::filters::table::table_oid::get_table_oid("public", "notes5", &mut conn).unwrap(); + + let claim_sub = uuid::Uuid::new_v4(); + let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); + + insert_into(subscription) + .values(( + subscription_id.eq(sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "postgres", + "email": "example@example.com", + "sub": claim_sub + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscriptions = subscription.load::(&mut conn).unwrap(); + + let old_note_id: i32 = 0; + + let wal2json_json = r#"{ + "action":"D", + "timestamp":"2022-07-07 14:52:58.092695+00", + "schema":"public", + "table":"notes5", + "identity":[ + {"name":"id","type":"integer","typeoid":23,"value":0} + ], + "pk":[ + {"name":"id","type":"integer","typeoid":23} + ] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_json).unwrap(); + + let walrus_output = crate::process_record( &rec, &subscriptions, "supabase_multiplayer", @@ -735,27 +838,24 @@ mod tests { ) .unwrap(); - assert!(res.len() == 1); - assert!(res[0].subscription_ids.len() == 1); - let expected = vec![realtime::WALRLS { wal: realtime::Data { schema: "public", - table: "notes", - r#type: realtime::Action::INSERT, - commit_timestamp: &ts, + table: "notes5", + r#type: realtime::Action::DELETE, + commit_timestamp: &rec.timestamp, columns: vec![realtime::Column { - name: "id", - type_: "integer", + name: "id".to_string(), + type_: "int4".to_string(), }], - record: HashMap::from([("id", json!(note_id))]), - old_record: None, + record: HashMap::new(), + old_record: Some(HashMap::from([("id", json!(old_note_id))])), }, is_rls_enabled: false, subscription_ids: vec![sub_id], errors: vec![], }]; - assert_eq!(res, expected); + assert_eq!(walrus_output, expected); } } diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index 39d1769..3c6964c 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -32,11 +32,11 @@ impl Action { } } -#[derive(Serialize, Clone, Debug, Eq, PartialEq)] -pub struct Column<'a> { - pub name: &'a str, +#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)] +pub struct Column { + pub name: String, #[serde(rename(serialize = "type", deserialize = "type"))] - pub type_: &'a str, + pub type_: String, } #[derive(Serialize, Clone, Debug, Eq, PartialEq)] @@ -46,7 +46,7 @@ pub struct Data<'a> { pub r#type: Action, #[serde(with = "crate::timestamp_fmt")] pub commit_timestamp: &'a DateTime, - pub columns: Vec>, + pub columns: Vec, pub record: HashMap<&'a str, serde_json::Value>, pub old_record: Option>, } From d12e8ce0c157a4d9f8847830365407b3debe35e3 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 7 Jul 2022 12:57:58 -0500 Subject: [PATCH 62/88] test unauthorized --- worker/walrus/src/main.rs | 89 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 88 insertions(+), 1 deletion(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index a876d46..73d9273 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -293,7 +293,7 @@ fn process_record<'a>( */ // If the role can not select any columns in the table, return - if action != realtime::Action::DELETE && columns.len() == 0 { + if columns.len() == 0 { let r = realtime::WALRLS { wal: realtime::Data { schema: rec.schema, @@ -858,4 +858,91 @@ mod tests { assert_eq!(walrus_output, expected); } + + #[test] + fn test_error_unauthorized() { + let mut conn = establish_connection(); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + create_auth_schema(&mut conn); + truncate("realtime", "subscription", &mut conn); + create_publication_for_all_tables("supabase_multiplayer", &mut conn); + + drop_table("public", "notes6", &mut conn); + diesel::sql_query("create table if not exists public.notes6(id int primary key);") + .execute(&mut conn) + .unwrap(); + + create_role("authenticated", &mut conn); + + diesel::sql_query("create table if not exists public.notes6(id int primary key);") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("revoke select on public.notes6 from authenticated;") + .execute(&mut conn) + .unwrap(); + + let notes_oid: u32 = + crate::filters::table::table_oid::get_table_oid("public", "notes6", &mut conn).unwrap(); + + let claim_sub = uuid::Uuid::new_v4(); + let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); + + insert_into(subscription) + .values(( + subscription_id.eq(sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "authenticated", + "email": "example@example.com", + "sub": claim_sub + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscriptions = subscription.load::(&mut conn).unwrap(); + + let wal2json_json = r#"{ + "action":"D", + "timestamp":"2022-07-07 14:52:58.092695+00", + "schema":"public", + "table":"notes6", + "identity":[ + {"name":"id","type":"integer","typeoid":23,"value":0} + ], + "pk":[ + {"name":"id","type":"integer","typeoid":23} + ] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_json).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes6", + r#type: realtime::Action::DELETE, + commit_timestamp: &rec.timestamp, + columns: vec![], + record: HashMap::new(), + old_record: None, + }, + is_rls_enabled: false, + subscription_ids: vec![sub_id], + errors: vec!["Error 401: Unauthorized"], + }]; + + assert_eq!(walrus_output, expected); + } } From 0926e1f5f022fd73859bd75f0f99ea846262acb0 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 7 Jul 2022 15:33:30 -0500 Subject: [PATCH 63/88] test quoted tables, schemas, types --- .../2022-06-01-154923_helper_functions/up.sql | 2 +- worker/walrus/src/main.rs | 117 +++++++++++++++++- 2 files changed, 117 insertions(+), 2 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index a6b2d87..4ab0562 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -124,7 +124,7 @@ $$ case when maybe_quoted_name like '"%"' then substring( maybe_quoted_name, - 1, + 2, character_length(maybe_quoted_name)-2 ) else maybe_quoted_name diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 73d9273..01a1edf 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -529,9 +529,16 @@ mod tests { } fn grant_all_on_schema(schema: &str, role: &str, conn: &mut PgConnection) { - diesel::sql_query(format!("grant all on schema {} to {};", schema, role)) + diesel::sql_query(format!("grant all on schema \"{}\" to {};", schema, role)) .execute(conn) .unwrap(); + + diesel::sql_query(format!( + "grant select on all tables in schema \"{}\" to {};", + schema, role + )) + .execute(conn) + .unwrap(); } fn truncate(schema: &str, table: &str, conn: &mut PgConnection) { @@ -945,4 +952,112 @@ mod tests { assert_eq!(walrus_output, expected); } + + #[test] + fn test_quoted_type_schema_and_table() { + // TODO: Enable RLS to make sure that works + // TODO: Add user defined filter to make sure delegate to sql works + + let mut conn = establish_connection(); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + create_auth_schema(&mut conn); + truncate("realtime", "subscription", &mut conn); + create_publication_for_all_tables("supabase_multiplayer", &mut conn); + + diesel::sql_query("create schema if not exists \"dEv\";") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("drop type if exists \"Color\" cascade;") + .execute(&mut conn) + .unwrap(); + diesel::sql_query("create type \"Color\" as enum ('RED', 'YELLOW', 'GREEN');") + .execute(&mut conn) + .unwrap(); + + drop_table("dEv", "Notes7", &mut conn); + diesel::sql_query( + "create table if not exists \"dEv\".\"Notes7\"(id \"Color\" primary key);", + ) + .execute(&mut conn) + .unwrap(); + + use diesel::dsl::sql; + + let type_oid = + sql::("select oid from pg_type where typname = 'Color' limit 1") + .get_result::(&mut conn) + .unwrap(); + + create_role("authenticated", &mut conn); + grant_all_on_schema("dEv", "authenticated", &mut conn); + + let notes_oid: u32 = + crate::filters::table::table_oid::get_table_oid("dEv", "Notes7", &mut conn).unwrap(); + + let claim_sub = uuid::Uuid::new_v4(); + let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); + + insert_into(subscription) + .values(( + subscription_id.eq(sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "authenticated", + "email": "example@example.com", + "sub": claim_sub + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscriptions = subscription.load::(&mut conn).unwrap(); + + let wal2json_json = format!( + r#"{{ + "action":"I", + "timestamp":"2022-07-07 14:52:58.092695+00", + "schema":"dEv", + "table":"Notes7", + "columns":[ + {{"name":"id","type":"Color","typeoid":{type_oid},"value": "YELLOW"}} + ], + "pk":[ + {{"name":"id","type":"Color","typeoid":{type_oid}}} + ] + }}"# + ); + + let rec: wal2json::Record = serde_json::from_str(wal2json_json.as_ref()).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "dEv", + table: "Notes7", + r#type: realtime::Action::INSERT, + commit_timestamp: &rec.timestamp, + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "Color".to_string(), + }], + record: HashMap::from([("id", json!("YELLOW"))]), + old_record: None, + }, + is_rls_enabled: false, + subscription_ids: vec![sub_id], + errors: vec![], + }]; + + assert_eq!(walrus_output, expected); + } } From b3c1b0905abb73ae856a09d3fb0486dc9fd6dedc Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 8 Jul 2022 05:34:19 -0500 Subject: [PATCH 64/88] subscription filter speedup --- .../down.sql | 1 + .../up.sql | 62 +++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/down.sql create mode 100644 worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/up.sql diff --git a/worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/down.sql b/worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/down.sql new file mode 100644 index 0000000..291a97c --- /dev/null +++ b/worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/down.sql @@ -0,0 +1 @@ +-- This file should undo anything in `up.sql` \ No newline at end of file diff --git a/worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/up.sql b/worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/up.sql new file mode 100644 index 0000000..e704fd5 --- /dev/null +++ b/worker/walrus/migrations/2022-07-08-100736_subscriptions filter perf/up.sql @@ -0,0 +1,62 @@ +create or replace function realtime.subscription_check_filters() + returns trigger + language plpgsql +as $$ +/* +Validates that the user defined filters for a subscription: +- refer to valid columns that the claimed role may access +- values are coercable to the correct column type +*/ +declare + + col_names text[] = coalesce( + array_agg(pa.attname order by pa.attnum asc), + '{}'::text[] + ) + from + pg_catalog.pg_attribute pa + where + pa.attrelid = new.entity + and pa.attnum > 0 + and not pa.attisdropped + and pg_catalog.has_column_privilege( + (new.claims ->> 'role'), + new.entity, + pa.attname, + 'SELECT' + ); + filter realtime.user_defined_filter; + col_type regtype; +begin + for filter in select * from unnest(new.filters) loop + -- Filtered column is valid + if not filter.column_name = any(col_names) then + raise exception 'invalid column for filter %', filter.column_name; + end if; + + -- Type is sanitized and safe for string interpolation + col_type = ( + select atttypid::regtype + from pg_catalog.pg_attribute + where attrelid = new.entity + and attname = filter.column_name + ); + if col_type is null then + raise exception 'failed to lookup type for column %', filter.column_name; + end if; + -- raises an exception if value is not coercable to type + perform realtime.cast(filter.value, col_type); + end loop; + + -- Apply consistent order to filters so the unique constraint on + -- (subscription_id, entity, filters) can't be tricked by a different filter order + new.filters = coalesce( + array_agg(f order by f.column_name, f.op, f.value), + '{}' + ) from unnest(new.filters) f; + + return new; +end; +$$; + + From d8b8f1c8fb0b4a835556caed1a9ad757276521fd Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 8 Jul 2022 05:50:25 -0500 Subject: [PATCH 65/88] sql functions to table_oid where practical --- .../2022-06-01-154923_helper_functions/up.sql | 31 ++++--------------- .../src/filters/record/row_level_security.rs | 8 ++--- .../src/filters/table/row_level_security.rs | 12 +++---- worker/walrus/src/main.rs | 12 +++---- 4 files changed, 18 insertions(+), 45 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 4ab0562..0d38637 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -11,18 +11,16 @@ as $$ select 1 from - pg_publication pp - left join pg_publication_tables ppt - on pp.pubname = ppt.pubname + pg_publication_tables ppt where - pp.pubname = publication_name + ppt.pubname = publication_name and ppt.schemaname = schema_name and ppt.tablename = table_name limit 1 ) $$; -create function realtime.is_rls_enabled(schema_name text, table_name text) +create function realtime.is_rls_enabled(table_oid oid) returns bool language sql as $$ @@ -31,26 +29,10 @@ as $$ from pg_class where - oid = format('%I.%I', schema_name, table_name)::regclass + oid = table_oid limit 1; $$; -create function realtime.subscribed_roles( - schema_name text, - table_name text -) - returns text[] - language sql -as $$ - select - coalesce(array_agg(distinct claims_role), '{}') - from - realtime.subscription s - where - s.entity = format('%I.%I', schema_name, table_name)::regclass - limit 1 -$$; - create function realtime.selectable_columns( table_oid oid, @@ -191,8 +173,7 @@ $$; create function realtime.is_visible_through_rls( - schema_name text, - table_name text, + table_oid oid, columns jsonb, ids int8[] ) @@ -200,7 +181,7 @@ create function realtime.is_visible_through_rls( language plpgsql as $$ declare - entity_ regclass = format('%I.%I', schema_name, table_name)::regclass; + entity_ regclass = table_oid::regclass; cols realtime.wal_column[]; visible_to_subscription_ids int8[] = '{}'; subscription_id int8; diff --git a/worker/walrus/src/filters/record/row_level_security.rs b/worker/walrus/src/filters/record/row_level_security.rs index 5821fa8..dbafad3 100644 --- a/worker/walrus/src/filters/record/row_level_security.rs +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -7,20 +7,18 @@ pub mod sql { sql_function! { #[sql_name = "realtime.is_visible_through_rls"] - fn is_visible_through_rls(schema_name: Text, table_name: Text, columns: Jsonb, ids: Array) -> Array + fn is_visible_through_rls(table_oid: Oid, columns: Jsonb, ids: Array) -> Array } } pub fn is_visible_through_rls( - schema_name: &str, - table_name: &str, + table_oid: u32, columns: &Vec, ids: &Vec, conn: &mut PgConnection, ) -> Result, String> { select(sql::is_visible_through_rls( - schema_name, - table_name, + table_oid, serde_json::to_value(columns).unwrap(), ids, )) diff --git a/worker/walrus/src/filters/table/row_level_security.rs b/worker/walrus/src/filters/table/row_level_security.rs index 6ca6e18..278892b 100644 --- a/worker/walrus/src/filters/table/row_level_security.rs +++ b/worker/walrus/src/filters/table/row_level_security.rs @@ -8,22 +8,18 @@ pub mod sql { sql_function! { #[sql_name = "realtime.is_rls_enabled"] - fn is_rls_enabled(schema_name: Text, table_name: Text) -> Bool; + fn is_rls_enabled(table_oid: Oid) -> Bool; } } #[cached( type = "TimedSizedCache>", create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", - convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, + convert = r#"{ format!("{}", table_oid) }"#, sync_writes = true )] -pub fn is_rls_enabled( - schema_name: &str, - table_name: &str, - conn: &mut PgConnection, -) -> Result { - select(sql::is_rls_enabled(schema_name, table_name)) +pub fn is_rls_enabled(table_oid: u32, conn: &mut PgConnection) -> Result { + select(sql::is_rls_enabled(table_oid)) .first(conn) .map_err(|x| format!("{}", x)) } diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 01a1edf..f4a1696 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -200,6 +200,9 @@ fn process_record<'a>( * Table Level Filters */ + // Will not be necessary after replacing wal2json + let table_oid = crate::filters::table::table_oid::get_table_oid(rec.schema, rec.table, conn)?; + let is_in_publication = filters::table::publication::is_in_publication(&rec.schema, &rec.table, publication, conn)?; @@ -212,8 +215,7 @@ fn process_record<'a>( .collect(); let is_subscribed_to = entity_subscriptions.len() > 0; - let is_rls_enabled = - filters::table::row_level_security::is_rls_enabled(&rec.schema, &rec.table, conn)?; + let is_rls_enabled = filters::table::row_level_security::is_rls_enabled(table_oid, conn)?; debug!( "Processing record: {}.{} inpub: {}, entity_subs {}, rls_on {}", @@ -278,9 +280,6 @@ fn process_record<'a>( .map(|x| *x) .collect(); - let table_oid = - crate::filters::table::table_oid::get_table_oid(rec.schema, rec.table, conn)?; - // realtime columns to output with correct type name conversion (e.g. interger -> int4) // and filtered to columns selectable by the current user let columns: Vec = @@ -413,8 +412,7 @@ fn process_record<'a>( false => visible_through_filters, true => { match filters::record::row_level_security::is_visible_through_rls( - &rec.schema, - &rec.table, + table_oid, &walcols, &visible_through_filters.iter().map(|x| x.id).collect(), conn, From 9a5b2f3a3da6fa5810f3f36c8d9eff388bf672c5 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 8 Jul 2022 06:47:33 -0500 Subject: [PATCH 66/88] explicit error enum --- worker/walrus/src/errors.rs | 52 +++++++++++++++++++ .../src/filters/record/row_level_security.rs | 5 +- .../walrus/src/filters/record/user_defined.rs | 25 ++++++--- .../src/filters/table/column_security.rs | 14 +++-- .../walrus/src/filters/table/publication.rs | 7 +-- .../src/filters/table/row_level_security.rs | 7 +-- worker/walrus/src/filters/table/table_oid.rs | 7 +-- worker/walrus/src/main.rs | 22 ++++---- 8 files changed, 108 insertions(+), 31 deletions(-) create mode 100644 worker/walrus/src/errors.rs diff --git a/worker/walrus/src/errors.rs b/worker/walrus/src/errors.rs new file mode 100644 index 0000000..be74b84 --- /dev/null +++ b/worker/walrus/src/errors.rs @@ -0,0 +1,52 @@ +use std::fmt; + +#[derive(Debug, Clone)] +pub enum Error { + // can't reach the database + PostgresConnectionError(String), + + // pg_recvlogical subprocess error + PgRecvLogicalError(String), + + // issue loading subscriptions + Subscriptions(String), + + // generic business logic error + Walrus(String), + + // Unexpected result from SQL function + SQLFunction(String), +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let msg = match self { + Self::PostgresConnectionError(x) => x, + Self::PgRecvLogicalError(x) => x, + Self::Subscriptions(x) => x, + Self::Walrus(x) => x, + Self::SQLFunction(x) => x, + }; + + write!(f, "{}", msg) + } +} + +pub enum FilterError { + // unrecoverable. trash the subscription + FilterParsing(String), + + // Filter too difficult to handle locally. Delgate it to SQL + DelegateToSQL(String), +} + +impl fmt::Display for FilterError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let msg = match self { + Self::FilterParsing(x) => x, + Self::DelegateToSQL(x) => x, + }; + + write!(f, "{}", msg) + } +} diff --git a/worker/walrus/src/filters/record/row_level_security.rs b/worker/walrus/src/filters/record/row_level_security.rs index dbafad3..6d2d8dd 100644 --- a/worker/walrus/src/filters/record/row_level_security.rs +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -1,3 +1,4 @@ +use crate::errors; use crate::models::walrus; use diesel::*; @@ -16,12 +17,12 @@ pub fn is_visible_through_rls( columns: &Vec, ids: &Vec, conn: &mut PgConnection, -) -> Result, String> { +) -> Result, errors::Error> { select(sql::is_visible_through_rls( table_oid, serde_json::to_value(columns).unwrap(), ids, )) .first(conn) - .map_err(|x| format!("{}", x)) + .map_err(|x| errors::Error::SQLFunction(format!("{}", x))) } diff --git a/worker/walrus/src/filters/record/user_defined.rs b/worker/walrus/src/filters/record/user_defined.rs index 6001daf..8e36789 100644 --- a/worker/walrus/src/filters/record/user_defined.rs +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -1,3 +1,4 @@ +use crate::errors; use crate::models::walrus; use crate::models::{realtime, wal2json}; use diesel::*; @@ -34,13 +35,13 @@ fn is_null(v: &serde_json::Value) -> bool { pub fn visible_through_filters( filters: &Vec, columns: &Vec, -) -> Result { +) -> Result { use realtime::Op; for filter in filters { let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { Ok(v) => v, - Err(err) => return Err(format!("{}", err)), + Err(err) => return Err(errors::FilterError::FilterParsing(format!("{}", err))), }; let column = match columns @@ -90,9 +91,17 @@ pub fn visible_through_filters( return Ok(false); }; } - _ => return Err("could not handle filter op for allowed types".to_string()), + _ => { + return Err(errors::FilterError::DelegateToSQL( + "could not handle filter op for allowed types".to_string(), + )) + } }, - _ => return Err("Could not handle type. Delegate comparison to SQL".to_string()), + _ => { + return Err(errors::FilterError::DelegateToSQL( + "Could not handle type. Delegate comparison to SQL".to_string(), + )) + } }; } Ok(true) @@ -102,7 +111,7 @@ pub fn visible_through_filters( fn get_valid_ops( a: &serde_json::Value, b: &serde_json::Value, -) -> Result, String> { +) -> Result, errors::FilterError> { use serde_json::Value; match (a, b) { @@ -112,7 +121,11 @@ fn get_valid_ops( (Value::String(a_), Value::String(b_)) => Ok(get_matching_ops(a_, b_)), // Array possible // Object possible - _ => return Err("non-scalar or mismatched json value types".to_string()), + _ => { + return Err(errors::FilterError::DelegateToSQL( + "non-scalar or mismatched json value types".to_string(), + )); + } } } diff --git a/worker/walrus/src/filters/table/column_security.rs b/worker/walrus/src/filters/table/column_security.rs index 6c8d37a..231d5ee 100644 --- a/worker/walrus/src/filters/table/column_security.rs +++ b/worker/walrus/src/filters/table/column_security.rs @@ -1,3 +1,4 @@ +use crate::errors; use crate::models::realtime; use cached::proc_macro::cached; use cached::TimedSizedCache; @@ -14,7 +15,7 @@ pub mod sql { } #[cached( - type = "TimedSizedCache, String>>", + type = "TimedSizedCache, errors::Error>>", create = "{ TimedSizedCache::with_size_and_lifespan(500, 1)}", convert = r#"{ format!("{}-{}", table_oid, role_name) }"#, sync_writes = true @@ -23,15 +24,18 @@ pub fn selectable_columns( table_oid: u32, role_name: &str, conn: &mut PgConnection, -) -> Result, String> { +) -> Result, errors::Error> { let cols_as_json: Vec = select(sql::selectable_columns(table_oid, role_name)) .first(conn) - .map_err(|x| format!("{}", x))?; + .map_err(|x| errors::Error::SQLFunction(format!("{}", x)))?; - let r: Result, String> = cols_as_json + let r: Result, errors::Error> = cols_as_json .into_iter() - .map(|col_json| serde_json::from_value(col_json).map_err(|x| format!("{}", x))) + .map(|col_json| { + serde_json::from_value(col_json) + .map_err(|x| errors::Error::SQLFunction(format!("{}", x))) + }) .collect(); r } diff --git a/worker/walrus/src/filters/table/publication.rs b/worker/walrus/src/filters/table/publication.rs index 190ed4e..1ad2d3b 100644 --- a/worker/walrus/src/filters/table/publication.rs +++ b/worker/walrus/src/filters/table/publication.rs @@ -1,3 +1,4 @@ +use crate::errors; use cached::proc_macro::cached; use cached::TimedSizedCache; use diesel::*; @@ -13,7 +14,7 @@ pub mod sql { } #[cached( - type = "TimedSizedCache>", + type = "TimedSizedCache>", create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", convert = r#"{ format!("{}.{}-{}", schema_name, table_name, publication_name) }"#, sync_writes = true @@ -23,12 +24,12 @@ pub fn is_in_publication( table_name: &str, publication_name: &str, conn: &mut PgConnection, -) -> Result { +) -> Result { select(sql::is_in_publication( schema_name, table_name, publication_name, )) .first(conn) - .map_err(|x| format!("{}", x)) + .map_err(|x| errors::Error::SQLFunction(format!("{}", x))) } diff --git a/worker/walrus/src/filters/table/row_level_security.rs b/worker/walrus/src/filters/table/row_level_security.rs index 278892b..2ec3d4d 100644 --- a/worker/walrus/src/filters/table/row_level_security.rs +++ b/worker/walrus/src/filters/table/row_level_security.rs @@ -1,3 +1,4 @@ +use crate::errors; use cached::proc_macro::cached; use cached::TimedSizedCache; use diesel::*; @@ -13,13 +14,13 @@ pub mod sql { } #[cached( - type = "TimedSizedCache>", + type = "TimedSizedCache>", create = "{ TimedSizedCache::with_size_and_lifespan(250, 1)}", convert = r#"{ format!("{}", table_oid) }"#, sync_writes = true )] -pub fn is_rls_enabled(table_oid: u32, conn: &mut PgConnection) -> Result { +pub fn is_rls_enabled(table_oid: u32, conn: &mut PgConnection) -> Result { select(sql::is_rls_enabled(table_oid)) .first(conn) - .map_err(|x| format!("{}", x)) + .map_err(|x| errors::Error::SQLFunction(format!("{}", x))) } diff --git a/worker/walrus/src/filters/table/table_oid.rs b/worker/walrus/src/filters/table/table_oid.rs index 43474a3..ced5580 100644 --- a/worker/walrus/src/filters/table/table_oid.rs +++ b/worker/walrus/src/filters/table/table_oid.rs @@ -1,3 +1,4 @@ +use crate::errors; use cached::proc_macro::cached; use cached::TimedSizedCache; use diesel::*; @@ -13,7 +14,7 @@ pub mod sql { } #[cached( - type = "TimedSizedCache>", + type = "TimedSizedCache>", create = "{ TimedSizedCache::with_size_and_lifespan(10000, 1)}", convert = r#"{ format!("{}.{}", schema_name, table_name) }"#, sync_writes = true @@ -22,8 +23,8 @@ pub fn get_table_oid( schema_name: &str, table_name: &str, conn: &mut PgConnection, -) -> Result { +) -> Result { select(sql::get_table_oid(schema_name, table_name)) .first(conn) - .map_err(|x| format!("{}", x)) + .map_err(|x| errors::Error::SQLFunction(format!("{}", x))) } diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index f4a1696..cd045a3 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -15,6 +15,7 @@ use std::process::{Command, Stdio}; use std::thread::sleep; use std::time; +mod errors; mod filters; mod models; mod sql; @@ -63,15 +64,15 @@ fn main() { } } -fn run(args: &Args) -> Result<(), String> { +fn run(args: &Args) -> Result<(), errors::Error> { // Connect to Postgres let conn_result = &mut PgConnection::establish(&args.connection); let publication = &args.publication; let conn = match conn_result { Ok(c) => c, - Err(_) => { - return Err("failed to make postgres connection".to_string()); + Err(err) => { + return Err(errors::Error::PostgresConnectionError(format!("{}", err))); } }; @@ -107,7 +108,7 @@ fn run(args: &Args) -> Result<(), String> { .spawn(); match cmd { - Err(err) => Err(format!("{}", err)), + Err(err) => Err(errors::Error::PgRecvLogicalError(format!("{}", err))), Ok(mut cmd) => { info!("pg_recvlogical started"); // Reading from stdin @@ -122,7 +123,7 @@ fn run(args: &Args) -> Result<(), String> { Err(err) => { cmd.kill().unwrap(); error!("Error loading subscriptions: {}", err); - return Err("Error loading subscriptions".to_string()); + return Err(errors::Error::Subscriptions(format!("{}", err))); } }; info!("Snapshot of subscriptions loaded"); @@ -171,7 +172,7 @@ fn run(args: &Args) -> Result<(), String> { Err(err) => { cmd.kill().unwrap(); error!("WALRUS Error: {}", err); - return Err("walrus error".to_string()); + return Err(errors::Error::Walrus(format!("{}", err))); } } } @@ -183,7 +184,7 @@ fn run(args: &Args) -> Result<(), String> { } match cmd.wait() { Ok(_) => Ok(()), - Err(err) => Err(format!("{}", err)), + Err(err) => Err(errors::Error::PgRecvLogicalError(format!("{}", err))), } } } @@ -195,7 +196,7 @@ fn process_record<'a>( publication: &str, max_record_bytes: usize, conn: &mut PgConnection, -) -> Result>, String> { +) -> Result>, errors::Error> { /* * Table Level Filters */ @@ -374,13 +375,16 @@ fn process_record<'a>( } Ok(false) => (), // delegate to SQL when we can't handle the comparison in rust - Err(_) => { + Err(errors::FilterError::DelegateToSQL(_)) => { //debug!( // "Filters delegated to SQL: {:?}. Error: {}", // &sub.filters, err //); delegate_to_sql_filters.push(sub); } + Err(err) => { + error!("Failed to apply filters: {}", err) + } } } From fced27d596469672ccc47db5fa4bbd3e2efe34a4 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 8 Jul 2022 08:18:38 -0500 Subject: [PATCH 67/88] core rust-side filter logic tests --- worker/walrus/src/errors.rs | 1 + .../walrus/src/filters/record/user_defined.rs | 57 ++++++++++++++++++- worker/walrus/src/models/realtime.rs | 4 +- worker/walrus/src/sql/types.rs | 1 + 4 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 worker/walrus/src/sql/types.rs diff --git a/worker/walrus/src/errors.rs b/worker/walrus/src/errors.rs index be74b84..668a508 100644 --- a/worker/walrus/src/errors.rs +++ b/worker/walrus/src/errors.rs @@ -32,6 +32,7 @@ impl fmt::Display for Error { } } +#[derive(Debug)] pub enum FilterError { // unrecoverable. trash the subscription FilterParsing(String), diff --git a/worker/walrus/src/filters/record/user_defined.rs b/worker/walrus/src/filters/record/user_defined.rs index 8e36789..750144c 100644 --- a/worker/walrus/src/filters/record/user_defined.rs +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -68,8 +68,12 @@ pub fn visible_through_filters( match column.type_.as_ref() { // All operations supported + // regtype names (provided by wal2json) "boolean" | "smallint" | "integer" | "bigint" | "serial" | "bigserial" | "numeric" - | "double precision" | "character" | "character varying" | "text" => match &filter.op { + | "double precision" | "character" | "character varying" | "text" + // Type 2 (from pg_type.typname) + | "bool" | "char" | "int2" | "int4" | "int8" | "float4" | "float8" | "varchar" + => match &filter.op { // List is currently exhastive but that may be possible if types/ops expand Op::Equal | Op::NotEqual @@ -154,5 +158,56 @@ where if a >= b { ops.push(Op::GreaterThanOrEqual); }; + ops.sort(); ops } + +#[cfg(test)] +mod tests { + + use crate::filters::record::user_defined::get_valid_ops; + use crate::models::realtime::Op; + use serde_json::json; + + #[test] + fn test_get_valid_ops_eq() { + let a: i32 = 1; + let a = json!(a); + + let eq_ops = get_valid_ops(&a, &a).unwrap(); + assert_eq!( + eq_ops, + vec![Op::Equal, Op::LessThanOrEqual, Op::GreaterThanOrEqual] + ); + } + + #[test] + fn test_get_valid_ops_lt() { + let a: i32 = 1; + let a = json!(a); + + let b: i32 = 2; + let b = json!(b); + + let eq_ops = get_valid_ops(&a, &b).unwrap(); + assert_eq!( + eq_ops, + vec![Op::NotEqual, Op::LessThan, Op::LessThanOrEqual] + ); + } + + #[test] + fn test_get_valid_ops_gt() { + let a: i32 = 2; + let a = json!(a); + + let b: i32 = 1; + let b = json!(b); + + let eq_ops = get_valid_ops(&a, &b).unwrap(); + assert_eq!( + eq_ops, + vec![Op::NotEqual, Op::GreaterThan, Op::GreaterThanOrEqual] + ); + } +} diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index 3c6964c..87728d5 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -186,7 +186,9 @@ pub fn update_subscriptions( #[diesel(postgres_type(schema = "realtime", name = "equality_op"))] pub struct OpType; -#[derive(Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq)] +#[derive( + Debug, PartialEq, FromSqlRow, AsExpression, Clone, Deserialize, Serialize, Eq, Ord, PartialOrd, +)] #[diesel(sql_type = OpType)] pub enum Op { #[serde(alias = "eq")] diff --git a/worker/walrus/src/sql/types.rs b/worker/walrus/src/sql/types.rs new file mode 100644 index 0000000..2024c71 --- /dev/null +++ b/worker/walrus/src/sql/types.rs @@ -0,0 +1 @@ +pub fn oid_to_name(type_oid: u32) -> String {} From d4651c1fa0d877dc9001ef12135756aee448dec1 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 8 Jul 2022 09:28:29 -0500 Subject: [PATCH 68/88] rls on quoted test --- .../2022-06-01-154923_helper_functions/up.sql | 1 + .../src/filters/record/row_level_security.rs | 2 + worker/walrus/src/main.rs | 56 ++++++++++++++----- worker/walrus/src/models/realtime.rs | 2 +- 4 files changed, 45 insertions(+), 16 deletions(-) diff --git a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql index 0d38637..69e15b6 100644 --- a/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -208,6 +208,7 @@ begin if (select 1 from pg_prepared_statements where name = 'walrus_rls_stmt' limit 1) > 0 then deallocate walrus_rls_stmt; end if; + execute realtime.build_prepared_statement_sql('walrus_rls_stmt', entity_, cols); for subscription_id, claims in ( diff --git a/worker/walrus/src/filters/record/row_level_security.rs b/worker/walrus/src/filters/record/row_level_security.rs index 6d2d8dd..f46c8dd 100644 --- a/worker/walrus/src/filters/record/row_level_security.rs +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -18,6 +18,8 @@ pub fn is_visible_through_rls( ids: &Vec, conn: &mut PgConnection, ) -> Result, errors::Error> { + println!("columns {:?}", columns); + println!("ids {:?}", ids); select(sql::is_visible_through_rls( table_oid, serde_json::to_value(columns).unwrap(), diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index cd045a3..375ca1d 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -408,13 +408,15 @@ fn process_record<'a>( } } + println!("through filters {:?}", visible_through_filters); + // Row Level Security - let subscriptions_to_notify: Vec<&realtime::Subscription> = match is_rls_enabled - && visible_through_filters.len() > 0 - && !vec![realtime::Action::DELETE, realtime::Action::TRUNCATE].contains(&action) - { - false => visible_through_filters, - true => { + let subscriptions_to_notify: Vec<&realtime::Subscription> = match ( + is_rls_enabled && visible_through_filters.len() > 0, + vec![realtime::Action::DELETE, realtime::Action::TRUNCATE].contains(&action), + ) { + (false, _) | (true, true) => visible_through_filters, + _ => { match filters::record::row_level_security::is_visible_through_rls( table_oid, &walcols, @@ -434,6 +436,8 @@ fn process_record<'a>( } }; + println!("notify {:?}", subscriptions_to_notify); + let r = realtime::WALRLS { wal: realtime::Data { schema: rec.schema, @@ -466,7 +470,7 @@ fn process_record<'a>( mod tests { extern crate diesel; use crate::models::{realtime, wal2json}; - use crate::realtime::Subscription; + use crate::realtime::{Subscription, UserDefinedFilter}; use crate::sql::schema::realtime::subscription::dsl::*; use chrono::{TimeZone, Utc}; use diesel::prelude::*; @@ -957,7 +961,6 @@ mod tests { #[test] fn test_quoted_type_schema_and_table() { - // TODO: Enable RLS to make sure that works // TODO: Add user defined filter to make sure delegate to sql works let mut conn = establish_connection(); @@ -974,6 +977,7 @@ mod tests { diesel::sql_query("drop type if exists \"Color\" cascade;") .execute(&mut conn) .unwrap(); + diesel::sql_query("create type \"Color\" as enum ('RED', 'YELLOW', 'GREEN');") .execute(&mut conn) .unwrap(); @@ -985,16 +989,28 @@ mod tests { .execute(&mut conn) .unwrap(); - use diesel::dsl::sql; - - let type_oid = - sql::("select oid from pg_type where typname = 'Color' limit 1") - .get_result::(&mut conn) - .unwrap(); + let type_oid = diesel::dsl::sql::( + "select oid from pg_type where typname = 'Color' limit 1", + ) + .get_result::(&mut conn) + .unwrap(); create_role("authenticated", &mut conn); grant_all_on_schema("dEv", "authenticated", &mut conn); + diesel::sql_query( + "create policy rls_note_select + on \"dEv\".\"Notes7\" + to authenticated + using (true);", + ) + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("alter table \"dEv\".\"Notes7\" enable row level security;") + .execute(&mut conn) + .unwrap(); + let notes_oid: u32 = crate::filters::table::table_oid::get_table_oid("dEv", "Notes7", &mut conn).unwrap(); @@ -1010,12 +1026,22 @@ mod tests { "email": "example@example.com", "sub": claim_sub })), + /*filters.eq(vec![UserDefinedFilter { + column_name: "id".to_string(), + op: realtime::Op::Equal, + value: "YELLOW".to_string(), + }]), + */ )) .execute(&mut conn) .unwrap(); let subscriptions = subscription.load::(&mut conn).unwrap(); + diesel::sql_query("insert into \"dEv\".\"Notes7\"(id) values ('YELLOW');") + .execute(&mut conn) + .unwrap(); + let wal2json_json = format!( r#"{{ "action":"I", @@ -1055,7 +1081,7 @@ mod tests { record: HashMap::from([("id", json!("YELLOW"))]), old_record: None, }, - is_rls_enabled: false, + is_rls_enabled: true, subscription_ids: vec![sub_id], errors: vec![], }]; diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index 87728d5..e96cb66 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -233,7 +233,7 @@ impl FromSql for Op { } } -#[derive(SqlType, PartialEq)] +#[derive(SqlType, PartialEq, QueryId)] #[diesel(postgres_type(schema = "realtime", name = "user_defined_filter"))] pub struct UserDefinedFilterType; From 733cccc24e15252f0c0003474f123a3b812d61ef Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Fri, 8 Jul 2022 11:48:16 -0500 Subject: [PATCH 69/88] filter in quoted user defined type test --- worker/walrus/src/errors.rs | 4 - .../walrus/src/filters/record/user_defined.rs | 3 +- worker/walrus/src/main.rs | 10 +- worker/walrus/src/models/walrus.rs | 122 ++++++++++++++++++ 4 files changed, 129 insertions(+), 10 deletions(-) diff --git a/worker/walrus/src/errors.rs b/worker/walrus/src/errors.rs index 668a508..8b8fa21 100644 --- a/worker/walrus/src/errors.rs +++ b/worker/walrus/src/errors.rs @@ -34,9 +34,6 @@ impl fmt::Display for Error { #[derive(Debug)] pub enum FilterError { - // unrecoverable. trash the subscription - FilterParsing(String), - // Filter too difficult to handle locally. Delgate it to SQL DelegateToSQL(String), } @@ -44,7 +41,6 @@ pub enum FilterError { impl fmt::Display for FilterError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let msg = match self { - Self::FilterParsing(x) => x, Self::DelegateToSQL(x) => x, }; diff --git a/worker/walrus/src/filters/record/user_defined.rs b/worker/walrus/src/filters/record/user_defined.rs index 750144c..238830a 100644 --- a/worker/walrus/src/filters/record/user_defined.rs +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -41,7 +41,8 @@ pub fn visible_through_filters( for filter in filters { let filter_value: serde_json::Value = match serde_json::from_str(&filter.value) { Ok(v) => v, - Err(err) => return Err(errors::FilterError::FilterParsing(format!("{}", err))), + // Composite types are not parsable as json + Err(err) => return Err(errors::FilterError::DelegateToSQL(format!("{}", err))), }; let column = match columns diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 375ca1d..50123cb 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -365,6 +365,7 @@ fn process_record<'a>( let mut delegate_to_sql_filters = vec![]; for sub in entity_role_subscriptions { + println!("sub: {:?}", sub); match filters::record::user_defined::visible_through_filters( &sub.filters, rec.columns.as_ref().unwrap_or(&vec![]), @@ -974,17 +975,17 @@ mod tests { .execute(&mut conn) .unwrap(); - diesel::sql_query("drop type if exists \"Color\" cascade;") + diesel::sql_query("drop type if exists \"dEv\".\"Color\" cascade;") .execute(&mut conn) .unwrap(); - diesel::sql_query("create type \"Color\" as enum ('RED', 'YELLOW', 'GREEN');") + diesel::sql_query("create type \"dEv\".\"Color\" as enum ('RED', 'YELLOW', 'GREEN');") .execute(&mut conn) .unwrap(); drop_table("dEv", "Notes7", &mut conn); diesel::sql_query( - "create table if not exists \"dEv\".\"Notes7\"(id \"Color\" primary key);", + "create table if not exists \"dEv\".\"Notes7\"(id \"dEv\".\"Color\" primary key);", ) .execute(&mut conn) .unwrap(); @@ -1026,12 +1027,11 @@ mod tests { "email": "example@example.com", "sub": claim_sub })), - /*filters.eq(vec![UserDefinedFilter { + filters.eq(vec![UserDefinedFilter { column_name: "id".to_string(), op: realtime::Op::Equal, value: "YELLOW".to_string(), }]), - */ )) .execute(&mut conn) .unwrap(); diff --git a/worker/walrus/src/models/walrus.rs b/worker/walrus/src/models/walrus.rs index 7b0f663..75cccb3 100644 --- a/worker/walrus/src/models/walrus.rs +++ b/worker/walrus/src/models/walrus.rs @@ -11,3 +11,125 @@ pub struct Column<'a> { pub is_pkey: bool, pub is_selectable: bool, } + +#[cfg(test)] +mod tests { + extern crate diesel; + use crate::models::realtime; + use chrono::{DateTime, Utc}; + use pretty_assertions::assert_eq; + use serde_json::json; + use std::collections::HashMap; + + #[test] + fn test_serialize_insert() { + let id: i32 = 2; + let sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); + let record = realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes", + r#type: realtime::Action::INSERT, + + commit_timestamp: &DateTime::parse_from_rfc3339("2020-04-12T22:10:57.002456+02:00") + .unwrap() + .with_timezone(&Utc), + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + record: HashMap::from([("id", json!(id))]), + old_record: None, + }, + is_rls_enabled: true, + subscription_ids: vec![sub_id, sub_id], + errors: vec!["sample error"], + }; + + let ser = serde_json::to_string_pretty(&record).unwrap(); + + assert_eq!( + ser, + r#"{ + "wal": { + "schema": "public", + "table": "notes", + "type": "INSERT", + "commit_timestamp": "2020-04-12T20:10:57.002456Z", + "columns": [ + { + "name": "id", + "type": "int4" + } + ], + "record": { + "id": 2 + }, + "old_record": null + }, + "is_rls_enabled": true, + "subscription_ids": [ + "37c7e506-9eca-4671-8c48-526d404660ce", + "37c7e506-9eca-4671-8c48-526d404660ce" + ], + "errors": [ + "sample error" + ] +}"# + ) + } + + #[test] + fn test_serialize_delete() { + let id: i32 = 2; + let record = realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes6", + r#type: realtime::Action::DELETE, + + commit_timestamp: &DateTime::parse_from_rfc3339("2020-04-12T22:10:57.002456+02:00") + .unwrap() + .with_timezone(&Utc), + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + record: HashMap::new(), + old_record: Some(HashMap::from([("id", json!(id))])), + }, + is_rls_enabled: false, + subscription_ids: vec![], + errors: vec!["sample error"], + }; + + let ser = serde_json::to_string_pretty(&record).unwrap(); + + assert_eq!( + ser, + r#"{ + "wal": { + "schema": "public", + "table": "notes6", + "type": "DELETE", + "commit_timestamp": "2020-04-12T20:10:57.002456Z", + "columns": [ + { + "name": "id", + "type": "int4" + } + ], + "record": {}, + "old_record": { + "id": 2 + } + }, + "is_rls_enabled": false, + "subscription_ids": [], + "errors": [ + "sample error" + ] +}"# + ) + } +} From 9e22f4ff959caa742a045a4ed1e7a176aa206d44 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 11 Jul 2022 11:52:24 -0500 Subject: [PATCH 70/88] reconnecting ws --- worker/realtime/Cargo.toml | 4 + worker/realtime/src/main.rs | 177 +++++++++++++++++++++++++++++++----- 2 files changed, 158 insertions(+), 23 deletions(-) diff --git a/worker/realtime/Cargo.toml b/worker/realtime/Cargo.toml index b63d157..6d1644c 100644 --- a/worker/realtime/Cargo.toml +++ b/worker/realtime/Cargo.toml @@ -12,9 +12,13 @@ uuid = { version = "1.0", features = ["serde"] } tokio = { version = "1", features = ["full"] } tokio-util = { version="0.7.2", features=["codec"] } tokio-tungstenite = { verison = "0.17.1", features=["native-tls"] } +tungstenite = "0.17.1" futures-util = "0.3.21" url = "2.2.2" futures-channel = "0.3.21" futures = "0.3.21" log = "0.4.17" env_logger = "0.9.0" +stream-reconnect = { version = "0.3", default-features = true } +stubborn-io = "0.3" + diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 171e8a7..9805bfe 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -1,7 +1,8 @@ use clap::Parser; use env_logger; use futures::stream::SplitSink; -use futures_util::SinkExt; +use futures::{Sink, SinkExt, Stream}; +//use futures_util::SinkExt; use futures_util::{future, pin_mut, StreamExt}; use log::{error, info, warn}; use serde::Serialize; @@ -70,7 +71,6 @@ async fn main() { // url let args = Args::parse(); let addr = build_url(&args.url, &args.header); - let url = url::Url::parse(&addr).expect("invalid URL"); let topic = args.topic.to_string(); info!("{:?}", args); @@ -78,10 +78,20 @@ async fn main() { // enable logger env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); - loop { - // websocket - info!("Connecting to websocket"); - let ws_connection = connect_async(&url).await; + info!("Connecting to websocket"); + let mut ws_stream: ReconnectWs = ReconnectWs::connect(addr).await.unwrap(); + + ws_stream + .send(Message::text(String::from("hello world!"))) + .await + .unwrap(); + + // websocket + //let ws_connection: String = connect_async(&url).await; + //Result<(WebSocketStream>, Response<()>), tokio_tungstenite::tungstenite::Error> + //let mut tcp_stream = StubbornTcpStream::connect(addr).await.unwrap(); + + /* match ws_connection { Err(msg) => { @@ -92,33 +102,42 @@ async fn main() { } Ok((ws_stream, _)) => { info!("WebSocket handshake successful"); - let (mut write, read) = ws_stream.split(); - write = join_topic(write, topic.to_string()).await; + let (mut write, read) = ws_stream.split(); + join_topic(&mut write, topic.to_string()).await; // Futures channel let (tx, rx) = futures_channel::mpsc::unbounded(); let heartbeat_tx = tx.clone(); - tokio::spawn(read_stdin(tx, topic.to_string())); + tokio::spawn(heartbeat(heartbeat_tx)); // Map let tx_to_ws = rx.map(Ok).forward(write); + let ws_to_stdout = { read.for_each(|message| async { match message { Ok(msg) => match msg.into_text() { Ok(msg_text) => info!("{}", msg_text), - Err(err) => error!( - "Failed to parse message from realtime service: Error: {}", - err - ), + Err(err) => { + error!( + "Failed to parse message from realtime service: Error: {}", + err + ); + // Attempt to rejoin + //write = join_topic(write, topic.to_string()).await; + } }, - Err(err) => error!( - "Failed to read message from realtime service: Error: {}", - err - ), + Err(err) => { + error!( + "Failed to read message from realtime service: Error: {}", + err + ); + // Attempt to rejoin + //write = join_topic(write, topic.to_string()).await; + } }; }) }; @@ -127,7 +146,7 @@ async fn main() { future::select(tx_to_ws, ws_to_stdout).await; } } - } + */ } pub fn build_url(url: &str, params: &Vec<(String, String)>) -> String { @@ -140,12 +159,14 @@ pub fn build_url(url: &str, params: &Vec<(String, String)>) -> String { } async fn join_topic( - mut writer: SplitSink< + writer: &mut SplitSink< WebSocketStream>, Message, >, topic_name: String, -) -> SplitSink>, Message> { +) +//-> SplitSink>, Message> { +{ // join channel let join_message = PhoenixMessage { event: PhoenixMessageEvent::Join, @@ -157,7 +178,26 @@ async fn join_topic( let join_message = serde_json::to_string(&join_message).unwrap(); let msg = Message::Text(join_message); writer.send(msg).await.unwrap(); - writer + //writer +} + +fn join_topic_msg(tx: futures_channel::mpsc::UnboundedSender, topic_name: String) { + // join channel + let join_message = PhoenixMessage { + event: PhoenixMessageEvent::Join, + payload: serde_json::json!({}), + reference: None, + topic: topic_name, + }; + + let msg = Message::Text(serde_json::to_string(&join_message).unwrap()); + // push to output stream + match tx.unbounded_send(msg) { + Ok(()) => (), + Err(err) => { + error!("Error sending message: {}", err); + } + }; } // Our helper method which will read data from stdin and send it along the @@ -188,7 +228,12 @@ async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); // push to output stream - tx.unbounded_send(msg).unwrap(); + match tx.unbounded_send(msg) { + Ok(()) => (), + Err(err) => { + error!("Error sending message: {}", err); + } + }; } Err(err) => { warn!( @@ -224,7 +269,93 @@ async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender) { // push to futures stream match tx.unbounded_send(msg) { Ok(()) => (), - Err(err) => error!("Error sending heatbeat: {}", err), + Err(err) => { + error!("Error sending heatbeat: {}", err); + } }; } } + +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; +use stream_reconnect::{ReconnectStream, UnderlyingStream}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::error::Error as WsError; +use tokio_tungstenite::MaybeTlsStream; + +struct MyWs(WebSocketStream>); + +impl Stream for MyWs { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_next(cx) + } +} + +impl Sink for MyWs { + type Error = WsError; + + fn poll_ready( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.0).poll_ready(_cx) + } + + fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + Pin::new(&mut self.0).start_send(item) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_flush(cx) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.0).poll_close(cx) + } +} + +// implement Stream & Sink for MyWs + +impl UnderlyingStream, WsError> for MyWs { + // Establishes connection. + // Additionally, this will be used when reconnect tries are attempted. + fn establish(addr: String) -> Pin> + Send>> { + Box::pin(async move { + // In this case, we are trying to connect to the WebSocket endpoint + let ws_connection = connect_async(addr).await.unwrap().0; + Ok(MyWs(ws_connection)) + }) + } + + // The following errors are considered disconnect errors. + fn is_write_disconnect_error(&self, err: &WsError) -> bool { + matches!( + err, + WsError::ConnectionClosed + | WsError::AlreadyClosed + | WsError::Io(_) + | WsError::Tls(_) + | WsError::Protocol(_) + ) + } + + // If an `Err` is read, then there might be an disconnection. + fn is_read_disconnect_error(&self, item: &Result) -> bool { + if let Err(e) = item { + self.is_write_disconnect_error(e) + } else { + false + } + } + + // Return "Exhausted" if all retry attempts are failed. + fn exhaust_err() -> WsError { + WsError::Io(io::Error::new(io::ErrorKind::Other, "Exhausted")) + } +} + +type ReconnectWs = ReconnectStream, WsError>; From 29dfadbd118d4ecf8e31fe686de5cb279885377c Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 11 Jul 2022 12:06:12 -0500 Subject: [PATCH 71/88] reconnecting ws --- worker/realtime/src/main.rs | 149 ++++++++++++++---------------------- 1 file changed, 58 insertions(+), 91 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index 9805bfe..e32da55 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -79,74 +79,56 @@ async fn main() { env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); info!("Connecting to websocket"); - let mut ws_stream: ReconnectWs = ReconnectWs::connect(addr).await.unwrap(); - - ws_stream - .send(Message::text(String::from("hello world!"))) - .await - .unwrap(); + let ws_stream: ReconnectWs = ReconnectWs::connect(addr).await.unwrap(); // websocket //let ws_connection: String = connect_async(&url).await; //Result<(WebSocketStream>, Response<()>), tokio_tungstenite::tungstenite::Error> //let mut tcp_stream = StubbornTcpStream::connect(addr).await.unwrap(); - /* - - match ws_connection { - Err(msg) => { - let n_seconds = 3; - error!("Failed to connect to websocket. Error: {}", msg); - info!("Attempting websocket reconnect in {} seconds", n_seconds); - sleep(Duration::from_secs(n_seconds)).await; - } - Ok((ws_stream, _)) => { - info!("WebSocket handshake successful"); - - let (mut write, read) = ws_stream.split(); - join_topic(&mut write, topic.to_string()).await; - - // Futures channel - let (tx, rx) = futures_channel::mpsc::unbounded(); - let heartbeat_tx = tx.clone(); - tokio::spawn(read_stdin(tx, topic.to_string())); - - tokio::spawn(heartbeat(heartbeat_tx)); - - // Map - let tx_to_ws = rx.map(Ok).forward(write); - - let ws_to_stdout = { - read.for_each(|message| async { - match message { - Ok(msg) => match msg.into_text() { - Ok(msg_text) => info!("{}", msg_text), - Err(err) => { - error!( - "Failed to parse message from realtime service: Error: {}", - err - ); - // Attempt to rejoin - //write = join_topic(write, topic.to_string()).await; - } - }, - Err(err) => { - error!( - "Failed to read message from realtime service: Error: {}", - err - ); - // Attempt to rejoin - //write = join_topic(write, topic.to_string()).await; - } - }; - }) - }; + info!("WebSocket handshake successful"); + + let (mut write, read) = ws_stream.split(); + join_topic(&mut write, topic.to_string()).await; + + // Futures channel + let (tx, rx) = futures_channel::mpsc::unbounded(); + let heartbeat_tx = tx.clone(); + tokio::spawn(read_stdin(tx, topic.to_string())); + + tokio::spawn(heartbeat(heartbeat_tx)); + + // Map + let tx_to_ws = rx.map(Ok).forward(write); + + let ws_to_stdout = { + read.for_each(|message| async { + match message { + Ok(msg) => match msg.into_text() { + Ok(msg_text) => info!("{}", msg_text), + Err(err) => { + error!( + "Failed to parse message from realtime service: Error: {}", + err + ); + // Attempt to rejoin + //write = join_topic(write, topic.to_string()).await; + } + }, + Err(err) => { + error!( + "Failed to read message from realtime service: Error: {}", + err + ); + // Attempt to rejoin + //write = join_topic(write, topic.to_string()).await; + } + }; + }) + }; - pin_mut!(tx_to_ws, ws_to_stdout); - future::select(tx_to_ws, ws_to_stdout).await; - } - } - */ + pin_mut!(tx_to_ws, ws_to_stdout); + future::select(tx_to_ws, ws_to_stdout).await; } pub fn build_url(url: &str, params: &Vec<(String, String)>) -> String { @@ -160,13 +142,16 @@ pub fn build_url(url: &str, params: &Vec<(String, String)>) -> String { async fn join_topic( writer: &mut SplitSink< - WebSocketStream>, + ReconnectStream< + MyWs, + String, + Result, + tokio_tungstenite::tungstenite::Error, + >, Message, >, topic_name: String, -) -//-> SplitSink>, Message> { -{ +) { // join channel let join_message = PhoenixMessage { event: PhoenixMessageEvent::Join, @@ -178,26 +163,6 @@ async fn join_topic( let join_message = serde_json::to_string(&join_message).unwrap(); let msg = Message::Text(join_message); writer.send(msg).await.unwrap(); - //writer -} - -fn join_topic_msg(tx: futures_channel::mpsc::UnboundedSender, topic_name: String) { - // join channel - let join_message = PhoenixMessage { - event: PhoenixMessageEvent::Join, - payload: serde_json::json!({}), - reference: None, - topic: topic_name, - }; - - let msg = Message::Text(serde_json::to_string(&join_message).unwrap()); - // push to output stream - match tx.unbounded_send(msg) { - Ok(()) => (), - Err(err) => { - error!("Error sending message: {}", err); - } - }; } // Our helper method which will read data from stdin and send it along the @@ -285,9 +250,11 @@ use tokio::net::TcpStream; use tokio_tungstenite::tungstenite::error::Error as WsError; use tokio_tungstenite::MaybeTlsStream; -struct MyWs(WebSocketStream>); +// A websocket to communicate with a Phoenix server +// It reconnects and re-subscribes to a topic if disconnected +struct PhoenixWs(WebSocketStream>); -impl Stream for MyWs { +impl Stream for PhoenixWs { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { @@ -295,7 +262,7 @@ impl Stream for MyWs { } } -impl Sink for MyWs { +impl Sink for PhoenixWs { type Error = WsError; fn poll_ready( @@ -320,14 +287,14 @@ impl Sink for MyWs { // implement Stream & Sink for MyWs -impl UnderlyingStream, WsError> for MyWs { +impl UnderlyingStream, WsError> for PhoenixWs { // Establishes connection. // Additionally, this will be used when reconnect tries are attempted. fn establish(addr: String) -> Pin> + Send>> { Box::pin(async move { // In this case, we are trying to connect to the WebSocket endpoint let ws_connection = connect_async(addr).await.unwrap().0; - Ok(MyWs(ws_connection)) + Ok(PhoenixWs(ws_connection)) }) } @@ -358,4 +325,4 @@ impl UnderlyingStream, WsError> for MyWs { } } -type ReconnectWs = ReconnectStream, WsError>; +type ReconnectWs = ReconnectStream, WsError>; From c6e7c8a4d570381b1fe7a07feef3a5b663be8162 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 11 Jul 2022 12:22:32 -0500 Subject: [PATCH 72/88] reconnects function --- worker/realtime/Cargo.toml | 2 -- worker/realtime/src/main.rs | 46 +++++++++++-------------------------- 2 files changed, 13 insertions(+), 35 deletions(-) diff --git a/worker/realtime/Cargo.toml b/worker/realtime/Cargo.toml index 6d1644c..950b6a2 100644 --- a/worker/realtime/Cargo.toml +++ b/worker/realtime/Cargo.toml @@ -20,5 +20,3 @@ futures = "0.3.21" log = "0.4.17" env_logger = "0.9.0" stream-reconnect = { version = "0.3", default-features = true } -stubborn-io = "0.3" - diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index e32da55..d9b7abf 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -81,15 +81,9 @@ async fn main() { info!("Connecting to websocket"); let ws_stream: ReconnectWs = ReconnectWs::connect(addr).await.unwrap(); - // websocket - //let ws_connection: String = connect_async(&url).await; - //Result<(WebSocketStream>, Response<()>), tokio_tungstenite::tungstenite::Error> - //let mut tcp_stream = StubbornTcpStream::connect(addr).await.unwrap(); - info!("WebSocket handshake successful"); - let (mut write, read) = ws_stream.split(); - join_topic(&mut write, topic.to_string()).await; + let (write, read) = ws_stream.split(); // Futures channel let (tx, rx) = futures_channel::mpsc::unbounded(); @@ -140,31 +134,6 @@ pub fn build_url(url: &str, params: &Vec<(String, String)>) -> String { addr } -async fn join_topic( - writer: &mut SplitSink< - ReconnectStream< - MyWs, - String, - Result, - tokio_tungstenite::tungstenite::Error, - >, - Message, - >, - topic_name: String, -) { - // join channel - let join_message = PhoenixMessage { - event: PhoenixMessageEvent::Join, - payload: serde_json::json!({}), - reference: None, - topic: topic_name, - }; - - let join_message = serde_json::to_string(&join_message).unwrap(); - let msg = Message::Text(join_message); - writer.send(msg).await.unwrap(); -} - // Our helper method which will read data from stdin and send it along the // sender provided. async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: String) { @@ -293,7 +262,18 @@ impl UnderlyingStream, WsError> for PhoenixWs { fn establish(addr: String) -> Pin> + Send>> { Box::pin(async move { // In this case, we are trying to connect to the WebSocket endpoint - let ws_connection = connect_async(addr).await.unwrap().0; + let mut ws_connection = connect_async(addr).await.unwrap().0; + + // (re)Join Topic + let join_message = PhoenixMessage { + event: PhoenixMessageEvent::Join, + payload: serde_json::json!({}), + reference: None, + topic: "walrus:bcd".to_string(), + }; + let join_message = serde_json::to_string(&join_message).unwrap(); + let msg = Message::Text(join_message); + ws_connection.send(msg).await.unwrap(); Ok(PhoenixWs(ws_connection)) }) } From 01ef7ab5758da16a71229cae585c47533e48a569 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 11 Jul 2022 13:37:42 -0500 Subject: [PATCH 73/88] e2e reconnect handle --- worker/realtime/src/main.rs | 45 +++++++++++++++++++++++-------------- worker/walrus/src/main.rs | 14 +++--------- 2 files changed, 31 insertions(+), 28 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index d9b7abf..a6a07a2 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -1,8 +1,6 @@ use clap::Parser; use env_logger; -use futures::stream::SplitSink; use futures::{Sink, SinkExt, Stream}; -//use futures_util::SinkExt; use futures_util::{future, pin_mut, StreamExt}; use log::{error, info, warn}; use serde::Serialize; @@ -11,7 +9,6 @@ use std::str; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::time::{sleep, Duration}; use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream}; -use url; /// reads JSON from stdin and forwards it to supabase realtime #[derive(Parser, Debug)] @@ -70,17 +67,20 @@ struct PhoenixMessage { async fn main() { // url let args = Args::parse(); - let addr = build_url(&args.url, &args.header); - let topic = args.topic.to_string(); + + let config = PhoenixWsConfig { + addr: build_url(&args.url, &args.header), + topic_name: args.topic.to_string(), + }; + + let topic = config.topic_name.to_string(); info!("{:?}", args); // enable logger - env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init(); - - info!("Connecting to websocket"); - let ws_stream: ReconnectWs = ReconnectWs::connect(addr).await.unwrap(); + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("debug")).init(); + let ws_stream: ReconnectWs = ReconnectWs::connect(config).await.unwrap(); info!("WebSocket handshake successful"); let (write, read) = ws_stream.split(); @@ -221,7 +221,13 @@ use tokio_tungstenite::MaybeTlsStream; // A websocket to communicate with a Phoenix server // It reconnects and re-subscribes to a topic if disconnected -struct PhoenixWs(WebSocketStream>); +struct PhoenixWs(WebSocketStream>, PhoenixWsConfig); + +#[derive(Clone)] +struct PhoenixWsConfig { + addr: String, + topic_name: String, +} impl Stream for PhoenixWs { type Item = Result; @@ -256,25 +262,30 @@ impl Sink for PhoenixWs { // implement Stream & Sink for MyWs -impl UnderlyingStream, WsError> for PhoenixWs { +impl UnderlyingStream, WsError> for PhoenixWs { // Establishes connection. // Additionally, this will be used when reconnect tries are attempted. - fn establish(addr: String) -> Pin> + Send>> { + fn establish( + config: PhoenixWsConfig, + ) -> Pin> + Send>> { Box::pin(async move { // In this case, we are trying to connect to the WebSocket endpoint - let mut ws_connection = connect_async(addr).await.unwrap().0; + info!("Connecting to Realtime"); + + let mut ws_connection = connect_async(config.addr.clone()).await?.0; // (re)Join Topic let join_message = PhoenixMessage { event: PhoenixMessageEvent::Join, payload: serde_json::json!({}), reference: None, - topic: "walrus:bcd".to_string(), + topic: config.topic_name.to_string(), }; let join_message = serde_json::to_string(&join_message).unwrap(); let msg = Message::Text(join_message); - ws_connection.send(msg).await.unwrap(); - Ok(PhoenixWs(ws_connection)) + ws_connection.send(msg).await?; + info!("Joining topic"); + Ok(PhoenixWs(ws_connection, config.clone())) }) } @@ -305,4 +316,4 @@ impl UnderlyingStream, WsError> for PhoenixWs { } } -type ReconnectWs = ReconnectStream, WsError>; +type ReconnectWs = ReconnectStream, WsError>; diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 50123cb..a610d02 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -223,11 +223,6 @@ fn process_record<'a>( &rec.schema, &rec.table, is_in_publication, is_subscribed_to, is_rls_enabled ); - println!( - "Processing record: {}.{} inpub: {}, entity_subs {}, rls_on {}", - &rec.schema, &rec.table, is_in_publication, is_subscribed_to, is_rls_enabled - ); - let exceeds_max_size = serde_json::json!(rec).to_string().len() > max_record_bytes; let action = realtime::Action::from_wal2json(&rec.action); @@ -365,7 +360,7 @@ fn process_record<'a>( let mut delegate_to_sql_filters = vec![]; for sub in entity_role_subscriptions { - println!("sub: {:?}", sub); + //println!("sub: {:?}", sub); match filters::record::user_defined::visible_through_filters( &sub.filters, rec.columns.as_ref().unwrap_or(&vec![]), @@ -383,9 +378,6 @@ fn process_record<'a>( //); delegate_to_sql_filters.push(sub); } - Err(err) => { - error!("Failed to apply filters: {}", err) - } } } @@ -409,7 +401,7 @@ fn process_record<'a>( } } - println!("through filters {:?}", visible_through_filters); + //println!("through filters {:?}", visible_through_filters); // Row Level Security let subscriptions_to_notify: Vec<&realtime::Subscription> = match ( @@ -437,7 +429,7 @@ fn process_record<'a>( } }; - println!("notify {:?}", subscriptions_to_notify); + //println!("notify {:?}", subscriptions_to_notify); let r = realtime::WALRLS { wal: realtime::Data { From 554dc8fdd9a8ebc2eb70f104f7f81be760ed8ea4 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 11 Jul 2022 13:45:13 -0500 Subject: [PATCH 74/88] display messages in stream logical --- worker/stream_logical/src/main.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/worker/stream_logical/src/main.rs b/worker/stream_logical/src/main.rs index 64137a3..660c16a 100644 --- a/worker/stream_logical/src/main.rs +++ b/worker/stream_logical/src/main.rs @@ -125,6 +125,7 @@ async fn main() { let mut last_keepalive = Instant::now(); let mut inserts: Vec = vec![]; let mut deletes: Vec = vec![]; + let mut lsn: u64; // TODO(olirice) track table defs mapping // TODO(olirice) send keepalive every 20 seconds @@ -136,27 +137,28 @@ async fn main() { let msg_res = match msg { Some(Ok(XLogData(xlog_data))) => match xlog_data.data() { - Begin(_) => { - println!("begin"); + Begin(begin) => { + println!("{:?}", begin) } Insert(insert) => { - println!("insert"); + println!("{:?}", insert) } Update(update) => { - println!("update"); + println!("{:?}", update) } Delete(delete) => { - println!("delete"); + println!("{:?}", delete) } Commit(commit) => { - println!("commit"); + println!("{:?}", commit); + lsn = commit.commit_lsn(); } Relation(relation) => { - println!("relation"); + println!("{:?}", relation) } Origin(_) | Type(_) => {} Truncate(truncate) => { - println!("truncate"); + println!("{:?}", truncate) } _ => println!("unknown logical replication message type"), }, From b993104965178dfdbdd76977f9669cbe045c6272 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Mon, 11 Jul 2022 15:21:52 -0500 Subject: [PATCH 75/88] rm filter printlines --- worker/walrus/src/filters/record/row_level_security.rs | 2 -- worker/walrus/src/main.rs | 6 ------ 2 files changed, 8 deletions(-) diff --git a/worker/walrus/src/filters/record/row_level_security.rs b/worker/walrus/src/filters/record/row_level_security.rs index f46c8dd..6d2d8dd 100644 --- a/worker/walrus/src/filters/record/row_level_security.rs +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -18,8 +18,6 @@ pub fn is_visible_through_rls( ids: &Vec, conn: &mut PgConnection, ) -> Result, errors::Error> { - println!("columns {:?}", columns); - println!("ids {:?}", ids); select(sql::is_visible_through_rls( table_oid, serde_json::to_value(columns).unwrap(), diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index a610d02..da1308d 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -135,7 +135,6 @@ fn run(args: &Args) -> Result<(), errors::Error> { let result_record = serde_json::from_str::(&line); match result_record { Ok(wal2json_record) => { - //println!("rec {:?}", wal2json_record); // Update subscriptions if needed realtime::update_subscriptions( &wal2json_record, @@ -360,7 +359,6 @@ fn process_record<'a>( let mut delegate_to_sql_filters = vec![]; for sub in entity_role_subscriptions { - //println!("sub: {:?}", sub); match filters::record::user_defined::visible_through_filters( &sub.filters, rec.columns.as_ref().unwrap_or(&vec![]), @@ -401,8 +399,6 @@ fn process_record<'a>( } } - //println!("through filters {:?}", visible_through_filters); - // Row Level Security let subscriptions_to_notify: Vec<&realtime::Subscription> = match ( is_rls_enabled && visible_through_filters.len() > 0, @@ -429,8 +425,6 @@ fn process_record<'a>( } }; - //println!("notify {:?}", subscriptions_to_notify); - let r = realtime::WALRLS { wal: realtime::Data { schema: rec.schema, From 5b15a1ffe7a55bb67ede95987c8dfe9ec3c4d2b3 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 08:29:13 -0500 Subject: [PATCH 76/88] remove (incomplete) stream_logical from branch --- worker/stream_logical/Cargo.toml | 12 --- worker/stream_logical/README.md | 44 -------- worker/stream_logical/src/main.rs | 174 ------------------------------ 3 files changed, 230 deletions(-) delete mode 100644 worker/stream_logical/Cargo.toml delete mode 100644 worker/stream_logical/README.md delete mode 100644 worker/stream_logical/src/main.rs diff --git a/worker/stream_logical/Cargo.toml b/worker/stream_logical/Cargo.toml deleted file mode 100644 index a05fa29..0000000 --- a/worker/stream_logical/Cargo.toml +++ /dev/null @@ -1,12 +0,0 @@ -[package] -name = "stream_logical" -version = "0.1.0" -edition = "2021" - -[dependencies] -tokio = { version = "1", features = ["full"] } -serde = { version = "1" } -futures = "0.3" -bytes = "1.0" -tokio-postgres = { git = "https://github.com/MaterializeInc/rust-postgres" } -postgres-protocol = { git = "https://github.com/MaterializeInc/rust-postgres" } diff --git a/worker/stream_logical/README.md b/worker/stream_logical/README.md deleted file mode 100644 index e27b3e3..0000000 --- a/worker/stream_logical/README.md +++ /dev/null @@ -1,44 +0,0 @@ -# stream_logical - -stream_logical is an option for listening to the logical replication stream from postgres without the existing dependencies on: -- pg_recvlogical (local dependency) -- wal2json (server dependency) - -The dependency on pg_recvlogical requires that postgres is installed on the same machine as the walrus server. The wal2json dependency is not installed by default on many postgres setups which reduces compatibility. - -This solution would enable walrus to run against unmodified postgres. - -### Current State - -- Connects to postgres (from parent dir docker-compose) -- Creates a publication and temporary replication slot -- Creates a small amount of WAL -- Receives the WAL messages and prints the message type - -### Known TODOS - -- A keepalive message needs to be sent on a schedule -- Tracking table state -- Convert logical replication stream to a useable format -- Serialize messages - - -### Key Commands / Notes - -```shell -cargo run --bin stream_logical -``` - -Output -```text -lsn 0/173C578 -keep alive -begin -relation -insert -commit -keep alive -keep alive -``` - - diff --git a/worker/stream_logical/src/main.rs b/worker/stream_logical/src/main.rs deleted file mode 100644 index 660c16a..0000000 --- a/worker/stream_logical/src/main.rs +++ /dev/null @@ -1,174 +0,0 @@ -use futures::StreamExt; -use postgres_protocol::message::backend::LogicalReplicationMessage::*; -use postgres_protocol::message::backend::ReplicationMessage::*; -use postgres_protocol::message::backend::{LogicalReplicationMessage, ReplicationMessage}; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::time::Instant; -use tokio_postgres::replication::LogicalReplicationStream; -use tokio_postgres::types::PgLsn; -use tokio_postgres::NoTls; -use tokio_postgres::SimpleQueryMessage::Row; - -/// Describes a table in a PostgreSQL database. -#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] -pub struct PostgresTableDesc { - /// The OID of the table. - pub oid: u32, - /// The name of the schema that the table belongs to. - pub namespace: String, - /// The name of the table. - pub name: String, - /// The description of each column, in order. - pub columns: Vec, -} - -/// Describes a column in a [`PostgresTableDesc`]. -#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] -pub struct PostgresColumnDesc { - /// The name of the column. - pub name: String, - /// The OID of the column's type. - pub type_oid: u32, - /// The modifier for the column's type. - pub type_mod: i32, - /// True if the column lacks a `NOT NULL` constraint. - pub nullable: bool, - /// Whether the column is part of the table's primary key. - /// TODO The order of the columns in the primary key matters too. - pub primary_key: bool, -} - -struct PostgresTaskInfo { - /// Our cursor into the WAL - lsn: PgLsn, - source_tables: HashMap, -} - -#[tokio::main] -async fn main() { - let conninfo = "host=127.0.0.1 port=5501 user=postgres password=password replication=database"; - let (client, connection) = tokio_postgres::connect(conninfo, NoTls).await.unwrap(); - - tokio::spawn(async move { - if let Err(e) = connection.await { - eprintln!("connection error: {}", e); - } - }); - - client - .simple_query("DROP TABLE IF EXISTS test_logical_replication") - .await - .unwrap(); - client - .simple_query("CREATE TABLE test_logical_replication(i int)") - .await - .unwrap(); - let res = client - .simple_query("SELECT 'test_logical_replication'::regclass::oid") - .await - .unwrap(); - - let rel_id: u32 = if let Row(row) = &res[0] { - row.get("oid").unwrap().parse().unwrap() - } else { - panic!("unexpeced query message"); - }; - - client - .simple_query("DROP PUBLICATION IF EXISTS test_pub") - .await - .unwrap(); - client - .simple_query("CREATE PUBLICATION test_pub FOR ALL TABLES") - .await - .unwrap(); - - let slot = "test_logical_slot"; - - let query = format!( - r#"CREATE_REPLICATION_SLOT {:?} TEMPORARY LOGICAL "pgoutput""#, - slot - ); - let slot_query = client.simple_query(&query).await.unwrap(); - let lsn = if let Row(row) = &slot_query[0] { - row.get("consistent_point").unwrap() - } else { - panic!("unexpeced query message"); - }; - - println!("lsn {}", lsn); - - // issue a query that will appear in the slot's stream since it happened after its creation - client - .simple_query("INSERT INTO test_logical_replication VALUES (42)") - .await - .unwrap(); - - let options = r#"("proto_version" '1', "publication_names" 'test_pub')"#; - let query = format!( - r#"START_REPLICATION SLOT {:?} LOGICAL {} {}"#, - slot, lsn, options - ); - let copy_stream = client - .copy_both_simple::(&query) - .await - .unwrap(); - - let stream = LogicalReplicationStream::new(copy_stream); - tokio::pin!(stream); - - let task = PostgresTaskInfo { - lsn: 0.into(), - source_tables: HashMap::new(), - }; - let mut last_keepalive = Instant::now(); - let mut inserts: Vec = vec![]; - let mut deletes: Vec = vec![]; - let mut lsn: u64; - - // TODO(olirice) track table defs mapping - // TODO(olirice) send keepalive every 20 seconds - - loop { - let msg: Option< - Result, tokio_postgres::Error>, - > = stream.next().await; - - let msg_res = match msg { - Some(Ok(XLogData(xlog_data))) => match xlog_data.data() { - Begin(begin) => { - println!("{:?}", begin) - } - Insert(insert) => { - println!("{:?}", insert) - } - Update(update) => { - println!("{:?}", update) - } - Delete(delete) => { - println!("{:?}", delete) - } - Commit(commit) => { - println!("{:?}", commit); - lsn = commit.commit_lsn(); - } - Relation(relation) => { - println!("{:?}", relation) - } - Origin(_) | Type(_) => {} - Truncate(truncate) => { - println!("{:?}", truncate) - } - _ => println!("unknown logical replication message type"), - }, - Some(Err(_)) => panic!("unexpected replication stream error"), - None => panic!("unexpected replication stream end"), - Some(Ok(PrimaryKeepAlive(_))) => { - println!("keep alive") - } - Some(Ok(_)) => (), - _ => println!("Unexpected replication message"), - }; - } -} From ec55a96a68f96feb999f4162602ca1f03bde2318 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 11:14:01 -0500 Subject: [PATCH 77/88] insert, update, delete, truncate integration test w/ filters and rls --- worker/Cargo.toml | 1 - worker/walrus/Cargo.toml | 2 +- worker/walrus/src/main.rs | 588 +++++++++++++++++++++++++++++++++----- 3 files changed, 524 insertions(+), 67 deletions(-) diff --git a/worker/Cargo.toml b/worker/Cargo.toml index d5db734..cd4533e 100644 --- a/worker/Cargo.toml +++ b/worker/Cargo.toml @@ -2,5 +2,4 @@ members = [ "walrus", "realtime", - "stream_logical", ] diff --git a/worker/walrus/Cargo.toml b/worker/walrus/Cargo.toml index 1897253..a123321 100644 --- a/worker/walrus/Cargo.toml +++ b/worker/walrus/Cargo.toml @@ -8,7 +8,7 @@ clap = { version = "3.1.12", features = ["derive"] } diesel = { version = "2.0.0-rc.0", features = ["postgres", "serde_json", "uuid", "chrono"] } diesel_migrations = { version = "2.0.0-rc.0", features = ["postgres"] } dotenv = "0.15.0" -serde_json = "1.0" +serde_json = { version="1.0", features = ["preserve_order"] } serde = "1.0" uuid = { version = "1.0", features = ["serde", "v4"] } log = "0.4.17" diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index da1308d..0c4a8ac 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -355,75 +355,84 @@ fn process_record<'a>( }); // User Defined Filters - let mut visible_through_filters = vec![]; - let mut delegate_to_sql_filters = vec![]; - - for sub in entity_role_subscriptions { - match filters::record::user_defined::visible_through_filters( - &sub.filters, - rec.columns.as_ref().unwrap_or(&vec![]), - ) { - Ok(true) => { - //debug!("Filters handled in rust: {:?}", &sub.filters); - visible_through_filters.push(sub); - } - Ok(false) => (), - // delegate to SQL when we can't handle the comparison in rust - Err(errors::FilterError::DelegateToSQL(_)) => { - //debug!( - // "Filters delegated to SQL: {:?}. Error: {}", - // &sub.filters, err - //); - delegate_to_sql_filters.push(sub); - } - } - } + let subscriptions_to_notify: Vec<&realtime::Subscription>; - if delegate_to_sql_filters.len() > 0 { - match filters::record::user_defined::is_visible_through_filters_sql( - &walcols, - &delegate_to_sql_filters.iter().map(|x| x.id).collect(), - conn, - ) { - Ok(sub_ids) => { - for sub in delegate_to_sql_filters - .iter() - .filter(|x| sub_ids.contains(&x.id)) - { - visible_through_filters.push(sub) - } - } - Err(err) => { - error!("Failed to deletegate some filters to SQL: {}", err) - } + match vec![realtime::Action::DELETE, realtime::Action::TRUNCATE].contains(&action) { + true => { + subscriptions_to_notify = entity_role_subscriptions; } - } + false => { + let mut visible_through_filters = vec![]; + let mut delegate_to_sql_filters = vec![]; + + for sub in entity_role_subscriptions { + match filters::record::user_defined::visible_through_filters( + &sub.filters, + rec.columns.as_ref().unwrap_or(&vec![]), + ) { + Ok(true) => { + //debug!("Filters handled in rust: {:?}", &sub.filters); + visible_through_filters.push(sub); + } + Ok(false) => (), + // delegate to SQL when we can't handle the comparison in rust + Err(errors::FilterError::DelegateToSQL(_)) => { + //debug!( + // "Filters delegated to SQL: {:?}. Error: {}", + // &sub.filters, err + //); + delegate_to_sql_filters.push(sub); + } + } - // Row Level Security - let subscriptions_to_notify: Vec<&realtime::Subscription> = match ( - is_rls_enabled && visible_through_filters.len() > 0, - vec![realtime::Action::DELETE, realtime::Action::TRUNCATE].contains(&action), - ) { - (false, _) | (true, true) => visible_through_filters, - _ => { - match filters::record::row_level_security::is_visible_through_rls( - table_oid, - &walcols, - &visible_through_filters.iter().map(|x| x.id).collect(), - conn, - ) { - Ok(sub_ids) => visible_through_filters - .iter() - .filter(|x| sub_ids.contains(&x.id)) - .map(|x| *x) - .collect(), - Err(err) => { - error!("Failed to delegate RLS to SQL: {}", err); - vec![] + if delegate_to_sql_filters.len() > 0 { + match filters::record::user_defined::is_visible_through_filters_sql( + &walcols, + &delegate_to_sql_filters.iter().map(|x| x.id).collect(), + conn, + ) { + Ok(sub_ids) => { + for sub in delegate_to_sql_filters + .iter() + .filter(|x| sub_ids.contains(&x.id)) + { + visible_through_filters.push(sub) + } + } + Err(err) => { + error!("Failed to deletegate some filters to SQL: {}", err) + } + } } } + + // Row Level Security + subscriptions_to_notify = match is_rls_enabled { + false => visible_through_filters, + true => match visible_through_filters.len() > 0 { + true => { + match filters::record::row_level_security::is_visible_through_rls( + table_oid, + &walcols, + &visible_through_filters.iter().map(|x| x.id).collect(), + conn, + ) { + Ok(sub_ids) => visible_through_filters + .iter() + .filter(|x| sub_ids.contains(&x.id)) + .map(|x| *x) + .collect(), + Err(err) => { + error!("Failed to delegate RLS to SQL: {}", err); + vec![] + } + } + } + false => vec![], + }, + }; } - }; + } let r = realtime::WALRLS { wal: realtime::Data { @@ -948,8 +957,6 @@ mod tests { #[test] fn test_quoted_type_schema_and_table() { - // TODO: Add user defined filter to make sure delegate to sql works - let mut conn = establish_connection(); crate::sql::migrations::run_migrations(&mut conn) .expect("Pending migrations failed to execute"); @@ -1074,4 +1081,455 @@ mod tests { assert_eq!(walrus_output, expected); } + + #[test] + fn test_integration() { + let mut conn = establish_connection(); + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + create_auth_schema(&mut conn); + truncate("realtime", "subscription", &mut conn); + create_publication_for_all_tables("supabase_multiplayer", &mut conn); + + create_role("authenticated", &mut conn); + + drop_table("public", "notes_integ", &mut conn); + diesel::sql_query( + " + create table if not exists public.notes_integ( + id int primary key, + body text + );", + ) + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("alter table public.notes_integ replica identity full;") + .execute(&mut conn) + .unwrap(); + + let rls_excluded_auth_uid = uuid::uuid!("876286ab-4a4e-47df-9c08-a91054c87e1d"); + + diesel::sql_query(format!( + "create policy rls_note_integ_select + on public.notes_integ + to authenticated + using (auth.uid() <> '{rls_excluded_auth_uid}');" + )) + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("alter table public.notes_integ enable row level security;") + .execute(&mut conn) + .unwrap(); + + let notes_oid = + crate::filters::table::table_oid::get_table_oid("public", "notes_integ", &mut conn) + .unwrap(); + + grant_all_on_schema("public", "authenticated", &mut conn); + + diesel::sql_query("revoke select on public.notes_integ from authenticated;") + .execute(&mut conn) + .unwrap(); + + diesel::sql_query("grant select (id) on public.notes_integ to authenticated;") + .execute(&mut conn) + .unwrap(); + + // Superuser + let postgres_sub_id = uuid::uuid!("37c7e506-9eca-4671-8c48-526d404660ce"); + insert_into(subscription) + .values(( + subscription_id.eq(postgres_sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "postgres", + "email": "example@example.com", + "sub": uuid::Uuid::new_v4(), + })), + )) + .execute(&mut conn) + .unwrap(); + + // Authenticated: No filter + let authenticated_no_filter_sub_id = uuid::uuid!("6257552a-76c6-433c-8777-44cfb00851f9"); + insert_into(subscription) + .values(( + subscription_id.eq(authenticated_no_filter_sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "authenticated", + "email": "example@example.com", + "sub": uuid::Uuid::new_v4(), + })), + )) + .execute(&mut conn) + .unwrap(); + + // Authenticated: Filter match + let authenticated_filter_match_sub_id = uuid::uuid!("3e31b5e1-828a-404b-9f06-4b6e0826b027"); + insert_into(subscription) + .values(( + subscription_id.eq(authenticated_filter_match_sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "authenticated", + "email": "example@example.com", + "sub": uuid::Uuid::new_v4(), + })), + filters.eq(vec![UserDefinedFilter { + column_name: "id".to_string(), + op: realtime::Op::Equal, + value: "1".to_string(), + }]), + )) + .execute(&mut conn) + .unwrap(); + + // Authenticated: Filter no match + let authenticated_filter_no_match_sub_id = + uuid::uuid!("086145fd-1598-451c-9f00-a38b0d9c7a7e"); + insert_into(subscription) + .values(( + subscription_id.eq(authenticated_filter_no_match_sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "authenticated", + "email": "example@example.com", + "sub": uuid::Uuid::new_v4(), + })), + filters.eq(vec![UserDefinedFilter { + column_name: "id".to_string(), + op: realtime::Op::Equal, + value: "999".to_string(), + }]), + )) + .execute(&mut conn) + .unwrap(); + + // Authenticated: User excluded by RLS policy + let authenticated_filter_rls_exclude_sub_id = + uuid::uuid!("54249b4a-98ca-4941-8af7-0154123df504"); + insert_into(subscription) + .values(( + subscription_id.eq(authenticated_filter_rls_exclude_sub_id), + entity.eq(notes_oid), + claims.eq(json!({ + "role": "authenticated", + "email": "example@example.com", + "sub": rls_excluded_auth_uid, + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscriptions = subscription.load::(&mut conn).unwrap(); + + let note_id: i32 = 1; + + { + // INSERT + + diesel::sql_query(format!( + "insert into public.notes_integ( id, body) values (1, 'starting body');" + )) + .execute(&mut conn) + .unwrap(); + + let wal2json_insert = r#"{ + "action":"I", + "timestamp":"2022-07-13 13:46:35.021414+00", + "schema":"public", + "table":"notes_integ", + "columns":[ + {"name":"id","type":"integer","typeoid":23,"value":1}, + {"name":"body","type":"text","typeoid":25,"value":"starting body"} + ], + "pk":[{"name":"id","type":"integer","typeoid":23}] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_insert).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![ + // Record for postgres role + realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes_integ", + r#type: realtime::Action::INSERT, + commit_timestamp: &rec.timestamp, + columns: vec![ + realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }, + realtime::Column { + name: "body".to_string(), + type_: "text".to_string(), + }, + ], + record: HashMap::from([ + ("id", json!(note_id)), + ("body", json!("starting body")), + ]), + old_record: None, + }, + is_rls_enabled: true, + subscription_ids: vec![postgres_sub_id], + errors: vec![], + }, + // Record for authenticated role + realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes_integ", + r#type: realtime::Action::INSERT, + commit_timestamp: &rec.timestamp, + // body key was excluded + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + record: HashMap::from([("id", json!(note_id))]), + old_record: None, + }, + is_rls_enabled: true, + subscription_ids: vec![ + authenticated_no_filter_sub_id, + authenticated_filter_match_sub_id, + ], + errors: vec![], + }, + ]; + + assert_eq!(walrus_output, expected); + } + + { + // Update + + diesel::sql_query(format!( + "update public.notes_integ set body = 'updated body';" + )) + .execute(&mut conn) + .unwrap(); + + let wal2json_update = r#"{ + "action":"U", + "timestamp":"2022-07-13 13:46:35.021414+00", + "schema":"public", + "table":"notes_integ", + "columns":[ + {"name":"id","type":"integer","typeoid":23,"value":1}, + {"name":"body","type":"text","typeoid":25,"value":"updated body"} + ], + "identity":[ + {"name":"id","type":"integer","typeoid":23,"value":1}, + {"name":"body","type":"text","typeoid":25,"value":"starting body"} + ], + "pk":[ + {"name":"id","type":"integer","typeoid":23} + ] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_update).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![ + // Record for postgres role + realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes_integ", + r#type: realtime::Action::UPDATE, + commit_timestamp: &rec.timestamp, + columns: vec![ + realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }, + realtime::Column { + name: "body".to_string(), + type_: "text".to_string(), + }, + ], + record: HashMap::from([ + ("id", json!(note_id)), + ("body", json!("updated body")), + ]), + old_record: Some(HashMap::from([ + ("id", json!(note_id)), + ("body", json!("starting body")), + ])), + }, + is_rls_enabled: true, + subscription_ids: vec![postgres_sub_id], + errors: vec![], + }, + // Record for authenticated role + realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes_integ", + r#type: realtime::Action::UPDATE, + commit_timestamp: &rec.timestamp, + // body key was excluded + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + record: HashMap::from([("id", json!(note_id))]), + old_record: Some(HashMap::from([("id", json!(note_id))])), + }, + is_rls_enabled: true, + subscription_ids: vec![ + authenticated_no_filter_sub_id, + authenticated_filter_match_sub_id, + ], + errors: vec![], + }, + ]; + + assert_eq!(walrus_output, expected); + } + + { + // Delete + + diesel::sql_query(format!("delete from public.notes_integ;")) + .execute(&mut conn) + .unwrap(); + + let wal2json_delete = r#"{ + "action":"D", + "timestamp":"2022-07-13 13:46:35.021414+00", + "schema":"public", + "table":"notes_integ", + "identity":[ + {"name":"id","type":"integer","typeoid":23,"value":1}, + {"name":"body","type":"text","typeoid":25,"value":"updated body"} + ], + "pk":[ + {"name":"id","type":"integer","typeoid":23} + ] + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_delete).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![ + // Record for postgres role + realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes_integ", + r#type: realtime::Action::DELETE, + commit_timestamp: &rec.timestamp, + columns: vec![ + realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }, + realtime::Column { + name: "body".to_string(), + type_: "text".to_string(), + }, + ], + record: HashMap::new(), + old_record: Some(HashMap::from([ + ("id", json!(note_id)), + ("body", json!("updated body")), + ])), + }, + is_rls_enabled: true, + subscription_ids: vec![postgres_sub_id], + errors: vec![], + }, + // Record for authenticated role + realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes_integ", + r#type: realtime::Action::DELETE, + commit_timestamp: &rec.timestamp, + // body key was excluded + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + record: HashMap::new(), + old_record: Some(HashMap::from([("id", json!(note_id))])), + }, + is_rls_enabled: true, + subscription_ids: vec![ + authenticated_no_filter_sub_id, + authenticated_filter_match_sub_id, + // the non-matching filter subscriber is included + authenticated_filter_no_match_sub_id, + // the rls excluded subscriber is also included + authenticated_filter_rls_exclude_sub_id, + ], + errors: vec![], + }, + ]; + + assert_eq!(walrus_output, expected); + } + + { + // Truncate + diesel::sql_query(format!("truncate table public.notes_integ;")) + .execute(&mut conn) + .unwrap(); + + let wal2json_truncate = r#"{ + "action":"T", + "timestamp":"2022-07-13 13:46:35.021414+00", + "schema":"public", + "table":"notes_integ" + }"#; + + let rec: wal2json::Record = serde_json::from_str(wal2json_truncate).unwrap(); + + let walrus_output = crate::process_record( + &rec, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + // truncates do not broadcast + let expected = vec![]; + + assert_eq!(walrus_output, expected); + } + } } From fd40c907a135ca769b2e68db7fb02ca0a8b75275 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 11:31:22 -0500 Subject: [PATCH 78/88] walrus worker ci --- .github/workflows/worker_test.yaml | 34 ++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 .github/workflows/worker_test.yaml diff --git a/.github/workflows/worker_test.yaml b/.github/workflows/worker_test.yaml new file mode 100644 index 0000000..8a4cce3 --- /dev/null +++ b/.github/workflows/worker_test.yaml @@ -0,0 +1,34 @@ +name: WALRUS worker tests + +on: + pull_request: + branches: [master] + push: + branches: [master] + +jobs: + build: + name: worker tests + runs-on: ubuntu-latest + timeout-minutes: 15 + + services: + postgres: + image: supabase/postgres:lastest + env: + POSTGRES_DB: postgres + POSTGRES_HOST: localhost + POSTGRES_USER: postgres + POSTGRES_PASSWORD: password + ports: + - 5501:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + steps: + - name: run tests + run: | + cd worker + cargo test --bin walrus -- --test-threads=1 From c6b996bed0167d08e7ff1855cbe23003a52de559 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 11:33:17 -0500 Subject: [PATCH 79/88] docker image typo --- .github/workflows/worker_test.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/worker_test.yaml b/.github/workflows/worker_test.yaml index 8a4cce3..3682fa7 100644 --- a/.github/workflows/worker_test.yaml +++ b/.github/workflows/worker_test.yaml @@ -14,7 +14,7 @@ jobs: services: postgres: - image: supabase/postgres:lastest + image: supabase/postgres:latest env: POSTGRES_DB: postgres POSTGRES_HOST: localhost From 4158b0337c9843b427da12e4c9820d1c6d8ede20 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 11:36:53 -0500 Subject: [PATCH 80/88] check out the code in ci... --- .github/workflows/worker_test.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/worker_test.yaml b/.github/workflows/worker_test.yaml index 3682fa7..3beea32 100644 --- a/.github/workflows/worker_test.yaml +++ b/.github/workflows/worker_test.yaml @@ -28,6 +28,8 @@ jobs: --health-timeout 5s --health-retries 5 steps: + - uses: actions/checkout@v3 + - name: run tests run: | cd worker From c65c7e36c0855d195f58ced686c315148a2834e8 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 12:58:13 -0500 Subject: [PATCH 81/88] test subscription manager --- worker/walrus/src/main.rs | 6 +- worker/walrus/src/models/realtime.rs | 230 ++++++++++++++++++++++++++- 2 files changed, 232 insertions(+), 4 deletions(-) diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs index 0c4a8ac..9218668 100644 --- a/worker/walrus/src/main.rs +++ b/worker/walrus/src/main.rs @@ -468,7 +468,7 @@ mod tests { use crate::models::{realtime, wal2json}; use crate::realtime::{Subscription, UserDefinedFilter}; use crate::sql::schema::realtime::subscription::dsl::*; - use chrono::{TimeZone, Utc}; + use chrono::Utc; use diesel::prelude::*; use diesel::*; use pretty_assertions::assert_eq; @@ -481,7 +481,7 @@ mod tests { //const INT8OID: i32 = 20; //const TEXTOID: i32 = 25; - fn establish_connection() -> PgConnection { + pub fn establish_connection() -> PgConnection { let database_url = "postgresql://postgres:password@localhost:5501/postgres"; PgConnection::establish(&database_url).unwrap() } @@ -543,7 +543,7 @@ mod tests { .unwrap(); } - fn truncate(schema: &str, table: &str, conn: &mut PgConnection) { + pub fn truncate(schema: &str, table: &str, conn: &mut PgConnection) { diesel::sql_query(format!("truncate table \"{}\".\"{}\";", schema, table)) .execute(conn) .unwrap(); diff --git a/worker/walrus/src/models/realtime.rs b/worker/walrus/src/models/realtime.rs index e96cb66..f1f3326 100644 --- a/worker/walrus/src/models/realtime.rs +++ b/worker/walrus/src/models/realtime.rs @@ -141,7 +141,9 @@ pub fn update_subscriptions( subscriptions.push(new_sub); debug!("Subscription inserted. Total {}", subscriptions.len()); } - Err(err) => error!("No subscription found: id={}, Error: {} ", id_val, err), + Err(err) => { + error!("No subscription found: id={}, Error: {} ", id_val, err); + } }; } wal2json::Action::U => { @@ -265,3 +267,229 @@ impl FromSql for UserDefinedFilter { }) } } + +#[cfg(test)] +mod tests { + extern crate diesel; + use crate::models::realtime::*; + use crate::models::{realtime, wal2json}; + use crate::tests::*; + use chrono::Utc; + use pretty_assertions::assert_eq; + use serde_json::json; + use uuid; + + #[test] + fn test_update_subscriptions_insert() { + let mut conn = establish_connection(); + + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + + crate::tests::truncate("realtime", "subscription", &mut conn); + + let subscriptions_table_oid = + crate::filters::table::table_oid::get_table_oid("realtime", "subscription", &mut conn) + .unwrap(); + + let mut subscriptions = vec![]; + + let sub_id = uuid::uuid!("54249b4a-98ca-4941-8af7-0154123df504"); + + insert_into(subscription) + .values(( + subscription_id.eq(sub_id), + entity.eq(subscriptions_table_oid), + claims.eq(json!({ + "role": "postgres", + "email": "example@example.com", + "sub": sub_id, + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscription_row = subscription + .first::(&mut conn) + .unwrap(); + let subscription_row_id = subscription_row.id; + + let wal2json = format!( + r#"{{ + "action":"I", + "timestamp":"2022-07-13 17:04:40.784361+00", + "schema":"realtime", + "table":"subscription", + "columns":[ + {{ + "name":"id", + "type":"bigint", + "typeoid":20, + "value":{subscription_row_id} + }} + ], + "pk":[{{"name":"id","type":"bigint","typeoid":20}}] + }}"# + ); + + let rec: wal2json::Record = serde_json::from_str(&wal2json).unwrap(); + update_subscriptions(&rec, &mut subscriptions, &mut conn); + + // Subscription was added + assert_eq!(subscriptions, vec![subscription_row]); + } + + #[test] + fn test_update_subscriptions_update() { + let mut conn = establish_connection(); + + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + + crate::tests::truncate("realtime", "subscription", &mut conn); + + let subscriptions_table_oid = + crate::filters::table::table_oid::get_table_oid("realtime", "subscription", &mut conn) + .unwrap(); + + let sub_id = uuid::uuid!("54249b4a-98ca-4941-8af7-0154123df504"); + + insert_into(subscription) + .values(( + subscription_id.eq(sub_id), + entity.eq(subscriptions_table_oid), + claims.eq(json!({ + "role": "postgres", + "email": "example@example.com", + "sub": sub_id, + })), + )) + .execute(&mut conn) + .unwrap(); + + let subscription_row = subscription + .first::(&mut conn) + .unwrap(); + let subscription_row_id = subscription_row.id; + + let mut subscriptions = vec![subscription_row.clone()]; + + // Update the subscription id + let updated_sub_id = uuid::uuid!("54249b4a-98ca-4941-8af7-0154123df504"); + diesel::sql_query(format!( + "update realtime.subscription set subscription_id = '{updated_sub_id}'" + )) + .execute(&mut conn) + .unwrap(); + + let wal2json = format!( + r#"{{ + "action":"U", + "timestamp":"2022-07-13 17:04:40.784361+00", + "schema":"realtime", + "table":"subscription", + "columns":[ + {{ + "name":"id", + "type":"bigint", + "typeoid":20, + "value":{subscription_row_id} + }} + ], + "pk":[{{"name":"id","type":"bigint","typeoid":20}}] + }}"# + ); + + let rec: wal2json::Record = serde_json::from_str(&wal2json).unwrap(); + update_subscriptions(&rec, &mut subscriptions, &mut conn); + + assert_eq!(subscriptions[0].subscription_id, updated_sub_id); + } + + #[test] + fn test_update_subscriptions_delete() { + let mut conn = establish_connection(); + + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + + crate::tests::truncate("realtime", "subscription", &mut conn); + + let subscription_row_id = 1; + + let mut subscriptions = vec![realtime::Subscription { + id: subscription_row_id, + subscription_id: uuid::uuid!("54249b4a-98ca-4941-8af7-0154123df504"), + entity: 999, + filters: vec![], + claims: json!({}), + claims_role: 999, + schema_name: "abc".to_string(), + table_name: "abc".to_string(), + claims_role_name: "abc".to_string(), + created_at: Utc::now().naive_utc(), + }]; + + let wal2json = format!( + r#"{{ + "action":"D", + "timestamp":"2022-07-13 17:04:40.784361+00", + "schema":"realtime", + "table":"subscription", + "identity":[ + {{ + "name":"id", + "type":"bigint", + "typeoid":20, + "value":{subscription_row_id} + }} + ], + "pk":[{{"name":"id","type":"bigint","typeoid":20}}] + }}"# + ); + + let rec: wal2json::Record = serde_json::from_str(&wal2json).unwrap(); + update_subscriptions(&rec, &mut subscriptions, &mut conn); + + assert_eq!(subscriptions, vec![]); + } + + #[test] + fn test_update_subscriptions_truncate() { + let mut conn = establish_connection(); + + crate::sql::migrations::run_migrations(&mut conn) + .expect("Pending migrations failed to execute"); + + crate::tests::truncate("realtime", "subscription", &mut conn); + + let subscription_row_id = 1; + + let mut subscriptions = vec![realtime::Subscription { + id: subscription_row_id, + subscription_id: uuid::uuid!("54249b4a-98ca-4941-8af7-0154123df504"), + entity: 999, + filters: vec![], + claims: json!({}), + claims_role: 999, + schema_name: "abc".to_string(), + table_name: "abc".to_string(), + claims_role_name: "abc".to_string(), + created_at: Utc::now().naive_utc(), + }]; + + let wal2json = format!( + r#"{{ + "action":"T", + "timestamp":"2022-07-13 17:04:40.784361+00", + "schema":"realtime", + "table":"subscription" + }}"# + ); + + let rec: wal2json::Record = serde_json::from_str(&wal2json).unwrap(); + update_subscriptions(&rec, &mut subscriptions, &mut conn); + + assert_eq!(subscriptions, vec![]); + } +} From fa7ea98a7a82bd3aff427f81dda9fa05df10fe23 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 13:00:10 -0500 Subject: [PATCH 82/88] reduce heartbeat to every 30 seconds --- worker/realtime/src/main.rs | 2 +- worker/walrus/src/filters/record/user_defined.rs | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs index a6a07a2..94c512f 100644 --- a/worker/realtime/src/main.rs +++ b/worker/realtime/src/main.rs @@ -191,7 +191,7 @@ async fn read_stdin(tx: futures_channel::mpsc::UnboundedSender, topic: async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender) { loop { - sleep(Duration::from_secs(20)).await; + sleep(Duration::from_secs(30)).await; let phoenix_msg = PhoenixMessage { event: PhoenixMessageEvent::Heartbeat, payload: serde_json::json!({"msg": "ping"}), diff --git a/worker/walrus/src/filters/record/user_defined.rs b/worker/walrus/src/filters/record/user_defined.rs index 238830a..9589e76 100644 --- a/worker/walrus/src/filters/record/user_defined.rs +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -17,7 +17,6 @@ pub mod sql_functions { pub fn is_visible_through_filters_sql( columns: &Vec, ids: &Vec, - // TODO: convert this to use subscription_ids to reduce n calls conn: &mut PgConnection, ) -> Result, String> { select(sql_functions::is_visible_through_filters( From b3ba019d6af11287759488a11adb6d96c9cb684b Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Wed, 13 Jul 2022 13:32:07 -0500 Subject: [PATCH 83/88] drop unused functions --- .../migrations/2022-07-13-181325_drop unused functions/down.sql | 1 + .../migrations/2022-07-13-181325_drop unused functions/up.sql | 2 ++ 2 files changed, 3 insertions(+) create mode 100644 worker/walrus/migrations/2022-07-13-181325_drop unused functions/down.sql create mode 100644 worker/walrus/migrations/2022-07-13-181325_drop unused functions/up.sql diff --git a/worker/walrus/migrations/2022-07-13-181325_drop unused functions/down.sql b/worker/walrus/migrations/2022-07-13-181325_drop unused functions/down.sql new file mode 100644 index 0000000..291a97c --- /dev/null +++ b/worker/walrus/migrations/2022-07-13-181325_drop unused functions/down.sql @@ -0,0 +1 @@ +-- This file should undo anything in `up.sql` \ No newline at end of file diff --git a/worker/walrus/migrations/2022-07-13-181325_drop unused functions/up.sql b/worker/walrus/migrations/2022-07-13-181325_drop unused functions/up.sql new file mode 100644 index 0000000..f15b788 --- /dev/null +++ b/worker/walrus/migrations/2022-07-13-181325_drop unused functions/up.sql @@ -0,0 +1,2 @@ +drop function realtime.apply_rls(jsonb, integer); +drop function realtime.quote_wal2json(regclass); From c71134c934a0c957d823b7e3e7891e0fb68490f0 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 21 Jul 2022 13:42:20 -0500 Subject: [PATCH 84/88] document publication usage in readmes --- worker/README.md | 6 ++++-- worker/walrus/README.md | 12 ++++++------ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/worker/README.md b/worker/README.md index f1d8283..0ff1c67 100644 --- a/worker/README.md +++ b/worker/README.md @@ -30,7 +30,7 @@ docker-compose up Run the `walrus` worker, piping its output to `realtime` transport ```sh cargo run --bin walrus -- \ - --connection=postgresql://postgres:password@localhost:5501/postgres | + cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres --publication=walrus_pub | cargo run --bin realtime -- \ --url=wss://sendwal.fly.dev/socket \ --header=apikey= @@ -47,6 +47,8 @@ create table book( title text ); +create publication walrus_pub for all tables; + -- Create a dummy subscription to our new table insert into realtime.subscription(subscription_id, entity, claims) select @@ -61,4 +63,4 @@ select -- Create a record insert into book(id, title) values (1, 'Foo'); -``` \ No newline at end of file +``` diff --git a/worker/walrus/README.md b/worker/walrus/README.md index 5571ef4..f716bd5 100644 --- a/worker/walrus/README.md +++ b/worker/walrus/README.md @@ -12,6 +12,7 @@ USAGE: OPTIONS: --connection [default: postgresql://postgres@localhost:5432/postgres] -h, --help Print help information + --pubilcation [default: supabase_multiplayer] --slot [default: realtime] -V, --version Print version information ``` @@ -59,10 +60,10 @@ docker-compose up Run the walrus worker ```sh cd walrus -cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres +cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres --publication=walrus_pub # Note: if you have jq installed, the output is more readable with -# cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres | jq +# cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres --publication=walrus_pub | jq ``` Connect to the database at `postgresql://postgres:password@localhost:5501/postgres` @@ -76,6 +77,8 @@ create table book( title text ); +create publication walrus_pub for all tables; + -- Create a dummy subscription to our new table insert into realtime.subscription(subscription_id, entity, claims) select @@ -95,7 +98,7 @@ values (1, 'Foo'); Now, looking back out the output from the `cargo run` command, you see the following printed to stdout ```json -> cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres +> cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres --publication=walrus_pub ... { "wal":{ @@ -131,6 +134,3 @@ Now, looking back out the output from the `cargo run` command, you see the follo ## Possible Enhancements - Rate limiting -- Some work that is currently being performed in SQL could be shuffled out rust for a performance improvement - - Reshaping the WAL records - - Any state we want to track between calls to `realtime.apply_rls` From 962db205322bdd6224c14d0c435cfce16dc28019 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 21 Jul 2022 13:48:10 -0500 Subject: [PATCH 85/88] use modern module declarations --- worker/walrus/src/{filters.rs => filters/mod.rs} | 0 worker/walrus/src/{models.rs => models/mod.rs} | 0 worker/walrus/src/{sql.rs => sql/mod.rs} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename worker/walrus/src/{filters.rs => filters/mod.rs} (100%) rename worker/walrus/src/{models.rs => models/mod.rs} (100%) rename worker/walrus/src/{sql.rs => sql/mod.rs} (100%) diff --git a/worker/walrus/src/filters.rs b/worker/walrus/src/filters/mod.rs similarity index 100% rename from worker/walrus/src/filters.rs rename to worker/walrus/src/filters/mod.rs diff --git a/worker/walrus/src/models.rs b/worker/walrus/src/models/mod.rs similarity index 100% rename from worker/walrus/src/models.rs rename to worker/walrus/src/models/mod.rs diff --git a/worker/walrus/src/sql.rs b/worker/walrus/src/sql/mod.rs similarity index 100% rename from worker/walrus/src/sql.rs rename to worker/walrus/src/sql/mod.rs From 5cf69bea2d29f4544b9f2ea87367152426609d81 Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 21 Jul 2022 13:51:02 -0500 Subject: [PATCH 86/88] document running walrus tests --- worker/walrus/README.md | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/worker/walrus/README.md b/worker/walrus/README.md index f716bd5..e7aad8d 100644 --- a/worker/walrus/README.md +++ b/worker/walrus/README.md @@ -48,7 +48,6 @@ Clone and Navigate ```sh git clone https://github.com/supabase/walrus.git cd walrus -git checkout worker cd worker ``` @@ -134,3 +133,12 @@ Now, looking back out the output from the `cargo run` command, you see the follo ## Possible Enhancements - Rate limiting + + +## Testing + +``` +cd worker +docker-compose up +cargo test -- --test-threads=1 +``` From c1dc68af43795186fdd92b98d1a793af0e1a14a1 Mon Sep 17 00:00:00 2001 From: Qiao Han Date: Fri, 22 Jul 2022 12:52:47 +0800 Subject: [PATCH 87/88] Move remaining modules to mod.rs --- worker/walrus/src/filters/{record.rs => record/mod.rs} | 0 worker/walrus/src/filters/{table.rs => table/mod.rs} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename worker/walrus/src/filters/{record.rs => record/mod.rs} (100%) rename worker/walrus/src/filters/{table.rs => table/mod.rs} (100%) diff --git a/worker/walrus/src/filters/record.rs b/worker/walrus/src/filters/record/mod.rs similarity index 100% rename from worker/walrus/src/filters/record.rs rename to worker/walrus/src/filters/record/mod.rs diff --git a/worker/walrus/src/filters/table.rs b/worker/walrus/src/filters/table/mod.rs similarity index 100% rename from worker/walrus/src/filters/table.rs rename to worker/walrus/src/filters/table/mod.rs From d8ed48095a77a803104ed007ea5677400da28beb Mon Sep 17 00:00:00 2001 From: Oliver Rice Date: Thu, 17 Aug 2023 13:27:13 -0500 Subject: [PATCH 88/88] Update README.md --- worker/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/worker/README.md b/worker/README.md index 0ff1c67..8506941 100644 --- a/worker/README.md +++ b/worker/README.md @@ -30,7 +30,7 @@ docker-compose up Run the `walrus` worker, piping its output to `realtime` transport ```sh cargo run --bin walrus -- \ - cargo run -- --connection=postgresql://postgres:password@localhost:5501/postgres --publication=walrus_pub | + --connection=postgresql://postgres:password@localhost:5501/postgres --publication=walrus_pub | cargo run --bin realtime -- \ --url=wss://sendwal.fly.dev/socket \ --header=apikey=