[FEAT] (compile-time error): prevent stub entities from having fields beyond their key fields

ballerina/Dependencies.toml

@@ -59,7 +59,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "data.jsondata"
-version = "1.1.2"
+version = "1.1.3"
 dependencies = [
 	{org = "ballerina", name = "jballerina.java"},
 	{org = "ballerina", name = "lang.object"}
@@ -168,7 +168,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "jwt"
-version = "2.15.0"
+version = "2.15.1"
 dependencies = [
 	{org = "ballerina", name = "cache"},
 	{org = "ballerina", name = "crypto"},
@@ -334,7 +334,7 @@ modules = [
 [[package]]
 org = "ballerina"
 name = "os"
-version = "1.10.0"
+version = "1.10.1"
 dependencies = [
 	{org = "ballerina", name = "io"},
 	{org = "ballerina", name = "jballerina.java"}

compiler-plugin-tests/src/test/java/io/ballerina/stdlib/graphql/compiler/ServiceValidationTest.java

@@ -206,6 +206,21 @@ public void testDynamicallyAttachingService() {
         Assert.assertEquals(diagnosticResult.errorCount(), 0);
     }
 
+    @Test(groups = "valid")
+    public void testValidEntityWithKeyOnly() {
+        String packagePath = "88_valid_entity_with_key_only";
+        DiagnosticResult result = getDiagnosticResult(packagePath);
+        Assert.assertEquals(result.errorCount(), 0, "Expected no errors for key-only entity definition.");
+
+    }
+
+    @Test(groups = "valid")
+    public void testValidEntityWithKeyAndResolveReference() {
+        String packagePath = "89_valid_entity_with_key_and_resolve_reference";
+        DiagnosticResult result = getDiagnosticResult(packagePath);
+        Assert.assertEquals(result.errorCount(), 0, "Expected no errors for entity with resolveReference.");
+    }
+
     @Test(groups = "invalid")
     public void testMultipleListenersOnSameService() {
         String packagePath = "24_multiple_listeners_on_same_service";
@@ -1051,6 +1066,16 @@ public void testInvalidUseOfReservedFederatedResolverNames() {
         assertErrorMessage(diagnostic, message, 34, 21);
     }
 
+    @Test(groups = "invalid")
+    public void testInvalidEntityMissingResolveReference() {
+        String packagePath = "90_invalid_entity_missing_resolve_reference";
+        DiagnosticResult result = getDiagnosticResult(packagePath);
+        Assert.assertEquals(result.errorCount(), 1, "Expected one validation error when resolveReference is missing.");
+
+        Diagnostic diagnostic = result.errors().iterator().next();
+        String expectedMessage = getErrorMessage(CompilationDiagnostic.INVALID_ENTITY_FIELD, "User");
+        assertErrorMessage(diagnostic, expectedMessage, 8, 5);
+    }
     @Test(groups = "invalid")
     public void testInvalidUseOfReservedFederatedTypeNames() {
         String packagePath = "60_invalid_use_of_reserved_federation_type_names";

compiler-plugin-tests/src/test/resources/ballerina_sources/validator_tests/88_valid_entity_with_key_only/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "graphql_test"
+name = "test_package"
+version = "0.1.0"

compiler-plugin-tests/src/test/resources/ballerina_sources/validator_tests/88_valid_entity_with_key_only/service.bal

@@ -0,0 +1,29 @@
+// Copyright (c) 2024 WSO2 LLC. (http://www.wso2.com) All Rights Reserved.
+//
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import ballerina/graphql as gql;
+
+@gql:ServiceConfig {
+    entities: ["Product"]
+   }
+
+service /graphql on new gql:Listener(9090) {
+
+    @gql:Type
+    type Product @key(fields: "id") record {|
+        readonly string id;
+    |};
+}

compiler-plugin-tests/src/test/resources/ballerina_sources/validator_tests/89_valid_entity_with_key_and_resolve_reference/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "graphql_test"
+name = "test_package"
+version = "0.1.0"

compiler-plugin-tests/src/test/resources/ballerina_sources/validator_tests/89_valid_entity_with_key_and_resolve_reference/service.bal

@@ -0,0 +1,33 @@
+// Copyright (c) 2024 WSO2 LLC. (http://www.wso2.com) All Rights Reserved.
+//
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import ballerina/graphql;
+
+@graphql:ServiceConfig {
+    entities: ["User"]
+}
+service /graphql on new graphql:Listener(9090) {
+
+    @graphql:Type
+    type User @key(fields: "id") record {|
+        readonly int id;
+        string name?;
+    |};
+
+    isolated function resolveReference(User ref) returns User {
+        return {id: ref.id, name: "John Doe"};
+    }
+}

compiler-plugin-tests/src/test/resources/ballerina_sources/validator_tests/90_invalid_entity_missing_resolve_reference/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "graphql_test"
+name = "test_package"
+version = "0.1.0"

compiler-plugin-tests/src/test/resources/ballerina_sources/validator_tests/90_invalid_entity_missing_resolve_reference/service.bal

@@ -0,0 +1,29 @@
+// Copyright (c) 2024 WSO2 LLC. (http://www.wso2.com) All Rights Reserved.
+//
+// WSO2 LLC. licenses this file to you under the Apache License,
+// Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import ballerina/graphql;
+
+@graphql:ServiceConfig {
+    entities: ["User"]
+}
+service /graphql on new graphql:Listener(9090) {
+
+    @graphql:Type
+    type User @key(fields: "id") record {|
+        readonly int id;
+        string name?;
+    |};
+}

compiler-plugin/src/main/java/io/ballerina/stdlib/graphql/compiler/diagnostics/CompilationDiagnostic.java

@@ -89,6 +89,7 @@ public enum CompilationDiagnostic {
     INVALID_EMPTY_RECORD_INPUT_TYPE(DiagnosticMessage.ERROR_147, DiagnosticCode.GRAPHQL_147, DiagnosticSeverity.ERROR),
     INVALID_MODIFICATION_OF_SERVICE_CONFIG_FIELD(DiagnosticMessage.ERROR_148, DiagnosticCode.GRAPHQL_148,
                                                  DiagnosticSeverity.ERROR),
+    INVALID_ENTITY_FIELD(DiagnosticMessage.ERROR_149, DiagnosticCode.GRAPHQL_149, DiagnosticSeverity.ERROR),
 
     // Warnings
     UNSUPPORTED_INPUT_FIELD_DEPRECATION(DiagnosticMessage.WARNING_201, DiagnosticCode.GRAPHQL_201,

compiler-plugin/src/main/java/io/ballerina/stdlib/graphql/compiler/diagnostics/DiagnosticCode.java

@@ -70,6 +70,7 @@ public enum DiagnosticCode {
     GRAPHQL_146,
     GRAPHQL_147,
     GRAPHQL_148,
+    GRAPHQL_149,
     GRAPHQL_201,
     GRAPHQL_202,
     GRAPHQL_203,

compiler-plugin/src/main/java/io/ballerina/stdlib/graphql/compiler/diagnostics/DiagnosticMessage.java

@@ -98,6 +98,7 @@ public enum DiagnosticMessage {
     ERROR_147("invalid empty record type ''{0}'' found for GraphQL input object type at field ''{1}''"),
     ERROR_148("field ''{0}'' in ServiceConfig is not allowed to be modified. "
                       + "The value will be generated automatically by the GraphQL module"),
+    ERROR_149("field ''{0}'' is not allowed in stub entity." + " Only key fields can be included"),
 
     WARNING_201("invalid usage of @deprecated directive found in ''{0}''. Input object field(s) deprecation "
                         + "is not supported by the current GraphQL spec"),

compiler-plugin/src/main/java/io/ballerina/stdlib/graphql/compiler/service/validator/ServiceValidator.java

@@ -18,10 +18,12 @@
 
 package io.ballerina.stdlib.graphql.compiler.service.validator;
 
+import io.ballerina.compiler.api.SemanticModel;
 import io.ballerina.compiler.api.symbols.AnnotationSymbol;
 import io.ballerina.compiler.api.symbols.ArrayTypeSymbol;
 import io.ballerina.compiler.api.symbols.ClassSymbol;
 import io.ballerina.compiler.api.symbols.EnumSymbol;
+import io.ballerina.compiler.api.symbols.FunctionSymbol;
 import io.ballerina.compiler.api.symbols.FunctionTypeSymbol;
 import io.ballerina.compiler.api.symbols.IntersectionTypeSymbol;
 import io.ballerina.compiler.api.symbols.MapTypeSymbol;
@@ -76,6 +78,7 @@
 import io.ballerina.tools.diagnostics.Location;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -189,6 +192,7 @@ private void validateEntityAnnotation(AnnotationNode entityAnnotation) {
         if (entityAnnotation.annotValue().isEmpty()) {
             return;
         }
+        List<String> keyFields = new ArrayList<>();
         for (MappingFieldNode fieldNode : entityAnnotation.annotValue().get().fields()) {
             if (fieldNode.kind() != SPECIFIC_FIELD) {
                 addDiagnostic(CompilationDiagnostic.PROVIDE_KEY_VALUE_PAIR_FOR_ENTITY_ANNOTATION, fieldNode.location());
@@ -203,8 +207,159 @@ private void validateEntityAnnotation(AnnotationNode entityAnnotation) {
             String fieldName = fieldNameToken.text().trim();
             if (KEY.equals(fieldName)) {
                 validateKeyField(specificFieldNode);
+                keyFields = extractKeyFields(specificFieldNode);
+
+            }
+
+            validateFieldsAgainstKeys(keyFields, entityAnnotation.location());
+
+        }
+    }
+
+    private void validateFieldsAgainstKeys(List<String> keyFields, Location location) {
+
+        if (keyFields == null || keyFields.isEmpty()) {
+            return;
+
+        }
+
+        Set<String> keyFieldSet = new HashSet<>(keyFields);
+        List<RecordFieldSymbol> recordFields = getRecordFields();
+
+        boolean hasExtraFields = false;
+
+        for (RecordFieldSymbol recordField : recordFields) {
+            String fieldName = recordField.getName().orElse(null);
+
+            if (fieldName == null) {
+                continue;
+            }
+
+            if (!keyFieldSet.contains(fieldName)) {
+                hasExtraFields = true;
+                addDiagnostic(CompilationDiagnostic.INVALID_ENTITY_FIELD,
+                        recordField.location(), fieldName);
+            }
+        }
+        Symbol currentEntity = getCurrentEntitySymbol();
+        if (hasExtraFields && !hasResolveReferenceFunctionForEntity(currentEntity)) {
+            addDiagnostic(CompilationDiagnostic.INVALID_ENTITY_FIELD, location);
+        }
+    }
+
+    private Symbol getCurrentEntitySymbol() {
+
+        for (Map.Entry<String, Symbol> entry : this.interfaceEntityFinder.getEntities().entrySet()) {
+            AnnotationSymbol entityAnnotation = getEntityAnnotationSymbol(entry.getValue());
+            if (entityAnnotation != null) {
+                return entry.getValue();
+            }
+        }
+        return null;
+    }
+
+    private boolean hasResolveReferenceFunctionForEntity(Symbol entitySymbol) {
+        if (entitySymbol == null || context == null) {
+            return false;
+        }
+
+        String entityName = entitySymbol.getName().orElse(null);
+        if (entityName == null) {
+            return false;
+        }
+
+        // Get semantic model from context
+        SemanticModel semanticModel = context.semanticModel();
+
+        // Get module symbols and search for resolveReference function
+        List<Symbol> moduleSymbols = new ArrayList<>();
+        semanticModel.moduleSymbols().forEach(moduleSymbols::add);
+
+        for (Symbol symbol : moduleSymbols) {
+            if (symbol.kind() != SymbolKind.FUNCTION) {
+                continue;
+            }
+
+            FunctionSymbol functionSymbol = (FunctionSymbol) symbol;
+            if (!"resolveReference".equals(functionSymbol.getName().orElse(""))) {
+                continue;
+            }
+
+            Optional<List<ParameterSymbol>> paramsOpt = functionSymbol.typeDescriptor().params();
+            if (paramsOpt.isEmpty() || paramsOpt.get().isEmpty()) {
+                continue;
+            }
+
+            ParameterSymbol firstParam = paramsOpt.get().get(0);
+            TypeSymbol paramType = firstParam.typeDescriptor();
+
+
+            // Check if the parameter type matches the entity
+            if (paramType != null) {
+                // Try matching by name
+                if (entityName.equals(paramType.getName().orElse(null))) {
+                    return true;
+                }
+
+                // For type reference types, also check the original type
+                if (paramType.typeKind() == TypeDescKind.TYPE_REFERENCE) {
+                    TypeReferenceTypeSymbol typeRef = (TypeReferenceTypeSymbol) paramType;
+                    if (entityName.equals(typeRef.getName().orElse(null))) {
+                        return true;
+                    }
+                }
+            }
+        }
+
+        return false;
+    }
+    private List<String> extractKeyFields(SpecificFieldNode specificFieldNode) {
+
+        List<String> keyFields = new ArrayList<>();
+        if (specificFieldNode.valueExpr().isEmpty()) {
+            return keyFields;
+
+        }
+
+        ExpressionNode valueNode = specificFieldNode.valueExpr().get();
+
+        if (valueNode.kind() == SyntaxKind.STRING_LITERAL) {
+            keyFields.add(valueNode.toString().replace("\"", "").trim());
+        } else if (valueNode.kind() == SyntaxKind.LIST_CONSTRUCTOR) {
+            for (Node keyNode : ((ListConstructorExpressionNode) valueNode).expressions()) {
+                if (keyNode.kind() == SyntaxKind.STRING_LITERAL) {
+                    keyFields.add(keyNode.toString().replace("\"", "").trim());
+                }
+            }
+        } else {
+            addDiagnostic(CompilationDiagnostic.PROVIDE_A_STRING_LITERAL_OR_AN_ARRAY_OF_STRING_LITERALS_FOR_KEY_FIELD,
+                    valueNode.location());
+
+        }
+
+        return keyFields;
+
+    }
+
+    private List<RecordFieldSymbol> getRecordFields() {
+
+        for (Map.Entry<String, Symbol> entry : this.interfaceEntityFinder.getEntities().entrySet()) {
+            AnnotationSymbol entityAnnotation = getEntityAnnotationSymbol(entry.getValue());
+            if (entityAnnotation != null) {
+                Symbol entitySymbol = entry.getValue();
+
+                if (entitySymbol instanceof TypeDefinitionSymbol) {
+                    TypeDefinitionSymbol typeDefSymbol = (TypeDefinitionSymbol) entitySymbol;
+                    TypeSymbol typeDescriptor = typeDefSymbol.typeDescriptor();
+
+                    if (typeDescriptor instanceof RecordTypeSymbol) {
+                        RecordTypeSymbol recordType = (RecordTypeSymbol) typeDescriptor;
+                        return new ArrayList<>(recordType.fieldDescriptors().values());
+                    }
+                }
             }
         }
+        return Collections.emptyList();
     }
 
     private void validateKeyField(SpecificFieldNode specificFieldNode) {