Tenant-safe queries

What you will be able to do

Make it hard to write a query that forgets the organization. Then make the database hide another organization’s rows, even when a query forgets it anyway.

Before you startWhat to know and to have readyAbout 55 minutes

What you need to know

• The four checks of Authorization in the API, and the composite foreign keys of Users and organizations.
• Transactions, from Transactions; the connection pool, from Connection management; query plans, from Read a query plan.

Your project

• ~/projects/ticket-api at the snapshot end-of-m04, from the last lesson: 17 migrations, 218 passed, 1 deselected. The provider is not needed in this lesson.

A hotel room has two locks

Think of a hotel. At check-in, the front desk checks your passport and gives you a card for room 214. That is the first lock: a person checks who you are and which room is yours.

The door of room 214 has its own lock. It opens only for a card made for 214. If the front desk makes a mistake, or a guest walks down the wrong corridor, the door still stays shut. That is the second lock.

But the hotel also has a master key. The manager’s card opens every door, and the door locks do not stop it. A second lock helps only if the people who walk the corridors do not carry the master key.

The ticket API now has a front desk: the four checks. This lesson adds the door locks in the database, and it finds a master key in your own project.

The problem: one forgotten WHERE

The API has dozens of queries. Every one of them must say tenant_id = …. Now imagine one new query, written on a busy day, that forgets it:

SELECT count(*) FROM tickets;

In your database, it counts 240 tickets: 200 of Larkfield and 40 of Bramble Books. A list or a search built on that query mixes the two shops, as the attempts of Module 1 did. The four checks of the last lesson do not help here. They let Ines in to Bramble Books, and then the query gives her everything.

The mistake is easy to make and hard to see. It passes every test that uses only one shop. You need protection at two levels: in the code that writes queries, and in the database that runs them.

Lock 1: the organization is a required argument

Open ticket_api/repository.py. Every function that reads or writes an organization’s records has the same parameter:

def get_ticket(conn: psycopg.Connection, ticket_id: str, *, tenant_id: str) -> dict:

The * makes tenant_id keyword-only: a caller must write tenant_id=… by name. It has no default value, so a caller cannot leave it out. Try to forget it:

python - <<'EOF'
import psycopg
from psycopg.rows import dict_row
from ticket_api import repository
from ticket_api.config import load_env, read_secret
load_env()
with psycopg.connect(read_secret("DATABASE_URL"), row_factory=dict_row) as conn:
    try:
        repository.get_ticket(conn, "T-40001")
    except TypeError as e:
        print("TypeError:", e)
    try:
        repository.get_ticket(conn, "T-40001", tenant_id="larkfield")
    except Exception as e:
        print(type(e).__name__, e)
    print(repository.get_ticket(conn, "T-40001", tenant_id="bramble")["subject"])
EOF

TypeError: get_ticket() missing 1 required keyword-only argument: 'tenant_id'
NotFound Ticket T-40001 does not exist.
Charged twice for order BB-310729

PowerShell has no <<'EOF': save the Python lines as lock1.py in the project folder, and run python lock1.py.

Three calls, three results. Without the organization, Python refuses the call before any SQL runs. With Larkfield, Bramble’s ticket “does not exist”, which is exactly the 404 of the last lesson. With Bramble Books, the ticket comes back. A forgotten organization is now a loud error in development, not a quiet leak in production.

Inside each function, the organization is the first condition of every statement. The ticket list starts like this:

where, params = ["t.tenant_id = %(tenant_id)s"], {"limit": limit + 1, "tenant_id": tenant_id}

The filters for status, team and customer come after it. A test, test_every_function_that_reads_an_organizations_records_needs_the_organization, checks the signature of nine functions. If someone adds a default value, the test fails.

The foreign keys of Module 1 guard the writes. A message points to its ticket by both columns, (tenant_id, ticket_id). So the database refuses a Larkfield message on Bramble’s T-40001 with a ForeignKeyViolation, as you saw in Users and organizations. add_message() writes the caller’s organization into the new row, never a value from the body.

Lock 1 depends on people. It works as long as every query goes through these functions, and every function follows the rule. A script, a report or a quick fix in psql can still forget.

Lock 2: row-level security

PostgreSQL can filter rows itself. A table with row-level security has policies: rules that decide, for each row, whether the current user may see or write it. The database adds the rule to every query on the table, written or forgotten. It is the door lock of the hotel.

Migration 017_row_level_security.sql adds one policy to each of the seven tables that hold an organization’s records:

ALTER TABLE tickets ENABLE ROW LEVEL SECURITY;
CREATE POLICY tenant_isolation ON tickets
    USING (tenant_id = current_setting('app.tenant_id', true))
    WITH CHECK (tenant_id = current_setting('app.tenant_id', true));

USING filters the rows that a query reads. WITH CHECK refuses rows that a query writes. Both compare the row’s tenant_id with a setting, app.tenant_id. The API sets it at the start of each transaction, from the checked caller.

A tiny example, step by step

The script scripts/rls_demo.py runs the forgotten query four times, with different settings. It changes nothing: each step ends with a rollback.

python -m scripts.rls_demo

SELECT count(*) FROM tickets;
  1. as the Compose user, no organization              user tickets    app.tenant_id (empty)    -> 240 tickets
  2. as the Compose user, organization bramble         user tickets    app.tenant_id bramble    -> 240 tickets
  3. as ticket_app, no organization                    user ticket_app app.tenant_id (empty)    -> 0 tickets
  4. as ticket_app, organization bramble               user ticket_app app.tenant_id bramble    -> 40 tickets
  5. as ticket_app for bramble, insert a message for Larkfield's T-30002:
     InsufficientPrivilege: new row violates row-level security policy for table "messages"

The same query that forgot the organization, four times. Steps 1 and 2 run as the Compose user tickets; steps 3 and 4 as the role ticket_app. Steps 2 and 4 set app.tenant_id to bramble.

Text description of the chart

A column chart with four columns. Step 1, the Compose user without an organization: 240. Step 2, the Compose user with the organization bramble: still 240. Step 3, the role ticket_app without an organization: 0. Step 4, ticket_app with bramble: 40, Bramble Books’ tickets only.

Step 2 is the surprise. The organization is set, and the policy exists, but the count is still 240. Look at the two roles in psql. Open it with docker compose exec db psql -U tickets -d tickets:

SELECT rolname, rolsuper, rolbypassrls FROM pg_roles WHERE rolname IN ('tickets', 'ticket_app');

  rolname   | rolsuper | rolbypassrls
------------+----------+--------------
 tickets    | t        | t
 ticket_app | f        | f
(2 rows)

The user tickets comes from compose.yaml. The PostgreSQL image makes it a superuser, and policies never apply to a superuser. It is the master key of the hotel. The table’s owner also skips its policies, unless the table has FORCE ROW LEVEL SECURITY. So the migration creates a normal role, ticket_app, without these powers, and the API switches to it in every transaction. Steps 3 and 4 show the result. Without an organization, ticket_app sees nothing at all. With Bramble, it sees Bramble’s 40 tickets.

Step 5 tests WITH CHECK. The transaction is for Bramble, and the new message names Larkfield, so the database refuses it.

Two locks for one rule. The application refuses a call without the organization; the database filters a query that forgot it, but only when the API is not a superuser.

Text description of the diagram

On the left, lock 1: a call to get_ticket without tenant_id fails in Python with a TypeError, before any SQL runs. On the right, lock 2: the API opens a transaction for Bramble Books, switches to the role ticket_app and sets app.tenant_id. A query that forgot WHERE tenant_id passes through the policy tenant_isolation and returns 40 tickets, only Bramble’s. The same query as the superuser tickets skips the policy and returns 240.

How the API sets it

Open ticket_api/db.py. When DB_ROW_SECURITY=true, which is the default, each transaction for an organization starts with one statement:

conn.execute(
    "SELECT set_config('role', 'ticket_app', true),"
    " set_config('app.tenant_id', %s, true)",
    (tenant_id,),
)

The routes ask for it with db.connection(caller.tenant_id), as you saw in tickets.py. The last argument, true, means local to the transaction. Both settings end with the transaction, before the connection goes back to the pool. This detail matters, as the failure cases below show.

Try one transaction by hand in psql. Predict the two counts first:

BEGIN;
SET LOCAL ROLE ticket_app;
SELECT set_config('app.tenant_id', 'larkfield', true);
SELECT count(*) FROM tickets;
SELECT count(*) FROM messages;
ROLLBACK;

The counts are 200 and 331: Larkfield’s tickets and messages, as in the load table of Module 1. After ROLLBACK, SELECT current_user; shows tickets again. The tables tenants, users and memberships have no policy, so SELECT count(*) FROM tenants; inside the transaction still shows 2. The membership check of the last lesson must read them before any organization is set.

Some programs still connect as the superuser, on purpose. They are the migrations, the loader scripts.load, the tests’ conn fixture, and the background worker’s claim of a job in the next module. They work across organizations, or they create the tables. Each of them is a master key, so keep the list short and known.

What it costs: measured

A second lock costs something. The API sends one more statement per request. The bench script of the databases course shows it, with Ines’s token on Bramble’s list. It starts its own copy of the API on port 8765, so that port must be free:

python -m scripts.bench --tenant bramble --user usr-ines
python -m scripts.bench --tenant bramble --user usr-ines --row-security off

The last line of each run says errors: 0 {200: 200} | database queries per request: 3, then … per request: 2. The three queries are the membership check, the settings and the page. The times are left out here: on your computer they will differ.

The course measured the time on the large data: 300,000 Larkfield tickets and Bramble’s 40. Each case ran 5 passes of 200 requests, 4 at a time. Little else ran on the computer during the passes.

The ticket list endpoint with row-level security off and on, on the large data: five passes per case. With the lock on, each request runs one more statement; the passes of off and on still overlap.

Text description of the chart

Four rows of five dots each, one dot per pass, in milliseconds:

• Larkfield, row-level security off: 4.3, 4.2, 4.3, 4.5 and 4.2; median 4.3.
• Larkfield, on: 4.4, 4.4, 4.5, 4.4 and 5.0; median 4.4.
• Bramble Books, off: 4.3, 4.0, 4.6, 5.0 and 4.6; median 4.6.
• Bramble Books, on: 5.2, 5.0, 5.0, 7.8 and 4.8; median 5.0.

The dots of off and on overlap in both organizations.

Read the dots, not only the medians. With the lock on, Larkfield’s median was 4.4 ms against 4.3, and Bramble’s was 5.0 against 4.6. That is 0.1 to 0.4 ms per request: the cost of the extra statement that sets the organization. The dots of off and on overlap, so the cost is at the edge of the noise. Inside the database, the list query for Larkfield took 0.100 ms, without policies and as ticket_app. The plan shows the policy as a One-Time Filter, checked once per query:

One-Time Filter: (current_setting('app.tenant_id'::text, true) = 'larkfield'::text)

Migration 016_tenant_indexes.sql changed the queue index to start with tenant_id, because every list query now starts with it. Check the plan for Bramble:

python -m scripts.explain queue --tenant bramble

->  Bitmap Index Scan on tickets_tenant_queue_idx  (cost=0.00..4.29 rows=2 width=0) (actual time=0.067..0.068 rows=3.00 loops=1)
      Index Cond: ((tenant_id = 'bramble'::text) AND (status = 'open'::text) AND (team = 'billing'::text))

This is two lines of a longer plan. One honest note: the course expected a “noisy neighbour” without this index. That means a small organization whose list reads through the big organization’s rows first. On the large data it did not show. With the databases course’s old index, Bramble’s list took 0.035 ms instead of 0.026 ms. The planner used another index that starts with tenant_id, the key (tenant_id, ticket_id). Its filter removed 37 rows, all of them Bramble’s own tickets, none of Larkfield’s. The new index is still the right shape, but at this size it is not dramatic.

When to use it, and when not

Use lock 1 always: it is free, and it catches the mistake where it happens. Add row-level security when many people write queries, when scripts and reports touch the same tables, or when one leak would be very costly. In this course it is on by default, because it costs one statement per request, a fraction of a millisecond.

Do not treat it as the only lock. The policies are far from the code, so a developer can forget that they exist, and a superuser skips them without any message. It also needs care with pools: the setting must end with the transaction. And it does not replace the four checks. The database knows the organization only because the API set it from a checked membership.

Common mistakes

Common mistake

Row-level security is on, and the query still returns 240 rows

Symptom: you enabled the policies and set app.tenant_id, but every count still covers both shops. This is step 2 of rls_demo.

Cause: the connection’s user is a superuser, or the table’s owner. Policies do not apply to them, and PostgreSQL gives no warning.

Diagnose: run SELECT current_user; inside the transaction, and check rolsuper in pg_roles.

Fix: switch to a normal role in each transaction, as db.py does, or connect as a normal role from the start. For a table owner, add ALTER TABLE … FORCE ROW LEVEL SECURITY.

Common mistake

A setting outlives its request

Symptom: a request without an organization sees another organization’s rows. It happens only sometimes, and only under load.

Cause: the code made the setting for the whole session (set_config(…, false) or SET without LOCAL). The pool gives the same connection to the next request, setting included. This constructed script shows it. Save it as setting_leak.py in the project folder:

# Constructed for the course: a setting that outlives its request.
from psycopg_pool import ConnectionPool

from ticket_api.config import load_env, read_secret

load_env()
for local in (False, True):
    with ConnectionPool(read_secret("DATABASE_URL"), min_size=1, max_size=1) as pool:
        with pool.connection() as conn:  # request 1: Ines, Bramble Books
            conn.execute("SELECT set_config('app.tenant_id', 'bramble', %s)", (local,))
        with pool.connection() as conn:  # request 2: the next caller, the same connection
            value = conn.execute("SELECT current_setting('app.tenant_id', true)").fetchone()[0]
        print(f"is_local={local}: request 2 sees app.tenant_id = {value!r}")

python setting_leak.py

is_local=False: request 2 sees app.tenant_id = 'bramble'
is_local=True: request 2 sees app.tenant_id = ''

Fix: always make both settings local to the transaction (true, or SET LOCAL). The test test_the_settings_end_with_the_transaction checks it. Delete setting_leak.py after the test.

Common mistake

TypeError: missing 1 required keyword-only argument: 'tenant_id'

Symptom: a new route or script stops with this error.

Cause: lock 1 works: the call has no organization.

Fix: pass the organization of the checked caller, tenant_id=caller.tenant_id. Never add a default value to make the error go away.

Reset. rls_demo and the psql transaction change nothing. If you inserted rows by hand, run python -m scripts.load --reset.

See it end to endAdd sign-in and roles to an existing APICase study

The case study gives every query of a data API a required department argument. Without it, a parks member got 94 transport requests in the first 100. With it, none.

Knowledge check

Check your understanding

Explain why setting app.tenant_id to bramble did not change the count for the Compose user tickets, and what the API does about it.

Show the answer

The user tickets is a superuser (rolsuper and rolbypassrls are true), and PostgreSQL never applies row-level security policies to a superuser. So the count stays 240. The API switches to the normal role ticket_app at the start of each transaction, with set_config('role', 'ticket_app', true), and sets app.tenant_id in the same statement. Then the policies apply.

Check your understanding

Predict: in a transaction as ticket_app with app.tenant_id = 'bramble', what do these return? (a) SELECT count(*) FROM messages; (b) SELECT count(*) FROM tickets WHERE ticket_id = 'T-30002'; (c) an INSERT of a message with tenant_id 'bramble' for T-40001.

Show the answer

(a) 58: Bramble’s messages only, as in the load table. (b) 0: T-30002 is Larkfield’s, and the policy hides it, as if it did not exist. (c) The insert works: the row’s tenant_id matches the setting, so WITH CHECK passes, and the composite foreign key finds Bramble’s T-40001. Roll back so that the extra message does not stay.

Check your understanding

Practical check: run python -m scripts.rls_demo. Which step shows the master key, and which line shows that a write into another organization is refused?

Show the answer

Step 2: the Compose user tickets, with app.tenant_id set to bramble, still counts 240 tickets. Step 5 ends with InsufficientPrivilege: new row violates row-level security policy for table "messages".

Recap

• Lock 1: tenant_id is a required keyword argument of every repository function, and the first condition of every statement. A forgotten organization is a TypeError. Composite foreign keys refuse cross-organization links.
• Lock 2: row-level security policies filter and check every row by app.tenant_id. A forgotten WHERE returns 40 tickets for Bramble, not 240.
• Policies never apply to a superuser. The Compose user tickets is one, so the API switches to ticket_app in every transaction. Both settings are local to the transaction.
• Cost: one more statement per request (3 instead of 2): 0.1 to 0.4 ms per request in the course’s measurement, at the edge of the noise.

Next, the organization follows the data out of the tables: into file links, the similar-ticket search, document citations and the cache.

Next: File and retrieval access