ROLLBACK;
BEGIN transaction;

-- setup a role which cannot bypass rls
CREATE role authenticated;
grant usage ON schema PUBLIC TO authenticated;

ALTER DEFAULT privileges IN schema PUBLIC grant SELECT ON tables TO authenticated;

-- helper function to get user id
CREATE FUNCTION getUserId()
    returns int
    stable strict parallel safe leakproof
BEGIN atomic;
    RETURN current_setting('app.user_id')::int;
END;

-- users table
CREATE TABLE users (
    id int primary key generated always AS identity,
    name text
);

-- permission table
CREATE TABLE permissions (
    id int primary key generated always AS identity,
    user_id int references users(id),
    can_access_assets BOOLEAN
);
CREATE INDEX ON permissions (user_id);
CREATE INDEX ON permissions (user_id, can_access_assets);

-- assets table
CREATE TABLE assets (
    id int primary key generated always AS identity,
    owner_id int references users(id),
    url text
);
CREATE INDEX ON assets (owner_id);

-- setup rls for assets
ALTER TABLE assets enable ROW LEVEL security;
CREATE policy assets_policy ON assets FOR ALL using (
    EXISTS (
        SELECT 1
        FROM permissions p
        WHERE p.user_id = getUserId() AND p.can_access_assets
    ) AND (
        owner_id IS NULL OR
        owner_id = getUserId()
    )
);

-- insert example users
INSERT INTO users (name) VALUES
    ('Alice can access assets'),
    ('Bob cannot access assets');

-- insert example assets using generate_series
INSERT INTO assets (owner_id, url)
SELECT
    FLOOR(random() * 2 + 1),
    'http://example.com/asset/' || i
FROM generate_series(1, 1000000) AS i;

-- set user id (eg. from session)
SET local app.user_id = '1';

-- query as superuser with bypassrls
explain analyze
SELECT COUNT(*) FROM assets
WHERE (
    EXISTS (
        SELECT 1
        FROM permissions p
        WHERE p.user_id = 1 AND p.can_access_assets
    ) AND (
        owner_id IS NULL OR
        owner_id = 1
    )
);
/*
Aggregate  (cost=9029.74..9029.75 rows=1 width=8) (actual time=0.016..0.017 rows=1 loops=1)
  InitPlan 1 (returns $0)
    ->  Index Only Scan using permissions_user_id_can_access_assets_idx on permissions p  (cost=0.15..24.27 rows=6 width=0) (actual time=0.003..0.003 rows=0 loops=1)
          Index Cond: ((user_id = 1) AND (can_access_assets = true))
          Heap Fetches: 0
  ->  Result  (cost=112.84..9000.63 rows=9976 width=0) (actual time=0.004..0.004 rows=0 loops=1)
        One-Time Filter: $0
        ->  Bitmap Heap Scan on assets  (cost=112.84..9000.63 rows=9976 width=0) (never executed)
              Recheck Cond: ((owner_id IS NULL) OR (owner_id = 1))
              ->  BitmapOr  (cost=112.84..112.84 rows=10001 width=0) (never executed)
                    ->  Bitmap Index Scan on assets_owner_id_idx  (cost=0.00..53.92 rows=5000 width=0) (never executed)
                          Index Cond: (owner_id IS NULL)
                    ->  Bitmap Index Scan on assets_owner_id_idx  (cost=0.00..53.92 rows=5000 width=0) (never executed)
                          Index Cond: (owner_id = 1)
Planning Time: 0.206 ms
Execution Time: 0.038 ms
*/

-- perform query as authenticated user with rls
SET local role authenticated;

explain analyze
SELECT COUNT(*) FROM assets;
/*
Aggregate  (cost=9089.80..9089.81 rows=1 width=8) (actual time=34.791..34.794 rows=1 loops=1)
  InitPlan 1 (returns $0)
    ->  Index Only Scan using permissions_user_id_can_access_assets_idx on permissions p  (cost=0.16..24.28 rows=6 width=0) (actual time=0.004..0.004 rows=0 loops=1)
          Index Cond: ((user_id = (current_setting('app.user_id'::text))::integer) AND (can_access_assets = true))
          Heap Fetches: 0
  ->  Bitmap Heap Scan on assets  (cost=110.35..9073.15 rows=4988 width=0) (actual time=34.787..34.789 rows=0 loops=1)
        Recheck Cond: ((owner_id IS NULL) OR (owner_id = (current_setting('app.user_id'::text))::integer))
        Filter: $0
        Rows Removed by Filter: 499840
        Heap Blocks: exact=8334
        ->  BitmapOr  (cost=110.35..110.35 rows=10001 width=0) (actual time=7.306..7.308 rows=0 loops=1)
              ->  Bitmap Index Scan on assets_owner_id_idx  (cost=0.00..53.92 rows=5000 width=0) (actual time=0.005..0.005 rows=0 loops=1)
                    Index Cond: (owner_id IS NULL)
              ->  Bitmap Index Scan on assets_owner_id_idx  (cost=0.00..53.93 rows=5000 width=0) (actual time=7.301..7.301 rows=499840 loops=1)
                    Index Cond: (owner_id = (current_setting('app.user_id'::text))::integer)
Planning Time: 0.080 ms
Execution Time: 34.815 ms
*/

ROLLBACK;