feat: JSON type replacement for Prisma TypedSQL queries

src/handler/typedsql.ts

@@ -0,0 +1,101 @@
+import type { SqlQueryOutput } from '@prisma/generator';
+import ts from 'typescript';
+import type { PrismaJsonTypesGeneratorConfig } from '../util/config';
+import { createType } from '../util/create-signature';
+import type { DeclarationWriter } from '../util/declaration-writer';
+
+/** A map from column name to field documentation for JSON fields. */
+export type ColumnAnnotationMap = Map<string, string | undefined>;
+
+/**
+ * Handles a TypedSQL query file by replacing JSON column types with annotated types.
+ *
+ * TypedSQL files are TypeScript modules generated by Prisma that represent typed SQL
+ * queries. Each file exports a factory function and a namespace with `Parameters` and
+ * `Result` types. This handler replaces `$runtime.JsonValue` types in the `Result` type
+ * with user-annotated types from the Prisma schema.
+ */
+export function handleTypedSqlFile(
+  tsSource: ts.SourceFile,
+  writer: DeclarationWriter,
+  query: SqlQueryOutput,
+  columnAnnotationMap: ColumnAnnotationMap,
+  config: PrismaJsonTypesGeneratorConfig
+) {
+  // Only process json and json-array typed columns
+  const jsonColumns = query.resultColumns.filter(
+    (col) => col.typ === 'json' || col.typ === 'json-array'
+  );
+
+  if (jsonColumns.length === 0) return;
+
+  // Build a map of column name → replacement info
+  const columnsToReplace = new Map<
+    string,
+    { newType: string; nullable: boolean; isArray: boolean }
+  >();
+
+  for (const col of jsonColumns) {
+    const documentation = columnAnnotationMap.get(col.name);
+
+    // No annotation and allowAny is set — skip replacing
+    if (!documentation && config.allowAny) continue;
+
+    const newType = createType(documentation, config);
+    columnsToReplace.set(col.name, {
+      newType,
+      nullable: col.nullable,
+      isArray: col.typ === 'json-array'
+    });
+  }
+
+  if (columnsToReplace.size === 0) return;
+
+  // Traverse the top-level namespace declaration matching the query name
+  tsSource.forEachChild((child) => {
+    if (child.kind !== ts.SyntaxKind.ModuleDeclaration) return;
+
+    const ns = child as ts.ModuleDeclaration;
+    if (ns.name.getText() !== query.name) return;
+
+    const body = ns.body;
+    if (!body || body.kind !== ts.SyntaxKind.ModuleBlock) return;
+
+    for (const stmt of (body as ts.ModuleBlock).statements) {
+      if (stmt.kind !== ts.SyntaxKind.TypeAliasDeclaration) continue;
+
+      const typeAlias = stmt as ts.TypeAliasDeclaration;
+      if (typeAlias.name.getText() !== 'Result') continue;
+      if (typeAlias.type.kind !== ts.SyntaxKind.TypeLiteral) continue;
+
+      const typeLiteral = typeAlias.type as ts.TypeLiteralNode;
+
+      for (const member of typeLiteral.members) {
+        if (member.kind !== ts.SyntaxKind.PropertySignature) continue;
+
+        const prop = member as ts.PropertySignature;
+        const propName = prop.name?.getText();
+
+        if (!propName || !columnsToReplace.has(propName)) continue;
+
+        const { newType, nullable, isArray } = columnsToReplace.get(propName)!;
+        const typeNode = prop.type;
+
+        if (!typeNode) continue;
+
+        let replacement: string;
+        if (isArray) {
+          replacement = `${newType}[]`;
+        } else if (nullable) {
+          replacement = `${newType} | null`;
+        } else {
+          replacement = newType;
+        }
+
+        // Use getStart() (not pos) to exclude leading trivia so the
+        // space between the colon and the type is preserved.
+        writer.replace(typeNode.getStart(tsSource), typeNode.end, replacement);
+      }
+    }
+  });
+}

src/helpers/dmmf.ts

@@ -1,4 +1,5 @@
 import type DMMF from '@prisma/dmmf';
+import type { ColumnAnnotationMap } from '../handler/typedsql';
 import { createRegexForType, generateTypeNamesFromName } from './regex';
 import { parseTypeSyntax } from './type-parser';
 
@@ -63,3 +64,32 @@ export function extractPrismaModels(dmmf: DMMF.Document): {
   }
   return { typeToNameMap, modelMap, knownNoOps };
 }
+
+/**
+ * Builds a map from database column names to their documentation for JSON fields.
+ *
+ * This is used when processing TypedSQL query result columns — each column name is
+ * matched against the Prisma schema field (or its `@map`-ed database name) to find
+ * the documentation annotation that drives type replacement.
+ *
+ * When the same column name exists in multiple models, the first model encountered wins.
+ */
+export function buildTypedSqlColumnAnnotationMap(dmmf: DMMF.Document): ColumnAnnotationMap {
+  const map: ColumnAnnotationMap = new Map();
+
+  for (const model of dmmf.datamodel.models) {
+    for (const field of model.fields) {
+      if (field.type !== 'Json') continue;
+
+      // Use the @map column name if provided, otherwise the field name
+      const columnName = field.dbName || field.name;
+
+      // First model with this column name wins to avoid ambiguity
+      if (!map.has(columnName)) {
+        map.set(columnName, field.documentation);
+      }
+    }
+  }
+
+  return map;
+}

src/on-generate.ts

@@ -1,10 +1,11 @@
 import fs from 'node:fs/promises';
 import { join } from 'node:path';
-import type { GeneratorOptions } from '@prisma/generator';
+import type { GeneratorOptions, SqlQueryOutput } from '@prisma/generator';
 import ts from 'typescript';
 import { handlePrismaModule } from './handler/module';
 import { handleStatement } from './handler/statement';
-import { extractPrismaModels } from './helpers/dmmf';
+import { type ColumnAnnotationMap, handleTypedSqlFile } from './handler/typedsql';
+import { buildTypedSqlColumnAnnotationMap, extractPrismaModels } from './helpers/dmmf';
 import { type PrismaJsonTypesGeneratorConfig, parseConfig } from './util/config';
 import { DeclarationWriter, getNamespacePrelude } from './util/declaration-writer';
 import { findPrismaClientGenerators, type GeneratorWithOutput } from './util/prisma-generator';
@@ -53,6 +54,9 @@ async function generateClient(
         })
       );
 
+      // Process TypedSQL queries after pjtg.ts is in place
+      await handleTypedSqlQueries(prismaClient, config, options, ext);
+
       return;
     }
   }
@@ -62,6 +66,84 @@ async function generateClient(
     : buildTypesFilePath(prismaClient.output.value, config.clientOutput, options.schemaPath);
 
   await handleDeclarationFile(clientOutput, config, options, ext, false);
+
+  // For TypedSQL with old-client (or new-client single-file mode), ensure pjtg.ts exists
+  if (options.typedSql?.length) {
+    if (!isNewClient) {
+      const pjtgPath = join(prismaClient.output.value, 'pjtg.ts');
+      if (!(await fs.stat(pjtgPath).catch(() => null))) {
+        await fs.writeFile(
+          pjtgPath,
+          await getNamespacePrelude({
+            namespace: config.namespace,
+            isNewClient: false,
+            dotExt: ext ? `.${ext}` : ''
+          })
+        );
+      }
+    }
+    await handleTypedSqlQueries(prismaClient, config, options, ext);
+  }
+}
+
+async function handleTypedSqlQueries(
+  prismaClient: GeneratorWithOutput,
+  config: PrismaJsonTypesGeneratorConfig,
+  options: GeneratorOptions,
+  importFileExtension: string | undefined
+) {
+  const typedSql = options.typedSql;
+  if (!typedSql?.length) return;
+
+  const sqlDir = join(prismaClient.output.value, 'sql');
+  const sqlDirStat = await fs.stat(sqlDir).catch(() => null);
+  if (!sqlDirStat?.isDirectory()) return;
+
+  const columnAnnotationMap = buildTypedSqlColumnAnnotationMap(options.dmmf);
+
+  for (const query of typedSql) {
+    const filePath = join(sqlDir, `${query.name}.ts`);
+    const fileStat = await fs.stat(filePath).catch(() => null);
+    if (!fileStat?.isFile()) continue;
+
+    await handleTypedSqlDeclarationFile(
+      filePath,
+      query,
+      columnAnnotationMap,
+      config,
+      importFileExtension
+    );
+  }
+}
+
+async function handleTypedSqlDeclarationFile(
+  filepath: string,
+  query: SqlQueryOutput,
+  columnAnnotationMap: ColumnAnnotationMap,
+  config: PrismaJsonTypesGeneratorConfig,
+  importFileExtension: string | undefined
+) {
+  // TypedSQL files live in sql/ (one level below client output), same as model files in
+  // the multifile layout — so multifile=true causes the writer to emit
+  // `import type * as PJTG from '../pjtg'`, which is the correct relative path.
+  const writer = new DeclarationWriter(filepath, config, true, importFileExtension);
+  await writer.load();
+
+  const tsSource = ts.createSourceFile(
+    writer.filepath,
+    writer.content,
+    ts.ScriptTarget.ESNext,
+    true,
+    ts.ScriptKind.TS
+  );
+
+  try {
+    handleTypedSqlFile(tsSource, writer, query, columnAnnotationMap, config);
+  } catch (error) {
+    console.error(error);
+  }
+
+  await writer.save();
 }
 
 async function handleDeclarationFile(

test/schemas/typedsql.prisma

@@ -0,0 +1,24 @@
+generator client {
+  provider        = "prisma-client"
+  output          = "../target/typedSql"
+  previewFeatures = ["typedSql"]
+}
+
+generator json {
+  provider  = "node ./index.js"
+  namespace = "PTypedSqlJson"
+}
+
+datasource db {
+  provider = "postgresql"
+}
+
+model Model {
+  id Int @id
+
+  /// ![number]
+  field Json
+
+  /// [OptionalField]
+  optField Json?
+}

test/types/typedsql.test-d.ts

@@ -0,0 +1,48 @@
+import { expectAssignable, expectNotAssignable } from 'tsd';
+import type { Model, Prisma } from '../target/typedSql/client';
+
+declare global {
+  export namespace PTypedSqlJson {
+    export type OptionalField = string;
+  }
+}
+
+// field is typed as (number) from the /// ![number] annotation
+expectAssignable<Model>({
+  id: 0,
+  field: 42,
+  optField: null
+});
+
+// optField is PTypedSqlJson.OptionalField | null (string | null here)
+expectAssignable<Model>({
+  id: 0,
+  field: 0,
+  optField: 'hello'
+});
+
+// field should NOT accept a string (it's (number))
+expectNotAssignable<Model>({
+  id: 0,
+  field: 'not-a-number',
+  optField: null
+});
+
+// optField should NOT accept a number (it's OptionalField = string)
+expectNotAssignable<Model>({
+  id: 0,
+  field: 0,
+  optField: 42
+});
+
+// CreateInput should reflect the same type constraints
+expectAssignable<Prisma.ModelCreateInput>({
+  id: 1,
+  field: 1,
+  optField: 'value'
+});
+
+expectNotAssignable<Prisma.ModelCreateInput>({
+  id: 1,
+  field: 'not-a-number'
+});