Merge pull request #1745 from ballerina-platform/bal-ext

Enable ballerina type extension for OpenAPI generation through CLI

ballerina-to-openapi/src/main/java/io/ballerina/openapi/service/mapper/ServiceToOpenAPIMapper.java

@@ -17,7 +17,9 @@
  */
 package io.ballerina.openapi.service.mapper;
 
+import io.ballerina.compiler.api.ModuleID;
 import io.ballerina.compiler.api.SemanticModel;
+import io.ballerina.compiler.api.symbols.ModuleSymbol;
 import io.ballerina.compiler.api.symbols.ServiceDeclarationSymbol;
 import io.ballerina.compiler.api.symbols.Symbol;
 import io.ballerina.compiler.syntax.tree.FunctionDefinitionNode;
@@ -95,6 +97,24 @@ public static List<OASResult> generateOAS3Definition(Project project, SyntaxTree
                                                          SemanticModel semanticModel,
                                                          String serviceName, Boolean needJson,
                                                          Path inputPath) {
+        return generateOAS3Definition(project, syntaxTree, semanticModel, serviceName, needJson, inputPath, false);
+    }
+
+    /**
+     * This method will generate  openapi definition Map lists with ballerina code.
+     *
+     * @param syntaxTree    - Syntax tree the related to ballerina service
+     * @param semanticModel - Semantic model related to ballerina module
+     * @param serviceName   - Service name that need to generate the openAPI specification
+     * @param needJson      - Flag for enabling the generated file format with json or YAML
+     * @param inputPath     - Input file path for resolve the annotation details
+     * @param ballerinaExtension - Flag to enable ballerina type extension
+     * @return - {@link java.util.Map} with openAPI definitions for service nodes
+     */
+    public static List<OASResult> generateOAS3Definition(Project project, SyntaxTree syntaxTree,
+                                                         SemanticModel semanticModel,
+                                                         String serviceName, Boolean needJson,
+                                                         Path inputPath, Boolean ballerinaExtension) {
         Map<String, ServiceNode> servicesToGenerate = new HashMap<>();
         List<String> availableService = new ArrayList<>();
         List<OpenAPIMapperDiagnostic> diagnostics = new ArrayList<>();
@@ -115,7 +135,7 @@ public static List<OASResult> generateOAS3Definition(Project project, SyntaxTree
             for (Map.Entry<String, ServiceNode> serviceNode : servicesToGenerate.entrySet()) {
                 String openApiName = getOpenApiFileName(syntaxTree.filePath(), serviceNode.getKey(), needJson);
                 OASResult oasDefinition = generateOasFroServiceNode(project, openApiName,
-                        semanticModel, inputPath, serviceNode.getValue());
+                        semanticModel, inputPath, serviceNode.getValue(), ballerinaExtension);
                 outputs.add(oasDefinition);
             }
         }
@@ -127,13 +147,15 @@ public static List<OASResult> generateOAS3Definition(Project project, SyntaxTree
     }
 
     public static OASResult generateOasFroServiceNode(Project project, String openApiName, SemanticModel semanticModel,
-                                                       Path inputPath, ServiceNode serviceNode) {
+                                                      Path inputPath, ServiceNode serviceNode,
+                                                      Boolean ballerinaExtension) {
         OASGenerationMetaInfo.OASGenerationMetaInfoBuilder builder =
                 new OASGenerationMetaInfo.OASGenerationMetaInfoBuilder();
         builder.setServiceNode(serviceNode)
                 .setSemanticModel(semanticModel)
                 .setOpenApiFileName(openApiName)
                 .setBallerinaFilePath(inputPath)
+                .setBallerinaExtension(ballerinaExtension)
                 .setProject(project);
         OASGenerationMetaInfo oasGenerationMetaInfo = builder.build();
         OASResult oasDefinition = generateOAS(oasGenerationMetaInfo);
@@ -275,7 +297,12 @@ public static OASResult generateOAS(OASGenerationMetaInfo oasGenerationMetaInfo)
                     openapi.setComponents(null);
                 }
                 // Remove ballerina extensions
-                BallerinaTypeExtensioner.removeExtensions(openapi);
+                Optional<ModuleID> serviceModuleId = getServiceModuleId(serviceDefinition, semanticModel);
+                if (oasGenerationMetaInfo.getBallerinaExtension() && serviceModuleId.isPresent()) {
+                    BallerinaTypeExtensioner.removeCurrentModuleTypeExtensions(openapi, serviceModuleId.get());
+                } else {
+                    BallerinaTypeExtensioner.removeExtensions(openapi);
+                }
                 return new OASResult(openapi, diagnostics);
             } else {
                 return new OASResult(openapi, oasResult.getDiagnostics());
@@ -331,6 +358,11 @@ private static String extractBasePath(OpenAPI openApiFromServiceContract) {
         return servers.get(0).getUrl().split("\\{server}:\\{port}").length == 2 ? parts[1] : "";
     }
 
+    private static Optional<ModuleID> getServiceModuleId(ServiceNode serviceNode, SemanticModel semanticModel) {
+        return semanticModel.symbol(serviceNode.getInternalNode())
+                .flatMap(symbol -> symbol.getModule().map(ModuleSymbol::id));
+    }
+
     /**
      * Travers every syntax tree and collect all the listener nodes.
      *

ballerina-to-openapi/src/main/java/io/ballerina/openapi/service/mapper/example/OpenAPIExampleMapperImpl.java

@@ -149,7 +149,7 @@ private Optional<TypeDefinitionSymbol> getTypeDefinitionNode(String name, Schema
         Optional<Symbol> ballerinaType;
         if (ballerinaExt.isPresent()) {
             BallerinaPackage ballerinaPkg = ballerinaExt.get();
-            String typeName = ballerinaPkg.name().orElse(name);
+            String typeName = Objects.isNull(ballerinaPkg.name()) ? name : ballerinaPkg.name();
             ballerinaType = semanticModel.types().getTypeByName(ballerinaPkg.orgName(), ballerinaPkg.moduleName(),
                     ballerinaPkg.version(), typeName);
         } else {

ballerina-to-openapi/src/main/java/io/ballerina/openapi/service/mapper/model/OASGenerationMetaInfo.java

@@ -36,6 +36,7 @@ public class OASGenerationMetaInfo {
     private final SemanticModel semanticModel;
     private final ServiceNode serviceNode;
     private final Project project;
+    private final Boolean ballerinaExtensionLevel;
 
     public OASGenerationMetaInfo(OASGenerationMetaInfoBuilder builder) {
         this.openApiFileName = builder.openApiFileName;
@@ -47,6 +48,7 @@ public OASGenerationMetaInfo(OASGenerationMetaInfoBuilder builder) {
         }
         this.serviceNode = serviceNodeFromBuilder;
         this.project = builder.project;
+        this.ballerinaExtensionLevel = builder.ballerinaExtension;
     }
 
     public String getOpenApiFileName() {
@@ -69,6 +71,10 @@ public Project getProject() {
         return project;
     }
 
+    public Boolean getBallerinaExtension() {
+        return ballerinaExtensionLevel;
+    }
+
     /**
      * This method is used to create a new {@link OASGenerationMetaInfoBuilder} instance.
      */
@@ -80,6 +86,7 @@ public static class OASGenerationMetaInfoBuilder {
         private ServiceDeclarationNode serviceDeclarationNode;
         private ServiceNode serviceNode;
         private Project project;
+        private Boolean ballerinaExtension = false;
 
         public OASGenerationMetaInfoBuilder setBallerinaFilePath(Path ballerinaFilePath) {
             this.ballerinaFilePath = ballerinaFilePath;
@@ -111,6 +118,11 @@ public OASGenerationMetaInfoBuilder setOpenApiFileName(String openApiFileName) {
             return this;
         }
 
+        public OASGenerationMetaInfoBuilder setBallerinaExtension(Boolean balExt) {
+            this.ballerinaExtension = balExt;
+            return this;
+        }
+
         public void setProject(Project project) {
             this.project = project;
         }

ballerina-to-openapi/src/main/java/io/ballerina/openapi/service/mapper/model/ServiceDeclaration.java

@@ -236,7 +236,7 @@ private Optional<String> getOpenAPISpecFromResources(Package pkg, SemanticModel
                 return Optional.empty();
             }
             OASResult oasResult = generateOasFroServiceNode(pkg.project(), serviceName.get(), semanticModel, null,
-                    serviceContract.get());
+                    serviceContract.get(), false);
             if (oasResult.getDiagnostics().stream()
                     .anyMatch(diagnostic -> diagnostic.getDiagnosticSeverity().equals(DiagnosticSeverity.ERROR))) {
                 diagnostics.add(new ExceptionDiagnostic(DiagnosticMessages.OAS_CONVERTOR_136, serviceName.get()));

ballerina-to-openapi/src/main/java/io/ballerina/openapi/service/mapper/type/extension/BallerinaPackage.java

@@ -17,8 +17,6 @@
  */
 package io.ballerina.openapi.service.mapper.type.extension;
 
-import java.util.Optional;
-
 /**
  * This {@link BallerinaPackage} record represents the Ballerina package details.
  * @param orgName organization name
@@ -30,5 +28,5 @@
  * @since 2.1.0
  */
 public record BallerinaPackage(String orgName, String pkgName, String moduleName, String version,
-                               String modulePrefix, Optional<String> name) {
+                               String modulePrefix, String name) {
 }

ballerina-to-openapi/src/main/java/io/ballerina/openapi/service/mapper/type/extension/BallerinaTypeExtensioner.java

@@ -46,6 +46,18 @@ public static void addExtension(Schema schema, TypeSymbol typeSymbol) {
     }
 
     public static void removeExtensions(OpenAPI openAPI) {
+        removeExtensionsFromSchemas(openAPI, (extensions, orgName, moduleName) -> true);
+    }
+
+    public static void removeCurrentModuleTypeExtensions(OpenAPI openAPI, ModuleID moduleID) {
+        String orgName = moduleID.orgName();
+        String moduleName = moduleID.moduleName();
+
+        removeExtensionsFromSchemas(openAPI, (extensions, org, mod) -> fromSameModule(extensions, orgName,
+                moduleName));
+    }
+
+    private static void removeExtensionsFromSchemas(OpenAPI openAPI, ExtensionRemovalCondition condition) {
         Components components = openAPI.getComponents();
         if (Objects.isNull(components)) {
             return;
@@ -58,12 +70,22 @@ public static void removeExtensions(OpenAPI openAPI) {
 
         schemas.forEach((key, schema) -> {
             Map<?, ?> extensions = schema.getExtensions();
-            if (Objects.nonNull(extensions)) {
+            if (Objects.nonNull(extensions) && condition.shouldRemove(extensions, null, null)) {
                 extensions.remove(X_BALLERINA_TYPE);
             }
         });
     }
 
+    @FunctionalInterface
+    private interface ExtensionRemovalCondition {
+        boolean shouldRemove(Map<?, ?> extensions, String orgName, String moduleName);
+    }
+
+    private static boolean fromSameModule(Map<?, ?> extensions, String orgName, String moduleName) {
+        return extensions.get(X_BALLERINA_TYPE) instanceof BallerinaPackage ballerinaPkg &&
+                orgName.equals(ballerinaPkg.orgName()) && moduleName.equals(ballerinaPkg.moduleName());
+    }
+
     public static Optional<BallerinaPackage> getExtension(Schema schema) {
         Map<?, ?> extensions = schema.getExtensions();
         if (Objects.isNull(extensions)) {
@@ -84,6 +106,6 @@ static Optional<BallerinaPackage> getBallerinaPackage(TypeSymbol typeSymbol) {
         }
         ModuleID moduleID = module.get().id();
         return Optional.of(new BallerinaPackage(moduleID.orgName(), moduleID.packageName(), moduleID.moduleName(),
-                moduleID.version(), moduleID.modulePrefix(), typeSymbol.getName()));
+                moduleID.version(), moduleID.modulePrefix(), typeSymbol.getName().orElse(null)));
     }
 }

openapi-cli/src/main/java/io/ballerina/openapi/cmd/BaseCmd.java

@@ -54,8 +54,7 @@ public class BaseCmd {
     @CommandLine.Option(names = {"--status-code-binding"}, description = "Generate the client methods with " +
             "status code response binding")
     public boolean statusCodeBinding;
-
-
+
     @CommandLine.Option(names = {"--mock"}, hidden = true,
             description = "Generate mock client with given response example")
     public boolean mock;

openapi-cli/src/main/java/io/ballerina/openapi/cmd/OASContractGenerator.java

@@ -43,6 +43,7 @@
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 
 import static io.ballerina.openapi.service.mapper.utils.CodegenUtils.resolveContractFileName;
@@ -59,6 +60,7 @@ public class OASContractGenerator {
     private Project project;
     private List<OpenAPIMapperDiagnostic> diagnostics = new ArrayList<>();
     private PrintStream outStream = System.out;
+    private Boolean ballerinaExtension = false;
 
     /**
      * Initialize constructor.
@@ -67,6 +69,12 @@ public OASContractGenerator() {
 
     }
 
+    public void setBallerinaExtension(Boolean ballerinaExtension) {
+        if (Objects.nonNull(ballerinaExtension)) {
+            this.ballerinaExtension = ballerinaExtension;
+        }
+    }
+
     public List<OpenAPIMapperDiagnostic> getDiagnostics() {
         return diagnostics;
     }
@@ -118,7 +126,7 @@ public void generateOAS3DefinitionsAllService(Path servicePath, Path outPath, St
         }
         semanticModel = compilation.getSemanticModel(docId.moduleId());
         List<OASResult> openAPIDefinitions = ServiceToOpenAPIMapper.generateOAS3Definition(project, syntaxTree,
-                semanticModel, serviceName, needJson, inputPath);
+                semanticModel, serviceName, needJson, inputPath, ballerinaExtension);
 
         if (!openAPIDefinitions.isEmpty()) {
             List<String> fileNames = new ArrayList<>();

openapi-cli/src/main/java/io/ballerina/openapi/cmd/OpenApiCmd.java

@@ -126,6 +126,9 @@ public class OpenApiCmd implements BLauncherCmd {
             description = "Generate service without data binding")
     private boolean generateWithoutDataBinding;
 
+    @CommandLine.Option(names = {"--with-bal-ext"}, hidden = true, description = "Generate ballerina type extensions")
+    private boolean addBallerinaExtension;
+
 
     @CommandLine.Parameters
     private List<String> argList;
@@ -223,7 +226,7 @@ public void execute() {
 
             if (baseCmd.statusCodeBinding) {
                 if (baseCmd.mode != null && baseCmd.mode.equals(SERVICE)) {
-                    outStream.println("ERROR: the '--status-code-binding' option is only available in client " +
+                    outStream.println("the '--status-code-binding' option is only available in client " +
                             "generation mode.");
                     exitError(this.exitWhenFinish);
                 }
@@ -238,6 +241,12 @@ public void execute() {
                 }
             }
 
+            if (addBallerinaExtension) {
+                outStream.println("'--with-bal-ext' option is only available in OpenAPI specification " +
+                        "generation mode.");
+                exitError(this.exitWhenFinish);
+            }
+
             try {
                 openApiToBallerina(fileName, filter);
             } catch (IOException e) {
@@ -289,6 +298,7 @@ private void ballerinaToOpenApi(String fileName) {
         getTargetOutputPath();
         // Check service name it is mandatory
         OASContractGenerator openApiConverter = new OASContractGenerator();
+        openApiConverter.setBallerinaExtension(addBallerinaExtension);
         openApiConverter.generateOAS3DefinitionsAllService(balFilePath, targetOutputPath, service,
                 generatedFileType);
         mapperDiagnostics.addAll(openApiConverter.getDiagnostics());

openapi-client-native/ballerina-tests/Dependencies.toml

@@ -75,7 +75,7 @@ dependencies = [
 [[package]]
 org = "ballerina"
 name = "http"
-version = "2.11.2"
+version = "2.11.1"
 dependencies = [
 	{org = "ballerina", name = "auth"},
 	{org = "ballerina", name = "cache"},