implement requested changes from review

Azzerty23 · Azzerty23 · commit be2fcd5e97b5 · 2025-07-26T19:45:36.000+02:00
diff --git a/packages/schema/src/language-server/validator/attribute-application-validator.ts b/packages/schema/src/language-server/validator/attribute-application-validator.ts
@@ -1,4 +1,5 @@
 import {
+    AliasDecl,
     ArrayExpr,
     Attribute,
     AttributeArg,
@@ -300,14 +301,9 @@ function assignableToAttributeParam(arg: AttributeArg, param: AttributeParam, at
         }
     }
 
-    // alias expression is compared to corresponding expression resolved shape
+    // Handle alias expressions by comparing to their resolved shape
     if (isAliasDecl(arg.$resolvedType?.decl)) {
-        // TODO: what is context type? Passed to true to avoid error, to be fixed later
-        if (dstType === 'ContextType') return true;
-
-        const alias = arg.$resolvedType.decl;
-        const mappedAliasResolvedType = mappedRawExpressionTypeToResolvedShape(alias.expression.$type);
-        return dstType === mappedAliasResolvedType || dstType === 'Any' || mappedAliasResolvedType === 'Any';
+        return isAliasAssignableToType(arg.$resolvedType.decl, dstType ?? 'Any', attr);
     }
 
     // destination is field reference or transitive field reference, check if
@@ -422,6 +418,34 @@ function isValidAttributeTarget(attrDecl: Attribute, targetDecl: DataModelField)
     return allowed;
 }
 
+function isAliasAssignableToType(alias: AliasDecl, dstType: string, attr: AttributeApplication): boolean {
+    const effectiveDstType = resolveEffectiveDestinationType(dstType, attr);
+    if (effectiveDstType === null) {
+        return false;
+    }
+
+    const mappedAliasResolvedType = mappedRawExpressionTypeToResolvedShape(alias.expression.$type);
+    return (
+        effectiveDstType === mappedAliasResolvedType || effectiveDstType === 'Any' || mappedAliasResolvedType === 'Any'
+    );
+}
+
+function resolveEffectiveDestinationType(dstType: string, attr: AttributeApplication): string | null {
+    if (dstType !== 'ContextType') {
+        return dstType;
+    }
+
+    // ContextType is inferred from the attribute's container's type
+    if (isDataModelField(attr.$container)) {
+        if (!attr.$container?.type?.type) {
+            return null;
+        }
+        return mapBuiltinTypeToExpressionType(attr.$container.type.type);
+    }
+
+    return 'Any';
+}
+
 export function validateAttributeApplication(attr: AttributeApplication, accept: ValidationAcceptor) {
     new AttributeApplicationValidator().validate(attr, accept);
 }
diff --git a/packages/schema/src/language-server/validator/function-invocation-validator.ts b/packages/schema/src/language-server/validator/function-invocation-validator.ts
@@ -110,12 +110,6 @@ export default class FunctionInvocationValidator implements AstValidator<Express
 
                 // first argument must refer to a model field
                 const firstArg = expr.args?.[0]?.value;
-                const callableDecl = expr.function.ref;
-                if (!callableDecl) {
-                    accept('error', 'function or rule cannot be resolved', { node: expr });
-                    return;
-                }
-
                 if (firstArg) {
                     if (!getFieldReference(firstArg)) {
                         accept('error', 'first argument must be a field reference', { node: firstArg });
diff --git a/packages/schema/src/language-server/zmodel-scope.ts b/packages/schema/src/language-server/zmodel-scope.ts
@@ -282,6 +282,22 @@ export class ZModelScopeProvider extends DefaultScopeProvider {
         const globalScope = this.getGlobalScope(referenceType, context);
         const node = context.container as MemberAccessExpr;
         
+        // Handle auth() invocation in alias
+        if (isInvocationExpr(node.operand)) {
+            if (isAuthInvocation(node.operand)) {
+                return this.createScopeForAuth(node, globalScope);
+            }
+            if (isFutureInvocation(node.operand)) {
+                return this.createScopeForContainingModel(node, globalScope);
+            }
+            return EMPTY_SCOPE;
+        }
+        
+        // Handle this expression in alias
+        if (isThisExpr(node.operand)) {
+            return this.createScopeForContainingModel(node, globalScope);
+        }
+        
         // For member access in aliases, we need to check all possible contexts
         if (isReferenceExpr(node.operand)) {
             const operandName = node.operand.$cstNode?.text;
@@ -333,6 +349,17 @@ export class ZModelScopeProvider extends DefaultScopeProvider {
             }
         }
         
+        // Handle nested member access (e.g., some.nested.field)
+        if (isMemberAccessExpr(node.operand)) {
+            const ref = node.operand.member.ref;
+            if (isDataModelField(ref) && !ref.type.array) {
+                return this.createScopeForContainer(ref.type.reference?.ref, globalScope, true);
+            }
+            if (isTypeDefField(ref) && !ref.type.array) {
+                return this.createScopeForContainer(ref.type.reference?.ref, globalScope, true);
+            }
+        }
+        
         return EMPTY_SCOPE;
     }
     
diff --git a/packages/schema/src/plugins/enhancer/policy/constraint-transformer.ts b/packages/schema/src/plugins/enhancer/policy/constraint-transformer.ts
@@ -32,7 +32,12 @@ import {
 } from '@zenstackhq/sdk/ast';
 import { P, match } from 'ts-pattern';
 import { name } from '..';
-import { isCheckInvocation } from '../../../utils/ast-utils';
+import {
+    getAliasDeclaration,
+    getFieldsFromAliasExpression,
+    isAliasInvocation,
+    isCheckInvocation,
+} from '../../../utils/ast-utils';
 
 /**
  * Options for {@link ConstraintTransformer}.
@@ -381,6 +386,19 @@ export class ConstraintTransformer {
         if (isMemberAccessExpr(expr)) {
             return isThisExpr(expr.operand) ? { name: expr.member.$refText } : undefined;
         }
+        if (isAliasInvocation(expr)) {
+            // Resolve alias to actual fields - for constraint transformation,
+            // we need to return a single field name representation
+            const aliasDecl = getAliasDeclaration(expr);
+            if (aliasDecl) {
+                const fields = getFieldsFromAliasExpression(aliasDecl);
+                if (fields.length > 0) {
+                    // For constraint purposes, use the alias name itself as the field identifier
+                    // The actual field resolution will be handled by the expression transformer
+                    return { name: aliasDecl.name };
+                }
+            }
+        }
         return undefined;
     }
 
diff --git a/packages/schema/src/plugins/enhancer/policy/policy-guard-generator.ts b/packages/schema/src/plugins/enhancer/policy/policy-guard-generator.ts
@@ -36,7 +36,12 @@ import { lowerCaseFirst } from '@zenstackhq/runtime/local-helpers';
 import { streamAst } from 'langium';
 import path from 'path';
 import { FunctionDeclarationStructure, OptionalKind, Project, SourceFile, VariableDeclarationKind } from 'ts-morph';
-import { isCheckInvocation } from '../../../utils/ast-utils';
+import {
+    getAliasDeclaration,
+    getFieldsFromAliasExpression,
+    isAliasInvocation,
+    isCheckInvocation,
+} from '../../../utils/ast-utils';
 import { ConstraintTransformer } from './constraint-transformer';
 import {
     generateConstantQueryGuardFunction,
@@ -224,6 +229,25 @@ export class PolicyGenerator {
                     }
                 }
 
+                // Handle alias invocations
+                if (isAliasInvocation(expr)) {
+                    const aliasDecl = getAliasDeclaration(expr);
+                    if (aliasDecl) {
+                        const referencedFields = getFieldsFromAliasExpression(aliasDecl);
+                        // Check if any referenced field would prevent create input checking
+                        for (const field of referencedFields) {
+                            if (field.$container === model && hasAttribute(field, '@default')) {
+                                // Alias references field with default value
+                                return false;
+                            }
+                            if (isForeignKeyField(field)) {
+                                // Alias references foreign key field
+                                return false;
+                            }
+                        }
+                    }
+                }
+
                 return true;
             });
         });
diff --git a/packages/schema/src/plugins/enhancer/policy/utils.ts b/packages/schema/src/plugins/enhancer/policy/utils.ts
@@ -44,6 +44,8 @@ import { getContainerOfType, streamAllContents, streamAst, streamContents } from
 import { FunctionDeclarationStructure, OptionalKind } from 'ts-morph';
 import { name } from '..';
 import {
+    getAliasDeclaration,
+    getFieldsFromAliasExpression,
     isAliasInvocation,
     isCheckInvocation,
     isCollectionPredicate,
@@ -318,8 +320,8 @@ export function generateQueryGuardFunction(
                 isFutureExpr(child) ||
                 // field reference
                 (isReferenceExpr(child) && isDataModelField(child.target.ref)) ||
-                // TODO: field access from alias expression
-                isAliasInvocation(child)
+                // field access from alias expression - resolve to actual fields
+                (isAliasInvocation(child) && isExpression(child) && hasFieldAccessInAlias(child))
         )
     );
 
@@ -598,3 +600,17 @@ function getSourceModelOfFieldAccess(expr: Expression) {
 
     return undefined;
 }
+
+/**
+ * Checks if an alias invocation resolves to actual field accesses
+ */
+function hasFieldAccessInAlias(aliasInvocation: Expression): boolean {
+    const aliasDecl = getAliasDeclaration(aliasInvocation);
+    if (!aliasDecl) {
+        return false;
+    }
+    
+    // Get all fields referenced in the alias expression
+    const fields = getFieldsFromAliasExpression(aliasDecl);
+    return fields.length > 0;
+}
diff --git a/packages/schema/src/utils/ast-utils.ts b/packages/schema/src/utils/ast-utils.ts
@@ -9,8 +9,10 @@ import {
     isDataModel,
     isDataModelField,
     isInvocationExpr,
+    isMemberAccessExpr,
     isModel,
     isReferenceExpr,
+    isThisExpr,
     isAliasDecl,
     isTypeDef,
     Model,
@@ -166,11 +168,63 @@ export function isCheckInvocation(node: AstNode) {
 }
 
 export function isAliasInvocation(node: AstNode) {
-    // check if a matching alias exists
+    if (!isInvocationExpr(node)) {
+        return false;
+    }
+
+    // Check if the resolved reference is an alias declaration
+    if (node.function.ref && isAliasDecl(node.function.ref)) {
+        return true;
+    }
+
+    // Fallback: check by name in the current model
     const allAlias = getContainerOfType(node, isModel)?.declarations.filter(isAliasDecl) ?? [];
-    // const aliasDecls = getAllLoadedAlias(this.langiumDocuments());
-    return isInvocationExpr(node) && allAlias.some((alias) => alias.name === node.function.$refText);
-    // (!node.function.ref || !isFromStdlib(node.function.ref)) /*  && isAliasDecl(node.function.ref) */
+    return allAlias.some((alias) => alias.name === node.function.$refText);
+}
+
+/**
+ * Gets the alias declaration for a given alias invocation
+ */
+export function getAliasDeclaration(node: AstNode): AliasDecl | undefined {
+    if (!isInvocationExpr(node)) {
+        return undefined;
+    }
+
+    const allAlias = getContainerOfType(node, isModel)?.declarations.filter(isAliasDecl) ?? [];
+    return allAlias.find((alias) => alias.name === node.function.$refText);
+}
+
+/**
+ * Extracts all DataModelField references from an alias expression
+ */
+export function getFieldsFromAliasExpression(alias: AliasDecl): DataModelField[] {
+    const fields: DataModelField[] = [];
+
+    function extractFields(expr: Expression): void {
+        if (isReferenceExpr(expr) && isDataModelField(expr.target.ref)) {
+            fields.push(expr.target.ref);
+        } else if (isMemberAccessExpr(expr)) {
+            // Handle this.fieldName
+            if (isThisExpr(expr.operand) && expr.member.ref && isDataModelField(expr.member.ref)) {
+                fields.push(expr.member.ref);
+            }
+        } else if (isInvocationExpr(expr)) {
+            // Handle nested alias invocations
+            const nestedAlias = getAliasDeclaration(expr);
+            if (nestedAlias) {
+                fields.push(...getFieldsFromAliasExpression(nestedAlias));
+            }
+            // Also check arguments for field references
+            expr.args.forEach((arg) => extractFields(arg.value));
+        } else if (isBinaryExpr(expr)) {
+            extractFields(expr.left);
+            extractFields(expr.right);
+        }
+        // Add more expression types as needed
+    }
+
+    extractFields(alias.expression);
+    return [...new Set(fields)]; // Remove duplicates
 }
 
 export function resolveImportUri(imp: ModelImport): URI | undefined {
diff --git a/packages/sdk/src/typescript-expression-transformer.ts b/packages/sdk/src/typescript-expression-transformer.ts
@@ -89,7 +89,6 @@ export class TypeScriptExpressionTransformer {
                 return this.this(expr as ThisExpr);
 
             case ReferenceExpr:
-                // TODO: ensure referenceExpr from alias is resolved
                 return this.reference(expr as ReferenceExpr);
 
             case InvocationExpr:
@@ -147,7 +146,10 @@ export class TypeScriptExpressionTransformer {
 
         if (isAlias) {
             // if the function invocation comes from an alias, we transform its expression
-            return this.transform(expr.function.ref.expression!, normalizeUndefined);
+            if (!expr.function.ref.expression) {
+                throw new TypeScriptExpressionTransformerError(`Unresolved alias expression`);
+            }
+            return this.transform(expr.function.ref.expression, normalizeUndefined);
         }
 
         if (!isStdFunc) {
diff --git a/tests/integration/tests/plugins/policy.test.ts b/tests/integration/tests/plugins/policy.test.ts
