com.google.cloud.bigquery.Field Java Examples

@Test
public void shouldTestConvertToSchemaForRepeatedFields() {
    ProtoField protoField = new ProtoField(new ArrayList<ProtoField>() {{
        add(new ProtoField("field1_map",
                DescriptorProtos.FieldDescriptorProto.Type.TYPE_INT32,
                DescriptorProtos.FieldDescriptorProto.Label.LABEL_REPEATED));
        add(new ProtoField("field2_repeated",
                DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING,
                DescriptorProtos.FieldDescriptorProto.Label.LABEL_REPEATED));

    }});

    List<Field> fields = converter.generateBigquerySchema(protoField);

    assertEquals(protoField.getFields().size(), fields.size());
    assertBqField(protoField.getFields().get(0).getName(), LegacySQLTypeName.INTEGER, Field.Mode.REPEATED, fields.get(0));
    assertBqField(protoField.getFields().get(1).getName(), LegacySQLTypeName.STRING, Field.Mode.REPEATED, fields.get(1));
}
 

@Test
public void shouldCreatePartitionedTable() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(true);
    when(bqConfig.getBQTablePartitionKey()).thenReturn("timestamp_field");
    Schema bqSchema = Schema.of(
            Field.newBuilder("timestamp_field", LegacySQLTypeName.TIMESTAMP).build()
    );

    BQTableDefinition bqTableDefinition = new BQTableDefinition(bqConfig);
    StandardTableDefinition tableDefinition = bqTableDefinition.getTableDefinition(bqSchema);

    Schema returnedSchema = tableDefinition.getSchema();
    assertEquals(returnedSchema.getFields().size(), bqSchema.getFields().size());
    assertEquals(returnedSchema.getFields().get(0).getName(), bqSchema.getFields().get(0).getName());
    assertEquals(returnedSchema.getFields().get(0).getMode(), bqSchema.getFields().get(0).getMode());
    assertEquals(returnedSchema.getFields().get(0).getType(), bqSchema.getFields().get(0).getType());
    assertEquals("timestamp_field", tableDefinition.getTimePartitioning().getField());
}
 

static BigQueryColumnHandle toColumnHandle(Field field)
{
    FieldList subFields = field.getSubFields();
    List<BigQueryColumnHandle> subColumns = subFields == null ?
            Collections.emptyList() :
            subFields.stream()
                    .map(Conversions::toColumnHandle)
                    .collect(Collectors.toList());
    return new BigQueryColumnHandle(
            field.getName(),
            BigQueryType.valueOf(field.getType().name()),
            getMode(field),
            subColumns,
            field.getDescription(),
            false);
}
 

/**
 * This method gives us a chance to perform some additional type coercions in case the BigQuery
 * field type is different from the source data type. This should rarely happen, since only
 * validated payloads get through to this BQ sink path, but there are sets of probes with
 * heterogeneous types that appear as explicit fields in BQ, but are treated as loosely typed
 * maps at the validation phase; we need to catch these or they can cause the entire pipeline
 * to stall.
 *
 * <p>Returning {@link Optional#empty} here indicates that no coercion is defined and that the
 * field should be put to {@code additional_properties}.
 */
private Optional<JsonNode> coerceToBqType(JsonNode o, Field field) {
  if (o.isNull()) {
    // null is valid for any type, just not as an element of a list
    return Optional.of(o);
  } else if (field.getMode() == Field.Mode.REPEATED) {
    if (o.isArray()) {
      // We have not yet observed a case where an array type contains values that cannot be
      // coerced to appropriate values, but if it does this will throw NoSuchElementException
      // and prevent the message from being delivered to BigQuery in a form that could lead to
      // data being missed in additional_properties.
      return Optional.of(Json.createArrayNode()
          .addAll(Streams.stream(o).map(v -> coerceSingleValueToBqType(v, field))
              .map(Optional::get).collect(Collectors.toList())));
    } else {
      return Optional.empty();
    }
  } else {
    return coerceSingleValueToBqType(o, field);
  }
}
 

private Boolean addNewTableField(TableId tableId, TableRow row, String rowKey,
    List<Field> newFieldList, Map<String, LegacySQLTypeName> inputSchema) {
  // Call Get Schema and Extract New Field Type
  Field newField;

  if (inputSchema.containsKey(rowKey)) {
    newField = Field.of(rowKey, inputSchema.get(rowKey));
  } else {
    newField = Field.of(rowKey, LegacySQLTypeName.STRING);
  }

  newFieldList.add(newField);

  // Currently we always add new fields for each call
  // TODO: add an option to ignore new field and why boolean?
  return true;
}
 

/** Example of creating a table. */
// [TARGET create(TableInfo, TableOption...)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
// [VARIABLE "string_field"]
public Table createTable(String datasetName, String tableName, String fieldName) {
  // [START bigquery_create_table]
  TableId tableId = TableId.of(datasetName, tableName);
  // Table field definition
  Field field = Field.of(fieldName, LegacySQLTypeName.STRING);
  // Table schema definition
  Schema schema = Schema.of(field);
  TableDefinition tableDefinition = StandardTableDefinition.of(schema);
  TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();
  Table table = bigquery.create(tableInfo);
  // [END bigquery_create_table]
  return table;
}
 

private org.apache.avro.Schema inferSchemaField(Field field) {
    Field.Mode mode = field.getMode();

    // Get the "basic" type of the field.
    org.apache.avro.Schema fieldSchema = inferSchemaFieldWithoutMode(field);

    // BigQuery fields are NULLABLE by default.
    if (Field.Mode.NULLABLE == mode || mode == null) {
        fieldSchema = AvroUtils.wrapAsNullable(fieldSchema);
    } else if (Field.Mode.REPEATED == mode) {
        // Determine if the field is an array.
        // https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#array-type
        fieldSchema = SchemaBuilder.array().items(fieldSchema);
    }
    return fieldSchema;
}
 

@JsonCreator
public BigQueryColumnHandle(
        @JsonProperty("name") String name,
        @JsonProperty("bigQueryType") BigQueryType bigQueryType,
        @JsonProperty("mode") Field.Mode mode,
        @JsonProperty("subColumns") List<BigQueryColumnHandle> subColumns,
        @JsonProperty("description") String description,
        @JsonProperty("hidden") boolean hidden)
{
    this.name = requireNonNull(name, "column name cannot be null");
    this.bigQueryType = requireNonNull(bigQueryType, () -> format("column type cannot be null for column [%s]", name));
    this.mode = requireNonNull(mode, "Field mode cannot be null");
    this.subColumns = ImmutableList.copyOf(requireNonNull(subColumns, "subColumns is null"));
    this.description = description;
    this.hidden = hidden;
}
 

/**
 * Extracts and applies new column information to BigQuery by comparing the TableRow against the
 * BigQuery Table.
 *
 * @param tableId a TableId referencing the BigQuery table to be loaded to.
 * @param row a TableRow with the raw data to be loaded into BigQuery.
 * @param inputSchema The source schema lookup to be used in mapping.
 */
private void updateTableIfRequired(
    TableId tableId, TableRow row, Map<String, LegacySQLTypeName> inputSchema) {
  Table table = getOrCreateBigQueryTable(tableId);
  FieldList tableFields = table.getDefinition().getSchema().getFields();

  Set<String> rowKeys = row.keySet();
  Boolean tableWasUpdated = false;
  List<Field> newFieldList = new ArrayList<Field>();
  for (String rowKey : rowKeys) {
    // Check if rowKey (column from data) is in the BQ Table
    try {
      Field tableField = tableFields.get(rowKey);
    } catch (IllegalArgumentException e) {
      tableWasUpdated = addNewTableField(tableId, row, rowKey, newFieldList, inputSchema);
    }
  }

  if (tableWasUpdated) {
    LOG.info("Updating Table");
    updateBigQueryTable(tableId, table, tableFields, newFieldList);
  }
}
 

/**
 * Create a function that converts an Avro row with the given BigQuery schema to a Spark SQL row
 * <p>
 * The conversion is based on the BigQuery schema, not Avro Schema, because the Avro schema is
 * very painful to use.
 * <p>
 * Not guaranteed to be stable across all versions of Spark.
 */

private static StructField convert(Field field) {
    DataType dataType = getDataType(field);
    boolean nullable = true;

    if (field.getMode() == Field.Mode.REQUIRED) {
        nullable = false;
    } else if (field.getMode() == Field.Mode.REPEATED) {
        dataType = new ArrayType(dataType, true);
    }

    MetadataBuilder metadata = new MetadataBuilder();
    if (field.getDescription() != null) {
        metadata.putString("description", field.getDescription());
    }

    return new StructField(field.getName(), dataType, nullable, metadata.build());
}
 

@Test
public void testUnmapWithoutValue() throws Exception {
  List<Field> bqFields = ImmutableList.of(MAP_FIELD_WITHOUT_VALUE);
  Map<String, Object> expectedParent = ImmutableMap.of("map_field",
      ImmutableList.of(ImmutableMap.of("key", "foo"), ImmutableMap.of("key", "bar")));
  Map<String, Object> expectedAdditional = ImmutableMap.of("mapField",
      ImmutableMap.of("bar", 4, "foo", 3));
  for (boolean withAdditional : ImmutableList.of(true, false)) {
    ObjectNode additionalProperties = withAdditional ? Json.createObjectNode() : null;
    ObjectNode parent = Json.createObjectNode().set("mapField",
        Json.createObjectNode().put("foo", 3).put("bar", 4));
    TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
    assertEquals(expectedParent, Json.asMap(parent));
    if (withAdditional) {
      assertEquals(expectedAdditional, Json.asMap(additionalProperties));
    }
  }
}
 

@Test
public void testConvertTwoLevelsRecordColumn()
{
    BigQueryColumnHandle column = new BigQueryColumnHandle("rec", BigQueryType.RECORD, Field.Mode.NULLABLE, ImmutableList.of(
            new BigQueryColumnHandle("sub_rec", BigQueryType.RECORD, Field.Mode.NULLABLE, ImmutableList.of(
                    new BigQueryColumnHandle("sub_sub_s", BigQueryType.STRING, Field.Mode.NULLABLE, ImmutableList.of(), null),
                    new BigQueryColumnHandle("sub_sub_i", BigQueryType.INTEGER, Field.Mode.NULLABLE, ImmutableList.of(), null)
            ), null),
            new BigQueryColumnHandle("sub_s", BigQueryType.STRING, Field.Mode.NULLABLE, ImmutableList.of(), null),
            new BigQueryColumnHandle("sub_i", BigQueryType.INTEGER, Field.Mode.NULLABLE, ImmutableList.of(), null)
    ), null);
    ColumnMetadata metadata = column.getColumnMetadata();
    RowType targetType = RowType.rowType(
            RowType.field("sub_rec", RowType.rowType(
                    RowType.field("sub_sub_s", VarcharType.VARCHAR),
                    RowType.field("sub_sub_i", BigintType.BIGINT))),
            RowType.field("sub_s", VarcharType.VARCHAR),
            RowType.field("sub_i", BigintType.BIGINT));
    assertThat(metadata.getType()).isEqualTo(targetType);
}
 

public static Field mapToField(Map fMap) {
    String typeStr = fMap.get("type").toString();
    String nameStr = fMap.get("name").toString();
    String modeStr = fMap.get("mode").toString();
    LegacySQLTypeName type = null;

    if (typeStr.equals("BOOLEAN")) {
        type = LegacySQLTypeName.BOOLEAN;
    } else if (typeStr.equals("STRING")) {
        type = LegacySQLTypeName.STRING;
    } else if (typeStr.equals("BYTES")) {
        type = LegacySQLTypeName.BYTES;
    } else if (typeStr.equals("INTEGER")) {
        type = LegacySQLTypeName.INTEGER;
    } else if (typeStr.equals("FLOAT")) {
        type = LegacySQLTypeName.FLOAT;
    } else if (typeStr.equals("TIMESTAMP") || typeStr.equals("DATE")
            || typeStr.equals("TIME") || typeStr.equals("DATETIME")) {
        type = LegacySQLTypeName.TIMESTAMP;
    } else if (typeStr.equals("RECORD")) {
        type = LegacySQLTypeName.RECORD;
    }

    return Field.newBuilder(nameStr, type).setMode(Field.Mode.valueOf(modeStr)).build();
}
 

private static List<String> getFieldNames(List<Field> fields, String parentFieldName) {
    List<String> fieldNames = new ArrayList<>();
    if (fields == null) {
        return fieldNames;
    }

    fields.stream().forEach(field -> {
        String fieldName = field.getName();
        if (parentFieldName != null) {
            fieldNames.add(parentFieldName + "." + fieldName);
        } else {
            fieldNames.add(fieldName);
        }
        fieldNames.addAll(getFieldNames(field.getSubFields(), fieldName));
    });
    return fieldNames;
}
 

@Test
public void testListWithNulls() throws Exception {
  ObjectNode additionalProperties = Json.createObjectNode();
  ObjectNode parent = Json.readObjectNode(
      ("{\"modules\":[{\"base_addr\":\"0x1390000\"},null]}").getBytes(StandardCharsets.UTF_8));
  List<Field> bqFields = ImmutableList.of(Field
      .newBuilder("modules", LegacySQLTypeName.RECORD,
          Field.of("base_addr", LegacySQLTypeName.STRING)) //
      .setMode(Mode.REPEATED).build() //
  ); //
  Map<String, Object> expected = Json.readMap("{\"modules\":[{\"base_addr\":\"0x1390000\"},{}]}");
  TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
  assertEquals(expected, Json.asMap(parent));

  Map<String, Object> expectedAdditional = Json.readMap("{\"modules\":[{},null]}");
  assertEquals(expectedAdditional, Json.asMap(additionalProperties));
}
 

private Table createBigQueryTable(TableId tableId) {
  // Create Blank BigQuery Table
  LOG.info(String.format("Creating Table: %s", tableId.toString()));

  List<Field> fieldList = new ArrayList<Field>();
  Schema schema = Schema.of(fieldList);

  StandardTableDefinition.Builder tableDefinitionBuilder =
      StandardTableDefinition.newBuilder().setSchema(schema);
  if (dayPartitioning) {
    tableDefinitionBuilder.setTimePartitioning(
        TimePartitioning.newBuilder(TimePartitioning.Type.DAY).build());
  }
  TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinitionBuilder.build()).build();
  Table table = bigquery.create(tableInfo);

  return table;
}
 

@Test
public void testTupleIntoStructAdditionalProperties() throws Exception {
  ObjectNode additionalProperties = Json.createObjectNode();
  ObjectNode parent = Json.readObjectNode("{\n" //
      + "  \"payload\": [26,{\"a\":83,\"b\":44}]\n" //
      + "}\n");
  List<Field> bqFields = ImmutableList.of(Field.newBuilder("payload", LegacySQLTypeName.RECORD, //
      Field.of("f0_", LegacySQLTypeName.INTEGER), //
      Field.of("f1_", LegacySQLTypeName.RECORD, //
          Field.of("a", LegacySQLTypeName.INTEGER))) //
      .setMode(Mode.NULLABLE).build());
  Map<String, Object> expected = Json.readMap("{\"payload\":{\"f0_\":26,\"f1_\":{\"a\":83}}}");
  TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
  assertEquals(expected, Json.asMap(parent));

  Map<String, Object> expectedAdditional = Json.readMap("{\"payload\":[null,{\"b\":44}]}");
  assertEquals(expectedAdditional, Json.asMap(additionalProperties));
}
 

public static void main(String... args) throws InterruptedException, TimeoutException {
  BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
  TableId tableId = TableId.of("dataset", "table");
  Table table = bigquery.getTable(tableId);
  if (table == null) {
    System.out.println("Creating table " + tableId);
    Field integerField = Field.of("fieldName", LegacySQLTypeName.INTEGER);
    Schema schema = Schema.of(integerField);
    table = bigquery.create(TableInfo.of(tableId, StandardTableDefinition.of(schema)));
  }
  System.out.println("Loading data into table " + tableId);
  Job loadJob = table.load(FormatOptions.csv(), "gs://bucket/path");
  loadJob = loadJob.waitFor();
  if (loadJob.getStatus().getError() != null) {
    System.out.println("Job completed with errors");
  } else {
    System.out.println("Job succeeded");
  }
}
 

@Test
public void testDoublyNestedList() throws Exception {
  ObjectNode additionalProperties = Json.createObjectNode();
  ObjectNode parent = Json.readObjectNode("{\n" //
      + "  \"payload\": [[[0],[1]],[[2]]]\n" //
      + "}\n");
  List<Field> bqFields = ImmutableList.of(Field.newBuilder("payload", LegacySQLTypeName.RECORD, //
      Field.newBuilder("list", LegacySQLTypeName.RECORD, //
          Field.newBuilder("list", LegacySQLTypeName.INTEGER).setMode(Mode.REPEATED).build()) //
          .setMode(Mode.REPEATED).build() //
  ).setMode(Mode.REPEATED).build()); //
  Map<String, Object> expected = Json.readMap("{\"payload\":[" //
      + "{\"list\":[{\"list\":[0]},{\"list\":[1]}]}," //
      + "{\"list\":[{\"list\":[2]}]}" //
      + "]}");
  TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
  assertEquals(expected, Json.asMap(parent));
}
 

/** Example of listing table rows with schema. */
// [TARGET listTableData(TableId, Schema, TableDataListOption...)]
public FieldValueList listTableDataSchemaId() {
  // [START ]
  Schema schema =
      Schema.of(
          Field.of("word", LegacySQLTypeName.STRING),
          Field.of("word_count", LegacySQLTypeName.STRING),
          Field.of("corpus", LegacySQLTypeName.STRING),
          Field.of("corpus_date", LegacySQLTypeName.STRING));
  TableResult tableData =
      bigquery.listTableData(
          TableId.of("bigquery-public-data", "samples", "shakespeare"), schema);
  FieldValueList row = tableData.getValues().iterator().next();
  System.out.println(row.get("word").getStringValue());
  // [END ]
  return row;
}
 

private void createTable(String tableName) {
    TableId tableId = TableId.of(dataset.getDatasetId().getDataset(), tableName);

    // Table field definition
    Field id = Field.newBuilder("id", LegacySQLTypeName.INTEGER).setMode(Mode.REQUIRED).build();
    Field name = Field.newBuilder("name", LegacySQLTypeName.STRING).setMode(Mode.NULLABLE).build();
    Field alias = Field.newBuilder("alias", LegacySQLTypeName.STRING).setMode(Mode.REPEATED).build();

    Field zip = Field.newBuilder("zip", LegacySQLTypeName.STRING).setMode(Mode.NULLABLE).build();
    Field city = Field.newBuilder("city", LegacySQLTypeName.STRING).setMode(Mode.NULLABLE).build();
    Field addresses = Field.newBuilder("addresses", LegacySQLTypeName.RECORD, zip, city).setMode(Mode.REPEATED).build();

    Field position = Field.newBuilder("position", LegacySQLTypeName.STRING).setMode(Mode.NULLABLE).build();
    Field company = Field.newBuilder("company", LegacySQLTypeName.STRING).setMode(Mode.NULLABLE).build();
    Field job = Field.newBuilder("job", LegacySQLTypeName.RECORD, position, company).setMode(Mode.NULLABLE).build();

    // Table schema definition
    schema = Schema.of(id, name, alias, addresses, job);
    TableDefinition tableDefinition = StandardTableDefinition.of(schema);
    TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();

    // create table
    bigquery.create(tableInfo);
}
 

@Test
public void testAdditionalProperties() throws Exception {
  ObjectNode parent = Json.createObjectNode().set("outer", Json.createObjectNode()
      .put("otherfield", 3).set("mapField", Json.createObjectNode().put("foo", 3).put("bar", 4)));
  parent.put("clientId", "abc123");
  parent.put("otherStrangeIdField", 3);
  List<Field> bqFields = ImmutableList.of(Field.of("client_id", LegacySQLTypeName.STRING), //
      Field.of("outer", LegacySQLTypeName.RECORD, //
          MAP_FIELD));
  Map<String, Object> expected = ImmutableMap.of("client_id", "abc123", "outer",
      ImmutableMap.of("map_field", ImmutableList.of(ImmutableMap.of("key", "foo", "value", 3),
          ImmutableMap.of("key", "bar", "value", 4))));
  Map<String, Object> expectedAdditional = ImmutableMap.of("otherStrangeIdField", 3, "outer",
      ImmutableMap.of("otherfield", 3));
  ObjectNode additionalProperties = Json.createObjectNode();
  TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
  assertEquals(expected, Json.asMap(parent));
  assertEquals(expectedAdditional, Json.asMap(additionalProperties));
}
 

@Test
public void testUnmap() throws Exception {
  ObjectNode parent = Json.createObjectNode().set("mapField",
      Json.createObjectNode().put("foo", 3).put("bar", 4));
  ObjectNode additionalProperties = Json.createObjectNode();
  List<Field> bqFields = ImmutableList.of(MAP_FIELD);
  Map<String, Object> expected = ImmutableMap.of("map_field", ImmutableList
      .of(ImmutableMap.of("key", "foo", "value", 3), ImmutableMap.of("key", "bar", "value", 4)));
  TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
  assertEquals(expected, Json.asMap(parent));
}
 

@Test
public void shouldReturnTableDefinitionIfPartitionDisabled() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(false);
    Schema bqSchema = Schema.of(
            Field.newBuilder("int_field", LegacySQLTypeName.INTEGER).build()
    );

    BQTableDefinition bqTableDefinition = new BQTableDefinition(bqConfig);
    StandardTableDefinition tableDefinition = bqTableDefinition.getTableDefinition(bqSchema);
    Schema returnedSchema = tableDefinition.getSchema();
    assertEquals(returnedSchema.getFields().size(), bqSchema.getFields().size());
    assertEquals(returnedSchema.getFields().get(0).getName(), bqSchema.getFields().get(0).getName());
    assertEquals(returnedSchema.getFields().get(0).getMode(), bqSchema.getFields().get(0).getMode());
    assertEquals(returnedSchema.getFields().get(0).getType(), bqSchema.getFields().get(0).getType());
}
 

private TableDefinition getPartitionedTableDefinition(ArrayList<Field> bqSchemaFields) {
    TimePartitioning.Builder timePartitioningBuilder = TimePartitioning.newBuilder(TimePartitioning.Type.DAY);
    timePartitioningBuilder.setField(bqConfig.getBQTablePartitionKey())
            .setRequirePartitionFilter(true);

    Schema schema = Schema.of(bqSchemaFields);

    TableDefinition tableDefinition = StandardTableDefinition.newBuilder()
            .setSchema(schema)
            .setTimePartitioning(timePartitioningBuilder.build())
            .build();
    return tableDefinition;
}
 

@Test
public void shouldIgnoreExceptionIfDatasetAlreadyExists() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(true);
    when(bqConfig.getBQTablePartitionKey()).thenReturn("partition_column");
    when(bqConfig.getTable()).thenReturn("bq-table");
    when(bqConfig.getDataset()).thenReturn("bq-proto");
    bqClient = new BQClient(bigquery, bqConfig);

    ArrayList<Field> bqSchemaFields = new ArrayList<Field>() {{
        add(Field.newBuilder("test-1", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder("partition_column", LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.OFFSET_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TOPIC_COLUMN_NAME, LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.LOAD_TIME_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TIMESTAMP_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.PARTITION_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
    }};

    TableDefinition tableDefinition = getPartitionedTableDefinition(bqSchemaFields);
    TableId tableId = TableId.of(bqConfig.getDataset(), bqConfig.getTable());
    TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();

    when(bigquery.getDataset(tableId.getDataset())).thenReturn(dataset);
    when(dataset.exists()).thenReturn(false);
    when(table.exists()).thenReturn(false);
    when(bigquery.getTable(tableId)).thenReturn(table);
    when(bigquery.create(tableInfo)).thenReturn(table);

    bqClient.upsertTable(bqSchemaFields);
    verify(bigquery).create(DatasetInfo.of(tableId.getDataset()));
    verify(bigquery).create(tableInfo);
    verify(bigquery, never()).update(tableInfo);
}
 

@Test
public void shouldUpdateTableIfTableAlreadyExistsAndSchemaChanges() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(false);
    when(bqConfig.getTable()).thenReturn("bq-table");
    when(bqConfig.getDataset()).thenReturn("bq-proto");
    bqClient = new BQClient(bigquery, bqConfig);

    ArrayList<Field> bqSchemaFields = new ArrayList<Field>() {{
        add(Field.newBuilder("test-1", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder("test-2", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.OFFSET_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TOPIC_COLUMN_NAME, LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.LOAD_TIME_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TIMESTAMP_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.PARTITION_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
    }};

    TableDefinition tableDefinition = getNonPartitionedTableDefinition(bqSchemaFields);
    ArrayList<Field> updatedBQSchemaFields = new ArrayList<>();
    updatedBQSchemaFields.addAll(bqSchemaFields);
    updatedBQSchemaFields.add(Field.newBuilder("new-field", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
    TableDefinition updatedBQTableDefinition = getNonPartitionedTableDefinition(updatedBQSchemaFields);

    TableId tableId = TableId.of(bqConfig.getDataset(), bqConfig.getTable());
    TableInfo tableInfo = TableInfo.newBuilder(tableId, updatedBQTableDefinition).build();
    when(bigquery.getDataset(tableId.getDataset())).thenReturn(dataset);
    when(dataset.exists()).thenReturn(true);
    when(table.exists()).thenReturn(true);
    when(bigquery.getTable(tableId)).thenReturn(table);
    when(table.getDefinition()).thenReturn(mockTableDefinition);
    when(mockTableDefinition.getSchema()).thenReturn(tableDefinition.getSchema());
    when(bigquery.update(tableInfo)).thenReturn(table);

    bqClient.upsertTable(updatedBQSchemaFields);
    verify(bigquery, never()).create(tableInfo);
    verify(bigquery).update(tableInfo);
}
 

@Test
public void shouldCreateBigqueryTableWithPartition() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(true);
    when(bqConfig.getBQTablePartitionKey()).thenReturn("partition_column");
    when(bqConfig.getTable()).thenReturn("bq-table");
    when(bqConfig.getDataset()).thenReturn("bq-proto");
    bqClient = new BQClient(bigquery, bqConfig);

    ArrayList<Field> bqSchemaFields = new ArrayList<Field>() {{
        add(Field.newBuilder("test-1", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder("partition_column", LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.OFFSET_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TOPIC_COLUMN_NAME, LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.LOAD_TIME_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TIMESTAMP_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.PARTITION_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
    }};
    TableDefinition tableDefinition = getPartitionedTableDefinition(bqSchemaFields);
    TableId tableId = TableId.of(bqConfig.getDataset(), bqConfig.getTable());
    TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();
    when(bigquery.getDataset(tableId.getDataset())).thenReturn(dataset);
    when(dataset.exists()).thenReturn(true);
    when(table.exists()).thenReturn(false);
    when(bigquery.getTable(tableId)).thenReturn(table);
    when(bigquery.create(tableInfo)).thenReturn(table);

    bqClient.upsertTable(bqSchemaFields);
    verify(bigquery).create(tableInfo);
    verify(bigquery, never()).update(tableInfo);
}
 

@Test
public void shouldNotUpdateTableIfTableAlreadyExistsWithSameSchema() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(false);
    when(bqConfig.getTable()).thenReturn("bq-table");
    when(bqConfig.getDataset()).thenReturn("bq-proto");
    bqClient = new BQClient(bigquery, bqConfig);

    ArrayList<Field> bqSchemaFields = new ArrayList<Field>() {{
        add(Field.newBuilder("test-1", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder("test-2", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.OFFSET_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TOPIC_COLUMN_NAME, LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.LOAD_TIME_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.TIMESTAMP_COLUMN_NAME, LegacySQLTypeName.TIMESTAMP).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder(Constants.PARTITION_COLUMN_NAME, LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
    }};

    TableDefinition tableDefinition = getNonPartitionedTableDefinition(bqSchemaFields);

    TableId tableId = TableId.of(bqConfig.getDataset(), bqConfig.getTable());
    TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinition).build();
    when(bigquery.getDataset(tableId.getDataset())).thenReturn(dataset);
    when(dataset.exists()).thenReturn(true);
    when(table.exists()).thenReturn(true);
    when(bigquery.getTable(tableId)).thenReturn(table);
    when(table.getDefinition()).thenReturn(mockTableDefinition);
    when(mockTableDefinition.getSchema()).thenReturn(tableDefinition.getSchema());
    when(table.exists()).thenReturn(true);

    bqClient.upsertTable(bqSchemaFields);
    verify(bigquery, never()).create(tableInfo);
    verify(bigquery, never()).update(tableInfo);
}
 

@BeforeClass
public static void initDatasetAndTable() throws IOException {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    for (String dataset : datasets) {
        DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, dataset);
        bigquery.create(DatasetInfo.of(datasetId));
    }

    for (String table : tables) {
        TableDefinition tableDefinition =
                StandardTableDefinition.of(Schema.of(Field.of("test", LegacySQLTypeName.STRING)));
        TableId tableId = TableId.of(BigQueryTestConstants.PROJECT, datasets.get(0), table);
        bigquery.create(TableInfo.of(tableId, tableDefinition));
    }
}