Skip to content

Commit 1bd1b8f

Browse files
ymc9ymc9
authored
fix(json): support enums in type declarations (#1837)
1 parent 1d1fec0 commit 1bd1b8f

File tree

7 files changed

+267
-21
lines changed

7 files changed

+267
-21
lines changed

‎packages/language/src/generated/ast.ts

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,14 @@ export function isTypeDeclaration(item: unknown): item is TypeDeclaration {
9898
return reflection.isInstance(item, TypeDeclaration);
9999
}
100100

101+
export type TypeDefFieldTypes = Enum | TypeDef;
102+
103+
export const TypeDefFieldTypes = 'TypeDefFieldTypes';
104+
105+
export function isTypeDefFieldTypes(item: unknown): item is TypeDefFieldTypes {
106+
return reflection.isInstance(item, TypeDefFieldTypes);
107+
}
108+
101109
export interface Argument extends AstNode {
102110
readonly $container: InvocationExpr;
103111
readonly $type: 'Argument';
@@ -654,7 +662,7 @@ export interface TypeDefFieldType extends AstNode {
654662
readonly $type: 'TypeDefFieldType';
655663
array: boolean
656664
optional: boolean
657-
reference?: Reference<TypeDef>
665+
reference?: Reference<TypeDefFieldTypes>
658666
type?: BuiltinType
659667
}
660668

@@ -738,14 +746,15 @@ export type ZModelAstType = {
738746
TypeDef: TypeDef
739747
TypeDefField: TypeDefField
740748
TypeDefFieldType: TypeDefFieldType
749+
TypeDefFieldTypes: TypeDefFieldTypes
741750
UnaryExpr: UnaryExpr
742751
UnsupportedFieldType: UnsupportedFieldType
743752
}
744753

745754
export class ZModelAstReflection extends AbstractAstReflection {
746755

747756
getAllTypes(): string[] {
748-
return ['AbstractDeclaration', 'Argument', 'ArrayExpr', 'Attribute', 'AttributeArg', 'AttributeParam', 'AttributeParamType', 'BinaryExpr', 'BooleanLiteral', 'ConfigArrayExpr', 'ConfigExpr', 'ConfigField', 'ConfigInvocationArg', 'ConfigInvocationExpr', 'DataModel', 'DataModelAttribute', 'DataModelField', 'DataModelFieldAttribute', 'DataModelFieldType', 'DataSource', 'Enum', 'EnumField', 'Expression', 'FieldInitializer', 'FunctionDecl', 'FunctionParam', 'FunctionParamType', 'GeneratorDecl', 'InternalAttribute', 'InvocationExpr', 'LiteralExpr', 'MemberAccessExpr', 'Model', 'ModelImport', 'NullExpr', 'NumberLiteral', 'ObjectExpr', 'Plugin', 'PluginField', 'ReferenceArg', 'ReferenceExpr', 'ReferenceTarget', 'StringLiteral', 'ThisExpr', 'TypeDeclaration', 'TypeDef', 'TypeDefField', 'TypeDefFieldType', 'UnaryExpr', 'UnsupportedFieldType'];
757+
return ['AbstractDeclaration', 'Argument', 'ArrayExpr', 'Attribute', 'AttributeArg', 'AttributeParam', 'AttributeParamType', 'BinaryExpr', 'BooleanLiteral', 'ConfigArrayExpr', 'ConfigExpr', 'ConfigField', 'ConfigInvocationArg', 'ConfigInvocationExpr', 'DataModel', 'DataModelAttribute', 'DataModelField', 'DataModelFieldAttribute', 'DataModelFieldType', 'DataSource', 'Enum', 'EnumField', 'Expression', 'FieldInitializer', 'FunctionDecl', 'FunctionParam', 'FunctionParamType', 'GeneratorDecl', 'InternalAttribute', 'InvocationExpr', 'LiteralExpr', 'MemberAccessExpr', 'Model', 'ModelImport', 'NullExpr', 'NumberLiteral', 'ObjectExpr', 'Plugin', 'PluginField', 'ReferenceArg', 'ReferenceExpr', 'ReferenceTarget', 'StringLiteral', 'ThisExpr', 'TypeDeclaration', 'TypeDef', 'TypeDefField', 'TypeDefFieldType', 'TypeDefFieldTypes', 'UnaryExpr', 'UnsupportedFieldType'];
749758
}
750759

751760
protected override computeIsSubtype(subtype: string, supertype: string): boolean {
@@ -775,16 +784,18 @@ export class ZModelAstReflection extends AbstractAstReflection {
775784
case ConfigArrayExpr: {
776785
return this.isSubtype(ConfigExpr, supertype);
777786
}
778-
case DataModel:
779-
case Enum:
780-
case TypeDef: {
787+
case DataModel: {
781788
return this.isSubtype(AbstractDeclaration, supertype) || this.isSubtype(TypeDeclaration, supertype);
782789
}
783790
case DataModelField:
784791
case EnumField:
785792
case FunctionParam: {
786793
return this.isSubtype(ReferenceTarget, supertype);
787794
}
795+
case Enum:
796+
case TypeDef: {
797+
return this.isSubtype(AbstractDeclaration, supertype) || this.isSubtype(TypeDeclaration, supertype) || this.isSubtype(TypeDefFieldTypes, supertype);
798+
}
788799
case InvocationExpr:
789800
case LiteralExpr: {
790801
return this.isSubtype(ConfigExpr, supertype) || this.isSubtype(Expression, supertype);
@@ -821,7 +832,7 @@ export class ZModelAstReflection extends AbstractAstReflection {
821832
return ReferenceTarget;
822833
}
823834
case 'TypeDefFieldType:reference': {
824-
return TypeDef;
835+
return TypeDefFieldTypes;
825836
}
826837
default: {
827838
throw new Error(`${referenceId} is not a valid reference id.`);

‎packages/language/src/generated/grammar.ts

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2165,7 +2165,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
21652165
"terminal": {
21662166
"$type": "CrossReference",
21672167
"type": {
2168-
"$ref": "#/types@1"
2168+
"$ref": "#/types@2"
21692169
},
21702170
"terminal": {
21712171
"$type": "RuleCall",
@@ -2267,7 +2267,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
22672267
},
22682268
"arguments": []
22692269
},
2270-
"cardinality": "+"
2270+
"cardinality": "*"
22712271
},
22722272
{
22732273
"$type": "Keyword",
@@ -2375,7 +2375,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
23752375
"terminal": {
23762376
"$type": "CrossReference",
23772377
"type": {
2378-
"$ref": "#/rules@40"
2378+
"$ref": "#/types@1"
23792379
},
23802380
"terminal": {
23812381
"$type": "RuleCall",
@@ -2827,7 +2827,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
28272827
"terminal": {
28282828
"$type": "CrossReference",
28292829
"type": {
2830-
"$ref": "#/types@1"
2830+
"$ref": "#/types@2"
28312831
},
28322832
"terminal": {
28332833
"$type": "RuleCall",
@@ -3255,7 +3255,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
32553255
"terminal": {
32563256
"$type": "CrossReference",
32573257
"type": {
3258-
"$ref": "#/types@1"
3258+
"$ref": "#/types@2"
32593259
},
32603260
"terminal": {
32613261
"$type": "RuleCall",
@@ -3838,6 +3838,27 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
38383838
]
38393839
}
38403840
},
3841+
{
3842+
"$type": "Type",
3843+
"name": "TypeDefFieldTypes",
3844+
"type": {
3845+
"$type": "UnionType",
3846+
"types": [
3847+
{
3848+
"$type": "SimpleType",
3849+
"typeRef": {
3850+
"$ref": "#/rules@40"
3851+
}
3852+
},
3853+
{
3854+
"$type": "SimpleType",
3855+
"typeRef": {
3856+
"$ref": "#/rules@44"
3857+
}
3858+
}
3859+
]
3860+
}
3861+
},
38413862
{
38423863
"$type": "Type",
38433864
"name": "TypeDeclaration",

‎packages/language/src/zmodel.langium

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,15 +183,17 @@ TypeDef:
183183
(comments+=TRIPLE_SLASH_COMMENT)*
184184
'type' name=RegularID '{' (
185185
fields+=TypeDefField
186-
)+
186+
)*
187187
'}';
188188

189+
type TypeDefFieldTypes = TypeDef | Enum;
190+
189191
TypeDefField:
190-
(comments+=TRIPLE_SLASH_COMMENT)*
192+
(comments+=TRIPLE_SLASH_COMMENT)*
191193
name=RegularIDWithTypeNames type=TypeDefFieldType (attributes+=DataModelFieldAttribute)*;
192194

193195
TypeDefFieldType:
194-
(type=BuiltinType | reference=[TypeDef:RegularID]) (array?='[' ']')? (optional?='?')?;
196+
(type=BuiltinType | reference=[TypeDefFieldTypes:RegularID]) (array?='[' ']')? (optional?='?')?;
195197

196198
UnsupportedFieldType:
197199
'Unsupported' '(' (value=LiteralExpr) ')';

‎packages/schema/src/plugins/enhancer/enhance/model-typedef-generator.ts

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import { PluginError } from '@zenstackhq/sdk';
2-
import { BuiltinType, TypeDef, TypeDefFieldType } from '@zenstackhq/sdk/ast';
1+
import { getDataModels, PluginError } from '@zenstackhq/sdk';
2+
import { BuiltinType, Enum, isEnum, TypeDef, TypeDefFieldType } from '@zenstackhq/sdk/ast';
33
import { SourceFile } from 'ts-morph';
44
import { match } from 'ts-pattern';
55
import { name } from '..';
@@ -36,7 +36,11 @@ function zmodelTypeToTsType(type: TypeDefFieldType) {
3636
if (type.type) {
3737
result = builtinTypeToTsType(type.type);
3838
} else if (type.reference?.ref) {
39-
result = type.reference.ref.name;
39+
if (isEnum(type.reference.ref)) {
40+
result = makeEnumTypeReference(type.reference.ref);
41+
} else {
42+
result = type.reference.ref.name;
43+
}
4044
} else {
4145
throw new PluginError(name, `Unsupported field type: ${type}`);
4246
}
@@ -61,3 +65,17 @@ function builtinTypeToTsType(type: BuiltinType) {
6165
.with('Json', () => 'unknown')
6266
.exhaustive();
6367
}
68+
69+
function makeEnumTypeReference(enumDecl: Enum) {
70+
const zmodel = enumDecl.$container;
71+
const models = getDataModels(zmodel);
72+
73+
if (models.some((model) => model.fields.some((field) => field.type.reference?.ref === enumDecl))) {
74+
// if the enum is referenced by any data model, Prisma already generates its type,
75+
// we just need to reference it
76+
return enumDecl.name;
77+
} else {
78+
// otherwise, we need to inline the enum
79+
return enumDecl.fields.map((field) => `'${field.name}'`).join(' | ');
80+
}
81+
}

‎packages/schema/src/plugins/zod/transformer.ts

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
/* eslint-disable @typescript-eslint/ban-ts-comment */
22
import { indentString, isDiscriminatorField, type PluginOptions } from '@zenstackhq/sdk';
3-
import { DataModel, isDataModel, isTypeDef, type Model } from '@zenstackhq/sdk/ast';
3+
import { DataModel, Enum, isDataModel, isEnum, isTypeDef, type Model } from '@zenstackhq/sdk/ast';
44
import { checkModelHasModelRelation, findModelByName, isAggregateInputType } from '@zenstackhq/sdk/dmmf-helpers';
55
import { supportCreateMany, type DMMF as PrismaDMMF } from '@zenstackhq/sdk/prisma';
66
import path from 'path';
@@ -53,6 +53,9 @@ export default class Transformer {
5353
}
5454

5555
async generateEnumSchemas() {
56+
const generated: string[] = [];
57+
58+
// generate for enums in DMMF
5659
for (const enumType of this.enumTypes) {
5760
const name = upperCaseFirst(enumType.name);
5861
const filePath = path.join(Transformer.outputPath, `enums/${name}.schema.ts`);
@@ -61,14 +64,26 @@ export default class Transformer {
6164
`z.enum(${JSON.stringify(enumType.values)})`
6265
)}`;
6366
this.sourceFiles.push(this.project.createSourceFile(filePath, content, { overwrite: true }));
67+
generated.push(enumType.name);
68+
}
69+
70+
// enums not referenced by data models are not in DMMF, deal with them separately
71+
const extraEnums = this.zmodel.declarations.filter((d): d is Enum => isEnum(d) && !generated.includes(d.name));
72+
for (const enumDecl of extraEnums) {
73+
const name = upperCaseFirst(enumDecl.name);
74+
const filePath = path.join(Transformer.outputPath, `enums/${name}.schema.ts`);
75+
const content = `/* eslint-disable */\n${this.generateImportZodStatement()}\n${this.generateExportSchemaStatement(
76+
`${name}`,
77+
`z.enum(${JSON.stringify(enumDecl.fields.map((f) => f.name))})`
78+
)}`;
79+
this.sourceFiles.push(this.project.createSourceFile(filePath, content, { overwrite: true }));
80+
generated.push(enumDecl.name);
6481
}
6582

6683
this.sourceFiles.push(
6784
this.project.createSourceFile(
6885
path.join(Transformer.outputPath, `enums/index.ts`),
69-
this.enumTypes
70-
.map((enumType) => `export * from './${upperCaseFirst(enumType.name)}.schema';`)
71-
.join('\n'),
86+
generated.map((name) => `export * from './${upperCaseFirst(name)}.schema';`).join('\n'),
7287
{ overwrite: true }
7388
)
7489
);

‎tests/integration/tests/enhancements/json/crud.test.ts

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,83 @@ describe('Json field CRUD', () => {
191191
).toResolveTruthy();
192192
});
193193

194+
it('respects enums used by data models', async () => {
195+
const params = await loadSchema(
196+
`
197+
enum Role {
198+
USER
199+
ADMIN
200+
}
201+
202+
type Profile {
203+
role Role
204+
}
205+
206+
model User {
207+
id Int @id @default(autoincrement())
208+
profile Profile @json
209+
@@allow('all', true)
210+
}
211+
212+
model Foo {
213+
id Int @id @default(autoincrement())
214+
role Role
215+
}
216+
`,
217+
{
218+
provider: 'postgresql',
219+
dbUrl,
220+
}
221+
);
222+
223+
prisma = params.prisma;
224+
const db = params.enhance();
225+
226+
await expect(db.user.create({ data: { profile: { role: 'MANAGER' } } })).toBeRejectedByPolicy();
227+
await expect(db.user.create({ data: { profile: { role: 'ADMIN' } } })).resolves.toMatchObject({
228+
profile: { role: 'ADMIN' },
229+
});
230+
await expect(db.user.findFirst()).resolves.toMatchObject({
231+
profile: { role: 'ADMIN' },
232+
});
233+
});
234+
235+
it('respects enums unused by data models', async () => {
236+
const params = await loadSchema(
237+
`
238+
enum Role {
239+
USER
240+
ADMIN
241+
}
242+
243+
type Profile {
244+
role Role
245+
}
246+
247+
model User {
248+
id Int @id @default(autoincrement())
249+
profile Profile @json
250+
@@allow('all', true)
251+
}
252+
`,
253+
{
254+
provider: 'postgresql',
255+
dbUrl,
256+
}
257+
);
258+
259+
prisma = params.prisma;
260+
const db = params.enhance();
261+
262+
await expect(db.user.create({ data: { profile: { role: 'MANAGER' } } })).toBeRejectedByPolicy();
263+
await expect(db.user.create({ data: { profile: { role: 'ADMIN' } } })).resolves.toMatchObject({
264+
profile: { role: 'ADMIN' },
265+
});
266+
await expect(db.user.findFirst()).resolves.toMatchObject({
267+
profile: { role: 'ADMIN' },
268+
});
269+
});
270+
194271
it('respects @default', async () => {
195272
const params = await loadSchema(
196273
`

0 commit comments

Comments
 (0)