feat(postgres,sqlite): add typed temp table creation in transactions

Signed-off-by: paulwer <paul@wer-ner.de>

Author: paulwer

packages/3-extensions/postgres/src/runtime/postgres.ts

@@ -4,12 +4,20 @@ import type { Contract } from '@prisma-next/contract/types';
 import postgresDriver from '@prisma-next/driver-postgres/runtime';
 import { instantiateExecutionStack } from '@prisma-next/framework-components/execution';
 import { sql as sqlBuilder } from '@prisma-next/sql-builder/runtime';
-import type { Db } from '@prisma-next/sql-builder/types';
+import type {
+  Db,
+  QueryContext,
+  Scope,
+  ScopeField,
+  SelectQuery,
+  Subquery,
+} from '@prisma-next/sql-builder/types';
 import type { ExtractCodecTypes, SqlStorage } from '@prisma-next/sql-contract/types';
 import { orm as ormBuilder } from '@prisma-next/sql-orm-client';
+import { RawSqlExpr, TableSource } from '@prisma-next/sql-relational-core/ast';
 import type { CodecTypesBase, RawSqlTag } from '@prisma-next/sql-relational-core/expression';
 import { createRawSql } from '@prisma-next/sql-relational-core/expression';
-import type { SqlQueryPlan } from '@prisma-next/sql-relational-core/plan';
+import { planFromAst, type SqlQueryPlan } from '@prisma-next/sql-relational-core/plan';
 import type {
   BindSiteParams,
   Declaration,
@@ -43,11 +51,34 @@ import { PostgresRuntimeImpl } from './postgres-runtime';
 export type PostgresTargetId = 'postgres';
 type OrmClient<TContract extends Contract<SqlStorage>> = ReturnType<typeof ormBuilder<TContract>>;
 
+const TEMP_TABLE_NAME_RE = /^[A-Za-z_][A-Za-z0-9_]*$/;
+const TEMP_TABLE_NAME_MAX_LEN = 63;
+
+export interface TempTableCreateOptions {
+  readonly name?: string;
+}
+
+type TempTableJoinSource<Row extends Record<string, ScopeField>> = ReturnType<
+  SelectQuery<QueryContext, Scope, Row>['as']
+>;
+
+export type TempTableHandle<Row extends Record<string, ScopeField> = Record<string, ScopeField>> =
+  TempTableJoinSource<Row> & {
+    readonly name: string;
+    readonly fields: Row;
+    drop(): Promise<void>;
+  };
+
+export interface TempTableBuilder {
+  as<Row extends Record<string, ScopeField>>(query: Subquery<Row>): Promise<TempTableHandle<Row>>;
+}
+
 export interface PostgresTransactionContext<TContract extends Contract<SqlStorage>>
   extends TransactionContext {
   readonly sql: Db<TContract>;
   readonly orm: OrmClient<TContract>;
   readonly enums: NamespacedEnums<TContract>;
+  tempTable(options?: string | TempTableCreateOptions): TempTableBuilder;
 }
 
 export interface PostgresClient<TContract extends Contract<SqlStorage>> {
@@ -140,6 +171,106 @@ function toRuntimeBinding<TContract extends Contract<SqlStorage>>(
   } as const;
 }
 
+function quoteIdentifier(name: string): string {
+  return `"${name.replaceAll('"', '""')}"`;
+}
+
+function resolveTempTableName(options?: string | TempTableCreateOptions): string {
+  const requestedName = typeof options === 'string' ? options : options?.name;
+  if (requestedName !== undefined) {
+    const trimmed = requestedName.trim();
+    if (!TEMP_TABLE_NAME_RE.test(trimmed)) {
+      throw new Error(
+        'Invalid temp table name. Use only letters, numbers, and underscore, and start with a letter/underscore.',
+      );
+    }
+    if (trimmed.length > TEMP_TABLE_NAME_MAX_LEN) {
+      throw new Error(`Invalid temp table name. Maximum length is ${TEMP_TABLE_NAME_MAX_LEN}.`);
+    }
+    return trimmed;
+  }
+
+  const suffix = crypto.randomUUID().replaceAll('-', '').slice(0, 20);
+  return `pn_temp_${suffix}`;
+}
+
+async function executeRawSql(
+  txCtx: TransactionContext,
+  contract: Contract<SqlStorage>,
+  sql: string,
+): Promise<void> {
+  await txCtx.execute(planFromAst(RawSqlExpr.of([sql], []), contract, 'raw.temp-table')).toArray();
+}
+
+function createTempTableBuilder(
+  txCtx: TransactionContext,
+  contract: Contract<SqlStorage>,
+  stack: SqlExecutionStackWithDriver<PostgresTargetId>,
+  options?: string | TempTableCreateOptions,
+): TempTableBuilder {
+  const asJoinSource = <Row extends Record<string, ScopeField>>(
+    tableName: string,
+    alias: string,
+    rowFields: Row,
+  ): TempTableJoinSource<Row> => {
+    const source = {
+      getJoinOuterScope: () => ({
+        topLevel: rowFields,
+        namespaces: { [alias]: rowFields } as Record<string, Row>,
+      }),
+      buildAst: () => TableSource.named(tableName, alias),
+    };
+    return blindCast<TempTableJoinSource<Row>, 'source implements TempTableJoinSource duck-type'>(
+      source,
+    );
+  };
+
+  return {
+    async as<Row extends Record<string, ScopeField>>(
+      query: Subquery<Row>,
+    ): Promise<TempTableHandle<Row>> {
+      const tableName = resolveTempTableName(options);
+      const quotedTableName = quoteIdentifier(tableName);
+      const queryPlan = planFromAst(query.buildAst(), contract, 'dsl');
+      const adapter = instantiateExecutionStack(stack).adapter;
+      const lowered = adapter.lower(queryPlan.ast, {
+        contract,
+        params: queryPlan.params,
+      });
+      const params = lowered.params.map((slot) => {
+        if (slot.kind === 'literal') return slot.value;
+        throw new Error('tempTable.as(...) does not accept bind-site parameters.');
+      });
+
+      const createPlan = Object.freeze({
+        sql: `CREATE TEMP TABLE ${quoteIdentifier(tableName)} AS ${lowered.sql}`,
+        params,
+        ast: queryPlan.ast,
+        meta: queryPlan.meta,
+      });
+
+      await txCtx.execute(createPlan).toArray();
+
+      const rowFields = blindCast<Row, 'subquery row fields align with Subquery<Row> generic'>(
+        query.getRowFields(),
+      );
+      const defaultJoin = asJoinSource<Row>(tableName, tableName, rowFields);
+
+      return blindCast<
+        TempTableHandle<Row>,
+        'temp table handle created from Subquery<Row> preserves the same row field shape'
+      >({
+        ...defaultJoin,
+        name: tableName,
+        fields: rowFields,
+        async drop(): Promise<void> {
+          await executeRawSql(txCtx, contract, `DROP TABLE IF EXISTS ${quotedTableName}`);
+        },
+      });
+    },
+  };
+}
+
 /**
  * Creates a lazy Postgres client from either `contractJson` or a TypeScript-authored `contract`.
  * Static query surfaces are available immediately, while `runtime()` instantiates the driver/pool on first call.
@@ -337,7 +468,14 @@ export default function postgres<TContract extends Contract<SqlStorage>>(
         // Spreading would evaluate the getter once and freeze its value.
         const tx: PostgresTransactionContext<TContract> = Object.assign(
           Object.create(txCtx) as TransactionContext,
-          { sql: txSql, orm: txOrm, enums },
+          {
+            sql: txSql,
+            orm: txOrm,
+            enums,
+            tempTable(options?: string | TempTableCreateOptions): TempTableBuilder {
+              return createTempTableBuilder(txCtx, context.contract, stack, options);
+            },
+          },
         );
 
         return fn(tx);

packages/3-extensions/postgres/test/postgres.test.ts

@@ -1,7 +1,11 @@
 import type { Contract } from '@prisma-next/contract/types';
+import type { ScopeField, Subquery } from '@prisma-next/sql-builder/types';
 import type { SqlStorage } from '@prisma-next/sql-contract/types';
+import { ProjectionItem, SelectAst, TableSource } from '@prisma-next/sql-relational-core/ast';
+import { planFromAst } from '@prisma-next/sql-relational-core/plan';
 import type { Runtime } from '@prisma-next/sql-runtime';
 import { createContract } from '@prisma-next/test-utils';
+import { blindCast } from '@prisma-next/utils/casts';
 import { beforeEach, describe, expect, expectTypeOf, it, vi } from 'vitest';
 
 // Only mock the third-party pg boundary. Real drivers, adapters, and runtimes
@@ -378,6 +382,150 @@ describe('postgres', () => {
     expect(receivedTx!.orm).toBeDefined();
   });
 
+  it('transaction tempTable() creates and drops a typed temp table with generated name', async () => {
+    const pool = new Pool({ connectionString: 'postgres://localhost:5432/db' });
+    const fakeClient = {
+      query: vi.fn().mockResolvedValue({ rows: [], rowCount: 0 }),
+      release: vi.fn(),
+    };
+    (pool as unknown as { connect: typeof vi.fn }).connect = vi.fn().mockResolvedValue(fakeClient);
+
+    const db = postgres({ contract, pg: pool });
+    await db.connect();
+
+    const subquery = blindCast<
+      Subquery<{ id: ScopeField }>,
+      'test fixture for temp-table typed subquery'
+    >({
+      buildAst: () =>
+        SelectAst.from(TableSource.named('source_table')).withProjection([
+          ProjectionItem.of('id', db.raw`1`.returns('pg/int4@1').buildAst()),
+        ]),
+      getRowFields: () => ({ id: { codecId: 'pg/int4@1', nullable: false } }),
+    });
+
+    await db.transaction(async (tx) => {
+      const temp = await tx.tempTable().as(subquery);
+      expect(temp.name).toMatch(/^pn_temp_[a-f0-9]+$/);
+      expect(temp.fields['id']?.codecId).toBe('pg/int4@1');
+      expect('buildAst' in temp).toBe(true);
+      expect('getJoinOuterScope' in temp).toBe(true);
+      await temp.drop();
+    });
+
+    await db.close();
+
+    const issuedSql = fakeClient.query.mock.calls.map((call) => {
+      const arg = call[0] as string | { text?: string };
+      return typeof arg === 'string' ? arg : (arg.text ?? '');
+    });
+    expect(issuedSql.some((sql) => sql.startsWith('CREATE TEMP TABLE'))).toBe(true);
+    expect(issuedSql.some((sql) => sql.startsWith('DROP TABLE IF EXISTS'))).toBe(true);
+  });
+
+  it('transaction tempTable() accepts a manual table name', async () => {
+    const pool = new Pool({ connectionString: 'postgres://localhost:5432/db' });
+    const fakeClient = {
+      query: vi.fn().mockResolvedValue({ rows: [], rowCount: 0 }),
+      release: vi.fn(),
+    };
+    (pool as unknown as { connect: typeof vi.fn }).connect = vi.fn().mockResolvedValue(fakeClient);
+
+    const db = postgres({ contract, pg: pool });
+    await db.connect();
+
+    const subquery = blindCast<
+      Subquery<{ id: ScopeField; email: ScopeField }>,
+      'test fixture for temp-table typed subquery'
+    >({
+      buildAst: () =>
+        SelectAst.from(TableSource.named('source_table')).withProjection([
+          ProjectionItem.of('id', db.raw`1`.returns('pg/int4@1').buildAst()),
+          ProjectionItem.of('email', db.raw`'a@example.com'`.returns('pg/text@1').buildAst()),
+        ]),
+      getRowFields: () => ({
+        id: { codecId: 'pg/int4@1', nullable: false },
+        email: { codecId: 'pg/text@1', nullable: false },
+      }),
+    });
+
+    await db.transaction(async (tx) => {
+      const temp = await tx.tempTable({ name: 'recent_users' }).as(subquery);
+      expect(temp.name).toBe('recent_users');
+      expect(temp.fields).toEqual({
+        id: { codecId: 'pg/int4@1', nullable: false },
+        email: { codecId: 'pg/text@1', nullable: false },
+      });
+      await temp.drop();
+    });
+
+    await db.close();
+
+    const issuedSql = fakeClient.query.mock.calls.map((call) => {
+      const arg = call[0] as string | { text?: string };
+      return typeof arg === 'string' ? arg : (arg.text ?? '');
+    });
+    expect(issuedSql.some((sql) => sql.includes('CREATE TEMP TABLE "recent_users" AS'))).toBe(true);
+    expect(issuedSql.some((sql) => sql.includes('DROP TABLE IF EXISTS "recent_users"'))).toBe(true);
+  });
+
+  it('transaction tempTable() can be reused in tx.sql join composition within the same transaction', async () => {
+    const pool = new Pool({ connectionString: 'postgres://localhost:5432/db' });
+    const fakeClient = {
+      query: vi.fn().mockResolvedValue({ rows: [], rowCount: 0 }),
+      release: vi.fn(),
+    };
+    (pool as unknown as { connect: typeof vi.fn }).connect = vi.fn().mockResolvedValue(fakeClient);
+
+    const db = postgres({ contract, pg: pool });
+    await db.connect();
+
+    const subquery = blindCast<
+      Subquery<{ id: ScopeField; email: ScopeField }>,
+      'test fixture for temp-table typed subquery'
+    >({
+      buildAst: () =>
+        SelectAst.from(TableSource.named('source_table')).withProjection([
+          ProjectionItem.of('id', db.raw`1`.returns('pg/int4@1').buildAst()),
+          ProjectionItem.of('email', db.raw`'a@example.com'`.returns('pg/text@1').buildAst()),
+        ]),
+      getRowFields: () => ({
+        id: { codecId: 'pg/int4@1', nullable: false },
+        email: { codecId: 'pg/text@1', nullable: false },
+      }),
+    });
+
+    await db.transaction(async (tx) => {
+      const recentUsers = await tx.tempTable({ name: 'recent_users' }).as(subquery);
+
+      await tx
+        .execute(
+          planFromAst(
+            SelectAst.from(recentUsers.buildAst()).withProjection([
+              ProjectionItem.of('id', db.raw`1`.returns('pg/int4@1').buildAst()),
+              ProjectionItem.of('email', db.raw`'a@example.com'`.returns('pg/text@1').buildAst()),
+            ]),
+            contract,
+            'dsl',
+          ),
+        )
+        .toArray();
+
+      await recentUsers.drop();
+    });
+
+    await db.close();
+
+    const issuedSql = fakeClient.query.mock.calls.map((call) => {
+      const arg = call[0] as string | { text?: string };
+      return typeof arg === 'string' ? arg : (arg.text ?? '');
+    });
+    expect(issuedSql.some((sql) => sql.includes('CREATE TEMP TABLE "recent_users" AS'))).toBe(true);
+    expect(issuedSql.some((sql) => sql.includes('FROM "recent_users" AS "recent_users"'))).toBe(
+      true,
+    );
+  });
+
   it('transaction() lazily creates runtime before connect()', async () => {
     const pool = new Pool({ connectionString: 'postgres://localhost:5432/db' });
     const fakeClient = {

packages/3-extensions/postgres/test/transaction.types.test-d.ts

@@ -33,3 +33,20 @@ test('tx.orm has the same type as db.orm', () => {
   type TxOrm = PostgresTransactionContext<TestContract>['orm'];
   expectTypeOf<TxOrm>().toEqualTypeOf<DbOrm>();
 });
+
+test('transaction context exposes tempTable()', () => {
+  type HasTempTable = 'tempTable' extends keyof PostgresTransactionContext<TestContract>
+    ? true
+    : false;
+  expectTypeOf<HasTempTable>().toEqualTypeOf<true>();
+});
+
+test('tempTable().as returns a metadata-rich handle', () => {
+  type Builder = ReturnType<PostgresTransactionContext<TestContract>['tempTable']>;
+  type HandlePromise = ReturnType<Builder['as']>;
+  expectTypeOf<Awaited<HandlePromise>>().toMatchTypeOf<{
+    name: string;
+    fields: Record<string, { codecId: string; nullable: boolean }>;
+    drop(): Promise<void>;
+  }>();
+});