Add returning() to inserts on Postgres

CHANGELOG.md

@@ -10,6 +10,9 @@
   `command` and returns no rows.
 - Add `TableOptions.defaults` for the columns an insert may leave out. `defineTable` works them
   out from its column options and records `MATERIALIZED` / `ALIAS` columns as not insertable.
+- Add `returning` to an insert (Postgres): column names or a callback, as in `select`. `run`
+  returns the inserted rows decoded through the derived row schema; `CompiledQuery.returning`
+  lists the aliases. Add `DialectClauses.returning`, optional, absent meaning no.
 - Add `CompiledQuery.kind` (`"select"` or `"insert"`); `rawCompiledQuery` takes it as an option.
 - Add `ParamStyle.maxParameters`; Postgres sets 65535, and a statement over it fails to compile.
 - Add `@maple-dev/effect-orm/database`, opt-in (see `docs/database.md`): a `Database` over the

design/writes.md

@@ -1,7 +1,7 @@
 # Writes: INSERT
 
-Status: phase 1 built (`insertInto(...).values(...)`, see `docs/inserts.md`); phases 2 to 4 not
-started. Section 11 lists where the build differs from the plan. UPDATE and DELETE come later and
+Status: phases 1 and 2 built (`insertInto(...).values(...).returning(...)`, see
+`docs/inserts.md`); phases 3 and 4 not started. Section 11 lists where the build differs from the plan. UPDATE and DELETE come later and
 will reuse what this note sets up (the write-statement state, the `RETURNING` path, value
 encoding).
 
@@ -238,4 +238,10 @@ note, reusing `kind: "write"`, the returning path and the `set` record type from
   takes it as an option, so a handwritten INSERT can say so too.
 - **The bound-value limit is a dialect field**, `ParamStyle.maxParameters`, so it applies to
   queries as well and to a later dialect. The error is `InvalidArguments`; no new error code.
-- **An insert's row schema** is an empty struct with `rowSchemaSource: "derived"`.
+- **An insert's row schema** is an empty struct with `rowSchemaSource: "derived"`, or the one
+  derived from `returning`.
+- **`returning` on a dialect without it is a defect**, not a failure, like `.format()` on
+  Postgres: the dialect is chosen in source, not by a value. `DialectClauses.returning` is
+  optional so a dialect written against the earlier interface still compiles.
+- **`CompiledQuery.returning`** lists the RETURNING aliases; `run` reads it to choose between
+  the command path and the row path, so a precompiled insert runs correctly too.

docs/database.md

@@ -247,6 +247,7 @@ fails at BEGIN today: through the query path with a syntax error, and through `a
 
 ## Writes
 
-The builder compiles SELECTs and [INSERTs](./inserts.md). `run` runs an insert and returns no
-rows; on ClickHouse it goes through `command`, as `execute` does. Write UPDATE, DELETE and
-`INSERT ... RETURNING` with `sql`, as above, and read `RETURNING` with `query` and a schema.
+The builder compiles SELECTs and [INSERTs](./inserts.md). `run` runs an insert and returns its
+`returning` rows, decoded, or none without `returning`; an insert without it goes through
+`command`, as `execute` does. Write UPDATE and DELETE with `sql`, as above, and read `RETURNING`
+with `query` and a schema.

docs/inserts.md

@@ -31,8 +31,8 @@ const insertKey = CH.insertInto(ApiKeys).values({
 // yield* Db.run(insertKey, { id, orgId })
 ```
 
-`Database.run` compiles the insert for its database's dialect and runs it. An insert returns no
-rows; `RETURNING` is not built yet (see [`design/writes.md`](../design/writes.md)).
+`Database.run` compiles the insert for its database's dialect and runs it. Without
+[`returning`](#returning) an insert returns no rows.
 
 ## The row type
 
@@ -90,6 +90,25 @@ the codec rejects fails to compile with a `QueryBuilderError` that names the row
   several rows is bound once. A statement over 65535 bound values (Postgres's limit) fails to
   compile instead of being split: send fewer rows per statement.
 
+## Returning
+
+On Postgres, `returning` adds a RETURNING list and `Database.run` returns the inserted rows,
+decoded. It takes column names, or a callback building one expression per alias, as `select`
+does:
+
+```ts
+const created = CH.insertInto(ApiKeys)
+	.values({ id: CH.param.string("id"), org_id: CH.param.string("orgId"), name: "default" })
+	.returning(($) => ({ id: $.id, createdAt: $.created_at }))
+
+// const [row] = yield* Db.run(created, { id, orgId }) // { id: string; createdAt: DateTime.Utc }
+```
+
+The row schema is derived from the list, as it is from a SELECT: an untyped expression
+(`untypedExpr`) leaves the insert undecoded, with `rowSchemaSource: "none"` and the alias in
+`untypedColumns`. `CompiledQuery.returning` lists the aliases. ClickHouse has no RETURNING, so
+compiling an insert with `returning` for it is a `QueryBuilderDefect`.
+
 ## Tenant scope
 
 An insert has a `tenantScope` like a query, worked out the same way. On a table with a
@@ -107,6 +126,7 @@ without a tenant column gives `"untenanted"`.
 | A param with no value                             | `QueryBuilderError` `UnresolvedParam`    |
 | Over the dialect's bound-value limit              | `QueryBuilderError` `InvalidArguments`   |
 | Compiling without `values`                        | `QueryBuilderDefect`                     |
+| `returning` for a dialect without RETURNING       | `QueryBuilderDefect`                     |
 
 _(Backed by `src/ch/insert.test.ts`, `src/database/database.test.ts` and
 `tests/database.clickhouse.test.ts`.)_

docs/reference.md

@@ -62,7 +62,7 @@ Note `/sql` exports a `compile` (fragment → string) distinct from the root `co
 | `fromQuery` | `(query, alias) => CHQuery`               |
 | `fromUnion` | `(union, alias) => CHQuery`               |
 | `unionAll`  | `(...queries) => CHUnionQuery`            |
-| `insertInto` | `(table) => CHInsert`; `.values(row \| rows)` sets its rows. See [Inserting rows](./inserts.md) |
+| `insertInto` | `(table) => CHInsert`; `.values(row \| rows)` sets its rows, `.returning(...)` the RETURNING list (Postgres). See [Inserting rows](./inserts.md) |
 
 ### `CHQuery` methods
 

src/ch/compile.ts

@@ -11,7 +11,7 @@ import { custom, dateTime, dateTime64, type CHType, type ColumnDefs } from "./ty
 import type { CHQuery, CHQueryState } from "./query"
 import type { CHUnionQuery } from "./union"
 import { isInsert, type CHInsert } from "./insert"
-import { createQualifiedColumnAccessor, createJoinedColumnAccessor, sourceAlias } from "./query"
+import { createColumnAccessor, createQualifiedColumnAccessor, createJoinedColumnAccessor, sourceAlias } from "./query"
 import { aliased, columnTypeOf, isExprLike } from "./expr"
 import { raw, identPath, quoteIdent, quoteIdentPath, compile as compileSqlFragment } from "../sql/sql-fragment"
 import { splitTerminalClauses } from "../sql/terminal-clauses"
@@ -134,6 +134,11 @@ interface CompiledQueryBase<Output> {
 	 * through a client path that expects a result set.
 	 */
 	readonly kind: "select" | "insert"
+	/**
+	 * The aliases of an insert's RETURNING list, when it has one. An insert
+	 * without it sends back no rows; one with it is read like a query.
+	 */
+	readonly returning?: ReadonlyArray<string>
 	/**
 	 * The values a binding dialect sends beside `sql`, in placeholder order.
 	 *
@@ -343,6 +348,7 @@ const makeCompiledQuery = <Output, Route extends string | undefined>(
 	rowSchemaMismatch?: RowSchemaMismatch,
 	dialect?: string,
 	kind: "select" | "insert" = "select",
+	returning?: ReadonlyArray<string>,
 ): CompiledQuery<Output, Route> => {
 	let cachedDecodeRow: ((row: unknown) => Effect.Effect<Output, unknown, never>) | undefined
 	let decoderBuilt = false
@@ -405,6 +411,7 @@ const makeCompiledQuery = <Output, Route extends string | undefined>(
 	return {
 		sql,
 		kind,
+		...(returning !== undefined ? { returning } : undefined),
 		parameters,
 		tenantScope,
 		// Resolved eagerly only here, where the getter is already memoised by
@@ -1326,6 +1333,24 @@ const VALUE_PARAM = "$$v"
 
 const EMPTY_ROW = Schema.Struct({}) as unknown as CompiledQueryRowSchema<never>
 
+/** The RETURNING clause and its row schema, or none without `returning`. */
+const returningOf = (insert: CHInsert<any, any, any, any>) => {
+	const { table, returningFn } = insert._state
+	if (returningFn === undefined) return undefined
+	const dialect = currentDialect()
+	if (dialect.clauses.returning !== true) {
+		throw new QueryBuilderDefect({
+			message: `insertInto(${table.name}): returning() has no meaning for the ${dialect.name} dialect, which has no RETURNING clause`,
+		})
+	}
+	const exprs = returningFn(createColumnAccessor(table.columns))
+	const aliases = Object.keys(exprs)
+	if (aliases.length === 0) {
+		throw new QueryBuilderDefect({ message: `insertInto(${table.name}): returning() needs at least one column` })
+	}
+	return { exprs, aliases, derived: deriveRowSchema(exprs) }
+}
+
 /**
  * An INSERT ... VALUES. Columns are written in table order, so two rows with
  * their keys in different orders cannot swap values; a column some rows leave
@@ -1403,27 +1428,34 @@ function compileInsert(insert: CHInsert<any, any, any, any>, params: Record<stri
 	)
 
 	const dialect = currentDialect()
+	const returning = returningOf(insert)
+	const returningSql =
+		returning === undefined
+			? ""
+			: `\nRETURNING ${returning.aliases.map((alias) => compileSqlFragment(aliased(returning.exprs[alias]!, alias))).join(", ")}`
 	const rendered = renderParams(
-		`INSERT INTO ${quoteIdentPath(table.name)} (${columns.map(quoteIdent).join(", ")})\nVALUES ${tuples.join(", ")}`,
+		`INSERT INTO ${quoteIdentPath(table.name)} (${columns.map(quoteIdent).join(", ")})\nVALUES ${tuples.join(", ")}${returningSql}`,
 		values,
 		dialect,
 	)
+	const returnedSchema = returning !== undefined && "schema" in returning.derived ? returning.derived.schema : undefined
 	const tenantScope: TenantScope =
 		tenant === undefined ? "untenanted" : pinned && bounds.size === 1 ? "single-tenant" : "cross-tenant"
 
 	return withTenantBound(
-		makeCompiledQuery<never, undefined>(
+		makeCompiledQuery<any, undefined>(
 			rendered.sql,
 			rendered.parameters,
 			tenantScope,
-			"derived",
-			() => EMPTY_ROW,
+			returning === undefined || returnedSchema !== undefined ? "derived" : "none",
+			() => (returning === undefined ? EMPTY_ROW : returnedSchema),
 			undefined,
-			[],
+			returning !== undefined && "untyped" in returning.derived ? returning.derived.untyped : [],
 			undefined,
 			undefined,
 			dialect.name,
 			"insert",
+			returning?.aliases,
 		),
 		tenantScope === "single-tenant" ? [...bounds][0] : undefined,
 	)

src/ch/dialect.ts

@@ -59,6 +59,9 @@ export interface DialectClauses {
 	/** Whether each `UNION ALL` branch is wrapped in parentheses. Postgres needs
 	 *  it for a branch with its own WITH, ORDER BY or LIMIT. */
 	readonly parenthesizeUnionBranches: boolean
+	/** `RETURNING` after an INSERT. Absent means no: an insert with
+	 *  `.returning()` fails to compile for the dialect. */
+	readonly returning?: boolean
 }
 
 /** A transaction isolation level. `read uncommitted` is left out: Postgres runs it as `read committed`. */
@@ -138,7 +141,13 @@ export const clickhouseDialect: Dialect = {
 	literal: sqlLiteral,
 	dateTimeLiteral: (value) => quoteClickHouseString(chDateTimeLiteral(value)),
 	params: { _tag: "inline" },
-	clauses: { format: true, derivedTableAlias: false, groupByAlias: true, parenthesizeUnionBranches: false },
+	clauses: {
+		format: true,
+		derivedTableAlias: false,
+		groupByAlias: true,
+		parenthesizeUnionBranches: false,
+		returning: false,
+	},
 	transactions: noTransactions,
 }
 

src/ch/insert.test-d.ts

@@ -81,6 +81,19 @@ const Keys = CH.table("api_keys", { id: PG.uuid, created_at: PG.timestamptz }, {
 CH.insertInto(Keys).values({ id: "k" })
 CH.insertInto(Keys).values({ id: "k", created_at: new Date() as unknown as DateTime.Utc })
 
+// RETURNING: column names or a callback, as in select.
+const returningNames = CH.insertInto(Keys).values({ id: "k" }).returning("id", "created_at")
+expectTypeOf<RowOf<typeof returningNames>>().toEqualTypeOf<{ readonly id: string; readonly created_at: DateTime.Utc }>()
+const returningExprs = CH.insertInto(Keys)
+	.values({ id: "k" })
+	.returning(($) => ({ key: $.id, n: CH.rawExpr("1", PG.int4) }))
+expectTypeOf<RowOf<typeof returningExprs>>().toEqualTypeOf<{ readonly key: string; readonly n: number }>()
+expectTypeOf(PG.compileUnsafe(returningExprs)).toEqualTypeOf<
+	CH.CompiledQuery<{ readonly key: string; readonly n: number }, undefined>
+>()
+// @ts-expect-error not a column
+CH.insertInto(Keys).values({ id: "k" }).returning("nope")
+
 // An insert without RETURNING runs to no rows; compile gives a CompiledQuery.
 const insert = CH.insertInto(Plain).values({ A: "a", B: 1 })
 expectTypeOf<RowOf<typeof insert>>().toEqualTypeOf<never>()

src/ch/insert.test.ts

@@ -95,6 +95,55 @@ describe("insertInto", () => {
 		expect(compiled.sql).toBe("INSERT INTO notes (Body)\nVALUES ('\\x5F_PARAM_string_x__')")
 	})
 
+	describe("returning", () => {
+		it.effect("writes RETURNING and derives the row schema from it", () =>
+			Effect.gen(function* () {
+				const compiled = PG.compileUnsafe(
+					CH.insertInto(Keys)
+						.values({ id: "a", org_id: "o", name: "n", meta: {} })
+						.returning(($) => ({ id: $.id, createdAt: $.created_at })),
+				)
+				expect(compiled.sql).toBe(
+					'INSERT INTO "api_keys" ("id", "org_id", "name", "meta")\nVALUES ($1, $2, $3, $4)\n' +
+						'RETURNING "id" AS "id", "created_at" AS "createdAt"',
+				)
+				expect(compiled.returning).toEqual(["id", "createdAt"])
+				expect(compiled.rowSchemaSource).toBe("derived")
+				const at = "2026-01-02T03:04:05.000Z"
+				const [row] = yield* compiled.decodeRows([{ id: "a", createdAt: at }])
+				expect(row!.id).toBe("a")
+				expect(DateTime.formatIso(row!.createdAt)).toBe(at)
+			}),
+		)
+
+		it("takes column names, and calling it again replaces the list", () => {
+			const insert = CH.insertInto(Keys).values({ id: "a", org_id: "o", name: "n", meta: {} })
+			expect(PG.compileUnsafe(insert.returning("id", "revoked")).sql).toMatch(/RETURNING "id" AS "id", "revoked" AS "revoked"$/)
+			expect(PG.compileUnsafe(insert.returning("revoked").returning("id")).returning).toEqual(["id"])
+			expect(PG.compileUnsafe(insert).returning).toBeUndefined()
+		})
+
+		it("an untyped expression leaves the insert undecoded, and names the alias", () => {
+			const compiled = PG.compileUnsafe(
+				CH.insertInto(Keys)
+					.values({ id: "a", org_id: "o", name: "n", meta: {} })
+					.returning(() => ({ txid: CH.untypedExpr("pg_current_xact_id()::xid::text") })),
+			)
+			expect(compiled.sql).toMatch(/RETURNING pg_current_xact_id\(\)::xid::text AS "txid"$/)
+			expect(compiled.rowSchemaSource).toBe("none")
+			expect(compiled.untypedColumns).toEqual(["txid"])
+		})
+
+		it.effect("is a defect on ClickHouse, which has no RETURNING", () =>
+			Effect.gen(function* () {
+				const exit = yield* Effect.exit(
+					CH.compile(CH.insertInto(Events).values({ OrgId: "o", At: new Date(0), Attrs: {}, Tags: [] }).returning("Id")),
+				)
+				expect(failure(exit)).toBeInstanceOf(QueryBuilderDefect)
+			}),
+		)
+	})
+
 	describe("tenant scope", () => {
 		const scope = (rows: ReadonlyArray<Record<string, unknown>>, params: Record<string, unknown> = {}) =>
 			CH.compileUnsafe(CH.insertInto(Events).values(rows as any), params).tenantScope

src/ch/insert.ts

@@ -14,6 +14,7 @@
 //   yield* Database.run(insert, { id, orgId })
 
 import type { Comparable, Expr, Widen } from "./expr"
+import type { ColumnAccessor, InferOutput } from "./query"
 import type { Table } from "./table"
 import type { CHType, ColumnDefs, InferTS } from "./types"
 
@@ -69,6 +70,8 @@ export interface CHInsertState {
 	readonly table: Table<string, ColumnDefs>
 	/** Set by `values`. Compiling without it is a defect. */
 	readonly rows?: ReadonlyArray<Readonly<Record<string, unknown>>>
+	/** Set by `returning`: the RETURNING list, as a select callback. */
+	readonly returningFn?: ($: any) => Record<string, Expr<any>>
 }
 
 export interface CHInsert<
@@ -91,6 +94,20 @@ export interface CHInsert<
 	values(
 		rows: InsertRow<Cols, Defaulted, Computed> | ReadonlyArray<InsertRow<Cols, Defaulted, Computed>>,
 	): CHInsert<Cols, Defaulted, Computed, Output>
+
+	/**
+	 * Return the inserted rows: column names, or a callback building an
+	 * expression per alias, as in `select`. `Database.run` then decodes them
+	 * through the derived row schema. Postgres only; on a dialect without
+	 * RETURNING (ClickHouse) compiling is a defect. Calling it again replaces
+	 * the list.
+	 */
+	returning<K extends keyof Cols & string>(
+		...columns: K[]
+	): CHInsert<Cols, Defaulted, Computed, { readonly [P in K]: InferTS<Cols[P]> }>
+	returning<S extends Record<string, Expr<any>>>(
+		fn: ($: ColumnAccessor<Cols>) => S,
+	): CHInsert<Cols, Defaulted, Computed, InferOutput<S>>
 }
 
 const makeInsert = <Cols extends ColumnDefs, Defaulted extends string, Computed extends string, Output>(
@@ -104,6 +121,14 @@ const makeInsert = <Cols extends ColumnDefs, Defaulted extends string, Computed
 			// Copied, so a caller pushing to its array later does not change the insert.
 			rows: Array.isArray(rows) ? [...rows] : [rows as Readonly<Record<string, unknown>>],
 		}),
+	returning: ((...args: ReadonlyArray<unknown>) => {
+		const [first] = args
+		const returningFn =
+			typeof first === "function"
+				? (first as ($: any) => Record<string, Expr<any>>)
+				: ($: any) => Object.fromEntries((args as ReadonlyArray<string>).map((column) => [column, $[column]]))
+		return makeInsert({ ...state, returningFn })
+	}) as CHInsert<Cols, Defaulted, Computed, Output>["returning"],
 })
 
 /** Start an INSERT into `table`. Give its rows with `values`. */

src/database/database.test.ts

@@ -405,7 +405,15 @@ layer(Live, { excludeTestServices: true })("Database on PGlite", (it) => {
 				{ org: "o1" },
 			)
 			expect(inserted).toEqual([])
-			const rows = yield* Db.run(CH.from(Keys).select("id", "uses", "created_at", "revoked", "meta", "tags", "note").orderBy(["id", "asc"]))
+			const returned = yield* Db.run(
+				CH.insertInto(Keys)
+					.values({ id: id(3), org_id: "o1", tags: ["z"] })
+					.returning(($) => ({ id: $.id, uses: $.uses, createdAt: $.created_at, revoked: $.revoked, tags: $.tags })),
+			)
+			expect(returned).toHaveLength(1)
+			expect(returned[0]).toMatchObject({ id: id(3), uses: 0, revoked: false, tags: ["z"] })
+			expect(DateTime.isDateTime(returned[0]!.createdAt)).toBe(true)
+			const rows = yield* Db.run(CH.from(Keys).where(($) => [$.id.neq(id(3))]).select("id", "uses", "created_at", "revoked", "meta", "tags", "note").orderBy(["id", "asc"]))
 			expect(rows[0]).toEqual({ id: id(1), uses: 0, created_at: at, revoked: false, meta: { k: [1, 2] }, tags: ["a", "it's"], note: null })
 			expect(rows[1]).toMatchObject({ id: id(2), uses: 7, revoked: true, meta: null, tags: [], note: "n" })
 		}),

src/database/database.ts

@@ -85,7 +85,7 @@ export interface DatabaseApi {
 	readonly dialect: Dialect
 	/**
 	 * Compile a query for this database's dialect, run it, and decode its rows.
-	 * An insert returns no rows.
+	 * An insert returns its RETURNING rows, or none without `returning`.
 	 * `params` fills the query's `param.*` markers. A query compiled elsewhere
 	 * runs as it is, if it was compiled for this dialect.
 	 */
@@ -263,11 +263,12 @@ export const fromSqlClient = (sql: SqlClient.SqlClient, options: FromSqlClientOp
 		return compileCH(runnable as CHQuery<any, any, any, any>, params, { dialect })
 	}
 
-	// An insert sends back no rows, so it runs the way `execute` does: through
-	// `command`, which a ClickHouse client needs for a statement with no result set.
+	// An insert without RETURNING sends back no rows, so it runs the way `execute`
+	// does: through `command`, which a ClickHouse client needs for a statement
+	// with no result set.
 	const run: DatabaseApi["run"] = (runnable, params = {}) =>
 		Effect.flatMap(compileFor(runnable, params), (compiled) =>
-			compiled.kind === "insert"
+			compiled.kind === "insert" && compiled.returning === undefined
 				? Effect.as(execute(compiled), [])
 				: Effect.flatMap(rows(compiled), (wire) => compiled.decodeRows(wire)),
 		)