step.datapool.gsheet.GoogleSheetv4DataPool Maven / Gradle / Ivy

Go to download
/*******************************************************************************
 * Copyright (C) 2020, exense GmbH
 *  
 * This file is part of STEP
 *  
 * STEP is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *  
 * STEP is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *  
 * You should have received a copy of the GNU Affero General Public License
 * along with STEP.  If not, see .
 ******************************************************************************/
package step.datapool.gsheet;

import java.io.FileInputStream;
import java.io.IOException;
import java.security.GeneralSecurityException;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
import com.google.api.client.json.jackson2.JacksonFactory;
import com.google.api.services.drive.DriveScopes;
import com.google.api.services.sheets.v4.Sheets;
import com.google.api.services.sheets.v4.SheetsScopes;
import com.google.api.services.sheets.v4.model.ValueRange;

import step.attachments.FileResolver;
import step.core.variables.SimpleStringMap;
import step.datapool.DataSet;

/**
 * @author doriancransac
 * @author Jérôme Comte
 *
 */
public class GoogleSheetv4DataPool extends DataSet {

	FileResolver fileResolver;
	
	private String saKey;
	private String fileId;
	private String tabName;

	List> datapool;

	List headers;

	int cursor;

	public GoogleSheetv4DataPool(GoogleSheetv4DataPoolConfiguration configuration) {
		super(configuration);
		saKey = configuration.getServiceAccountKey().get();
		fileId = configuration.getFileId().get();
		tabName = configuration.getTabName().get();
	}

	Sheets service = null;

	@Override
	public void init() {
		super.init();
		fileResolver = context.getFileResolver();
		createDatapool(saKey, fileId, tabName);
		createHeaders();
		this.cursor = 1; // skip headers
	}

	private void createHeaders() {
		this.headers = new LinkedList<>();
		for(Object o : datapool.get(0)) // 0 = headers row
			this.headers.add((String)o);
	}

	private int getIndexForHeader(String header) {
		return this.headers.indexOf(header);
	}

	@Override
	public void reset() {
		init();
	}

	@Override
	public void close() {
		super.close();
		//nothing to do, apparently the Sheet service doesn't need to be closed explicitely
	}

	@Override
	public Object next_() {
		if(cursor >= this.datapool.size())
			return null;
		RowWrapper row = new RowWrapper(cursor);
		cursor++;
		return row;
	}

	@Override
	public void addRow(Object row) {
		throw new RuntimeException("Not implemented");
	}

	private class RowWrapper extends SimpleStringMap{

		private int cursor;

		public RowWrapper(int cursor){
			this.cursor = cursor;
		}

		@Override
		public int size() {
			return datapool.get(cursor).size();
		}

		@Override
		public boolean isEmpty() {
			return datapool.get(cursor).isEmpty();
		}

		//TODO: handle other types?
		@Override
		public String get(String key) {
			if(getIndexForHeader(key) < 0)
				throw new RuntimeException("Column unknown:" + key);
			return (String) datapool.get(cursor).get(getIndexForHeader(key));
		}


		//TODO:easy to implement by turning the datapool into List>
		@Override
		public String put(String key, String value) {
			throw new RuntimeException("Not implemented.");
		}

		@Override
		public Set keySet() {
			return new LinkedHashSet<>(headers);
		}

	}

	public void createDatapool(String saKey, String fileId, String tabName){
		try {
			this.service = buildService(getCredential(saKey));
			this.datapool = getValuesForRange(fileId, tabName);
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	public GoogleCredential getCredential(String saKey) throws Exception {
		Set all = new HashSet<>();
		all.addAll(SheetsScopes.all());
		all.addAll(DriveScopes.all());

		GoogleCredential gCred = null;
		FileInputStream keyFileIS = null;
		try{
			keyFileIS = new FileInputStream(this.fileResolver.resolve(saKey));
			gCred = GoogleCredential.fromStream(keyFileIS)
					.createScoped(all);
		}finally{
			if(keyFileIS != null)
				keyFileIS.close();
		}

		return gCred;
	}

	public Sheets buildService(GoogleCredential credential) throws IOException,
	GeneralSecurityException {
		return new Sheets.Builder(
				GoogleNetHttpTransport.newTrustedTransport(),
				JacksonFactory.getDefaultInstance(),
				credential)
				.setApplicationName("Sheets API Snippets")
				.build();
	}

	public List> getValuesForRange(String spreadsheetId, String range) throws IOException {
		ValueRange result = service.spreadsheets().values().get(spreadsheetId, range).execute();
		return result.getValues();
	}
}