diff --git a/.github/workflows/worker_test.yaml b/.github/workflows/worker_test.yaml new file mode 100644 index 0000000..3beea32 --- /dev/null +++ b/.github/workflows/worker_test.yaml @@ -0,0 +1,36 @@ +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:latest + 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: + - uses: actions/checkout@v3 + + - name: run tests + run: | + cd worker + cargo test --bin walrus -- --test-threads=1 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..8506941 --- /dev/null +++ b/worker/README.md @@ -0,0 +1,66 @@ +# 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 --publication=walrus_pub | +cargo run --bin realtime -- \ + --url=wss://sendwal.fly.dev/socket \ + --header=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 publication walrus_pub for all tables; + +-- 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'); +``` diff --git a/worker/docker-compose.yml b/worker/docker-compose.yml new file mode 100644 index 0000000..4f9296e --- /dev/null +++ b/worker/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/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..950b6a2 --- /dev/null +++ b/worker/realtime/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "realtime" +version = "0.1.0" +edition = "2021" + +[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"] } +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 } diff --git a/worker/realtime/README.md b/worker/realtime/README.md new file mode 100644 index 0000000..a07eccb --- /dev/null +++ b/worker/realtime/README.md @@ -0,0 +1,23 @@ +# 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] + +OPTIONS: + -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 diff --git a/worker/realtime/src/main.rs b/worker/realtime/src/main.rs new file mode 100644 index 0000000..94c512f --- /dev/null +++ b/worker/realtime/src/main.rs @@ -0,0 +1,319 @@ +use clap::Parser; +use env_logger; +use futures::{Sink, SinkExt, Stream}; +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, BufReader}; +use tokio::time::{sleep, Duration}; +use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream}; + +/// 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, + 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"))] + Join, + #[serde(rename(serialize = "changes"))] + 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 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("debug")).init(); + + let ws_stream: ReconnectWs = ReconnectWs::connect(config).await.unwrap(); + info!("WebSocket handshake successful"); + + let (write, read) = ws_stream.split(); + + // 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; +} + +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 +} + +// 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 stdin = tokio::io::stdin(); + let buf = BufReader::new(stdin); + let mut lines = buf.lines(); + + 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 + 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 + match tx.unbounded_send(msg) { + Ok(()) => (), + Err(err) => { + error!("Error sending message: {}", err); + } + }; + } + Err(err) => { + warn!( + "Error parsing stdin line to json: error={}, line={}", + err, line + ) + } + } + } + None => { + warn!("Received empty line from stdin"); + } + } + } + Err(err) => { + error!("Error reading line from stdin: {}", err); + } + } + } +} + +async fn heartbeat(tx: futures_channel::mpsc::UnboundedSender) { + loop { + sleep(Duration::from_secs(30)).await; + let phoenix_msg = PhoenixMessage { + event: PhoenixMessageEvent::Heartbeat, + payload: serde_json::json!({"msg": "ping"}), + reference: None, + topic: "phoenix".to_string(), + }; + // Wrap phoenix message in a websocket message + let msg = Message::Text(serde_json::to_string(&phoenix_msg).unwrap()); + // push to futures stream + match tx.unbounded_send(msg) { + Ok(()) => (), + 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; + +// A websocket to communicate with a Phoenix server +// It reconnects and re-subscribes to a topic if disconnected +struct PhoenixWs(WebSocketStream>, PhoenixWsConfig); + +#[derive(Clone)] +struct PhoenixWsConfig { + addr: String, + topic_name: String, +} + +impl Stream for PhoenixWs { + 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 PhoenixWs { + 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 PhoenixWs { + // Establishes connection. + // Additionally, this will be used when reconnect tries are attempted. + fn establish( + config: PhoenixWsConfig, + ) -> Pin> + Send>> { + Box::pin(async move { + // In this case, we are trying to connect to the WebSocket endpoint + 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: 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?; + info!("Joining topic"); + Ok(PhoenixWs(ws_connection, config.clone())) + }) + } + + // 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>; 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..a123321 --- /dev/null +++ b/worker/walrus/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "walrus" +version = "0.1.0" +edition = "2021" + +[dependencies] +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 = { version="1.0", features = ["preserve_order"] } +serde = "1.0" +uuid = { version = "1.0", features = ["serde", "v4"] } +log = "0.4.17" +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/README.md b/worker/walrus/README.md new file mode 100644 index 0000000..e7aad8d --- /dev/null +++ b/worker/walrus/README.md @@ -0,0 +1,144 @@ +# 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 + --pubilcation [default: supabase_multiplayer] + --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 + + +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 +cd worker +``` + +Start the DB +```sh +docker-compose up +``` + +Run the walrus worker +```sh +cd walrus +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 --publication=walrus_pub | 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 publication walrus_pub for all tables; + +-- 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 --publication=walrus_pub +... +{ + "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 + + +## Testing + +``` +cd worker +docker-compose up +cargo test -- --test-threads=1 +``` 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..667d54b --- /dev/null +++ b/worker/walrus/migrations/2022-04-29-132611_initial/up.sql @@ -0,0 +1,627 @@ +create schema if not exists realtime; + + +-- 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 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 +$$; + + +-- 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.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 + 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 +as $$ +declare + res jsonb; +begin + execute format('select to_jsonb(%L::'|| type_::text || ')', val) into res; + return res; +end +$$; + + + +create or replace function realtime.to_regrole(role_name text) + returns regrole + immutable + language sql + -- required to allow use in generated clause +as $$ select role_name::regrole $$; + + +-- 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); + + end if; + end +$$; + + +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; +$$; + +-- 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 + 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; +$$; + +drop type if exists realtime.wal_column cascade; + +-- 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, + 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 +$$; + +-- 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[] + ); + + end if; + 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; +$$; + + +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; +$$; + + +drop function realtime.type_exists; +drop function realtime.table_exists; 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; +$$; 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..69e15b6 --- /dev/null +++ b/worker/walrus/migrations/2022-06-01-154923_helper_functions/up.sql @@ -0,0 +1,261 @@ +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_tables ppt + where + ppt.pubname = publication_name + and ppt.schemaname = schema_name + and ppt.tablename = table_name + limit 1 + ) +$$; + +create function realtime.is_rls_enabled(table_oid oid) + returns bool + language sql +as $$ + select + relrowsecurity + from + pg_class + where + oid = table_oid + limit 1; +$$; + + +create function realtime.selectable_columns( + table_oid oid, + role_name text +) + returns jsonb[] + language sql +as $$ + select + coalesce( + array_agg( + jsonb_build_object( + 'name', pa.attname::text, + 'type', pt.typname::text + ) + order by pa.attnum asc + ), + 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 = 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 +$$; + + +create function realtime.to_table_name(regclass) + returns text + language sql + immutable +as +$$ + with x(maybe_quoted_name) as ( + select + coalesce(nullif(split_part($1::text, '.', 2), ''), $1::text) + ) + select + 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 + x +$$; + +create function realtime.to_schema_name(regclass) + returns text + language sql + immutable +as +$$ + with x(maybe_quoted_name) as ( + select + relnamespace::regnamespace::text + from pg_class + where oid = $1 + limit 1 + ) + select + case + when maybe_quoted_name like '"%"' then substring( + maybe_quoted_name, + 2, + character_length(maybe_quoted_name)-2 + ) + else maybe_quoted_name + end + from + x +$$; + + +create function realtime.is_visible_through_filters( + columns jsonb, + ids int8[] -- realtime.subscription.id +) + returns int8[] + language plpgsql +as $$ +declare + cols realtime.wal_column[]; + visible_to_subscription_ids int8[] = '{}'; + subscription_id int8; + 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.id, + subs.filters + from + realtime.subscription subs + where + subs.id = any(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; +$$; + + +create function realtime.is_visible_through_rls( + table_oid oid, + columns jsonb, + ids int8[] +) + returns int8[] + language plpgsql +as $$ +declare + entity_ regclass = table_oid::regclass; + cols realtime.wal_column[]; + visible_to_subscription_ids int8[] = '{}'; + subscription_id int8; + subscription_has_access bool; + claims jsonb; +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 + ); + + -- 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.id, + subs.claims + from + realtime.subscription subs + where + subs.id = any(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; +$$; + + +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; + +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; + +create index ix_realtime_subscription_subscription_id on realtime.subscription (subscription_id); 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; +$$; 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; +$$; + + 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); diff --git a/worker/walrus/src/errors.rs b/worker/walrus/src/errors.rs new file mode 100644 index 0000000..8b8fa21 --- /dev/null +++ b/worker/walrus/src/errors.rs @@ -0,0 +1,49 @@ +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) + } +} + +#[derive(Debug)] +pub enum FilterError { + // 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::DelegateToSQL(x) => x, + }; + + write!(f, "{}", msg) + } +} diff --git a/worker/walrus/src/filters/mod.rs b/worker/walrus/src/filters/mod.rs new file mode 100644 index 0000000..659f88f --- /dev/null +++ b/worker/walrus/src/filters/mod.rs @@ -0,0 +1,2 @@ +pub mod record; +pub mod table; diff --git a/worker/walrus/src/filters/record/mod.rs b/worker/walrus/src/filters/record/mod.rs new file mode 100644 index 0000000..b8bbaa5 --- /dev/null +++ b/worker/walrus/src/filters/record/mod.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..6d2d8dd --- /dev/null +++ b/worker/walrus/src/filters/record/row_level_security.rs @@ -0,0 +1,28 @@ +use crate::errors; +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(table_oid: Oid, columns: Jsonb, ids: Array) -> Array + } +} + +pub fn is_visible_through_rls( + table_oid: u32, + columns: &Vec, + ids: &Vec, + conn: &mut PgConnection, +) -> Result, errors::Error> { + select(sql::is_visible_through_rls( + table_oid, + serde_json::to_value(columns).unwrap(), + ids, + )) + .first(conn) + .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 new file mode 100644 index 0000000..9589e76 --- /dev/null +++ b/worker/walrus/src/filters/record/user_defined.rs @@ -0,0 +1,213 @@ +use crate::errors; +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, + 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, + // Composite types are not parsable as json + Err(err) => return Err(errors::FilterError::DelegateToSQL(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 + // regtype names (provided by wal2json) + "boolean" | "smallint" | "integer" | "bigint" | "serial" | "bigserial" | "numeric" + | "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 + | 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(errors::FilterError::DelegateToSQL( + "could not handle filter op for allowed types".to_string(), + )) + } + }, + _ => { + return Err(errors::FilterError::DelegateToSQL( + "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, errors::FilterError> { + 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(errors::FilterError::DelegateToSQL( + "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.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/filters/table/column_security.rs b/worker/walrus/src/filters/table/column_security.rs new file mode 100644 index 0000000..231d5ee --- /dev/null +++ b/worker/walrus/src/filters/table/column_security.rs @@ -0,0 +1,41 @@ +use crate::errors; +use crate::models::realtime; +use cached::proc_macro::cached; +use cached::TimedSizedCache; +use diesel::*; + +pub mod sql { + use diesel::sql_types::*; + use diesel::*; + + sql_function! { + #[sql_name = "realtime.selectable_columns"] + fn selectable_columns(table_oid: Oid, role_name: Text) -> Array; + } +} + +#[cached( + type = "TimedSizedCache, errors::Error>>", + create = "{ TimedSizedCache::with_size_and_lifespan(500, 1)}", + convert = r#"{ format!("{}-{}", table_oid, role_name) }"#, + sync_writes = true +)] +pub fn selectable_columns( + table_oid: u32, + role_name: &str, + conn: &mut PgConnection, +) -> Result, errors::Error> { + let cols_as_json: Vec = + select(sql::selectable_columns(table_oid, role_name)) + .first(conn) + .map_err(|x| errors::Error::SQLFunction(format!("{}", x)))?; + + let r: Result, errors::Error> = cols_as_json + .into_iter() + .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/mod.rs b/worker/walrus/src/filters/table/mod.rs new file mode 100644 index 0000000..f121281 --- /dev/null +++ b/worker/walrus/src/filters/table/mod.rs @@ -0,0 +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/publication.rs b/worker/walrus/src/filters/table/publication.rs new file mode 100644 index 0000000..1ad2d3b --- /dev/null +++ b/worker/walrus/src/filters/table/publication.rs @@ -0,0 +1,35 @@ +use crate::errors; +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| 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 new file mode 100644 index 0000000..2ec3d4d --- /dev/null +++ b/worker/walrus/src/filters/table/row_level_security.rs @@ -0,0 +1,26 @@ +use crate::errors; +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(table_oid: Oid) -> Bool; + } +} + +#[cached( + 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 { + select(sql::is_rls_enabled(table_oid)) + .first(conn) + .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 new file mode 100644 index 0000000..ced5580 --- /dev/null +++ b/worker/walrus/src/filters/table/table_oid.rs @@ -0,0 +1,30 @@ +use crate::errors; +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| errors::Error::SQLFunction(format!("{}", x))) +} diff --git a/worker/walrus/src/main.rs b/worker/walrus/src/main.rs new file mode 100644 index 0000000..9218668 --- /dev/null +++ b/worker/walrus/src/main.rs @@ -0,0 +1,1535 @@ +#[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 models::{realtime, wal2json, walrus}; +use serde_json; +use sql::schema::realtime::subscription::dsl::*; +use std::collections::HashMap; +use std::io::{self, BufRead}; +use std::process::{Command, Stdio}; +use std::thread::sleep; +use std::time; + +mod errors; +mod filters; +mod models; +mod sql; +mod timestamp_fmt; + +/// 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 { + #[clap(long, default_value = "realtime")] + slot: String, + + #[clap(long, default_value = "postgresql://postgres@localhost:5432/postgres")] + connection: String, + + #[clap(long, default_value = "supabase_multiplayer")] + publication: String, + + // Exit when no work remains + #[clap(long)] + exit_on_no_work: bool, +} + +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) => { + warn!("Error: {}", err); + + if args.exit_on_no_work { + return (); + } + } + _ => continue, + }; + info!("Stream interrupted. Restarting in 5 seconds"); + sleep(time::Duration::from_secs(5)); + } +} + +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(err) => { + return Err(errors::Error::PostgresConnectionError(format!("{}", err))); + } + }; + + // Run pending migrations + sql::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![ + "--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,truncate", + &format!("--slot={}", args.slot), + "--create-slot", + "--if-not-exists", + "--start", + "--no-loop", + ]) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn(); + + match cmd { + Err(err) => Err(errors::Error::PgRecvLogicalError(format!("{}", err))), + Ok(mut cmd) => { + info!("pg_recvlogical started"); + // Reading from stdin + let stdin = cmd.stdout.as_mut().unwrap(); + let stdin_reader = io::BufReader::new(stdin); + let stdin_lines = stdin_reader.lines(); + + // Load initial snapshot of subscriptions + info!("Snapshot of subscriptions loading"); + let mut subscriptions = match subscription.load::(conn) { + Ok(subscriptions) => subscriptions, + Err(err) => { + cmd.kill().unwrap(); + error!("Error loading subscriptions: {}", err); + return Err(errors::Error::Subscriptions(format!("{}", err))); + } + }; + info!("Snapshot of subscriptions loaded"); + + // Iterate input data + for input_line in stdin_lines { + match input_line { + Ok(line) => { + let result_record = serde_json::from_str::(&line); + match result_record { + Ok(wal2json_record) => { + // Update subscriptions if needed + realtime::update_subscriptions( + &wal2json_record, + &mut subscriptions, + conn, + ); + + // New + let walrus = process_record( + &wal2json_record, + &subscriptions, + publication, + 1024 * 1024, + conn, + ); + + match walrus { + Ok(rows) => { + for row in rows { + match serde_json::to_string(&row) { + 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: {}", + err + ) + } + } + } + } + Err(err) => { + cmd.kill().unwrap(); + error!("WALRUS Error: {}", err); + return Err(errors::Error::Walrus(format!("{}", err))); + } + } + } + Err(err) => error!("Failed to parse: {}", err), + } + } + Err(err) => error!("Error: {}", err), + } + } + match cmd.wait() { + Ok(_) => Ok(()), + Err(err) => Err(errors::Error::PgRecvLogicalError(format!("{}", err))), + } + } + } +} + +fn process_record<'a>( + rec: &'a wal2json::Record, + subscriptions: &Vec, + publication: &str, + max_record_bytes: usize, + conn: &mut PgConnection, +) -> Result>, errors::Error> { + /* + * 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)?; + + // Subscriptions to the current entity + let entity_subscriptions: Vec<&realtime::Subscription> = subscriptions + .iter() + .filter(|x| &x.schema_name == &rec.schema) + .filter(|x| &x.table_name == &rec.table) + .map(|x| x) + .collect(); + + let is_subscribed_to = entity_subscriptions.len() > 0; + let is_rls_enabled = filters::table::row_level_security::is_rls_enabled(table_oid, 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 = 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) { + debug!("Early exit. Not in pub or no one listening"); + return Ok(vec![]); + } + + // Postgres role names of subscribed users + let subscribed_roles: Vec<&String> = entity_subscriptions + .iter() + .map(|x| &x.claims_role_name) + .unique() + .collect(); + + let mut result: Vec = vec![]; + + // If the table has no primary key, return + if action != realtime::Action::DELETE && !rec.has_primary_key() { + let r = realtime::WALRLS { + wal: realtime::Data { + schema: rec.schema, + table: rec.table, + r#type: action.clone(), + commit_timestamp: &rec.timestamp, + columns: vec![], + record: HashMap::new(), + old_record: None, + }, + is_rls_enabled, + subscription_ids: subscriptions + .iter() + .map(|x| x.subscription_id.clone()) + .collect(), + errors: vec!["Error 400: Bad Request, no primary key"], + }; + result.push(r); + return Ok(result); + } + + for role in subscribed_roles { + /* + * Role Level Filters + */ + + // Subscriptions to current entity + role + let entity_role_subscriptions: Vec<&realtime::Subscription> = entity_subscriptions + .iter() + .filter(|x| &x.claims_role_name == role) + .map(|x| *x) + .collect(); + + // 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 columns.len() == 0 { + let r = realtime::WALRLS { + wal: realtime::Data { + schema: rec.schema, + table: rec.table, + r#type: action.clone(), + commit_timestamp: &rec.timestamp, + columns: vec![], + record: HashMap::new(), + old_record: None, + }, + is_rls_enabled, + subscription_ids: entity_role_subscriptions + .iter() + .map(|x| x.subscription_id.clone()) + .collect(), + errors: vec!["Error 401: Unauthorized"], + }; + result.push(r); + } else { + // Repr of columns for internal use + let walcols: Vec = rec + .columns + .as_ref() + .unwrap_or(&vec![]) + //.unwrap_or(rec.identity.as_ref().unwrap_or(&vec![])) + .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_column_names.contains(&col.name), + }) + .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_column_names.contains(&col.name), + }) + // 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() + }); + + // User Defined Filters + let subscriptions_to_notify: Vec<&realtime::Subscription>; + + 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); + } + } + + 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 { + schema: rec.schema, + table: rec.table, + r#type: action.clone(), + commit_timestamp: &rec.timestamp, + columns, + record: record_elem, + old_record: old_record_elem, + }, + is_rls_enabled, + 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"], + false => vec![], + }, + }; + result.push(r); + } + } + + Ok(result) +} + +#[cfg(test)] +mod tests { + extern crate diesel; + use crate::models::{realtime, wal2json}; + use crate::realtime::{Subscription, UserDefinedFilter}; + use crate::sql::schema::realtime::subscription::dsl::*; + use chrono::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 INTEGER_OID: u32 = 23; + //const INT8OID: i32 = 20; + //const TEXTOID: i32 = 25; + + pub fn establish_connection() -> PgConnection { + let database_url = "postgresql://postgres:password@localhost:5501/postgres"; + PgConnection::establish(&database_url).unwrap() + } + + 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 = '{}') THEN + RAISE NOTICE 'Role {} already exists. Skipping.'; + ELSE + 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(); + + 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(conn) + .unwrap(); + } + + 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(format!( + "grant select on all tables in schema \"{}\" to {};", + schema, role + )) + .execute(conn) + .unwrap(); + } + + pub fn truncate(schema: &str, table: &str, conn: &mut PgConnection) { + diesel::sql_query(format!("truncate table \"{}\".\"{}\";", schema, table)) + .execute(conn) + .unwrap(); + } + + fn drop_table(schema: &str, table: &str, conn: &mut PgConnection) { + diesel::sql_query(format!( + "drop table if exists \"{}\".\"{}\";", + schema, table + )) + .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(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: "notes1", + 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(); + + 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", "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", "notes2", &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 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, + &subscriptions, + "supabase_multiplayer", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes2", + r#type: realtime::Action::INSERT, + 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: None, + }, + is_rls_enabled: false, + subscription_ids: vec![sub_id], // A SUBSCRIBER EXISTS + errors: vec![], + }]; + + assert_eq!(res, expected); + } + + #[test] + fn test_simple_update() { + 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", "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", "notes4", &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 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 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![], + }]; + + 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", + 1024 * 1024, + &mut conn, + ) + .unwrap(); + + let expected = vec![realtime::WALRLS { + wal: realtime::Data { + schema: "public", + table: "notes5", + r#type: realtime::Action::DELETE, + commit_timestamp: &rec.timestamp, + columns: vec![realtime::Column { + name: "id".to_string(), + type_: "int4".to_string(), + }], + 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!(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); + } + + #[test] + fn test_quoted_type_schema_and_table() { + 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 \"dEv\".\"Color\" cascade;") + .execute(&mut conn) + .unwrap(); + + 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 \"dEv\".\"Color\" primary key);", + ) + .execute(&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(); + + 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 + })), + 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", + "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: true, + subscription_ids: vec![sub_id], + errors: vec![], + }]; + + 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); + } + } +} diff --git a/worker/walrus/src/models/mod.rs b/worker/walrus/src/models/mod.rs new file mode 100644 index 0000000..6c4d81f --- /dev/null +++ b/worker/walrus/src/models/mod.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 new file mode 100644 index 0000000..f1f3326 --- /dev/null +++ b/worker/walrus/src/models/realtime.rs @@ -0,0 +1,495 @@ +use crate::models::wal2json; +use crate::sql::schema::realtime::subscription::dsl::*; +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, +} + +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, Deserialize, Clone, Debug, Eq, PartialEq)] +pub struct Column { + pub name: String, + #[serde(rename(serialize = "type", deserialize = "type"))] + pub type_: String, +} + +#[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: u32, + // 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, Ord, PartialOrd, +)] +#[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, QueryId)] +#[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, + }) + } +} + +#[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![]); + } +} diff --git a/worker/walrus/src/models/wal2json.rs b/worker/walrus/src/models/wal2json.rs new file mode 100644 index 0000000..a2140d6 --- /dev/null +++ b/worker/walrus/src/models/wal2json.rs @@ -0,0 +1,55 @@ +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, +} + +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 + } +} diff --git a/worker/walrus/src/models/walrus.rs b/worker/walrus/src/models/walrus.rs new file mode 100644 index 0000000..75cccb3 --- /dev/null +++ b/worker/walrus/src/models/walrus.rs @@ -0,0 +1,135 @@ +use serde::Serialize; + +/// An internal representation of columns used when passing column data to SQL is required +/// (user defined filters) +#[derive(Serialize, Debug)] +pub struct Column<'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, +} + +#[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" + ] +}"# + ) + } +} diff --git a/worker/walrus/src/sql/migrations.rs b/worker/walrus/src/sql/migrations.rs new file mode 100644 index 0000000..4ed4657 --- /dev/null +++ b/worker/walrus/src/sql/migrations.rs @@ -0,0 +1,21 @@ +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/mod.rs b/worker/walrus/src/sql/mod.rs new file mode 100644 index 0000000..c80cbca --- /dev/null +++ b/worker/walrus/src/sql/mod.rs @@ -0,0 +1,2 @@ +pub mod migrations; +pub mod schema; diff --git a/worker/walrus/src/sql/schema.rs b/worker/walrus/src/sql/schema.rs new file mode 100644 index 0000000..f64ae3a --- /dev/null +++ b/worker/walrus/src/sql/schema.rs @@ -0,0 +1,17 @@ +pub mod realtime { + table! { + realtime.subscription (id) { + id -> Int8, + subscription_id -> Uuid, + entity -> Oid, + //filters -> Array>, + filters -> Array, + claims -> Jsonb, + claims_role -> Int4, + created_at -> Timestamp, + schema_name -> Text, + table_name -> Text, + claims_role_name -> Text, + } + } +} 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 {} diff --git a/worker/walrus/src/timestamp_fmt.rs b/worker/walrus/src/timestamp_fmt.rs new file mode 100644 index 0000000..47c3208 --- /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 = "%Y-%m-%dT%H:%M:%S%.fZ"; + +// 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) +}