Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(json): support enums in type declarations #1837

Merged
merged 2 commits into from
Nov 7, 2024
Merged

fix(json): support enums in type declarations #1837

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a5257ee
fix(json): support enums in type declarations
ymc9 Nov 7, 2024
2f8fbc1
fixes
ymc9 Nov 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions packages/language/src/generated/ast.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,14 @@ export function isTypeDeclaration(item: unknown): item is TypeDeclaration {
return reflection.isInstance(item, TypeDeclaration);
}

export type TypeDefFieldTypes = Enum | TypeDef;

export const TypeDefFieldTypes = 'TypeDefFieldTypes';

export function isTypeDefFieldTypes(item: unknown): item is TypeDefFieldTypes {
return reflection.isInstance(item, TypeDefFieldTypes);
}

export interface Argument extends AstNode {
readonly $container: InvocationExpr;
readonly $type: 'Argument';
Expand Down Expand Up @@ -654,7 +662,7 @@ export interface TypeDefFieldType extends AstNode {
readonly $type: 'TypeDefFieldType';
array: boolean
optional: boolean
reference?: Reference<TypeDef>
reference?: Reference<TypeDefFieldTypes>
type?: BuiltinType
}

Expand Down Expand Up @@ -738,14 +746,15 @@ export type ZModelAstType = {
TypeDef: TypeDef
TypeDefField: TypeDefField
TypeDefFieldType: TypeDefFieldType
TypeDefFieldTypes: TypeDefFieldTypes
UnaryExpr: UnaryExpr
UnsupportedFieldType: UnsupportedFieldType
}

export class ZModelAstReflection extends AbstractAstReflection {

getAllTypes(): string[] {
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'];
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'];
}

protected override computeIsSubtype(subtype: string, supertype: string): boolean {
Expand Down Expand Up @@ -775,16 +784,18 @@ export class ZModelAstReflection extends AbstractAstReflection {
case ConfigArrayExpr: {
return this.isSubtype(ConfigExpr, supertype);
}
case DataModel:
case Enum:
case TypeDef: {
case DataModel: {
return this.isSubtype(AbstractDeclaration, supertype) || this.isSubtype(TypeDeclaration, supertype);
}
case DataModelField:
case EnumField:
case FunctionParam: {
return this.isSubtype(ReferenceTarget, supertype);
}
case Enum:
case TypeDef: {
return this.isSubtype(AbstractDeclaration, supertype) || this.isSubtype(TypeDeclaration, supertype) || this.isSubtype(TypeDefFieldTypes, supertype);
}
case InvocationExpr:
case LiteralExpr: {
return this.isSubtype(ConfigExpr, supertype) || this.isSubtype(Expression, supertype);
Expand Down Expand Up @@ -821,7 +832,7 @@ export class ZModelAstReflection extends AbstractAstReflection {
return ReferenceTarget;
}
case 'TypeDefFieldType:reference': {
return TypeDef;
return TypeDefFieldTypes;
}
default: {
throw new Error(`${referenceId} is not a valid reference id.`);
Expand Down
31 changes: 26 additions & 5 deletions packages/language/src/generated/grammar.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2165,7 +2165,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
"terminal": {
"$type": "CrossReference",
"type": {
"$ref": "#/types@1"
"$ref": "#/types@2"
},
"terminal": {
"$type": "RuleCall",
Expand Down Expand Up @@ -2267,7 +2267,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
},
"arguments": []
},
"cardinality": "+"
"cardinality": "*"
},
{
"$type": "Keyword",
Expand Down Expand Up @@ -2375,7 +2375,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
"terminal": {
"$type": "CrossReference",
"type": {
"$ref": "#/rules@40"
"$ref": "#/types@1"
},
"terminal": {
"$type": "RuleCall",
Expand Down Expand Up @@ -2827,7 +2827,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
"terminal": {
"$type": "CrossReference",
"type": {
"$ref": "#/types@1"
"$ref": "#/types@2"
},
"terminal": {
"$type": "RuleCall",
Expand Down Expand Up @@ -3255,7 +3255,7 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
"terminal": {
"$type": "CrossReference",
"type": {
"$ref": "#/types@1"
"$ref": "#/types@2"
},
"terminal": {
"$type": "RuleCall",
Expand Down Expand Up @@ -3838,6 +3838,27 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
]
}
},
{
"$type": "Type",
"name": "TypeDefFieldTypes",
"type": {
"$type": "UnionType",
"types": [
{
"$type": "SimpleType",
"typeRef": {
"$ref": "#/rules@40"
}
},
{
"$type": "SimpleType",
"typeRef": {
"$ref": "#/rules@44"
}
}
]
}
},
{
"$type": "Type",
"name": "TypeDeclaration",
Expand Down
8 changes: 5 additions & 3 deletions packages/language/src/zmodel.langium
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,15 +183,17 @@ TypeDef:
(comments+=TRIPLE_SLASH_COMMENT)*
'type' name=RegularID '{' (
fields+=TypeDefField
)+
)*
'}';

type TypeDefFieldTypes = TypeDef | Enum;

TypeDefField:
(comments+=TRIPLE_SLASH_COMMENT)*
(comments+=TRIPLE_SLASH_COMMENT)*
name=RegularIDWithTypeNames type=TypeDefFieldType (attributes+=DataModelFieldAttribute)*;

TypeDefFieldType:
(type=BuiltinType | reference=[TypeDef:RegularID]) (array?='[' ']')? (optional?='?')?;
(type=BuiltinType | reference=[TypeDefFieldTypes:RegularID]) (array?='[' ']')? (optional?='?')?;

UnsupportedFieldType:
'Unsupported' '(' (value=LiteralExpr) ')';
Expand Down
24 changes: 21 additions & 3 deletions packages/schema/src/plugins/enhancer/enhance/model-typedef-generator.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { PluginError } from '@zenstackhq/sdk';
import { BuiltinType, TypeDef, TypeDefFieldType } from '@zenstackhq/sdk/ast';
import { getDataModels, PluginError } from '@zenstackhq/sdk';
import { BuiltinType, Enum, isEnum, TypeDef, TypeDefFieldType } from '@zenstackhq/sdk/ast';
import { SourceFile } from 'ts-morph';
import { match } from 'ts-pattern';
import { name } from '..';
Expand Down Expand Up @@ -36,7 +36,11 @@ function zmodelTypeToTsType(type: TypeDefFieldType) {
if (type.type) {
result = builtinTypeToTsType(type.type);
} else if (type.reference?.ref) {
result = type.reference.ref.name;
if (isEnum(type.reference.ref)) {
result = makeEnumTypeReference(type.reference.ref);
} else {
result = type.reference.ref.name;
}
} else {
throw new PluginError(name, `Unsupported field type: ${type}`);
}
Expand All @@ -61,3 +65,17 @@ function builtinTypeToTsType(type: BuiltinType) {
.with('Json', () => 'unknown')
.exhaustive();
}

function makeEnumTypeReference(enumDecl: Enum) {
const zmodel = enumDecl.$container;
const models = getDataModels(zmodel);

if (models.some((model) => model.fields.some((field) => field.type.reference?.ref === enumDecl))) {
// if the enum is referenced by any data model, Prisma already generates its type,
// we just need to reference it
return enumDecl.name;
} else {
// otherwise, we need to inline the enum
return enumDecl.fields.map((field) => `'${field.name}'`).join(' | ');
}
}
23 changes: 19 additions & 4 deletions packages/schema/src/plugins/zod/transformer.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/* eslint-disable @typescript-eslint/ban-ts-comment */
import { indentString, isDiscriminatorField, type PluginOptions } from '@zenstackhq/sdk';
import { DataModel, isDataModel, isTypeDef, type Model } from '@zenstackhq/sdk/ast';
import { DataModel, Enum, isDataModel, isEnum, isTypeDef, type Model } from '@zenstackhq/sdk/ast';
import { checkModelHasModelRelation, findModelByName, isAggregateInputType } from '@zenstackhq/sdk/dmmf-helpers';
import { supportCreateMany, type DMMF as PrismaDMMF } from '@zenstackhq/sdk/prisma';
import path from 'path';
Expand Down Expand Up @@ -53,6 +53,9 @@ export default class Transformer {
}

async generateEnumSchemas() {
const generated: string[] = [];

// generate for enums in DMMF
for (const enumType of this.enumTypes) {
const name = upperCaseFirst(enumType.name);
const filePath = path.join(Transformer.outputPath, `enums/${name}.schema.ts`);
Expand All @@ -61,14 +64,26 @@ export default class Transformer {
`z.enum(${JSON.stringify(enumType.values)})`
)}`;
this.sourceFiles.push(this.project.createSourceFile(filePath, content, { overwrite: true }));
generated.push(enumType.name);
}

// enums not referenced by data models are not in DMMF, deal with them separately
const extraEnums = this.zmodel.declarations.filter((d): d is Enum => isEnum(d) && !generated.includes(d.name));
for (const enumDecl of extraEnums) {
const name = upperCaseFirst(enumDecl.name);
const filePath = path.join(Transformer.outputPath, `enums/${name}.schema.ts`);
const content = `/* eslint-disable */\n${this.generateImportZodStatement()}\n${this.generateExportSchemaStatement(
`${name}`,
`z.enum(${JSON.stringify(enumDecl.fields.map((f) => f.name))})`
)}`;
this.sourceFiles.push(this.project.createSourceFile(filePath, content, { overwrite: true }));
generated.push(enumDecl.name);
}

this.sourceFiles.push(
this.project.createSourceFile(
path.join(Transformer.outputPath, `enums/index.ts`),
this.enumTypes
.map((enumType) => `export * from './${upperCaseFirst(enumType.name)}.schema';`)
.join('\n'),
generated.map((name) => `export * from './${upperCaseFirst(name)}.schema';`).join('\n'),
{ overwrite: true }
)
);
Expand Down
77 changes: 77 additions & 0 deletions tests/integration/tests/enhancements/json/crud.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,83 @@ describe('Json field CRUD', () => {
).toResolveTruthy();
});

it('respects enums used by data models', async () => {
const params = await loadSchema(
`
enum Role {
USER
ADMIN
}

type Profile {
role Role
}

model User {
id Int @id @default(autoincrement())
profile Profile @json
@@allow('all', true)
}

model Foo {
id Int @id @default(autoincrement())
role Role
}
`,
{
provider: 'postgresql',
dbUrl,
}
);

prisma = params.prisma;
const db = params.enhance();

await expect(db.user.create({ data: { profile: { role: 'MANAGER' } } })).toBeRejectedByPolicy();
await expect(db.user.create({ data: { profile: { role: 'ADMIN' } } })).resolves.toMatchObject({
profile: { role: 'ADMIN' },
});
await expect(db.user.findFirst()).resolves.toMatchObject({
profile: { role: 'ADMIN' },
});
});

it('respects enums unused by data models', async () => {
const params = await loadSchema(
`
enum Role {
USER
ADMIN
}

type Profile {
role Role
}

model User {
id Int @id @default(autoincrement())
profile Profile @json
@@allow('all', true)
}
`,
{
provider: 'postgresql',
dbUrl,
}
);

prisma = params.prisma;
const db = params.enhance();

await expect(db.user.create({ data: { profile: { role: 'MANAGER' } } })).toBeRejectedByPolicy();
await expect(db.user.create({ data: { profile: { role: 'ADMIN' } } })).resolves.toMatchObject({
profile: { role: 'ADMIN' },
});
await expect(db.user.findFirst()).resolves.toMatchObject({
profile: { role: 'ADMIN' },
});
});
ymc9 marked this conversation as resolved.
Show resolved Hide resolved

it('respects @default', async () => {
const params = await loadSchema(
`
Expand Down
Loading
Loading