fix(zmodel): unique attribute validation issues

Author: ymc9

packages/language/src/utils.ts

@@ -5,6 +5,7 @@ import path from 'node:path';
 import { fileURLToPath, pathToFileURL } from 'node:url';
 import { PLUGIN_MODULE_NAME, STD_LIB_MODULE_NAME, type ExpressionContext } from './constants';
 import {
+    InternalAttribute,
     isArrayExpr,
     isBinaryExpr,
     isConfigArrayExpr,
@@ -173,7 +174,7 @@ export function getRecursiveBases(
     bases.forEach((base) => {
         // avoid using .ref since this function can be called before linking
         const baseDecl = decl.$container.declarations.find(
-            (d): d is TypeDef | DataModel => isTypeDef(d) || (isDataModel(d) && d.name === base.$refText),
+            (d): d is TypeDef | DataModel => (isTypeDef(d) || isDataModel(d)) && d.name === base.$refText,
         );
         if (baseDecl) {
             if (!includeDelegate && isDelegateModel(baseDecl)) {
@@ -321,8 +322,15 @@ function getArray(expr: Expression | ConfigExpr | undefined) {
     return isArrayExpr(expr) || isConfigArrayExpr(expr) ? expr.items : undefined;
 }
 
+export function getAttributeArg(
+    attr: DataModelAttribute | DataFieldAttribute | InternalAttribute,
+    name: string,
+): Expression | undefined {
+    return attr.args.find((arg) => arg.$resolvedParam?.name === name)?.value;
+}
+
 export function getAttributeArgLiteral<T extends string | number | boolean>(
-    attr: DataModelAttribute | DataFieldAttribute,
+    attr: DataModelAttribute | DataFieldAttribute | InternalAttribute,
     name: string,
 ): T | undefined {
     for (const arg of attr.args) {

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

@@ -21,6 +21,7 @@ import {
 } from '../generated/ast';
 import {
     getAllAttributes,
+    getAttributeArg,
     getStringLiteral,
     hasAttribute,
     isAuthOrAuthMemberAccess,
@@ -291,7 +292,9 @@ export default class AttributeApplicationValidator implements AstValidator<Attri
     @check('@@index')
     @check('@@unique')
     private _checkConstraint(attr: AttributeApplication, accept: ValidationAcceptor) {
-        const fields = attr.args[0]?.value;
+        const fields = getAttributeArg(attr, 'fields');
+
+        // const fields = attr.args[0]?.value;
         const attrName = attr.decl.ref?.name;
         if (!fields) {
             accept('error', `expects an array of field references`, {

packages/sdk/src/ts-schema-generator.ts

@@ -35,7 +35,7 @@ import {
     UnaryExpr,
     type Model,
 } from '@zenstackhq/language/ast';
-import { getAllAttributes, getAllFields, isDataFieldReference } from '@zenstackhq/language/utils';
+import { getAllAttributes, getAllFields, getAttributeArg, isDataFieldReference } from '@zenstackhq/language/utils';
 import fs from 'node:fs';
 import path from 'node:path';
 import { match } from 'ts-pattern';
@@ -840,7 +840,11 @@ export class TsSchemaGenerator {
         const seenKeys = new Set<string>();
         for (const attr of allAttributes) {
             if (attr.decl.$refText === '@@id' || attr.decl.$refText === '@@unique') {
-                const fieldNames = this.getReferenceNames(attr.args[0]!.value);
+                const fieldsArg = getAttributeArg(attr, 'fields');
+                if (!fieldsArg) {
+                    continue;
+                }
+                const fieldNames = this.getReferenceNames(fieldsArg);
                 if (!fieldNames) {
                     continue;
                 }

packages/testtools/src/client.ts

@@ -38,6 +38,7 @@ export type CreateTestClientOptions<Schema extends SchemaDef> = Omit<ClientOptio
     extraSourceFiles?: Record<string, string>;
     workDir?: string;
     debug?: boolean;
+    dbFile?: string;
 };
 
 export async function createTestClient<Schema extends SchemaDef>(
@@ -57,9 +58,14 @@ export async function createTestClient<Schema extends SchemaDef>(
     let workDir = options?.workDir;
     let _schema: Schema;
     const provider = options?.provider ?? getTestDbProvider() ?? 'sqlite';
-
     const dbName = options?.dbName ?? getTestDbName(provider);
 
+    if (options?.dbFile) {
+        if (provider !== 'sqlite') {
+            throw new Error('dbFile option is only supported for sqlite provider');
+        }
+    }
+
     const dbUrl =
         provider === 'sqlite'
             ? `file:${dbName}`
@@ -108,35 +114,48 @@ export async function createTestClient<Schema extends SchemaDef>(
         console.log(`Work directory: ${workDir}`);
     }
 
+    // copy db file to workDir if specified
+    if (options?.dbFile) {
+        if (provider !== 'sqlite') {
+            throw new Error('dbFile option is only supported for sqlite provider');
+        }
+        fs.copyFileSync(options.dbFile, path.join(workDir, dbName));
+    }
+
     const { plugins, ...rest } = options ?? {};
     const _options: ClientOptions<Schema> = {
         ...rest,
     } as ClientOptions<Schema>;
 
-    if (options?.usePrismaPush) {
-        invariant(typeof schema === 'string' || schemaFile, 'a schema file must be provided when using prisma db push');
-        if (!model) {
-            const r = await loadDocumentWithPlugins(path.join(workDir, 'schema.zmodel'));
-            if (!r.success) {
-                throw new Error(r.errors.join('\n'));
+    if (!options?.dbFile) {
+        if (options?.usePrismaPush) {
+            invariant(
+                typeof schema === 'string' || schemaFile,
+                'a schema file must be provided when using prisma db push',
+            );
+            if (!model) {
+                const r = await loadDocumentWithPlugins(path.join(workDir, 'schema.zmodel'));
+                if (!r.success) {
+                    throw new Error(r.errors.join('\n'));
+                }
+                model = r.model;
+            }
+            const prismaSchema = new PrismaSchemaGenerator(model);
+            const prismaSchemaText = await prismaSchema.generate();
+            fs.writeFileSync(path.resolve(workDir!, 'schema.prisma'), prismaSchemaText);
+            execSync('npx prisma db push --schema ./schema.prisma --skip-generate --force-reset', {
+                cwd: workDir,
+                stdio: 'ignore',
+            });
+        } else {
+            if (provider === 'postgresql') {
+                invariant(dbName, 'dbName is required');
+                const pgClient = new PGClient(TEST_PG_CONFIG);
+                await pgClient.connect();
+                await pgClient.query(`DROP DATABASE IF EXISTS "${dbName}"`);
+                await pgClient.query(`CREATE DATABASE "${dbName}"`);
+                await pgClient.end();
             }
-            model = r.model;
-        }
-        const prismaSchema = new PrismaSchemaGenerator(model);
-        const prismaSchemaText = await prismaSchema.generate();
-        fs.writeFileSync(path.resolve(workDir!, 'schema.prisma'), prismaSchemaText);
-        execSync('npx prisma db push --schema ./schema.prisma --skip-generate --force-reset', {
-            cwd: workDir,
-            stdio: 'ignore',
-        });
-    } else {
-        if (provider === 'postgresql') {
-            invariant(dbName, 'dbName is required');
-            const pgClient = new PGClient(TEST_PG_CONFIG);
-            await pgClient.connect();
-            await pgClient.query(`DROP DATABASE IF EXISTS "${dbName}"`);
-            await pgClient.query(`CREATE DATABASE "${dbName}"`);
-            await pgClient.end();
         }
     }
 
@@ -155,7 +174,7 @@ export async function createTestClient<Schema extends SchemaDef>(
 
     let client = new ZenStackClient(_schema, _options);
 
-    if (!options?.usePrismaPush) {
+    if (!options?.usePrismaPush && !options?.dbFile) {
         await client.$pushSchema();
     }