BigQueryPullServiceImpl.java

/*
 * This file is a component of thundr, a software library from 3wks.
 * Read more: http://www.3wks.com.au/thundr
 * Copyright (C) 2013 3wks, <thundr@3wks.com.au>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.threewks.thundr.bigquery;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.Bigquery.Jobs.Insert;
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobConfiguration;
import com.google.api.services.bigquery.model.JobConfigurationQuery;
import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.QueryRequest;
import com.google.api.services.bigquery.model.QueryResponse;
import com.google.api.services.bigquery.model.TableCell;
import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
import com.threewks.thundr.logger.Logger;

public class BigQueryPullServiceImpl implements BigQueryPullService {

	private Bigquery bigQuery;
	private String projectId;

	public BigQueryPullServiceImpl(Bigquery bigQuery, String bigQueryProjectId) {
		this.bigQuery = bigQuery;
		this.projectId = bigQueryProjectId;
	}

	@Override
	public List<List<String>> query(String query, boolean includeColumnNames) throws IOException {

		Logger.debug("Running query %s", query);

		QueryRequest queryRequest = new QueryRequest().setQuery(query);
		QueryResponse response = bigQuery.jobs().query(projectId, queryRequest).execute();

		if (response.getJobComplete()) {
			JobReference jobReference = response.getJobReference();
			if (jobReference != null) {
				Logger.debug("Successfully ran query for job id %s", jobReference.getJobId());
			}
		}

		return buildQueryResults(response.getSchema(), response.getRows(), includeColumnNames);
	}

	@Override
	public String queryAsync(String query) throws IOException {
		Job job = new Job();
		JobConfiguration config = new JobConfiguration();
		JobConfigurationQuery queryConfig = new JobConfigurationQuery();
		config.setQuery(queryConfig);

		job.setConfiguration(config);
		queryConfig.setQuery(query);

		Insert insert = bigQuery.jobs().insert(projectId, job);
		insert.setProjectId(projectId);
		JobReference jobReference = insert.execute().getJobReference();

		Logger.debug("Submitted job %s to BigQuery", jobReference.getJobId());
		return jobReference.getJobId();
	}

	@Override
	public String getJobStatus(String jobId) throws IOException {
		Job job = bigQuery.jobs().get(projectId, jobId).execute();
		return job.getStatus().getState();
	}

	@Override
	public List<List<String>> getAsyncQueryResults(String jobId, boolean includeColumnNames) throws IOException {
		Job job = bigQuery.jobs().get(projectId, jobId).execute();
		String state = job.getStatus().getState();
		if (state.equals("DONE")) {
			GetQueryResultsResponse response = bigQuery.jobs().getQueryResults(projectId, job.getJobReference().getJobId()).execute();
			TableSchema schema = response.getSchema();
			List<TableRow> rows = response.getRows();
			return buildQueryResults(schema, rows, includeColumnNames);
		}

		// job is not done yet, return null
		return null;
	}

	/**
	 * Build a list of results.
	 * 
	 * @param schema the table schema for the results.
	 * @param rows the rows of results.
	 * @param includeColumnNames a boolean indicating whether the column names should be included in the results or not.
	 * @return the results as a list of list of strings.
	 */
	private List<List<String>> buildQueryResults(TableSchema schema, List<TableRow> rows, boolean includeColumnNames) {
		List<List<String>> results = new ArrayList<List<String>>();
		if (includeColumnNames) {
			results.add(getColumnNames(schema));
		}
		results.addAll(getResultData(rows));
		return results;
	}

	/**
	 * Get a list of the column names from a {@link TableSchema}.
	 * 
	 * @param schema the table schema to extract the column names from.
	 * @return a list of strings.
	 */
	private List<String> getColumnNames(TableSchema schema) {
		List<TableFieldSchema> tableFieldSchema = schema.getFields();
		List<String> columnNames = new ArrayList<String>();
		for (TableFieldSchema fieldSchema : tableFieldSchema) {
			columnNames.add(fieldSchema.getName());
		}
		return columnNames;
	}

	/**
	 * Extract the result data from a list of {@link TableRow} objects.
	 * 
	 * @param rows a list of table rows to extract the result data from.
	 * @return a list of rows (represented as lists of strings).
	 */
	private List<List<String>> getResultData(List<TableRow> rows) {
		List<List<String>> results = new ArrayList<List<String>>();
		if (rows != null) {
			for (TableRow tableRow : rows) {
				List<String> row = new ArrayList<String>();
				List<TableCell> fields = tableRow.getF();
				for (TableCell cell : fields) {
					row.add((String) cell.getV());
				}
				results.add(row);
			}
		}
		return results;
	}

}