dropped union merging; field its own type; cleaned up value mapping

airbytehq · Sep 13, 2024 · a0f7167 · a0f7167
1 parent 27f0858
commit a0f7167
Show file tree

Hide file tree

Showing 5 changed files with 203 additions and 202 deletions.
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/data/AirbyteType.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/data/AirbyteType.kt
@@ -34,6 +34,6 @@ data object ObjectTypeWithoutSchema : AirbyteType()
 
 data class UnionType(val options: List<AirbyteType>) : AirbyteType()
 
-data class FieldType(val type: AirbyteType, val nullable: Boolean) : AirbyteType()
-
 data class UnknownType(val what: String) : AirbyteType()
+
+data class FieldType(val type: AirbyteType, val nullable: Boolean)
diff --git a/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/data/JsonSchemaToAirbyteType.kt b/airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/data/JsonSchemaToAirbyteType.kt
@@ -28,12 +28,12 @@ class JsonSchemaToAirbyteType {
                             )
                     }
                 } else if (schemaType.isArray) {
-                    return fromUnionTypes(schemaType.map { it.asText() }, schema)
+                    return unionFromCombinedTypes(schemaType.map { it.asText() }, schema)
                 }
             } else {
                 val options = schema.get("oneOf") ?: schema.get("anyOf") ?: schema.get("allOf")
                 if (options != null) {
-                    return fromUnionOptions(options.map { convert(it as ObjectNode) })
+                    return UnionType(options.map { convert(it as ObjectNode) })
                 }
             }
         } catch (t: Throwable) {
@@ -77,84 +77,59 @@ class JsonSchemaToAirbyteType {
             if (items.isEmpty) {
                 return ArrayTypeWithoutSchema
             }
-            val itemOptions = fromUnionOptions(items.map { convert(it as ObjectNode) })
-            return ArrayType(itemOptions)
+            val itemOptions = UnionType(items.map { convert(it as ObjectNode) })
+            return ArrayType(fieldFromUnion(itemOptions))
         }
-        return ArrayType(maybeFlattenFieldType(items as ObjectNode))
+        return ArrayType(fieldFromSchema(items as ObjectNode))
     }
 
-    private fun objectFromProperties(schema: ObjectNode): AirbyteType {
+    private fun fromObject(schema: ObjectNode): AirbyteType {
+        val properties = schema.get("properties") ?: return ObjectTypeWithoutSchema
+        if (properties.isEmpty) {
+            return ObjectTypeWithEmptySchema
+        }
+        val requiredFields = schema.get("required")?.map { it.asText() } ?: emptyList()
+        return objectFromProperties(properties as ObjectNode, requiredFields)
+    }
+
+    private fun fieldFromSchema(fieldSchema: ObjectNode, onRequiredList: Boolean = false): FieldType {
+        val markedRequired = fieldSchema.get("required")?.asBoolean() ?: false
+        val nullable = !(onRequiredList || markedRequired)
+        val airbyteType = convert(fieldSchema)
+        if (airbyteType is UnionType) {
+            return fieldFromUnion(airbyteType, nullable)
+        } else {
+            return FieldType(airbyteType, nullable)
+        }
+    }
+
+    private fun fieldFromUnion(unionType: UnionType, nullable: Boolean = false): FieldType {
+        if (unionType.options.contains(NullType)) {
+            val filtered = unionType.options.filter { it != NullType }
+            return FieldType(UnionType(filtered), nullable = true)
+        }
+        return FieldType(unionType, nullable = nullable)
+    }
+
+    private fun objectFromProperties(schema: ObjectNode, requiredFields: List<String>): ObjectType {
         val properties =
             schema
                 .fields()
                 .asSequence()
-                .map { (name, node) -> name to maybeFlattenFieldType(node as ObjectNode) }
+                .map { (name, node) -> name to fieldFromSchema(node as ObjectNode, requiredFields.contains(name)) }
                 .toMap(LinkedHashMap())
         return ObjectType(properties)
     }
 
-    private fun fromObject(schema: ObjectNode): AirbyteType {
-        val properties = schema.get("properties") ?: return ObjectTypeWithoutSchema
-        if (properties.isEmpty) {
-            return ObjectTypeWithEmptySchema
-        }
-        return objectFromProperties(properties as ObjectNode)
-    }
-
-    private fun fromUnionTypes(options: List<String>, parentSchema: ObjectNode): FieldType {
+    private fun unionFromCombinedTypes(options: List<String>, parentSchema: ObjectNode): UnionType {
+        // Denormalize the properties across each type (the converter only checks what matters
+        // per type).
         val unionOptions =
             options.map {
                 val schema = parentSchema.deepCopy()
                 schema.put("type", it)
                 convert(schema)
             }
-        return fromUnionOptions(unionOptions)
-    }
-
-    private fun fromUnionOptions(options: List<AirbyteType>): FieldType {
-        if (options.size == 1) {
-            return FieldType(options.first(), nullable = false)
-        } else if (options.contains(NullType)) {
-            val field = fromUnionOptions(options.filter { it != NullType })
-            return FieldType(field.type, nullable = true)
-        } else {
-            // Merge like objects and arrays and unions
-            options.fold(mutableSetOf<AirbyteType>()) { acc, type ->
-                if (acc.contains(type)) {
-                    // Do nothing: type exists and is identical
-                } else if (type is ObjectType && acc.any { it is ObjectType }) {
-                    val existing = acc.find { it is ObjectType } as ObjectType
-                    val newProperties = LinkedHashMap(existing.properties)
-                    type.properties.forEach { (name, field) -> newProperties[name] = field }
-                    acc.remove(existing)
-                    acc.add(ObjectType(newProperties))
-                } else if (type is ArrayType && acc.any { it is ArrayType }) {
-                    val existing = acc.find { it is ArrayType } as ArrayType
-                    acc.remove(existing)
-                    acc.add(ArrayType(fromUnionOptions(listOf(existing.items, type.items))))
-                } else if (type is UnionType && acc.any { it is UnionType }) {
-                    val existing = acc.find { it is UnionType } as UnionType
-                    acc.remove(existing)
-                    val unique = (existing.options + type.options).distinct()
-                    acc.add(fromUnionOptions(unique))
-                } else {
-                    acc.add(type)
-                }
-                acc
-            }
-            return FieldType(UnionType(options), nullable = false)
-        }
-    }
-
-    private fun maybeFlattenFieldType(jsonSchema: ObjectNode): FieldType {
-        val subType = convert(jsonSchema)
-        val type =
-            if (subType is FieldType) {
-                subType.type
-            } else {
-                subType
-            }
-        val nullable = !(jsonSchema.get("required")?.asBoolean() ?: false)
-        return FieldType(type, nullable)
+        return UnionType(unionOptions)
     }
 }