From 9eccb58ba2db129c408a9bfaeea9f68aa87b371f Mon Sep 17 00:00:00 2001
From: Astrid Jahn <astrid.jahn@mobimeo.com>
Date: Wed, 28 Feb 2024 15:16:15 +0100
Subject: [PATCH] Issue 145 Add schema mapper for nullable

When creating the OpenAPI schema, fields annotated
with a @nullable trait should be given an attribute
"nullable" with a value of true in the resulting
schema. We add a new JSON mapper to the OpenAPI
extensions that will take care of this.
---
 .../alloy/openapi/AlloyOpenApiExtension.scala |  3 +-
 .../src/alloy/openapi/NullableMapper.scala    | 42 +++++++++++++++++++
 modules/openapi/test/resources/foo.json       |  5 +++
 modules/openapi/test/resources/foo.smithy     |  4 ++
 4 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 modules/openapi/src/alloy/openapi/NullableMapper.scala

diff --git a/modules/openapi/src/alloy/openapi/AlloyOpenApiExtension.scala b/modules/openapi/src/alloy/openapi/AlloyOpenApiExtension.scala
index 01af267..e17179a 100644
--- a/modules/openapi/src/alloy/openapi/AlloyOpenApiExtension.scala
+++ b/modules/openapi/src/alloy/openapi/AlloyOpenApiExtension.scala
@@ -48,7 +48,8 @@ final class AlloyOpenApiExtension() extends Smithy2OpenApiExtension {
     new DiscriminatedUnions(),
     new UntaggedUnions(),
     new DataExamplesMapper(),
-    new ExternalDocumentationMapperJsonSchema()
+    new ExternalDocumentationMapperJsonSchema(),
+    new NullableMapper()
   ).asJava
 
 }
diff --git a/modules/openapi/src/alloy/openapi/NullableMapper.scala b/modules/openapi/src/alloy/openapi/NullableMapper.scala
new file mode 100644
index 0000000..0b7b0f6
--- /dev/null
+++ b/modules/openapi/src/alloy/openapi/NullableMapper.scala
@@ -0,0 +1,42 @@
+/* Copyright 2022 Disney Streaming
+ *
+ * Licensed under the Tomorrow Open Source Technology License, Version 1.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    https://disneystreaming.github.io/TOST-1.0.txt
+ *
+ * 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 alloy.openapi
+
+import alloy.NullableTrait
+import software.amazon.smithy.jsonschema.{
+  JsonSchemaMapper,
+  JsonSchemaMapperContext,
+  Schema
+}
+import software.amazon.smithy.model.node.Node
+
+class NullableMapper extends JsonSchemaMapper {
+  override def updateSchema(
+      context: JsonSchemaMapperContext,
+      schemaBuilder: Schema.Builder
+  ): Schema.Builder = {
+    context.getShape
+      .getTrait(classOf[NullableTrait])
+      .asScala
+      .map { _ =>
+        schemaBuilder.putExtension(
+          "nullable",
+          Node.from(true)
+        )
+      }
+      .getOrElse(schemaBuilder)
+  }
+}
diff --git a/modules/openapi/test/resources/foo.json b/modules/openapi/test/resources/foo.json
index 19fd4c4..68ca3d6 100644
--- a/modules/openapi/test/resources/foo.json
+++ b/modules/openapi/test/resources/foo.json
@@ -272,6 +272,11 @@
         "properties": {
           "message": {
             "type": "string"
+          },
+          "count": {
+            "type": "integer",
+            "format": "int32",
+            "nullable": true
           }
         }
       },
diff --git a/modules/openapi/test/resources/foo.smithy b/modules/openapi/test/resources/foo.smithy
index fb173c3..79b7b38 100644
--- a/modules/openapi/test/resources/foo.smithy
+++ b/modules/openapi/test/resources/foo.smithy
@@ -2,6 +2,7 @@ namespace foo
 
 use alloy#simpleRestJson
 use alloy#discriminated
+use alloy#nullable
 use alloy#untagged
 use alloy#dataExamples
 
@@ -80,6 +81,9 @@ structure Greeting {
 @httpError(500)
 structure GeneralServerError {
   message: String,
+
+  @nullable
+  count: Integer
 }
 
 structure GetUnionResponse {