From 8995997850b5b6f54798a00c57fcf76e7a9ee661 Mon Sep 17 00:00:00 2001
From: Francisco Javier Tirado Sarti <ftirados@redhat.com>
Date: Mon, 17 Jul 2023 11:51:56 +0200
Subject: [PATCH] [KOGITO-9584] Downloading external references

---
 .../parser/schema/JsonSchemaImpl.java         | 27 +++++++++
 .../schema/OpenApiModelSchemaGenerator.java   |  5 +-
 .../workflow/parser/schema/RefSchemas.java    | 57 +++++++++++++++++++
 3 files changed, 88 insertions(+), 1 deletion(-)
 create mode 100644 quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/RefSchemas.java

diff --git a/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/JsonSchemaImpl.java b/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/JsonSchemaImpl.java
index 8cf2d1de5fb..a5435539638 100644
--- a/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/JsonSchemaImpl.java
+++ b/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/JsonSchemaImpl.java
@@ -16,16 +16,24 @@
 
 package org.kie.kogito.serverless.workflow.parser.schema;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import org.eclipse.microprofile.openapi.models.media.Schema;
+import org.kie.kogito.jackson.utils.ObjectMapperFactory;
+import org.kie.kogito.serverless.workflow.io.URIContentLoaderFactory;
 
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonSetter;
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
 
+import io.smallrye.openapi.api.constants.OpenApiConstants;
 import io.smallrye.openapi.api.models.media.SchemaImpl;
 
 /**
@@ -39,6 +47,25 @@ public class JsonSchemaImpl extends SchemaImpl {
     @JsonSetter("$ref")
     @Override
     public void setRef(String ref) {
+        if (ref != null && !ref.startsWith("#")) {
+            try (InputStream is = URIContentLoaderFactory.loader(new URI(ref), Optional.empty(), Optional.empty(), Optional.empty(), null).getInputStream()) {
+                JsonSchemaImpl schema = ObjectMapperFactory.get().readValue(is.readAllBytes(), JsonSchemaImpl.class);
+                String key;
+                boolean notTitle = !OpenApiModelSchemaGenerator.useTitle() || schema.getTitle() == null;
+                if (notTitle) {
+                    key = RefSchemas.getKey();
+                    schema.title(key);
+                } else {
+                    key = schema.getTitle();
+                }
+                if (key != null) {
+                    RefSchemas.get().put(key, schema);
+                }
+                ref = OpenApiConstants.REF_PREFIX_SCHEMA + key;
+            } catch (URISyntaxException | IOException e) {
+                // if not a valid uri, let super handle it
+            }
+        }
         super.setRef(ref);
     }
 
diff --git a/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/OpenApiModelSchemaGenerator.java b/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/OpenApiModelSchemaGenerator.java
index 4759f8b5436..a4aac0d2dad 100644
--- a/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/OpenApiModelSchemaGenerator.java
+++ b/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/OpenApiModelSchemaGenerator.java
@@ -58,6 +58,7 @@ private OpenApiModelSchemaGenerator() {
 
     public static void addOpenAPIModelSchema(KogitoWorkflowProcess workflow, Map<String, Schema> schemas) {
         if (workflow instanceof WorkflowProcess) {
+            RefSchemas.init(workflow.getId());
             WorkflowProcess workflowProcess = (WorkflowProcess) workflow;
             getSchema(workflowProcess.getInputValidator()).ifPresent(v -> {
                 String key = getSchemaName(workflow.getId(), INPUT_SUFFIX);
@@ -67,6 +68,8 @@ public static void addOpenAPIModelSchema(KogitoWorkflowProcess workflow, Map<Str
                 String key = getSchemaName(workflow.getId(), OUTPUT_SUFFIX);
                 schemas.put(key, createOutputSchema(schemaTitle(key, v)));
             });
+            schemas.putAll(RefSchemas.get());
+            RefSchemas.reset();
         }
     }
 
@@ -81,7 +84,7 @@ private static Schema schemaTitle(String key, Schema schema) {
         return schema;
     }
 
-    private static boolean useTitle() {
+    static boolean useTitle() {
         return ConfigProvider.getConfig().getOptionalValue("kogito.sw.schema.use_title", Boolean.class).orElse(true);
     }
 
diff --git a/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/RefSchemas.java b/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/RefSchemas.java
new file mode 100644
index 00000000000..56be95fbcfe
--- /dev/null
+++ b/quarkus/extensions/kogito-quarkus-serverless-workflow-extension/kogito-quarkus-serverless-workflow-deployment/src/main/java/org/kie/kogito/serverless/workflow/parser/schema/RefSchemas.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates.
+ *
+ * Licensed 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.
+ */
+
+package org.kie.kogito.serverless.workflow.parser.schema;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.eclipse.microprofile.openapi.models.media.Schema;
+
+class RefSchemas {
+
+    private RefSchemas() {
+    }
+
+    private static class ThreadInfo {
+        public int counter;
+        public String id;
+        public Map<String, Schema> map = new HashMap<>();
+
+        public ThreadInfo(String id) {
+            this.id = id;
+        }
+    }
+
+    private static ThreadLocal<ThreadInfo> threadInfo = new ThreadLocal<>();
+
+    public static void init(String id) {
+        threadInfo.set(new ThreadInfo(id));
+    }
+
+    public static Map<String, Schema> get() {
+        return threadInfo.get().map;
+    }
+
+    public static String getKey() {
+        ThreadInfo t = threadInfo.get();
+        return t.id + "_nested_" + t.counter++;
+    }
+
+    public static void reset() {
+        threadInfo.set(null);
+    }
+}