diff --git a/README.md b/README.md index 131ec7b..158a4f0 100644 --- a/README.md +++ b/README.md @@ -372,8 +372,9 @@ await bash.exec('semgrep -k 3 "delivery options" /docs') **Migrate** an existing deployment in two idempotent steps: 1. Re-run `setup(client)` — it creates `fs_blob_chunks` (no new extensions). - Drizzle users: re-run `drizzle-kit generate` (the table is in - `createSchema()`) plus `generateMigrationSQL()` for RLS and the blob FK. + Drizzle users: re-run `drizzle-kit generate` (the table, its blob FK and + its RLS policy are all in `createSchema()`) plus `generateMigrationSQL()` + for extensions and `FORCE ROW LEVEL SECURITY`. 2. Index pre-existing content: `await fs.backfillChunks()` — chunks the blobs visible at the handle's version; content-addressed, safe to re-run. diff --git a/app/src/app/docs/schema/page.tsx b/app/src/app/docs/schema/page.tsx index 3680823..77dc4c1 100644 --- a/app/src/app/docs/schema/page.tsx +++ b/app/src/app/docs/schema/page.tsx @@ -307,8 +307,19 @@ export const schema = createSchema({ generateMigrationSQL() {" "} - produces SQL for extensions and RLS policies that Drizzle can't - express. Paste it into a custom migration. + produces SQL for what Drizzle can't express: extensions and{" "} + + FORCE ROW LEVEL SECURITY + + . The{" "} + + workspace_isolation + {" "} + policies and the blob-chunks FK are declared by{" "} + createSchema(){" "} + itself, so drizzle-kit generates (and diffs) them like any other + schema object. Paste the output into a custom migration — every + statement is idempotent.

+

+ Upgrading a project whose migrations predate the in-schema + declarations: the first{" "} + + drizzle-kit generate + {" "} + after the upgrade emits a migration your bootstrapped database + largely already satisfies. Edit the generated file before applying — + keep the file itself, it updates the snapshot: +

+ +

+ + createSchema({"{ enableRLS: false }"}) + {" "} + omits only the policies — the FK and the index-form change still + show up in the diff, so the edit above applies either way. +

diff --git a/lib/adapters/drizzle/migration.ts b/lib/adapters/drizzle/migration.ts index d94ae8b..51e2334 100644 --- a/lib/adapters/drizzle/migration.ts +++ b/lib/adapters/drizzle/migration.ts @@ -15,11 +15,15 @@ const RLS_TABLES = [ /** * Returns SQL for a custom Drizzle migration covering what `createSchema()` - * cannot express: extensions and RLS policies. + * cannot express: extensions and `FORCE ROW LEVEL SECURITY`. * - * The tables and indexes are handled by the Drizzle schema (`createSchema()`), - * so `drizzle-kit generate` picks those up automatically. This function - * produces the SQL for everything else. + * Tables, indexes, the `fs_blob_chunks` FK and the `workspace_isolation` RLS + * policies are all declared by `createSchema()` nowadays, so `drizzle-kit + * generate` picks them up automatically (and `drizzle-kit push` no longer + * flags them as drift). Every statement here is idempotent, so projects whose + * migrations predate the in-schema declarations can keep this bootstrap + * as-is; it also remains the safety net for the FK and policies when the + * schema is used with `enableRLS: false`. * * @example * ```ts diff --git a/lib/adapters/drizzle/schema.ts b/lib/adapters/drizzle/schema.ts index ded13bc..470ddb6 100644 --- a/lib/adapters/drizzle/schema.ts +++ b/lib/adapters/drizzle/schema.ts @@ -8,6 +8,8 @@ import { customType, uniqueIndex, index, + foreignKey, + pgPolicy, primaryKey, vector, } from "drizzle-orm/pg-core"; @@ -29,6 +31,26 @@ export interface SchemaOptions { enableFullTextSearch?: boolean; enableVectorSearch?: boolean; embeddingDimensions?: number; + /** + * Declare the per-workspace `workspace_isolation` RLS policies on the + * schema objects themselves (default true, mirroring `setup()` and + * `generateMigrationSQL()`). With the policies in the schema, drizzle-kit + * sees them: `generate` emits them into migrations and `push` no longer + * proposes `DROP POLICY workspace_isolation` as spurious drift. + * `FORCE ROW LEVEL SECURITY` has no drizzle-orm API and still rides the + * `generateMigrationSQL()` custom migration. + */ + enableRLS?: boolean; +} + +/** The `workspace_isolation` policy every bash-gres table carries under RLS. */ +function workspaceIsolationPolicy() { + const expr = sql`workspace_id = current_setting('app.workspace_id', true)`; + return pgPolicy("workspace_isolation", { + for: "all", + using: expr, + withCheck: expr, + }); } export interface BashGresSchema { @@ -46,7 +68,7 @@ export interface BashGresSchemaWithVector extends BashGresSchema { fsChunkEmbeddings: ReturnType; } -function buildVersionRoots() { +function buildVersionRoots(enableRLS: boolean) { return pgTable( "fs_version_roots", { @@ -64,13 +86,14 @@ function buildVersionRoots() { ), index("idx_fs_version_roots_path_gist").using( "gist", - sql`${table.path} gist_ltree_ops(siglen=124)`, + table.path.op("gist_ltree_ops(siglen=124)"), ), + ...(enableRLS ? [workspaceIsolationPolicy()] : []), ], ); } -function buildVersions() { +function buildVersions(enableRLS: boolean) { return pgTable( "fs_versions", { @@ -96,11 +119,12 @@ function buildVersions() { table.versionRootId, table.parentVersionId, ), + ...(enableRLS ? [workspaceIsolationPolicy()] : []), ], ); } -function buildAncestors() { +function buildAncestors(enableRLS: boolean) { return pgTable( "version_ancestors", { @@ -123,11 +147,12 @@ function buildAncestors() { table.workspaceId, table.ancestorId, ), + ...(enableRLS ? [workspaceIsolationPolicy()] : []), ], ); } -function buildBlobs() { +function buildBlobs(enableRLS: boolean) { return pgTable( "fs_blobs", { @@ -140,15 +165,22 @@ function buildBlobs() { .notNull() .defaultNow(), }, - (table) => [primaryKey({ columns: [table.workspaceId, table.hash] })], + (table) => [ + primaryKey({ columns: [table.workspaceId, table.hash] }), + ...(enableRLS ? [workspaceIsolationPolicy()] : []), + ], ); } // Section-level slices of text blobs for chunk-granular search; content- // addressed like fs_blobs (see lib/core/setup.ts for the column semantics). -// The FK to fs_blobs is declared in the core DDL only — this schema mirrors -// tables for query typing, and no table here declares FKs. -function buildBlobChunks(options: SchemaOptions) { +// The FK to fs_blobs is declared here (same name as the core DDL, which stays +// idempotent) so drizzle-kit sees it instead of flagging it as drift. +function buildBlobChunks( + options: SchemaOptions, + fsBlobs: ReturnType, + enableRLS: boolean, +) { const { enableFullTextSearch = true } = options; return pgTable( "fs_blob_chunks", @@ -170,6 +202,11 @@ function buildBlobChunks(options: SchemaOptions) { primaryKey({ columns: [table.workspaceId, table.blobHash, table.chunkIndex], }), + foreignKey({ + name: "fs_blob_chunks_blob_fkey", + columns: [table.workspaceId, table.blobHash], + foreignColumns: [fsBlobs.workspaceId, fsBlobs.hash], + }).onDelete("cascade"), ]; if (enableFullTextSearch) { indexes.push( @@ -178,6 +215,9 @@ function buildBlobChunks(options: SchemaOptions) { .with({ text_config: "english" }), ); } + if (enableRLS) { + indexes.push(workspaceIsolationPolicy()); + } return indexes as ReturnType[]; }, ); @@ -186,7 +226,7 @@ function buildBlobChunks(options: SchemaOptions) { // Per-content embedding cache for chunk-level semantic search. Deliberately // no FK to fs_blob_chunks — the cache outlives its chunk rows (see // lib/core/setup.ts for the semantics). -function buildChunkEmbeddings(embeddingDimensions: number) { +function buildChunkEmbeddings(embeddingDimensions: number, enableRLS: boolean) { return pgTable( "fs_chunk_embeddings", { @@ -205,11 +245,12 @@ function buildChunkEmbeddings(embeddingDimensions: number) { "hnsw", table.embedding.op("vector_cosine_ops"), ), + ...(enableRLS ? [workspaceIsolationPolicy()] : []), ], ); } -function buildEntries() { +function buildEntries(enableRLS: boolean) { return pgTable( "fs_entries", { @@ -237,11 +278,12 @@ function buildEntries() { ), index("idx_fs_entries_path_gist").using( "gist", - sql`${table.path} gist_ltree_ops(siglen=124)`, + table.path.op("gist_ltree_ops(siglen=124)"), ), index("idx_fs_entries_blob_hash") .on(table.workspaceId, table.blobHash) .where(sql`${table.blobHash} IS NOT NULL`), + ...(enableRLS ? [workspaceIsolationPolicy()] : []), ], ); } @@ -253,7 +295,11 @@ export function createSchema(options?: SchemaOptions): BashGresSchema; export function createSchema( options: SchemaOptions = {}, ): BashGresSchema | BashGresSchemaWithVector { - const { enableVectorSearch = false, embeddingDimensions } = options; + const { + enableVectorSearch = false, + embeddingDimensions, + enableRLS = true, + } = options; if (enableVectorSearch && !embeddingDimensions) { throw new Error( @@ -261,16 +307,20 @@ export function createSchema( ); } + const fsBlobs = buildBlobs(enableRLS); const base = { - fsVersionRoots: buildVersionRoots(), - fsVersions: buildVersions(), - versionAncestors: buildAncestors(), - fsBlobs: buildBlobs(), - fsEntries: buildEntries(), - fsBlobChunks: buildBlobChunks(options), + fsVersionRoots: buildVersionRoots(enableRLS), + fsVersions: buildVersions(enableRLS), + versionAncestors: buildAncestors(enableRLS), + fsBlobs, + fsEntries: buildEntries(enableRLS), + fsBlobChunks: buildBlobChunks(options, fsBlobs, enableRLS), }; if (enableVectorSearch && embeddingDimensions) { - return { ...base, fsChunkEmbeddings: buildChunkEmbeddings(embeddingDimensions) }; + return { + ...base, + fsChunkEmbeddings: buildChunkEmbeddings(embeddingDimensions, enableRLS), + }; } return base; } diff --git a/package.json b/package.json index 88eded1..2697153 100644 --- a/package.json +++ b/package.json @@ -57,7 +57,7 @@ ], "license": "MIT", "peerDependencies": { - "drizzle-orm": ">=0.30.0", + "drizzle-orm": ">=0.36.0", "just-bash": ">=2.0.0", "pg": ">=8.0.0", "postgres": ">=3.0.0" diff --git a/tests/drizzle-schema.test.ts b/tests/drizzle-schema.test.ts index 6001670..4421de2 100644 --- a/tests/drizzle-schema.test.ts +++ b/tests/drizzle-schema.test.ts @@ -1,6 +1,7 @@ import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { drizzle } from "drizzle-orm/postgres-js"; import { eq } from "drizzle-orm"; +import { getTableConfig } from "drizzle-orm/pg-core"; import type postgres from "postgres"; import { createSchema } from "../lib/adapters/drizzle/schema.js"; import { ensureSetup } from "./global-setup.js"; @@ -50,4 +51,45 @@ describe("createSchema drizzle typing", () => { "embeddingDimensions is required", ); }); + + it("declares the workspace_isolation policy on every table by default", () => { + const schema = createSchema({ + enableVectorSearch: true, + embeddingDimensions: 3, + }); + for (const table of Object.values(schema)) { + const config = getTableConfig(table); + const names = config.policies.map((p) => p.name); + expect(names, config.name).toEqual(["workspace_isolation"]); + expect(config.policies[0]?.for).toBe("all"); + expect(config.policies[0]?.using).toBeDefined(); + expect(config.policies[0]?.withCheck).toBeDefined(); + } + }); + + it("omits the policies with enableRLS: false", () => { + const schema = createSchema({ enableRLS: false }); + for (const table of Object.values(schema)) { + expect(getTableConfig(table).policies).toEqual([]); + } + }); + + it("declares the fs_blob_chunks FK to fs_blobs with the core DDL's name", () => { + const schema = createSchema(); + const fks = getTableConfig(schema.fsBlobChunks).foreignKeys; + expect(fks).toHaveLength(1); + const fk = fks[0]!; + expect(fk.getName()).toBe("fs_blob_chunks_blob_fkey"); + expect(fk.onDelete).toBe("cascade"); + const ref = fk.reference(); + expect(ref.foreignTable).toBe(schema.fsBlobs); + expect(ref.columns.map((c) => c.name)).toEqual([ + "workspace_id", + "blob_hash", + ]); + expect(ref.foreignColumns.map((c) => c.name)).toEqual([ + "workspace_id", + "hash", + ]); + }); });