diff --git a/.github/workflows/test-pets-rollback.yml b/.github/workflows/test-pets-rollback.yml index 57d8744..388fc4f 100644 --- a/.github/workflows/test-pets-rollback.yml +++ b/.github/workflows/test-pets-rollback.yml @@ -28,7 +28,7 @@ jobs: services: pg_db: - image: ghcr.io/constructive-io/docker/postgres-plus:17 + image: ghcr.io/constructive-io/docker/postgres-plus:18 env: POSTGRES_USER: postgres POSTGRES_PASSWORD: password diff --git a/.github/workflows/test-pgpm.yml b/.github/workflows/test-pgpm.yml index 45233d5..6864e55 100644 --- a/.github/workflows/test-pgpm.yml +++ b/.github/workflows/test-pgpm.yml @@ -16,7 +16,7 @@ concurrency: cancel-in-progress: true env: - PGPM_VERSION: '2.7.9' + PGPM_VERSION: '5.30.6' jobs: test-pgpm: @@ -30,7 +30,7 @@ jobs: services: pg_db: - image: ghcr.io/constructive-io/docker/postgres-plus:17 + image: ghcr.io/constructive-io/docker/postgres-plus:18 env: POSTGRES_USER: postgres POSTGRES_PASSWORD: password @@ -49,7 +49,7 @@ jobs: - name: Setup Node.js uses: actions/setup-node@v4 with: - node-version: '20' + node-version: '22' - name: Cache pgpm CLI uses: actions/cache@v4 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4b65497..d4bf8e8 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -27,7 +27,7 @@ jobs: services: pg_db: - image: ghcr.io/constructive-io/docker/postgres-plus:17 + image: ghcr.io/constructive-io/docker/postgres-plus:18 env: POSTGRES_USER: postgres POSTGRES_PASSWORD: password diff --git a/README.md b/README.md index ee38102..f07024c 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ The Python counterpart to [`pgsql-test`](https://www.npmjs.com/package/pgsql-tes * **Instant test DBs** — each one seeded, isolated, and UUID-named * **Per-test rollback** — every test runs in its own transaction with savepoint-based rollback via `before_each()`/`after_each()` -* **RLS-friendly** — test with role-based auth via `set_context()` +* **RLS-friendly** — `db` is a real non-superuser connection (`app_user`); switch roles and JWT claims via `set_context()` * **pgpm integration** — run database migrations using [pgpm](https://pgpm.io) (PostgreSQL Package Manager) * **Flexible seeding** — run `.sql` files, programmatic seeds, pgpm modules, or combine multiple strategies * **Auto teardown** — no residue, no reboots, just clean exits @@ -104,6 +104,10 @@ def test_my_function(db): assert result['result'] == expected_value ``` +> `db` connects as the non-superuser `app_user` (see [`pg` vs `db`](#pg-vs-db)). Your migrations +> must `GRANT USAGE`/`EXECUTE` to `anonymous`/`authenticated`/`administrator` for `db` to reach +> them — or use `conn.pg` when you only want to assert that the deploy happened. + ### pgpm with Dependencies If your module depends on other pgpm packages (like `@pgpm/faker`), install them first: @@ -220,23 +224,69 @@ Without per-test rollback, tests can interfere with each other: With `before_each()`/`after_each()`, each test is completely isolated, making your test suite reliable and deterministic. +## `pg` vs `db` + +`get_connections()` returns two **different** clients, mirroring the TypeScript `pgsql-test`: + +| Client | Connects as | Use it for | +|--------|-------------|------------| +| `pg` | the superuser from `PGUSER` | DDL, seeding, grants, asserting ground truth (bypasses RLS) | +| `db` | `app_user` (non-superuser), default role `anonymous` | the code under test — grants and RLS policies are enforced | + +On every `get_connections()` call, pgsql-test creates the `app_user` LOGIN role (if missing), +grants it membership in the `anonymous`, `authenticated` and `administrator` roles, grants it +`CONNECT` on the test database, and opens `db` with those credentials. Every query on `db` runs +after `SET LOCAL ROLE ` (default `anonymous`), so tables created by `pg` are invisible to +`db` until you `GRANT` access — exactly as they would be in production. + +Override the credentials/roles via `connection_options`: + +```python +conn = get_connections(connection_options={ + 'connection': {'user': 'app_user', 'password': 'app_password', 'role': 'anonymous'}, + 'roles': {'anonymous': 'anonymous', 'authenticated': 'authenticated', 'administrator': 'administrator'}, +}) +``` + ## RLS Testing -Test Row Level Security policies by switching contexts: +Seed with `pg`, then exercise policies with `db`: ```python -def test_rls_policy(db): +@pytest.fixture(scope='module') +def conn(): + c = get_connections() + c.pg.query(""" + CREATE TABLE documents (id SERIAL PRIMARY KEY, owner_id TEXT, title TEXT); + ALTER TABLE documents ENABLE ROW LEVEL SECURITY; + GRANT SELECT ON documents TO authenticated; + CREATE POLICY owner_only ON documents FOR SELECT TO authenticated + USING (owner_id = current_setting('jwt.claims.user_id', true)); + """) + c.pg.commit() + yield c + c.teardown() + +def test_rls_policy(conn): + db = conn.db db.before_each() - - # Set the user context - db.set_context({'app.user_id': '123'}) - - # Now queries will be filtered by RLS policies - result = db.many('SELECT * FROM user_data') - + + # anonymous (the default role) has no grant at all + with pytest.raises(psycopg2.errors.InsufficientPrivilege): + db.query('SELECT * FROM documents') + + # authenticated sees only its own rows + db.set_context({'role': 'authenticated', 'jwt.claims.user_id': '123'}) + rows = db.many_or_none('SELECT * FROM documents') + + db.clear_context() # back to anonymous, claims nulled db.after_each() ``` +`set_context()` turns `role` into `SET LOCAL ROLE` and every other key into +`set_config(key, value, true)`. Both are transaction-local and are re-applied before every +query, so they persist for the whole `before_each()`/`after_each()` window. + ## Seeding Strategies ### pgpm Modules @@ -300,8 +350,8 @@ conn = get_connections( Creates a new isolated test database and returns connection objects. Returns a `ConnectionResult` with: -- `pg`: PgTestClient connected as superuser -- `db`: PgTestClient for testing (same as pg for now) +- `pg`: PgTestClient connected as the superuser (bypasses RLS; use for setup/assertions) +- `db`: PgTestClient connected as `app_user` with default role `anonymous` (RLS enforced) - `admin`: DbAdmin for database management - `manager`: PgTestConnector managing connections - `teardown()`: Function to clean up @@ -316,7 +366,9 @@ Returns a `ConnectionResult` with: - `execute(sql, params?)`: Execute and return affected row count - `before_each()`: Start test isolation (transaction + savepoint) - `after_each()`: End test isolation (rollback) -- `set_context(dict)`: Set session variables for RLS testing +- `set_context(dict)`: Set `role` (`SET LOCAL ROLE`) and GUCs (`set_config(..., true)`) for RLS testing +- `get_context()`: Return the current context dict +- `clear_context()`: Null every GUC and restore the default role ## GitHub Actions Example @@ -356,7 +408,7 @@ jobs: - uses: actions/setup-node@v4 with: - node-version: '20' + node-version: '22' - name: Install pgpm run: npm install -g pgpm diff --git a/WORKSPACE_SETUP.md b/WORKSPACE_SETUP.md index 6fe693a..baa35e2 100644 --- a/WORKSPACE_SETUP.md +++ b/WORKSPACE_SETUP.md @@ -547,7 +547,7 @@ jobs: - name: Setup Node.js uses: actions/setup-node@v4 with: - node-version: "20" + node-version: "22" - name: Install pgpm run: npm install -g pgpm diff --git a/pyproject.toml b/pyproject.toml index 48471f3..f180b6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pgsql-test" -version = "0.2.0" +version = "0.3.0" description = "PostgreSQL testing framework for Python - instant, isolated databases with automatic transaction rollback" authors = ["Constructive "] license = "MIT" diff --git a/src/pgsql_test/__init__.py b/src/pgsql_test/__init__.py index 88071f4..646e474 100644 --- a/src/pgsql_test/__init__.py +++ b/src/pgsql_test/__init__.py @@ -19,4 +19,4 @@ "seed", ] -__version__ = "0.2.0" +__version__ = "0.3.0" diff --git a/src/pgsql_test/admin.py b/src/pgsql_test/admin.py index cc3dde6..6f0ce6e 100644 --- a/src/pgsql_test/admin.py +++ b/src/pgsql_test/admin.py @@ -239,6 +239,66 @@ def create_role( ) logger.info(f"Created role: {role_name}") + def role_exists(self, role_name: str) -> bool: + """Check whether a role exists.""" + conn = self._get_admin_connection() + with conn.cursor() as cur: + cur.execute("SELECT 1 FROM pg_roles WHERE rolname = %s", (role_name,)) + return cur.fetchone() is not None + + def grant_role(self, role_name: str, member: str) -> None: + """Grant membership in `role_name` to `member` (idempotent).""" + conn = self._get_admin_connection() + with conn.cursor() as cur: + cur.execute( + """ + SELECT 1 + FROM pg_auth_members am + JOIN pg_roles r ON am.roleid = r.oid + JOIN pg_roles m ON am.member = m.oid + WHERE r.rolname = %s AND m.rolname = %s + """, + (role_name, member), + ) + if cur.fetchone() is not None: + return + cur.execute( + sql.SQL("GRANT {} TO {}").format( + sql.Identifier(role_name), + sql.Identifier(member), + ) + ) + logger.debug(f"Granted {role_name} to {member}") + + def create_user_role(self, user: str, password: str, roles: list[str]) -> None: + """ + Create the app-level LOGIN user used by the `db` client and grant it + membership in the given NOLOGIN roles (created if missing). + + Mirrors pgsql-test (TS) `DbAdmin.createUserRole`. Granting `administrator` + is for the test harness only - never do this for a production app user. + """ + conn = self._get_admin_connection() + with conn.cursor() as cur: + for role in roles: + if self.role_exists(role): + continue + try: + cur.execute(sql.SQL("CREATE ROLE {} NOLOGIN").format(sql.Identifier(role))) + except psycopg2.errors.DuplicateObject: + pass + if not self.role_exists(user): + try: + cur.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD %s").format(sql.Identifier(user)), + (password,), + ) + except psycopg2.errors.DuplicateObject: + pass + for role in roles: + self.grant_role(role, user) + logger.info(f"Created user role {user} with grants: {roles}") + def grant_connect(self, role_name: str, database: str) -> None: """ Grant CONNECT privilege on a database to a role. diff --git a/src/pgsql_test/client.py b/src/pgsql_test/client.py index 326ac73..c1f91d7 100644 --- a/src/pgsql_test/client.py +++ b/src/pgsql_test/client.py @@ -11,6 +11,7 @@ from typing import Any import psycopg2 +from psycopg2 import sql as pgsql from psycopg2.extras import RealDictCursor from pgsql_test.types import PgConfig, QueryResult @@ -39,18 +40,25 @@ class PgTestClient: client.close() """ - def __init__(self, config: PgConfig, enhanced_errors: bool = True) -> None: + def __init__( + self, + config: PgConfig, + enhanced_errors: bool = True, + default_role: str | None = None, + ) -> None: """ Initialize the test client. Args: config: PostgreSQL connection configuration enhanced_errors: Whether to enhance error messages with PG details + default_role: Role restored by clear_context() (None = SET ROLE NONE) """ self._config = config self._enhanced_errors = enhanced_errors + self._default_role = default_role self._conn: psycopg2.extensions.connection | None = None - self._context: dict[str, str] = {} + self._context: dict[str, str | None] = {} self._in_transaction = False self._savepoint_name = "pgsql_test_savepoint" @@ -111,6 +119,7 @@ def query(self, sql: str, params: tuple[Any, ...] | None = None) -> QueryResult: conn = self.connection try: with conn.cursor() as cur: + self._apply_context(cur) cur.execute(sql, params) # Check if this is a SELECT-like query that returns rows if cur.description is not None: @@ -185,30 +194,40 @@ def execute(self, sql: str, params: tuple[Any, ...] | None = None) -> int: result = self.query(sql, params) return result.row_count - def set_context(self, context: dict[str, str]) -> None: + def set_context(self, context: dict[str, str | None]) -> None: """ - Set PostgreSQL session context variables. + Set PostgreSQL context variables for RLS testing. - Useful for simulating RLS contexts in tests. + The `role` key becomes `SET LOCAL ROLE`; every other key becomes + `set_config(key, value, true)`. Both are transaction-local, so the + context is re-applied before every query (like the TS client) and + therefore persists across the transaction opened by before_each(). + A value of None resets the setting (`SET LOCAL ROLE NONE` / NULL). Args: - context: Dictionary of context variables to set - e.g., {"role": "authenticated", "jwt.claims.user_id": "123"} + context: e.g. {"role": "authenticated", "jwt.claims.user_id": "123"} """ self._context.update(context) - self._apply_context() - def clear_context(self) -> None: - """Clear all context variables.""" - self._context = {} + def get_context(self) -> dict[str, str | None]: + """Return a copy of the current context settings.""" + return dict(self._context) - def _apply_context(self) -> None: - """Apply context variables to the current session.""" - conn = self.connection - with conn.cursor() as cur: - for key, value in self._context.items(): - # Use SET LOCAL so it only applies within the current transaction - cur.execute(f"SET LOCAL {key} = %s", (value,)) + def clear_context(self) -> None: + """Null out every context variable and restore the default role.""" + self._context = {key: None for key in self._context} + self._context["role"] = self._default_role + + def _apply_context(self, cur: Any) -> None: + """Apply context variables on the given cursor (transaction-local).""" + for key, value in self._context.items(): + if key == "role": + if value is None: + cur.execute("SET LOCAL ROLE NONE") + else: + cur.execute(pgsql.SQL("SET LOCAL ROLE {}").format(pgsql.Identifier(value))) + else: + cur.execute("SELECT set_config(%s, %s, true)", (key, value)) def begin(self) -> None: """Begin a new transaction.""" @@ -256,8 +275,6 @@ def before_each(self) -> None: """ self.begin() self.savepoint() - if self._context: - self._apply_context() def after_each(self) -> None: """ diff --git a/src/pgsql_test/connect.py b/src/pgsql_test/connect.py index ea952a1..3916317 100644 --- a/src/pgsql_test/connect.py +++ b/src/pgsql_test/connect.py @@ -15,7 +15,15 @@ from pgsql_test.admin import DbAdmin from pgsql_test.client import PgTestClient from pgsql_test.manager import PgTestConnector, generate_test_db_name -from pgsql_test.types import ConnectionOptions, PgConfig, SeedContext +from pgsql_test.types import ( + DEFAULT_APP_CONNECTION, + DEFAULT_ROLES, + AppConnection, + ConnectionOptions, + PgConfig, + RoleMapping, + SeedContext, +) logger = logging.getLogger(__name__) @@ -120,6 +128,11 @@ def db(): # Get configuration config = pg_config or get_pg_config_from_env() options = connection_options or {} + app: AppConnection = {**DEFAULT_APP_CONNECTION, **options.get("connection", {})} + roles: RoleMapping = {**DEFAULT_ROLES, **options.get("roles", {})} + app_user = app["user"] + app_password = app["password"] + default_role = app["role"] # Generate unique database name prefix = options.get("prefix", "pgsql_test_") @@ -137,6 +150,13 @@ def db(): admin = DbAdmin(admin_config, verbose=False) + # Roles are cluster-wide, so the app user is created once against the root db + admin.create_user_role( + app_user, + app_password, + [roles["anonymous"], roles["authenticated"], roles["administrator"]], + ) + # Create the test database template = options.get("template") if template: @@ -149,6 +169,8 @@ def db(): if extensions: admin.install_extensions(extensions, test_db_name) + admin.grant_connect(app_user, test_db_name) + # Create configuration for the test database test_config: PgConfig = { "host": config.get("host", "localhost"), @@ -178,10 +200,13 @@ def db(): logger.error(f"Seed adapter failed: {e}") # Continue without teardown to allow debugging raise + # Seeds run on the non-autocommit pg connection; make them visible to db + pg.commit() - # For now, db is the same as pg (both superuser) - # In the future, we can add app-level user support - db = pg + # The app-level client: a real non-superuser connection so RLS policies apply + db_config: PgConfig = {**test_config, "user": app_user, "password": app_password} + db = manager.get_client(db_config, default_role=default_role) + db.set_context({"role": default_role}) # Create teardown function def teardown_fn() -> None: diff --git a/src/pgsql_test/manager.py b/src/pgsql_test/manager.py index a22f233..d9634e1 100644 --- a/src/pgsql_test/manager.py +++ b/src/pgsql_test/manager.py @@ -95,12 +95,13 @@ def begin_teardown(self) -> None: """Mark that teardown has begun (prevents new clients).""" self._shutting_down = True - def get_client(self, config: PgConfig) -> PgTestClient: + def get_client(self, config: PgConfig, default_role: str | None = None) -> PgTestClient: """ Get a new test client for the given configuration. Args: config: PostgreSQL configuration for the client + default_role: Role the client falls back to in clear_context() Returns: A new PgTestClient instance @@ -111,7 +112,7 @@ def get_client(self, config: PgConfig) -> PgTestClient: if self._shutting_down: raise RuntimeError("PgTestConnector is shutting down; no new clients allowed") - client = PgTestClient(config) + client = PgTestClient(config, default_role=default_role) client.connect() self._clients.add(client) diff --git a/src/pgsql_test/types.py b/src/pgsql_test/types.py index 7f3a0e3..aed17b7 100644 --- a/src/pgsql_test/types.py +++ b/src/pgsql_test/types.py @@ -14,6 +14,22 @@ class PgConfig(TypedDict, total=False): password: str +class AppConnection(TypedDict, total=False): + """Credentials for the app-level (non-superuser) `db` client.""" + + user: str # Login role used by the `db` client (default: app_user) + password: str # Password for that role (default: app_password) + role: str # Default role applied via set_context() (default: anonymous) + + +class RoleMapping(TypedDict, total=False): + """Names of the NOLOGIN roles the app user is granted membership in.""" + + anonymous: str + authenticated: str + administrator: str + + class ConnectionOptions(TypedDict, total=False): """Options for database connections.""" @@ -21,6 +37,21 @@ class ConnectionOptions(TypedDict, total=False): root_db: str # Root database for admin operations (default: postgres) extensions: list[str] # Extensions to install template: str | None # Template database to use + connection: AppConnection # Credentials/default role for the `db` client + roles: RoleMapping # Role names granted to the app user + + +DEFAULT_APP_CONNECTION: AppConnection = { + "user": "app_user", + "password": "app_password", + "role": "anonymous", +} + +DEFAULT_ROLES: RoleMapping = { + "anonymous": "anonymous", + "authenticated": "authenticated", + "administrator": "administrator", +} @dataclass diff --git a/tests/fixtures/pgpm-workspace/.gitignore b/tests/fixtures/pgpm-workspace/.gitignore new file mode 100644 index 0000000..0ef1842 --- /dev/null +++ b/tests/fixtures/pgpm-workspace/.gitignore @@ -0,0 +1 @@ +extensions/ diff --git a/tests/fixtures/pgpm-workspace/packages/test-module/deploy/schemas/test_app.sql b/tests/fixtures/pgpm-workspace/packages/test-module/deploy/schemas/test_app.sql index 8796001..38253df 100644 --- a/tests/fixtures/pgpm-workspace/packages/test-module/deploy/schemas/test_app.sql +++ b/tests/fixtures/pgpm-workspace/packages/test-module/deploy/schemas/test_app.sql @@ -4,4 +4,6 @@ BEGIN; CREATE SCHEMA test_app; +GRANT USAGE ON SCHEMA test_app TO anonymous, authenticated, administrator; + COMMIT; diff --git a/tests/test_basic.py b/tests/test_basic.py index 39db38c..b053e68 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -115,11 +115,13 @@ def test_execute_returns_row_count(self, db_connection): def test_transaction_rollback(self, db_connection): """Test that before_each/after_each provides transaction isolation.""" + pg = db_connection.pg db = db_connection.db - # Create a table - db.query("CREATE TABLE rollback_test (id INT)") - db.connection.commit() + # Create a table as superuser and expose it to the app role + pg.query("CREATE TABLE rollback_test (id INT)") + pg.query("GRANT ALL ON rollback_test TO anonymous") + pg.commit() # Start test isolation db.before_each() @@ -162,7 +164,7 @@ def test_sqlfile_loads_schema(self, pg_config, sql_dir): ) try: # Verify tables were created - result = conn.db.query( + result = conn.pg.query( """ SELECT table_name FROM information_schema.tables @@ -176,12 +178,12 @@ def test_sqlfile_loads_schema(self, pg_config, sql_dir): assert "users" in tables # Verify data was inserted - users = conn.db.many("SELECT * FROM users ORDER BY id") + users = conn.pg.many("SELECT * FROM users ORDER BY id") assert len(users) == 2 assert users[0]["name"] == "Alice" assert users[1]["name"] == "Bob" - posts = conn.db.many("SELECT * FROM posts ORDER BY id") + posts = conn.pg.many("SELECT * FROM posts ORDER BY id") assert len(posts) == 3 finally: conn.teardown() @@ -214,7 +216,7 @@ def my_seed(ctx): seed_adapters=[seed.fn(my_seed)], ) try: - result = conn.db.one("SELECT value FROM fn_test") + result = conn.pg.one("SELECT value FROM fn_test") assert result["value"] == "seeded" finally: conn.teardown() @@ -252,7 +254,7 @@ def third_seed(ctx): try: assert execution_order == ["first", "second", "third"] - result = conn.db.many("SELECT step FROM compose_test ORDER BY step") + result = conn.pg.many("SELECT step FROM compose_test ORDER BY step") assert [r["step"] for r in result] == [1, 2] finally: conn.teardown() diff --git a/tests/test_example.py b/tests/test_example.py index ba896ae..5d62f0e 100644 --- a/tests/test_example.py +++ b/tests/test_example.py @@ -25,7 +25,8 @@ def db(): id SERIAL PRIMARY KEY, name TEXT NOT NULL, email TEXT UNIQUE - ) + ); + GRANT ALL ON users, users_id_seq TO anonymous; """)) ] ) diff --git a/tests/test_pets_rollback.py b/tests/test_pets_rollback.py index 51a1e18..bd38b66 100644 --- a/tests/test_pets_rollback.py +++ b/tests/test_pets_rollback.py @@ -31,7 +31,8 @@ def pets_db(): name TEXT NOT NULL, species TEXT NOT NULL, age INTEGER - ) + ); + GRANT ALL ON pets, pets_id_seq TO anonymous; """)) ] ) diff --git a/tests/test_pgpm_integration.py b/tests/test_pgpm_integration.py index 47c31b7..a2743a0 100644 --- a/tests/test_pgpm_integration.py +++ b/tests/test_pgpm_integration.py @@ -33,16 +33,17 @@ def pgpm_db(): seed.pgpm(module_path=str(FIXTURE_PATH), package=PACKAGE_NAME) ] ) - db = conn.db - db.before_each() - yield db - db.after_each() + conn.pg.before_each() + conn.db.before_each() + yield conn + conn.db.after_each() + conn.pg.after_each() conn.teardown() def test_pgpm_creates_schema(pgpm_db): """Test that pgpm deploy creates the test_app schema.""" - result = pgpm_db.one(""" + result = pgpm_db.pg.one(""" SELECT schema_name FROM information_schema.schemata WHERE schema_name = 'test_app' @@ -50,6 +51,12 @@ def test_pgpm_creates_schema(pgpm_db): assert result["schema_name"] == "test_app" +def test_app_role_can_use_granted_schema(pgpm_db): + """The non-superuser db client sees the schema because the migration granted USAGE.""" + result = pgpm_db.db.one("SELECT has_schema_privilege('test_app', 'USAGE') AS ok") + assert result["ok"] is True + + def test_pgpm_faker_available(pgpm_db): """ Test that @pgpm/faker is available after pgpm deploy. @@ -58,7 +65,7 @@ def test_pgpm_faker_available(pgpm_db): and the faker schema/functions are available. """ # Check if faker schema exists (installed via pgpm install @pgpm/faker) - result = pgpm_db.one_or_none(""" + result = pgpm_db.pg.one_or_none(""" SELECT schema_name FROM information_schema.schemata WHERE schema_name = 'faker' @@ -81,7 +88,7 @@ def test_pgpm_faker_city_function(pgpm_db): 4. Test can use faker functions """ # Check if faker schema exists first - schema_exists = pgpm_db.one_or_none(""" + schema_exists = pgpm_db.pg.one_or_none(""" SELECT schema_name FROM information_schema.schemata WHERE schema_name = 'faker' @@ -91,7 +98,7 @@ def test_pgpm_faker_city_function(pgpm_db): pytest.skip("@pgpm/faker not installed - run: cd tests/fixtures/pgpm-workspace/packages/test-module && pgpm install @pgpm/faker") # Test faker.city() function with Michigan state code - result = pgpm_db.one("SELECT faker.city('MI') as city") + result = pgpm_db.pg.one("SELECT faker.city('MI') as city") assert result["city"] is not None assert isinstance(result["city"], str) assert len(result["city"]) > 0 diff --git a/tests/test_rls.py b/tests/test_rls.py new file mode 100644 index 0000000..8f838d1 --- /dev/null +++ b/tests/test_rls.py @@ -0,0 +1,144 @@ +""" +Row-level security tests. + +These prove that `db` is a genuinely separate, non-superuser connection +(app_user) that is subject to role grants and RLS policies, while `pg` +is the superuser that bypasses them. If `db` were ever aliased to `pg` +again, every test in this module would fail. +""" + +import psycopg2 +import pytest + +from pgsql_test import get_connections + +SCHEMA = """ +CREATE TABLE documents ( + id SERIAL PRIMARY KEY, + owner_id TEXT NOT NULL, + title TEXT NOT NULL +); + +INSERT INTO documents (owner_id, title) VALUES + ('alice', 'alice-doc-1'), + ('alice', 'alice-doc-2'), + ('bob', 'bob-doc-1'); + +ALTER TABLE documents ENABLE ROW LEVEL SECURITY; + +GRANT SELECT, INSERT ON documents TO authenticated; +GRANT USAGE ON SEQUENCE documents_id_seq TO authenticated; + +CREATE POLICY documents_owner_select ON documents + FOR SELECT TO authenticated + USING (owner_id = current_setting('jwt.claims.user_id', true)); + +CREATE POLICY documents_owner_insert ON documents + FOR INSERT TO authenticated + WITH CHECK (owner_id = current_setting('jwt.claims.user_id', true)); +""" + + +@pytest.fixture(scope="module") +def conn(): + connection = get_connections() + connection.pg.query(SCHEMA) + connection.pg.commit() + yield connection + connection.teardown() + + +@pytest.fixture +def db(conn): + conn.db.before_each() + yield conn.db + conn.db.clear_context() + conn.db.after_each() + + +@pytest.fixture +def pg(conn): + conn.pg.before_each() + yield conn.pg + conn.pg.after_each() + + +def test_db_and_pg_are_different_connections(conn): + assert conn.db is not conn.pg + assert conn.pg.config["user"] != conn.db.config["user"] + assert conn.db.config["user"] == "app_user" + + pg_who = conn.pg.one("SELECT current_user, session_user, usesuper FROM pg_user " + "WHERE usename = session_user") + db_who = conn.db.one("SELECT current_user, session_user, usesuper FROM pg_user " + "WHERE usename = session_user") + assert pg_who["usesuper"] is True + assert db_who["usesuper"] is False + assert db_who["session_user"] == "app_user" + # default context role is applied on every query + assert db_who["current_user"] == "anonymous" + + +def test_superuser_pg_bypasses_rls(pg): + rows = pg.many("SELECT title FROM documents ORDER BY id") + assert len(rows) == 3 + + +def test_anonymous_has_no_table_access(db): + with pytest.raises(psycopg2.errors.InsufficientPrivilege): + db.query("SELECT * FROM documents") + + +def test_authenticated_sees_only_own_rows(db): + db.set_context({"role": "authenticated", "jwt.claims.user_id": "alice"}) + rows = db.many("SELECT title FROM documents ORDER BY id") + assert [r["title"] for r in rows] == ["alice-doc-1", "alice-doc-2"] + + db.set_context({"jwt.claims.user_id": "bob"}) + rows = db.many("SELECT title FROM documents ORDER BY id") + assert [r["title"] for r in rows] == ["bob-doc-1"] + + +def test_authenticated_without_claims_sees_nothing(db): + db.set_context({"role": "authenticated"}) + result = db.query("SELECT title FROM documents") + assert result.rows == [] + + +def test_insert_policy_enforced(db): + db.set_context({"role": "authenticated", "jwt.claims.user_id": "alice"}) + + db.execute("INSERT INTO documents (owner_id, title) VALUES ('alice', 'alice-doc-3')") + rows = db.many("SELECT title FROM documents ORDER BY id") + assert len(rows) == 3 + + with pytest.raises(psycopg2.errors.InsufficientPrivilege): + db.execute("INSERT INTO documents (owner_id, title) VALUES ('bob', 'forged')") + + +def test_insert_rolled_back_between_tests(db): + db.set_context({"role": "authenticated", "jwt.claims.user_id": "alice"}) + rows = db.many("SELECT title FROM documents ORDER BY id") + assert [r["title"] for r in rows] == ["alice-doc-1", "alice-doc-2"] + + +def test_clear_context_restores_default_role(db): + db.set_context({"role": "authenticated", "jwt.claims.user_id": "alice"}) + assert db.one("SELECT current_user AS u")["u"] == "authenticated" + + db.clear_context() + assert db.one("SELECT current_user AS u")["u"] == "anonymous" + assert db.one("SELECT current_setting('jwt.claims.user_id', true) AS v")["v"] in (None, "") + + +def test_administrator_role_membership_is_granted(conn, db): + # Schema change must be committed before db (which holds an open tx) reads it. + conn.pg.query("GRANT SELECT ON documents TO administrator") + conn.pg.query( + "CREATE POLICY documents_admin_select ON documents FOR SELECT TO administrator USING (true)" + ) + conn.pg.commit() + + db.set_context({"role": "administrator"}) + rows = db.many("SELECT title FROM documents") + assert len(rows) == 3