Implement document-level access control for RAG

Introduction

Teams that ship RAG over a shared vector store quickly discover that a bare cosine similarity search will happily return the top-k most relevant documents to anyone who asks — including documents the requesting user has no right to see. When you skip per-row access control, a junior contractor's question about office policy can pull back confidential board minutes that happen to be semantically close, and the leak surfaces in audit logs long after the damage is done. By the end of this lesson you will be able to extend a pgvector documents table with ownership, group, and classification columns, build a SecureRetriever that folds access control into the same SQL query as the similarity search, and document the model so reviewers can verify exactly which rows each user can see.

Key Terminology

• Owner-based access: the user identifier stored on each document row (owner_id); the document's creator can always retrieve it regardless of group or classification.
• Access groups: a JSONB array of group names attached to each document; a user retrieves the row when any of their group memberships overlap with the document's groups (the ?| operator).
• Classification level: a four-step sensitivity label (public, internal, confidential, restricted) compared against the user's clearance rank so users only see documents at or below their cleared level.
• Row-level security (RLS): PostgreSQL policies that filter every SELECT against the documents table using session variables, enforcing access rules even if the application forgets to add the WHERE clause.

Concepts

The core idea is that access control belongs inside the retrieval query, not as a post-filter on the result list. If you fetch the top-k similar documents first and then drop the ones the user cannot read, you can end up returning fewer than k results — or, worse, returning a padded set that no longer reflects true similarity ranking.

Access control here is the union of three independent checks, evaluated per row:

1. Owner match — owner_id = user.user_id.
2. Group overlap — access_groups ?| user.groups (JSONB array intersection).
3. Public classification — classification_level = 'public'.

A row is visible if any of the three is true. On top of that, an AND-bound clearance check (classification_level = ANY(allowed_levels)) caps what a user can see by sensitivity rank, even for documents they would otherwise own or share a group with.

Defense in depth comes from layering: the SecureRetriever enforces the same predicates the RLS policy enforces. The application layer is the fast path; RLS is the safety net for the day someone writes a new query and forgets the filter.

Code Walkthrough

Extending the pgvector Schema

The pgvector table schema needs additional columns to support access control. The standard schema stores the embedding vector, document text, and basic metadata. The extended schema adds owner_id for document ownership, access_groups as a JSONB array of group names, and classification_level as an enumerated type with four levels: public, internal, confidential, and restricted.

The SQL migration below extends the existing pgvector documents table with access control columns and creates appropriate indexes. The access_groups column uses PostgreSQL JSONB type to store a flexible array of group names, enabling efficient containment queries with the ?| operator. The classification_level column uses a custom enum type that enforces valid values at the database level.

1CREATE TYPE doc_classification AS ENUM (
2    'public', 'internal', 'confidential', 'restricted'
3);
4
5ALTER TABLE documents ADD COLUMN owner_id VARCHAR(128) NOT NULL DEFAULT '';
6ALTER TABLE documents ADD COLUMN access_groups JSONB NOT NULL DEFAULT '[]';
7ALTER TABLE documents ADD COLUMN classification_level doc_classification
8    NOT NULL DEFAULT 'internal';
9
10CREATE INDEX idx_documents_owner ON documents(owner_id);
11CREATE INDEX idx_documents_classification
12    ON documents(classification_level);
13CREATE INDEX idx_documents_access_groups
14    ON documents USING GIN(access_groups);

Building the SecureRetriever

The SecureRetriever class wraps standard pgvector similarity search with access control enforcement. Instead of executing a bare cosine similarity query, it adds WHERE clauses that filter results based on the requesting user's identity, group memberships, and clearance level. The key design decision is that access control filtering happens in the same SQL query as the similarity search, not as a post-processing step. This ensures that the top-k results are the k most similar documents that the user is actually authorized to access.

The SecureRetriever class below accepts a database connection, a user context object containing the user's identity and permissions, and query parameters. The retrieve method constructs a parameterized SQL query that combines cosine similarity ordering with access control WHERE clauses. The access control logic allows retrieval when any of three conditions is met: the user owns the document, the user belongs to one of the document's access groups, or the document is classified as public.

Code snippetpython
1from dataclasses import dataclass
2from typing import Optional
3
4CLASSIFICATION_HIERARCHY = {
5    "public": 0,
6    "internal": 1,
7    "confidential": 2,
8    "restricted": 3,
9}
10
11@dataclass
12class UserContext:
13    """Represents the requesting user's access permissions."""
14    user_id: str
15    groups: list[str]
16    clearance_level: str = "internal"
17
18    @property
19    def clearance_rank(self) -> int:
20        return CLASSIFICATION_HIERARCHY.get(
21            self.clearance_level, 0
22        )
23
24@dataclass
25class RetrievalResult:
26    """A single document returned from secure retrieval."""
27    document_id: str
28    text: str
29    similarity_score: float
30    classification_level: str
31    owner_id: str
32
33class SecureRetriever:
34    """Wraps pgvector similarity search with access control."""
35
36    def __init__(self, db_connection, default_top_k: int = 5):
37        self.db = db_connection
38        self.default_top_k = default_top_k
39
40    def retrieve(
41        self,
42        query_embedding: list[float],
43        user: UserContext,
44        top_k: Optional[int] = None,
45    ) -> list[RetrievalResult]:
46        """Retrieve documents with access control enforcement."""
47        k = top_k or self.default_top_k
48        allowed_levels = [
49            level for level, rank
50            in CLASSIFICATION_HIERARCHY.items()
51            if rank <= user.clearance_rank
52        ]
53        query = """
54            SELECT id, content, 1 - (embedding <=> %s) as similarity,
55                   classification_level, owner_id
56            FROM documents
57            WHERE (
58                owner_id = %s
59                OR access_groups ?| %s
60                OR classification_level = 'public'
61            )
62            AND classification_level = ANY(%s)
63            ORDER BY embedding <=> %s
64            LIMIT %s
65        """
66        params = (
67            query_embedding, user.user_id,
68            user.groups, allowed_levels,
69            query_embedding, k,
70        )
71        return self._execute_and_map(query, params)

• CLASSIFICATION_HIERARCHY: maps classification names to numeric ranks (public=0 is lowest, restricted=3 is highest) so clearance comparisons are integer comparisons
• UserContext: carries the requesting user's ID, group memberships, and clearance level with a computed clearance_rank property for rank comparison
• RetrievalResult: returns the similarity score plus access control metadata so the caller can see why each document was accessible
• SecureRetriever.retrieve: builds allowed_levels from the user's clearance rank, then constructs a parameterized SQL query that combines cosine distance ordering (<=>) with a three-part access control WHERE clause — owner match OR group overlap OR public classification — AND a clearance ceiling. Parameters are bound as a tuple to prevent SQL injection.

Row-Level Security for Defense in Depth

PostgreSQL row-level security (RLS) provides an additional layer of protection at the database level. Even if application code has a bug that constructs queries without proper access control filters, RLS policies still prevent unauthorized access. The RLS implementation enables RLS on the documents table and creates a SELECT policy that mirrors the application-level logic — owner match OR group overlap OR public — by reading app.current_user_id and app.current_user_groups session variables via current_setting(..., true). The application sets those variables per transaction with SET LOCAL before issuing the similarity query, so RLS enforcement does not require a separate PostgreSQL role per user.

Discipline Application for

Do's and Don'ts

Do's

1. ✓Do embed the three-part access predicate (owner_id = %s OR access_groups ?| %s OR classification_level = 'public') and the classification_level = ANY(%s) clearance ceiling inside the same SQL query as the embedding <=> cosine distance ordering — SecureRetriever.retrieve does this so the database evaluates ranking and authorization over the full corpus together, guaranteeing that the returned top-k are the k most similar documents the requesting user is actually permitted to see, not a post-filtered subset of a larger unrestricted result.
2. ✓Do derive allowed_levels from user.clearance_rank using CLASSIFICATION_HIERARCHY before building the query — converting clearance names to integer ranks first means the permitted classification set is computed once in Python and passed as classification_level = ANY(%s), so a user with clearance_level = "internal" (rank 1) can never retrieve a "confidential" or "restricted" document even when they satisfy the ownership or access_groups predicate.
3. ✓Do create a GIN index on the access_groups JSONB column — the ?| containment operator that SecureRetriever uses to match a user's group list against each document's access_groups array is only efficient with a GIN index; without it, every similarity query degrades to a full sequential scan of access_groups regardless of how selective the cosine distance filter is, destroying retrieval latency at scale.

Don'ts

1. ✗Don't filter retrieved documents for authorization in Python after running an unrestricted embedding <=> similarity query — post-filtering discards unauthorized rows after ranking, so callers reliably receive fewer than top_k results, breaking downstream rerankers and response generators that expect a full-sized candidate set and causing the leak to appear in access logs before the filter ever runs.
2. ✗Don't concatenate user.user_id or user.groups into the SQL string — the parameterized form (owner_id = %s, access_groups ?| %s) prevents a crafted group name such as "finance' OR 1=1--" from subverting the access predicate; string interpolation turns the group list into a direct SQL injection surface that bypasses every other control inside SecureRetriever.
3. ✗Don't omit the classification_level = ANY(%s) clearance ceiling and rely solely on the owner/group/public predicate — a document classified "restricted" that lists the requesting user as owner_id or in access_groups satisfies the three-part condition alone; without the clearance ceiling derived from CLASSIFICATION_HIERARCHY, a user with clearance_level = "internal" silently retrieves "restricted" documents they own, defeating the entire classification hierarchy.