-
Notifications
You must be signed in to change notification settings - Fork 4.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Bulk Load CDK: Airbyte Type cleanup, Basic ID Schema Mapper (#47190)
- Loading branch information
1 parent
f5bf624
commit fb8063b
Showing
12 changed files
with
217 additions
and
45 deletions.
There are no files selected for viewing
55 changes: 55 additions & 0 deletions
55
...dk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/AirbyteSchemaIdentityMapper.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
/* | ||
* Copyright (c) 2024 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.cdk.load.data | ||
|
||
interface AirbyteSchemaIdentityMapper { | ||
fun map(schema: AirbyteType): AirbyteType = | ||
when (schema) { | ||
is NullType -> mapNull(schema) | ||
is StringType -> mapString(schema) | ||
is BooleanType -> mapBoolean(schema) | ||
is IntegerType -> mapInteger(schema) | ||
is NumberType -> mapNumber(schema) | ||
is ArrayType -> mapArray(schema) | ||
is ArrayTypeWithoutSchema -> mapArrayWithoutSchema(schema) | ||
is ObjectType -> mapObject(schema) | ||
is ObjectTypeWithoutSchema -> mapObjectWithoutSchema(schema) | ||
is ObjectTypeWithEmptySchema -> mapObjectWithEmptySchema(schema) | ||
is UnionType -> mapUnion(schema) | ||
is DateType -> mapDate(schema) | ||
is TimeTypeWithTimezone -> mapTimeTypeWithTimezone(schema) | ||
is TimeTypeWithoutTimezone -> mapTimeTypeWithoutTimezone(schema) | ||
is TimestampTypeWithTimezone -> mapTimestampTypeWithTimezone(schema) | ||
is TimestampTypeWithoutTimezone -> mapTimestampTypeWithoutTimezone(schema) | ||
is UnknownType -> mapUnknown(schema) | ||
} | ||
|
||
fun mapNull(schema: NullType): AirbyteType = schema | ||
fun mapString(schema: StringType): AirbyteType = schema | ||
fun mapBoolean(schema: BooleanType): AirbyteType = schema | ||
fun mapInteger(schema: IntegerType): AirbyteType = schema | ||
fun mapNumber(schema: NumberType): AirbyteType = schema | ||
fun mapArray(schema: ArrayType): AirbyteType { | ||
return ArrayType(mapField(schema.items)) | ||
} | ||
fun mapArrayWithoutSchema(schema: ArrayTypeWithoutSchema): AirbyteType = schema | ||
fun mapObject(schema: ObjectType): AirbyteType { | ||
val properties = LinkedHashMap<String, FieldType>() | ||
schema.properties.forEach { (name, field) -> properties[name] = mapField(field) } | ||
return ObjectType(properties) | ||
} | ||
fun mapObjectWithoutSchema(schema: ObjectTypeWithoutSchema): AirbyteType = schema | ||
fun mapObjectWithEmptySchema(schema: ObjectTypeWithEmptySchema): AirbyteType = schema | ||
fun mapUnion(schema: UnionType): AirbyteType { | ||
return UnionType(schema.options.map { map(it) }) | ||
} | ||
fun mapDate(schema: DateType): AirbyteType = schema | ||
fun mapTimeTypeWithTimezone(schema: TimeTypeWithTimezone): AirbyteType = schema | ||
fun mapTimeTypeWithoutTimezone(schema: TimeTypeWithoutTimezone): AirbyteType = schema | ||
fun mapTimestampTypeWithTimezone(schema: TimestampTypeWithTimezone): AirbyteType = schema | ||
fun mapTimestampTypeWithoutTimezone(schema: TimestampTypeWithoutTimezone): AirbyteType = schema | ||
fun mapUnknown(schema: UnknownType): AirbyteType = schema | ||
fun mapField(field: FieldType): FieldType = FieldType(map(field.type), field.nullable) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
41 changes: 41 additions & 0 deletions
41
...ulk/core/load/src/test/kotlin/io/airbyte/cdk/load/data/AirbyteSchemaIdentityMapperTest.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
/* | ||
* Copyright (c) 2024 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.cdk.load.data | ||
|
||
import io.airbyte.cdk.load.test.util.SchemaTestBuilder | ||
import org.junit.jupiter.api.Assertions | ||
import org.junit.jupiter.api.Test | ||
|
||
class AirbyteSchemaIdentityMapperTest { | ||
@Test | ||
fun testIdMapping() { | ||
val (inputSchema, expectedOutput) = | ||
SchemaTestBuilder() | ||
.with(DateType) | ||
.with(StringType) | ||
.with(IntegerType) | ||
.with(BooleanType) | ||
.with(NumberType) | ||
.with(NullType) | ||
.with(ArrayType(FieldType(StringType, true))) | ||
.with(UnionType(listOf(StringType, IntegerType, NullType))) | ||
.withRecord() | ||
.with(TimeTypeWithTimezone) | ||
.with(TimeTypeWithoutTimezone) | ||
.with(TimestampTypeWithTimezone) | ||
.with(TimestampTypeWithoutTimezone) | ||
.withRecord() | ||
.with(ObjectTypeWithoutSchema) | ||
.with(ObjectTypeWithEmptySchema) | ||
.with(ArrayTypeWithoutSchema) | ||
.endRecord() | ||
.endRecord() | ||
.with(NullType) | ||
.build() | ||
|
||
val mapper = object : AirbyteSchemaIdentityMapper {} | ||
Assertions.assertEquals(expectedOutput, mapper.map(inputSchema)) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
56 changes: 56 additions & 0 deletions
56
...yte-cdk/bulk/core/load/src/test/kotlin/io/airbyte/cdk/load/test/util/SchemaTestBuilder.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,56 @@ | ||
/* | ||
* Copyright (c) 2024 Airbyte, Inc., all rights reserved. | ||
*/ | ||
|
||
package io.airbyte.cdk.load.test.util | ||
|
||
import io.airbyte.cdk.load.data.AirbyteType | ||
import io.airbyte.cdk.load.data.FieldType | ||
import io.airbyte.cdk.load.data.ObjectType | ||
import java.util.* | ||
import kotlin.collections.LinkedHashMap | ||
|
||
class SchemaTestBuilder( | ||
val inputSchema: ObjectType = ObjectType(properties = LinkedHashMap()), | ||
val expectedSchema: ObjectType = ObjectType(properties = LinkedHashMap()), | ||
val parent: SchemaTestBuilder? = null | ||
) { | ||
|
||
fun with(given: FieldType, expected: FieldType = given): SchemaTestBuilder { | ||
val name = UUID.randomUUID().toString() | ||
inputSchema.properties[name] = given | ||
expectedSchema.properties[name] = expected | ||
return this | ||
} | ||
|
||
fun with(given: AirbyteType, expected: AirbyteType = given): SchemaTestBuilder { | ||
return with(FieldType(given, false), FieldType(expected, false)) | ||
} | ||
|
||
fun withRecord(nullable: Boolean = false): SchemaTestBuilder { | ||
val name = UUID.randomUUID().toString() | ||
val inputRecord = ObjectType(properties = LinkedHashMap()) | ||
val outputRecord = ObjectType(properties = LinkedHashMap()) | ||
inputSchema.properties[name] = FieldType(inputRecord, nullable = nullable) | ||
expectedSchema.properties[name] = FieldType(outputRecord, nullable = nullable) | ||
return SchemaTestBuilder( | ||
inputSchema = inputRecord, | ||
expectedSchema = outputRecord, | ||
parent = this | ||
) | ||
} | ||
|
||
fun endRecord(): SchemaTestBuilder { | ||
if (parent == null) { | ||
throw IllegalStateException("Cannot end record without parent") | ||
} | ||
return parent | ||
} | ||
|
||
fun build(): Pair<ObjectType, ObjectType> { | ||
if (parent != null) { | ||
throw IllegalStateException("Cannot build nested schema") | ||
} | ||
return Pair(inputSchema, expectedSchema) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters