Java Code Examples for com.google.cloud.bigquery.TableId#of()

@Test
public void shouldRunExtractJobInTestMode() throws Exception {
  final TableId srcTable = TableId.of("foo", "bar", "baz");

  final String destinationUri = "gs://foo/bar";

  final Task<String> task = Task.named("task")
      .ofType(String.class)
      .operator(BigQueryOperator.create())
      .process(bq -> bq.job(
          JobInfo.of(ExtractJobConfiguration.of(srcTable, destinationUri)))
          .success(response -> destinationUri));

  try (TestScope scope = FloTesting.scope()) {

    final String result = FloRunner.runTask(task).future()
        .get(30, SECONDS);

    assertThat(result, is(destinationUri));
  }
}
 

@Override
TableInfo parse(String... args) throws Exception {
  String message;
  if (args.length == 3) {
    String dataset = args[0];
    String table = args[1];
    String query = args[2];
    TableId tableId = TableId.of(dataset, table);
    return TableInfo.of(tableId, ViewDefinition.of(query));
  } else if (args.length < 3) {
    message = "Missing required dataset id, table id or query.";
  } else {
    message = "Too many arguments.";
  }
  throw new IllegalArgumentException(message);
}
 

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

/** Example of deleting a table. */
// [TARGET delete(String, String)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public boolean deleteTable(String datasetName, String tableName) {
  // [START ]
  TableId tableId = TableId.of(datasetName, tableName);
  boolean deleted = bigquery.delete(tableId);
  if (deleted) {
    // the table was deleted
  } else {
    // the table was not found
  }
  // [END ]
  return deleted;
}
 

@Ignore
@Test
public void shouldPushMessagesToBqActual() {
    TableId tableId = TableId.of("bqsinktest", "users");
    BqSink bqSink = new BqSink(authenticatedBQ(), tableId, new BQResponseParser(), gcsSinkHandler, bqRow);

    HashMap<String, Object> columns = new HashMap<>();
    columns.put("name", "someone_else");
    columns.put("age", 26);
    columns.put("location", 15123);
    columns.put("created_at", new DateTime(new Date()));

    HashMap<String, Object> route1 = new HashMap<>();
    route1.put("name", "route1");
    route1.put("position", "1");
    HashMap<String, Object> route2 = new HashMap<>();
    route2.put("name", "route2");
    route2.put("position", "2");

    columns.put("routes", Arrays.asList(route1, route2));

    Status push = bqSink.push(new Records(Arrays.asList(new Record(new OffsetInfo("default-topic", 0, 0, Instant.now().toEpochMilli()), columns))));

    assertTrue(push.isSuccess());
}
 

@Before
public void setUp() {
    tableId = TableId.of("test-dataset", "test-table");
    builder = InsertAllRequest.newBuilder(tableId);
    bqRow = new BQRowWithInsertId();
    BQErrorHandler errorHandlerInstance = new OOBErrorHandler(new DefaultLogWriter());
    sink = new BqSink(bigquery, tableId, new BQResponseParser(), errorHandlerInstance, bqRow);
    when(successfulResponse.hasErrors()).thenReturn(false);
    when(bigquery.insertAll(any())).thenReturn(successfulResponse);
    when(failureResponse.hasErrors()).thenReturn(true);
    insertErrors = new HashMap<>();
    List<BigQueryError> columnError = Arrays.asList(new BigQueryError("failed since type mismatched", "column location", "message"));
    insertErrors.put(0L, columnError);
    when(failureResponse.getInsertErrors()).thenReturn(insertErrors);
}
 

/**
 * Given a KV containing a destination and a message, return the message content as a {@link
 * TableRow} ready to pass to {@link org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO}.
 */
public TableRow kvToTableRow(KV<TableDestination, PubsubMessage> kv) {
  if (format == null) {
    format = createFormat();
  }
  final TableReference ref = kv.getKey().getTableReference();
  final TableId tableId = TableId.of(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
  final PubsubMessage message = kv.getValue();
  return Json.asTableRow(format.apply(tableId, message.getAttributeMap(), message.getPayload()));
}
 

@Test
public void canWriteViaFileLoads() throws Exception {
  String table = "my_test_table";
  TableId tableId = TableId.of(dataset, table);

  bigquery.create(DatasetInfo.newBuilder(dataset).build());
  bigquery
      .create(TableInfo
          .newBuilder(tableId,
              StandardTableDefinition
                  .of(Schema.of(Field.of("client_id", LegacySQLTypeName.STRING),
                      Field.of("type", LegacySQLTypeName.STRING),
                      Field.of("submission_timestamp", LegacySQLTypeName.TIMESTAMP)))
                  .toBuilder().setTimePartitioning(TIME_PARTITIONING).setClustering(CLUSTERING)
                  .build())
          .build());

  String input = Resources
      .getResource("testdata/bigquery-integration/input-varied-doctypes.ndjson").getPath();
  String output = String.format("%s:%s.%s", projectId, dataset, "${document_type}_table");
  String errorOutput = outputPath + "/error/out";

  PipelineResult result = Sink.run(new String[] { "--inputFileFormat=json", "--inputType=file",
      "--input=" + input, "--outputType=bigquery", "--output=" + output,
      "--bqWriteMethod=file_loads", "--errorOutputType=file",
      "--tempLocation=gs://gcp-ingestion-static-test-bucket/temp/bq-loads",
      "--schemasLocation=schemas.tar.gz", "--errorOutputFileCompression=UNCOMPRESSED",
      "--errorOutput=" + errorOutput });

  result.waitUntilFinish();

  String tableSpec = String.format("%s.%s", dataset, table);
  assertThat(stringValuesQueryWithRetries("SELECT client_id FROM " + tableSpec),
      matchesInAnyOrder(ImmutableList.of("abc123")));

  List<String> errorOutputLines = Lines.files(outputPath + "/error/out*.ndjson");
  assertThat(errorOutputLines, Matchers.hasSize(2));
}
 

/** Example of writing a local file to a table. */
// [TARGET writer(WriteChannelConfiguration)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
// [VARIABLE FileSystems.getDefault().getPath(".", "my-data.csv")]
// [VARIABLE "us"]
public long writeFileToTable(String datasetName, String tableName, Path csvPath, String location)
    throws IOException, InterruptedException, TimeoutException {
  // [START bigquery_load_from_file]
  TableId tableId = TableId.of(datasetName, tableName);
  WriteChannelConfiguration writeChannelConfiguration =
      WriteChannelConfiguration.newBuilder(tableId).setFormatOptions(FormatOptions.csv()).build();
  // Generally, location can be inferred based on the location of the referenced dataset.
  // However,
  // it can also be set explicitly to force job execution to be routed to a specific processing
  // location.  See https://cloud.google.com/bigquery/docs/locations for more info.
  JobId jobId = JobId.newBuilder().setLocation(location).build();
  TableDataWriteChannel writer = bigquery.writer(jobId, writeChannelConfiguration);
  // Write data to writer
  try (OutputStream stream = Channels.newOutputStream(writer)) {
    Files.copy(csvPath, stream);
  } finally {
    writer.close();
  }
  // Get load job
  Job job = writer.getJob();
  job = job.waitFor();
  LoadStatistics stats = job.getStatistics();
  return stats.getOutputRows();
  // [END bigquery_load_from_file]
}
 

@Ignore
@Test
public void shouldPushNestedMessage() {
    Instant now = Instant.now();
    long second = now.getEpochSecond();
    int nano = now.getNano();
    Timestamp createdAt = Timestamp.newBuilder().setSeconds(second).setNanos(nano).build();
    TestMessage testMessage = TestMessage.newBuilder()
            .setOrderNumber("order-1")
            .setOrderUrl("order-url")
            .setOrderDetails("order-details")
            .setCreatedAt(createdAt)
            .setStatus(com.gojek.beast.Status.COMPLETED)
            .build();
    ProtoParser protoParser = new ProtoParser(StencilClientFactory.getClient(), TestNestedMessage.class.getName());
    TestNestedMessage nestedMsg = TestNestedMessage.newBuilder()
            .setSingleMessage(testMessage)
            .setNestedId("nested-id")
            .build();
    TableId tableId = TableId.of("bqsinktest", "test_nested_messages");
    BqSink bqSink = new BqSink(authenticatedBQ(), tableId, new BQResponseParser(), gcsSinkHandler, bqRow);


    OffsetInfo offsetInfo = new OffsetInfo("topic", 1, 1, Instant.now().toEpochMilli());
    Map<String, Object> columns = new HashMap<>();
    HashMap<String, Object> nested = new HashMap<>();
    nested.put("order_number", nestedMsg.getSingleMessage().getOrderNumber());
    nested.put("order_url", nestedMsg.getSingleMessage().getOrderUrl());

    columns.put("id", nestedMsg.getNestedId());
    columns.put("msg", nested);

    Status push = bqSink.push(new Records(Arrays.asList(new Record(offsetInfo, columns))));
    assertTrue(push.isSuccess());
}
 

@Test
public void lookupLatestDailyShouldReturnTableIdForExistingTableInLookBackRange() throws Exception {
  final TableId expected = TableId.of("foo", "bar", "baz_20180105");
  final Task<TableId> lookup = BigQueryTasks.lookupLatestDaily(() -> {
    FloBigQueryClient bq = mock(FloBigQueryClient.class);
    when(bq.tableExists(expected)).thenReturn(true);
    return bq;
  }, "foo", "bar", "baz", Date.parse("2018-01-10"), 7);
  final TableId tableId = FloRunner.runTask(lookup)
      .future().get(30, TimeUnit.SECONDS);
  assertThat(tableId, is(expected));
}
 

/**
 * Get the schema by table name or query. This method also needed for read and write, because we can not get schema
 * from the ongoing data. BigQueryIO.Read return TableRow, which do not include schema in itself. So use BigQuery
 * client to get it before read and write.
 * 
 * @return
 */
@Override
public Schema getSchema() {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    com.google.cloud.bigquery.Schema bqRowSchema = null;
    switch (properties.sourceType.getValue()) {
    case TABLE_NAME: {
        TableId tableId = TableId.of(properties.getDatastoreProperties().projectName.getValue(),
                properties.bqDataset.getValue(), properties.tableName.getValue());
        Table table = bigquery.getTable(tableId);
        if (table == null) {
            ComponentException.build(CommonErrorCodes.UNEXPECTED_EXCEPTION).setAndThrow(
                    "Table not found:" + tableId.toString());
        }
        bqRowSchema = table.getDefinition().getSchema();
        break;
    }
    case QUERY: {
        QueryJobConfiguration queryRequest = QueryJobConfiguration
                .newBuilder(properties.query.getValue())
                .setUseLegacySql(properties.useLegacySql.getValue())
                .build();
        TableResult queryResponse =
                query(bigquery, queryRequest, properties.getDatastoreProperties().projectName.getValue());
        bqRowSchema = queryResponse.getSchema();
        break;
    }
    default:
        throw new RuntimeException("To be implemented: " + properties.sourceType.getValue());
    }
    return BigQueryAvroRegistry.get().inferSchema(bqRowSchema);
}
 

@Override
JobInfo parse(String... args) throws Exception {
  if (args.length >= 4) {
    String dataset = args[0];
    String table = args[1];
    String format = args[2];
    TableId tableId = TableId.of(dataset, table);
    ExtractJobConfiguration configuration =
        ExtractJobConfiguration.of(
            tableId, Arrays.asList(args).subList(3, args.length), format);
    return JobInfo.of(configuration);
  }
  throw new IllegalArgumentException("Missing required arguments.");
}
 

/** Inserts multiple rows of the same schema to a BigQuery table. */
public void insertAll(Collection<Map<String, ?>> rows, String table) {
  TableId tableId = TableId.of(projectId, dataset, table);

  InsertAllRequest.Builder builder = InsertAllRequest.newBuilder(tableId);

  for (Map<String, ?> row : rows) {
    builder.addRow(row);
  }

  InsertAllResponse response = client.insertAll(builder.build());
  handleBigQueryResponseExceptions(response);
}
 

@Test
public void lookupShouldReturnTableIdForExistingTable() throws Exception {
  final TableId expected = TableId.of("foo", "bar", "baz");
  final Task<TableId> lookup = BigQueryTasks.lookup(() -> {
    FloBigQueryClient bq = mock(FloBigQueryClient.class);
    when(bq.tableExists(expected)).thenReturn(true);
    return bq;
  }, expected);
  final TableId tableId = FloRunner.runTask(lookup)
      .future().get(30, TimeUnit.SECONDS);
  assertThat(tableId, is(expected));
}
 

private TableId getCopyTableId() {
  return TableId.of(COPY_DATASET_NAME, getCopyTableName());
}
 

@Override
public TableId getTableId(GenericRecord input) {
  return TableId.of(datasetProvider.get(), tableNameProvider.get());
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final String projectId = context.getProperty(PROJECT_ID).evaluateAttributeExpressions().getValue();
    final String dataset = context.getProperty(DATASET).evaluateAttributeExpressions(flowFile).getValue();
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();

    final TableId tableId;
    if (StringUtils.isEmpty(projectId)) {
        tableId = TableId.of(dataset, tableName);
    } else {
        tableId = TableId.of(projectId, dataset, tableName);
    }

    try {

        InsertAllRequest.Builder request = InsertAllRequest.newBuilder(tableId);
        int nbrecord = 0;

        try (final InputStream in = session.read(flowFile)) {
            final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
            try (final RecordReader reader = readerFactory.createRecordReader(flowFile, in, getLogger());) {
                Record currentRecord;
                while ((currentRecord = reader.nextRecord()) != null) {
                    request.addRow(convertMapRecord(currentRecord.toMap()));
                    nbrecord++;
                }
            }
        }

        request.setIgnoreUnknownValues(context.getProperty(IGNORE_UNKNOWN).evaluateAttributeExpressions(flowFile).asBoolean());
        request.setSkipInvalidRows(context.getProperty(SKIP_INVALID_ROWS).evaluateAttributeExpressions(flowFile).asBoolean());

        InsertAllResponse response = getCloudService().insertAll(request.build());

        final Map<String, String> attributes = new HashMap<>();

        if (response.hasErrors()) {
            getLogger().log(LogLevel.WARN, "Failed to insert {} of {} records into BigQuery {} table.", new Object[] { response.getInsertErrors().size(), nbrecord, tableName });
            if (getLogger().isDebugEnabled()) {
                for (long index : response.getInsertErrors().keySet()) {
                    for (BigQueryError e : response.getInsertErrors().get(index)) {
                        getLogger().log(LogLevel.DEBUG, "Failed to insert record #{}: {}", new Object[] { index, e.getMessage() });
                    }
                }
            }

            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord - response.getInsertErrors().size()));

            flowFile = session.penalize(flowFile);
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_FAILURE);
        } else {
            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord));
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_SUCCESS);
        }

    } catch (Exception ex) {
        getLogger().log(LogLevel.ERROR, ex.getMessage(), ex);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}
 

public BQClient(BigQuery bigquery, BQConfig bqConfig) {
    this.bigquery = bigquery;
    this.bqConfig = bqConfig;
    this.tableID = TableId.of(bqConfig.getDataset(), bqConfig.getTable());
    this.bqTableDefinition = new BQTableDefinition(bqConfig);
}
 

@Test
public void testCreateGetAndDeleteTable() throws InterruptedException {
  String tableName = "test_create_get_delete";
  String fieldName = "aField";
  Table table = bigquerySnippets.createTable(DATASET, tableName, fieldName);
  assertNotNull(table);
  TableId tableId = TableId.of(bigquery.getOptions().getProjectId(), DATASET, tableName);
  assertEquals(
      tableId, bigquerySnippets.getTable(tableId.getDataset(), tableId.getTable()).getTableId());
  assertNotNull(bigquerySnippets.updateTableDescription(DATASET, tableName, "new description"));
  table = bigquerySnippets.updateTableExpiration(DATASET, tableName);
  assertNotNull(table.getExpirationTime());
  assertEquals(
      "new description",
      bigquerySnippets
          .getTableFromId(tableId.getProject(), tableId.getDataset(), tableId.getTable())
          .getDescription());
  Set<TableId> tables =
      Sets.newHashSet(
          Iterators.transform(
              bigquerySnippets.listTables(DATASET).iterateAll().iterator(),
              TO_TABLE_ID_FUNCTION));
  while (!tables.contains(tableId)) {
    Thread.sleep(500);
    tables =
        Sets.newHashSet(
            Iterators.transform(
                bigquerySnippets.listTables(DATASET).iterateAll().iterator(),
                TO_TABLE_ID_FUNCTION));
  }
  tables =
      Sets.newHashSet(
          Iterators.transform(
              bigquerySnippets
                  .listTablesFromId(tableId.getProject(), DATASET)
                  .iterateAll()
                  .iterator(),
              TO_TABLE_ID_FUNCTION));
  while (!tables.contains(tableId)) {
    Thread.sleep(500);
    tables =
        Sets.newHashSet(
            Iterators.transform(
                bigquerySnippets
                    .listTablesFromId(tableId.getProject(), DATASET)
                    .iterateAll()
                    .iterator(),
                TO_TABLE_ID_FUNCTION));
  }
  assertTrue(bigquerySnippets.deleteTable(DATASET, tableName));
  assertFalse(bigquerySnippets.deleteTableFromId(tableId.getProject(), DATASET, tableName));
}