feat: primitive custom types (#2852)

Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: 3 people

packages/cli/test/prisma-schema-gen.test.ts

@@ -91,4 +91,47 @@ model User {
         expect(prismaSchemaText.includes('@ds.JsonB')).toBe(true);
         expect(prismaSchemaText.includes('@ds.ByteA')).toBe(true);
     });
+
+    it('renames primitive type defs to match the base type', async () => {
+        const model = await loadSchema(`
+model User {
+    id   String   @id
+    name UserName
+}
+
+type UserName with String {
+    this String
+}
+        `);
+
+        const generator = new PrismaSchemaGenerator(model);
+        const prismaSchemaText = await generator.generate();
+
+        expect(prismaSchemaText.includes('name UserName')).toBe(false);
+        expect(prismaSchemaText.includes('name String')).toBe(true);
+    });
+
+    it('renames primitive type defs to match the base type when used with lists', async () => {
+        const model = await loadSchema(`
+datasource db {
+    provider = 'postgresql'
+    url      = env('DATABASE_URL')
+}
+
+model User {
+    id   String   @id
+    name UserName[]
+}
+
+type UserName with String {
+    this String
+}
+        `);
+
+        const generator = new PrismaSchemaGenerator(model);
+        const prismaSchemaText = await generator.generate();
+
+        expect(prismaSchemaText.includes('name UserName[]')).toBe(false);
+        expect(prismaSchemaText.includes('name String[]')).toBe(true);
+    });
 });

packages/cli/test/ts-schema-gen.test.ts

@@ -1015,4 +1015,61 @@ type Profile with Strict {
             },
         });
     });
+
+    it('supports primitive type defs', async () => {
+        const { schema } = await generateTsSchema(`
+            model User {
+                id   Int      @id
+                name UserName
+            }
+
+            type UserName with String {
+                this String
+            }
+            `);
+
+        expect(schema).toMatchObject({
+            models: {
+                User: {
+                    name: 'User',
+                    fields: {
+                        id: {
+                            name: 'id',
+                            type: 'Int',
+                            id: true,
+                            attributes: [
+                                {
+                                    name: '@id',
+                                },
+                            ],
+                        },
+                        name: {
+                            name: 'name',
+                            type: 'UserName',
+                        },
+                    },
+                    idFields: ['id'],
+                    uniqueFields: {
+                        id: {
+                            type: 'Int',
+                        },
+                    },
+                },
+            },
+            typeDefs: {
+                UserName: {
+                    name: 'UserName',
+                    base: 'String',
+                    fields: {
+                        this: {
+                            name: 'this',
+                            type: 'String',
+                        },
+                    },
+                },
+            },
+            authType: 'User',
+            plugins: {},
+        });
+    });
 });

packages/language/src/generated/ast.ts

@@ -435,6 +435,7 @@ export interface DataModel extends langium.AstNode {
     readonly $container: Model;
     readonly $type: 'DataModel';
     attributes: Array<DataModelAttribute>;
+    base?: BuiltinType;
     baseModel?: langium.Reference<DataModel>;
     comments: Array<string>;
     fields: Array<DataField>;
@@ -448,6 +449,7 @@ export interface DataModel extends langium.AstNode {
 export const DataModel = {
     $type: 'DataModel',
     attributes: 'attributes',
+    base: 'base',
     baseModel: 'baseModel',
     comments: 'comments',
     fields: 'fields',
@@ -912,10 +914,10 @@ export function isReferenceTarget(item: unknown): item is ReferenceTarget {
     return reflection.isInstance(item, ReferenceTarget.$type);
 }
 
-export type RegularID = 'abstract' | 'attribute' | 'datasource' | 'enum' | 'import' | 'in' | 'model' | 'plugin' | 'type' | 'view' | string;
+export type RegularID = 'abstract' | 'attribute' | 'datasource' | 'enum' | 'import' | 'in' | 'model' | 'plugin' | 'this' | 'type' | 'view' | string;
 
 export function isRegularID(item: unknown): item is RegularID {
-    return item === 'model' || item === 'enum' || item === 'attribute' || item === 'datasource' || item === 'plugin' || item === 'abstract' || item === 'in' || item === 'view' || item === 'import' || item === 'type' || (typeof item === 'string' && (/[_a-zA-Z][\w_]*/.test(item)));
+    return item === 'model' || item === 'enum' || item === 'attribute' || item === 'datasource' || item === 'plugin' || item === 'abstract' || item === 'in' || item === 'view' || item === 'import' || item === 'type' || item === 'this' || (typeof item === 'string' && (/[_a-zA-Z][\w_]*/.test(item)));
 }
 
 export type RegularIDWithTypeNames = 'Any' | 'BigInt' | 'Boolean' | 'Bytes' | 'DateTime' | 'Decimal' | 'Float' | 'Int' | 'Json' | 'Null' | 'Object' | 'String' | 'Unsupported' | 'Void' | RegularID;
@@ -968,6 +970,7 @@ export interface TypeDef extends langium.AstNode {
     readonly $container: Model;
     readonly $type: 'TypeDef';
     attributes: Array<DataModelAttribute>;
+    base?: BuiltinType;
     comments: Array<string>;
     fields: Array<DataField>;
     mixins: Array<langium.Reference<TypeDef>>;
@@ -977,6 +980,7 @@ export interface TypeDef extends langium.AstNode {
 export const TypeDef = {
     $type: 'TypeDef',
     attributes: 'attributes',
+    base: 'base',
     comments: 'comments',
     fields: 'fields',
     mixins: 'mixins',
@@ -1354,6 +1358,9 @@ export class ZModelAstReflection extends langium.AbstractAstReflection {
                     name: DataModel.attributes,
                     defaultValue: []
                 },
+                base: {
+                    name: DataModel.base
+                },
                 baseModel: {
                     name: DataModel.baseModel,
                     referenceType: DataModel.$type
@@ -1765,6 +1772,9 @@ export class ZModelAstReflection extends langium.AbstractAstReflection {
                     name: TypeDef.attributes,
                     defaultValue: []
                 },
+                base: {
+                    name: TypeDef.base
+                },
                 comments: {
                     name: TypeDef.comments,
                     defaultValue: []

packages/language/src/generated/grammar.ts

@@ -2013,17 +2013,34 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
             "value": "with"
           },
           {
-            "$type": "Assignment",
-            "feature": "mixins",
-            "operator": "+=",
-            "terminal": {
-              "$type": "CrossReference",
-              "type": {
-                "$ref": "#/rules@44"
+            "$type": "Alternatives",
+            "elements": [
+              {
+                "$type": "Assignment",
+                "feature": "mixins",
+                "operator": "+=",
+                "terminal": {
+                  "$type": "CrossReference",
+                  "type": {
+                    "$ref": "#/rules@44"
+                  },
+                  "deprecatedSyntax": false,
+                  "isMulti": false
+                }
               },
-              "deprecatedSyntax": false,
-              "isMulti": false
-            }
+              {
+                "$type": "Assignment",
+                "feature": "base",
+                "operator": "=",
+                "terminal": {
+                  "$type": "RuleCall",
+                  "rule": {
+                    "$ref": "#/rules@64"
+                  },
+                  "arguments": []
+                }
+              }
+            ]
           },
           {
             "$type": "Group",
@@ -3087,6 +3104,10 @@ export const ZModelGrammar = (): Grammar => loadedZModelGrammar ?? (loadedZModel
           {
             "$type": "Keyword",
             "value": "type"
+          },
+          {
+            "$type": "Keyword",
+            "value": "this"
           }
         ]
       },

packages/language/src/utils.ts

@@ -249,6 +249,13 @@ export function isLiteAttribute(node: AstNode): node is Attribute {
     return isAttribute(node) && hasAttribute(node, '@@@lite');
 }
 
+/**
+ * Returns if the given node is primitive type def.
+ */
+export function isPrimitiveTypeDef(node: AstNode): node is TypeDef {
+    return isTypeDef(node) && !!node.base;
+}
+
 /**
  * Returns the datasource provider literal (e.g. `'postgresql'`) declared in the schema, or undefined
  * if no datasource is found or its provider is not a literal.
@@ -737,6 +744,10 @@ export function getAllFields(
     return fields;
 }
 
+export function getPrimitiveTypeDefThisField(td: TypeDef) {
+    return td.fields.find((f) => f.name === 'this');
+}
+
 /**
  * Gets all attributes of a data model or type def, including inherited attributes
  * from base models and mixins.

packages/language/src/validators/attribute-application-validator.ts

@@ -22,6 +22,7 @@ import {
     isModel,
     isReferenceExpr,
     isStringLiteral,
+    isThisExpr,
     isTypeDef,
 } from '../generated/ast';
 import {
@@ -30,6 +31,7 @@ import {
     getContainingDataModel,
     getDataSourceProvider,
     getNumberLiteral,
+    getPrimitiveTypeDefThisField,
     getStringLiteral,
     hasAttribute,
     isAuthOrAuthMemberAccess,
@@ -40,6 +42,7 @@ import {
     isDelegateModel,
     isManyToManyField,
     isNativeTypeMappingAttribute,
+    isPrimitiveTypeDef,
     isRelationshipField,
     mapBuiltinTypeToExpressionType,
     resolved,
@@ -142,6 +145,14 @@ export default class AttributeApplicationValidator implements AstValidator<Attri
         }
 
         const targetDecl = attr.$container;
+        if (
+            (isPrimitiveTypeDef(targetDecl) || isPrimitiveTypeDef(targetDecl.$container)) &&
+            !hasAttribute(attr.decl.ref!, '@@@validation')
+        ) {
+            accept('error', `attribute "${decl.name}" cannot be used with primitive type defs`, { node: attr });
+            return;
+        }
+
         if (decl.name === '@@@targetField' && !isAttribute(targetDecl)) {
             accept('error', `attribute "${decl.name}" can only be used on attribute declarations`, { node: attr });
             return;
@@ -381,13 +392,24 @@ export default class AttributeApplicationValidator implements AstValidator<Attri
     @check('@@validate')
     private _checkValidate(attr: AttributeApplication, accept: ValidationAcceptor) {
         const condition = attr.args[0]?.value;
-        if (
-            condition &&
-            AstUtils.streamAst(condition).some(
-                (node) => isDataFieldReference(node) && isDataModel(node.$resolvedType?.decl),
-            )
-        ) {
-            accept('error', `\`@@validate\` condition cannot use relation fields`, { node: condition });
+        if (condition) {
+            if (
+                AstUtils.streamAst(condition).some(
+                    (node) => isDataFieldReference(node) && isDataModel(node.$resolvedType?.decl),
+                )
+            ) {
+                accept('error', `\`@@validate\` condition cannot use relation fields`, { node: condition });
+                return;
+            }
+
+            if (isPrimitiveTypeDef(attr.$container) && AstUtils.streamAst(condition).some((node) => isThisExpr(node))) {
+                const thisField = getPrimitiveTypeDefThisField(attr.$container);
+                if (!thisField) {
+                    accept('error', 'primitive type def cannot reference `this` field without declaring it', {
+                        node: condition,
+                    });
+                }
+            }
         }
     }
 
@@ -664,10 +686,16 @@ function assignableToAttributeParam(
             // attribute parameter type is ContextType, need to infer type from
             // the attribute's container
             if (isDataField(attr.$container)) {
-                if (!attr.$container?.type?.type) {
-                    return genericError;
+                if (attr.$container.type.reference?.ref && isPrimitiveTypeDef(attr.$container.type.reference.ref)) {
+                    dstType = mapBuiltinTypeToExpressionType(attr.$container.type.reference.ref.base!);
+                } else {
+                    if (!attr.$container?.type?.type) {
+                        return genericError;
+                    }
+
+                    dstType = mapBuiltinTypeToExpressionType(attr.$container.type.type);
                 }
-                dstType = mapBuiltinTypeToExpressionType(attr.$container.type.type);
+
                 dstIsArray = attr.$container.type.array;
             } else {
                 dstType = 'Any';
@@ -708,34 +736,35 @@ function isValidAttributeTarget(attrDecl: Attribute, targetDecl: DataField) {
         .filter((name): name is string => !!name);
 
     let allowed = false;
+    const targetDeclType = targetDecl.$resolvedType?.decl ?? targetDecl.type.type;
     for (const allowedType of fieldTypes) {
         switch (allowedType) {
             case 'StringField':
-                allowed = allowed || targetDecl.type.type === 'String';
+                allowed = allowed || targetDeclType === 'String';
                 break;
             case 'IntField':
-                allowed = allowed || targetDecl.type.type === 'Int';
+                allowed = allowed || targetDeclType === 'Int';
                 break;
             case 'BigIntField':
-                allowed = allowed || targetDecl.type.type === 'BigInt';
+                allowed = allowed || targetDeclType === 'BigInt';
                 break;
             case 'FloatField':
-                allowed = allowed || targetDecl.type.type === 'Float';
+                allowed = allowed || targetDeclType === 'Float';
                 break;
             case 'DecimalField':
-                allowed = allowed || targetDecl.type.type === 'Decimal';
+                allowed = allowed || targetDeclType === 'Decimal';
                 break;
             case 'BooleanField':
-                allowed = allowed || targetDecl.type.type === 'Boolean';
+                allowed = allowed || targetDeclType === 'Boolean';
                 break;
             case 'DateTimeField':
-                allowed = allowed || targetDecl.type.type === 'DateTime';
+                allowed = allowed || targetDeclType === 'DateTime';
                 break;
             case 'JsonField':
-                allowed = allowed || targetDecl.type.type === 'Json';
+                allowed = allowed || targetDeclType === 'Json';
                 break;
             case 'BytesField':
-                allowed = allowed || targetDecl.type.type === 'Bytes';
+                allowed = allowed || targetDeclType === 'Bytes';
                 break;
             case 'ModelField':
                 allowed = allowed || isDataModel(targetDecl.type.reference?.ref);