|
| 1 | +/* |
| 2 | + * Copyright (c) 2023 Airbyte, Inc., all rights reserved. |
| 3 | + */ |
| 4 | + |
| 5 | +package io.airbyte.integrations.source.bigquery; |
| 6 | + |
| 7 | +import static io.airbyte.integrations.source.bigquery.BigQuerySource.CONFIG_DATASET_ID; |
| 8 | +import static org.junit.jupiter.api.Assertions.assertEquals; |
| 9 | +import static org.junit.jupiter.api.Assertions.assertNotNull; |
| 10 | + |
| 11 | +import com.fasterxml.jackson.databind.JsonNode; |
| 12 | +import com.fasterxml.jackson.databind.ObjectMapper; |
| 13 | +import io.airbyte.commons.util.MoreIterators; |
| 14 | +import io.airbyte.protocol.models.Field; |
| 15 | +import io.airbyte.protocol.models.JsonSchemaType; |
| 16 | +import io.airbyte.protocol.models.v0.AirbyteMessage; |
| 17 | +import io.airbyte.protocol.models.v0.CatalogHelpers; |
| 18 | +import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog; |
| 19 | +import java.sql.SQLException; |
| 20 | +import java.util.List; |
| 21 | +import org.junit.jupiter.api.Test; |
| 22 | + |
| 23 | +public class BigQuerySourceStructRepeatedTest extends AbstractBigQuerySourceTest { |
| 24 | + |
| 25 | + @Override |
| 26 | + public void createTable(String datasetId) throws SQLException { |
| 27 | + // create column name interval which should be escaped |
| 28 | + database.execute("CREATE TABLE " + datasetId + ".struct_repeated(id int64, key_value_pairs ARRAY<STRUCT<key STRING, value FLOAT64>>);"); |
| 29 | + database.execute("INSERT INTO " + datasetId + ".struct_repeated (id, key_value_pairs) VALUES (1, [('a', 0.7), ('b', 0.8), ('c', 1.2)]);"); |
| 30 | + } |
| 31 | + |
| 32 | + @Test |
| 33 | + public void testReadSuccess() throws Exception { |
| 34 | + final List<AirbyteMessage> actualMessages = MoreIterators.toList(new BigQuerySource().read(config, getConfiguredCatalog(), null)); |
| 35 | + |
| 36 | + ObjectMapper mapper = new ObjectMapper(); |
| 37 | + // JsonNode actualObj = mapper.readTree("{\"key_value_pairs\":[{ \"key\": \"a\",\"value\": \"0.7\"}, |
| 38 | + // {\"key\": \"b\",\"value\": \"0.8\"}, {\"key\": \"c\",\"value\": \"1.2\"}]}"); |
| 39 | + JsonNode actualObj = mapper.readTree("[{ \"key\": \"a\",\"value\": 0.7}, {\"key\": \"b\",\"value\": 0.8}, {\"key\": \"c\",\"value\": 1.2}]"); |
| 40 | + |
| 41 | + assertNotNull(actualMessages); |
| 42 | + assertEquals(1, actualMessages.size()); |
| 43 | + |
| 44 | + assertNotNull(actualMessages.get(0).getRecord().getData().get("id")); |
| 45 | + assertEquals(actualObj, actualMessages.get(0).getRecord().getData().get("key_value_pairs")); |
| 46 | + } |
| 47 | + |
| 48 | + protected ConfiguredAirbyteCatalog getConfiguredCatalog() { |
| 49 | + return CatalogHelpers.createConfiguredAirbyteCatalog( |
| 50 | + "struct_repeated", |
| 51 | + config.get(CONFIG_DATASET_ID).asText(), |
| 52 | + Field.of("id", JsonSchemaType.NUMBER), |
| 53 | + Field.of("key_value_pairs", JsonSchemaType.ARRAY)); |
| 54 | + } |
| 55 | + |
| 56 | +} |
0 commit comments