CassandraUpdateExecution.java

/*
 * JBoss, Home of Professional Open Source.
 * See the COPYRIGHT.txt file distributed with this work for information
 * regarding copyright ownership.  Some portions may be licensed
 * to Red Hat, Inc. under one or more contributor license agreements.
 * 
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * 
 * This library 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
 * Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301 USA.
 */
package org.teiid.translator.cassandra;

import java.nio.ByteBuffer;
import java.sql.Blob;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import org.teiid.core.types.BinaryType;
import org.teiid.language.BatchedCommand;
import org.teiid.language.BatchedUpdates;
import org.teiid.language.Command;
import org.teiid.logging.LogConstants;
import org.teiid.logging.LogManager;
import org.teiid.metadata.RuntimeMetadata;
import org.teiid.translator.DataNotAvailableException;
import org.teiid.translator.ExecutionContext;
import org.teiid.translator.TranslatorException;
import org.teiid.translator.UpdateExecution;

import com.datastax.driver.core.ResultSetFuture;
import com.google.common.util.concurrent.MoreExecutors;

public class CassandraUpdateExecution implements UpdateExecution {
	
	private CassandraConnection connection;
	private ExecutionContext executionContext;
	private RuntimeMetadata metadata;
	private Command command;
	private int updateCount = 1;
	private ResultSetFuture resultSetFuture;
	
	public CassandraUpdateExecution(Command command,
			ExecutionContext executionContext, RuntimeMetadata metadata,
			CassandraConnection connection) {
		this.command = command;
		this.executionContext = executionContext;
		this.metadata = metadata;
		this.connection = connection;
	}

	@Override
	public void close() {
		this.resultSetFuture = null;
	}

	@Override
	public void cancel() throws TranslatorException {
		if (this.resultSetFuture != null) {
			this.resultSetFuture.cancel(true);
		}
	}

	@Override
	public void execute() throws TranslatorException {
		internalExecute();
		resultSetFuture.addListener(new Runnable() {
			@Override
			public void run() {
				executionContext.dataAvailable();
			}
		}, MoreExecutors.sameThreadExecutor());
	}

	private void internalExecute() throws TranslatorException {
		if (this.command instanceof BatchedUpdates) {
			handleBatchedUpdates();
			return;
		}
		CassandraSQLVisitor visitor = new CassandraSQLVisitor();
		visitor.translateSQL(this.command);
		String cql = visitor.getTranslatedSQL();
		LogManager.logDetail(LogConstants.CTX_CONNECTOR, "Source-Query:", cql); //$NON-NLS-1$
		this.executionContext.logCommand(cql);
		if (this.command instanceof BatchedCommand) {
			BatchedCommand bc = (BatchedCommand)this.command;
			if (bc.getParameterValues() != null) {
				int count = 0;
				List<Object[]> newValues = new ArrayList<Object[]>();
				Iterator<? extends List<?>> values = bc.getParameterValues();
				while (values.hasNext()) {
					Object[] bindValues = values.next().toArray();
					for (int i = 0; i < bindValues.length; i++) {
						if (bindValues[i] instanceof Blob) {
							Blob blob = (Blob)bindValues[i];
							try {
								if (blob.length() > Integer.MAX_VALUE) {
									throw new AssertionError("Blob is too large"); //$NON-NLS-1$
								}
								byte[] bytes = ((Blob)bindValues[i]).getBytes(0, (int) blob.length());
								bindValues[i] = ByteBuffer.wrap(bytes);
							} catch (SQLException e) {
								throw new TranslatorException(e);
							}
						} else if (bindValues[i] instanceof BinaryType) {
							bindValues[i] = ByteBuffer.wrap(((BinaryType)bindValues[i]).getBytesDirect());
						}
					}
					newValues.add(bindValues);
					count++;
				}
				updateCount = count;
				resultSetFuture = connection.executeBatch(cql, newValues);
				return;
			}
		}
		resultSetFuture = connection.executeQuery(cql);
	}

	private void handleBatchedUpdates() {
		BatchedUpdates updates = (BatchedUpdates)this.command;
		List<String> cqlUpdates = new ArrayList<String>();
		for (Command update : updates.getUpdateCommands()) {
			CassandraSQLVisitor visitor = new CassandraSQLVisitor();
			visitor.translateSQL(update);
			String cql = visitor.getTranslatedSQL();
			cqlUpdates.add(cql);
		}
		this.updateCount = cqlUpdates.size();
		resultSetFuture = connection.executeBatch(cqlUpdates);
	}

	@Override
	public int[] getUpdateCounts() throws DataNotAvailableException,
			TranslatorException {
		if (!resultSetFuture.isDone()) {
			throw DataNotAvailableException.NO_POLLING;
		}
		resultSetFuture.getUninterruptibly();
		return new int[] {this.updateCount};
	}

}