com.google.cloud.bigquery.LegacySQLTypeName 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 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));
}
 

@Test
public void testCoerceMapValueToString() throws Exception {
  String mainPing = "{\"payload\":{\"processes\":{\"parent\":{\"scalars\":"
      + "{\"timestamps.first_paint\":5405}}}}}";
  ObjectNode parent = Json.readObjectNode(mainPing);
  ObjectNode additionalProperties = Json.createObjectNode();
  parent.put("64bit", true);
  parent.put("hi-fi", true);
  List<Field> bqFields = ImmutableList.of(Field.of("payload", LegacySQLTypeName.RECORD,
      Field.of("processes", LegacySQLTypeName.RECORD,
          Field.of("parent", LegacySQLTypeName.RECORD,
              Field.newBuilder("scalars", LegacySQLTypeName.RECORD, //
                  Field.of("key", LegacySQLTypeName.STRING), //
                  Field.of("value", LegacySQLTypeName.STRING)) //
                  .setMode(Mode.REPEATED).build()))));
  Map<String, Object> expected = ImmutableMap.of("payload",
      ImmutableMap.of("processes", ImmutableMap.of("parent", ImmutableMap.of("scalars",
          ImmutableList.of(ImmutableMap.of("key", "timestamps.first_paint", "value", "5405"))))));
  TRANSFORM.transformForBqSchema(parent, bqFields, additionalProperties);
  assertEquals(expected, Json.asMap(parent));
}
 

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 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));
}
 

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();
}
 

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

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

@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));
}
 

/**
 * Extracts and applies new column information to BigQuery by comparing the TableRow against the
 * BigQuery Table. Retries the supplied number of times before failing.
 *
 * @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.
 * @param retries Number of remaining retries before error is raised.
 */
private void applyMapperToTableRow(
    TableId tableId, TableRow row, Map<String, LegacySQLTypeName> inputSchema, int retries) {
  try {
    updateTableIfRequired(tableId, row, inputSchema);
  } catch (Exception e) {
    if (retries > 0) {
      LOG.info("RETRY TABLE UPDATE - enter: {}", String.valueOf(retries));
      try {
        Thread.sleep(2000);
      } catch (InterruptedException i) {
        throw e;
      }
      LOG.info("RETRY TABLE UPDATE - apply: {}", String.valueOf(retries));
      applyMapperToTableRow(tableId, row, inputSchema, retries - 1);
    } else {
      LOG.info("RETRY TABLE UPDATE - throw: {}", String.valueOf(retries));
      throw e;
    }
  }
}
 

/**
 * 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);
  }
}
 

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 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;
}
 

@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());
}
 

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

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

/** Example of creating a table in the dataset with schema and time partitioning. */
// [TARGET create(String, TableDefinition, TableOption...)]
// [VARIABLE “my_table”]
// [VARIABLE “my_field”]
public Table createTable(String tableName, String fieldName) {
  // [START ]
  Schema schema = Schema.of(Field.of(fieldName, LegacySQLTypeName.STRING));
  StandardTableDefinition definition =
      StandardTableDefinition.newBuilder()
          .setSchema(schema)
          .setTimePartitioning(TimePartitioning.of(TimePartitioning.Type.DAY))
          .build();
  Table table = dataset.create(tableName, definition);
  // [END ]
  return table;
}
 

@Override
public PCollection<OutputT> expand(PCollection<InputT> tableKVPCollection) {
  return tableKVPCollection.apply(
      "TableRowExtractDestination",
      MapElements.via(
          new SimpleFunction<InputT, OutputT>() {
            @Override
            public OutputT apply(InputT input) {
              /*
                  We run validation against every event to ensure all columns
                  exist in source.
                  If a column is in the event and not in BigQuery,
                  the column is added to the table before the event can continue.
              */
              setUp();
              TableId tableId = getTableId(input);
              TableRow row = getTableRow(input);
              Map<String, LegacySQLTypeName> inputSchema = getObjectSchema(input);
              int retries = getMapperRetries();

              applyMapperToTableRow(tableId, row, inputSchema, retries);
              return getOutputObject(input);
            }
          }));
}
 

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;
}
 

@Test
public void testReadBigQuerySchema() throws Exception {
  Schema schema = Json.readBigQuerySchema(
      "[{\"mode\":\"NULLABLE\",\"name\":\"document_id\",\"type\": \"STRING\"}]"
          .getBytes(StandardCharsets.UTF_8));
  assertEquals(LegacySQLTypeName.STRING, schema.getFields().get(0).getType());
}
 

@Test
public void testTupleIntoStructNested() throws Exception {
  ObjectNode additionalProperties = Json.createObjectNode();
  ObjectNode parent = Json.readObjectNode("{\n" //
      + "  \"payload\": [[1],[2],[3]]\n" //
      + "}\n");
  List<Field> bqFields = ImmutableList.of(Field.newBuilder("payload", LegacySQLTypeName.RECORD, //
      Field.of("f0_", LegacySQLTypeName.INTEGER)) //
      .setMode(Mode.REPEATED).build());
  Map<String, Object> expected = Json
      .readMap("{\"payload\":[{\"f0_\":1},{\"f0_\":2},{\"f0_\":3}]}");
  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());
}
 

@Test(expected = UnsupportedOperationException.class)
public void shouldThrowUnsupportedExceptionForRangePartition() {
    when(bqConfig.isBQTablePartitioningEnabled()).thenReturn(true);
    when(bqConfig.getBQTablePartitionKey()).thenReturn("int_field");

    Schema bqSchema = Schema.of(
            Field.newBuilder("int_field", LegacySQLTypeName.INTEGER).build()
    );

    BQTableDefinition bqTableDefinition = new BQTableDefinition(bqConfig);
    bqTableDefinition.getTableDefinition(bqSchema);
}
 

@Test
public void testInsertAllAndListTableData() throws IOException, InterruptedException {
  String tableName = "test_insert_all_and_list_table_data";
  String fieldName1 = "booleanField";
  String fieldName2 = "bytesField";
  String fieldName3 = "recordField";
  String fieldName4 = "stringField";
  TableId tableId = TableId.of(DATASET, tableName);
  Schema schema =
      Schema.of(
          Field.of(fieldName1, LegacySQLTypeName.BOOLEAN),
          Field.of(fieldName2, LegacySQLTypeName.BYTES),
          Field.of(
              fieldName3,
              LegacySQLTypeName.RECORD,
              Field.of(fieldName4, LegacySQLTypeName.STRING)));
  TableInfo table = TableInfo.of(tableId, StandardTableDefinition.of(schema));
  assertNotNull(bigquery.create(table));
  InsertAllResponse response = bigquerySnippets.insertAll(DATASET, tableName);
  assertFalse(response.hasErrors());
  assertTrue(response.getInsertErrors().isEmpty());
  Page<FieldValueList> listPage = bigquerySnippets.listTableDataFromId(DATASET, tableName);
  while (Iterators.size(listPage.iterateAll().iterator()) < 1) {
    Thread.sleep(500);
    listPage = bigquerySnippets.listTableDataFromId(DATASET, tableName);
  }
  FieldValueList row = listPage.getValues().iterator().next();
  assertEquals(true, row.get(0).getBooleanValue());
  assertArrayEquals(new byte[] {0xA, 0xD, 0xD, 0xE, 0xD}, row.get(1).getBytesValue());
  assertEquals("Hello, World!", row.get(2).getRecordValue().get(0).getStringValue());

  listPage = bigquerySnippets.listTableDataSchema(DATASET, tableName, schema, fieldName1);
  row = listPage.getValues().iterator().next();
  assertNotNull(row.get(fieldName1));
  assertArrayEquals(new byte[] {0xA, 0xD, 0xD, 0xE, 0xD}, row.get(fieldName2).getBytesValue());

  bigquerySnippets.listTableDataSchemaId();

  assertTrue(bigquerySnippets.deleteTable(DATASET, tableName));
}
 

@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);
}
 

private Map<String, LegacySQLTypeName> getObjectSchema(InputT input) {
  Map<String, LegacySQLTypeName> inputSchema = getInputSchema(input);
  if (this.defaultSchema != null) {
    inputSchema.putAll(this.defaultSchema);
  }

  return inputSchema;
}
 

@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 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(expected = BQTableUpdateFailure.class)
public void shouldThrowExceptionIfConverterFails() throws IOException {
    ProtoField returnedProtoField = new ProtoField();
    when(protoFieldFactory.getProtoField()).thenReturn(returnedProtoField);
    returnedProtoField.addField(new ProtoField("order_number", 1));
    returnedProtoField.addField(new ProtoField("order_url", 2));

    HashMap<String, DescriptorAndTypeName> descriptorsMap = new HashMap<String, DescriptorAndTypeName>() {{
        put(String.format("%s.%s", TestKey.class.getPackage(), TestKey.class.getName()), new DescriptorAndTypeName(TestKey.getDescriptor(), String.format(".%s.%s", TestKey.getDescriptor().getFile().getPackage(), TestKey.getDescriptor().getName())));
    }};
    when(protoMappingParser.parseFields(returnedProtoField, stencilConfig.getProtoSchema(), StencilUtils.getAllProtobufDescriptors(descriptorsMap), StencilUtils.getTypeNameToPackageNameMap(descriptorsMap))).thenReturn(returnedProtoField);
    ObjectNode objNode = JsonNodeFactory.instance.objectNode();
    objNode.put("1", "order_number");
    objNode.put("2", "order_url");
    String expectedProtoMapping = objectMapper.writeValueAsString(objNode);
    when(protoMappingConverter.generateColumnMappings(returnedProtoField.getFields())).thenReturn(expectedProtoMapping);

    ArrayList<Field> returnedSchemaFields = new ArrayList<Field>() {{
        add(Field.newBuilder("order_number", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder("order_url", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build());
    }};
    when(protoMappingConverter.generateBigquerySchema(returnedProtoField)).thenReturn(returnedSchemaFields);

    ArrayList<Field> bqSchemaFields = new ArrayList<Field>() {{
        add(Field.newBuilder("order_number", LegacySQLTypeName.INTEGER).setMode(Field.Mode.NULLABLE).build());
        add(Field.newBuilder("order_url", 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());
    }};
    doThrow(new BigQueryException(10, "bigquery mapping has failed")).when(bqInstance).upsertTable(bqSchemaFields);

    protoUpdateListener.onProtoUpdate(stencilConfig.getStencilUrl(), descriptorsMap);
}
 

/** Example of loading a newline-delimited-json file with textual fields from GCS to a table. */
// [TARGET create(JobInfo, JobOption...)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public Long writeRemoteFileToTable(String datasetName, String tableName)
    throws InterruptedException {
  // [START bigquery_load_table_gcs_json]
  String sourceUri = "gs://cloud-samples-data/bigquery/us-states/us-states.json";
  TableId tableId = TableId.of(datasetName, tableName);
  // Table field definition
  Field[] fields =
      new Field[] {
        Field.of("name", LegacySQLTypeName.STRING),
        Field.of("post_abbr", LegacySQLTypeName.STRING)
      };
  // Table schema definition
  Schema schema = Schema.of(fields);
  LoadJobConfiguration configuration =
      LoadJobConfiguration.builder(tableId, sourceUri)
          .setFormatOptions(FormatOptions.json())
          .setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
          .setSchema(schema)
          .build();
  // Load the table
  Job loadJob = bigquery.create(JobInfo.of(configuration));
  loadJob = loadJob.waitFor();
  // Check the table
  System.out.println("State: " + loadJob.getStatus().getState());
  return ((StandardTableDefinition) bigquery.getTable(tableId).getDefinition()).getNumRows();
  // [END bigquery_load_table_gcs_json]
}
 

@Test
public void compareBQSchemaFieldIfSchemaIsNotChanged() {
    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());
    }};
    boolean areEqual = BQUtils.compareBQSchemaFields(Schema.of(bqSchemaFields), Schema.of(bqSchemaFields));
    assertTrue(areEqual);
}