org.apache.calcite.avatica.QueryState Maven / Gradle / Ivy
The newest version!
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to you 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 org.apache.calcite.avatica;
import org.apache.calcite.avatica.proto.Common;
import org.apache.calcite.avatica.proto.Common.MetaDataOperationArgument;
import org.apache.calcite.avatica.proto.Common.MetaDataOperationArgument.ArgumentType;
import org.apache.calcite.avatica.remote.MetaDataOperation;
import com.fasterxml.jackson.annotation.JsonProperty;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Arrays;
import java.util.Objects;
/**
* A struct used to encapsulate the necessary information to reconstitute a ResultSet in the
* Avatica server.
*/
public class QueryState {
/**
* An enumeration that represents how a ResultSet was created.
*/
public static enum StateType {
SQL,
METADATA;
public Common.StateType toProto() {
switch (this) {
case SQL:
return Common.StateType.SQL;
case METADATA:
return Common.StateType.METADATA;
default:
return Common.StateType.UNRECOGNIZED;
}
}
public static StateType fromProto(Common.StateType protoType) {
switch (protoType) {
case SQL:
return StateType.SQL;
case METADATA:
return StateType.METADATA;
default:
throw new IllegalArgumentException("Unhandled StateType " + protoType);
}
}
}
@JsonProperty("type")
public final StateType type;
@JsonProperty("sql")
public final String sql;
@JsonProperty("metaDataOperation")
public final MetaDataOperation metaDataOperation;
@JsonProperty("operationArgs")
public final Object[] operationArgs;
/**
* Constructor encapsulating a SQL query used to create a result set.
*
* @param sql The SQL query.
*/
public QueryState(String sql) {
// This doesn't to be non-null
this.sql = sql;
this.type = StateType.SQL;
// Null out the members we don't use.
this.metaDataOperation = null;
this.operationArgs = null;
}
/**
* Constructor encapsulating a metadata operation's result set.
*
* @param op A pointer to the {@link DatabaseMetaData} operation being invoked.
* @param args The arguments to the method being invoked.
*/
public QueryState(MetaDataOperation op, Object... args) {
this.metaDataOperation = Objects.requireNonNull(op);
this.operationArgs = Arrays.copyOf(Objects.requireNonNull(args), args.length);
this.type = StateType.METADATA;
// Null out the members we won't use
this.sql = null;
}
/**
* Not intended for external use. For Jackson-databind only.
*/
public QueryState(StateType type, String sql, MetaDataOperation op, Object... args) {
this.type = Objects.requireNonNull(type);
switch (type) {
case SQL:
this.sql = Objects.requireNonNull(sql);
if (null != op) {
throw new IllegalArgumentException("Expected null MetaDataOperation, but got " + op);
}
this.metaDataOperation = null;
if (null != args) {
throw new IllegalArgumentException("Expected null arguments, but got "
+ Arrays.toString(args));
}
this.operationArgs = null;
break;
case METADATA:
this.metaDataOperation = Objects.requireNonNull(op);
this.operationArgs = Objects.requireNonNull(args);
if (null != sql) {
throw new IllegalArgumentException("Expected null SQl but got " + sql);
}
this.sql = null;
break;
default:
throw new IllegalArgumentException("Unable to handle StateType " + type);
}
}
/**
* Not intended for external use. For Jackson-databind only.
*/
public QueryState() {
this.sql = null;
this.metaDataOperation = null;
this.type = null;
this.operationArgs = null;
}
/**
* @return The {@link StateType} for this encapsulated state.
*/
public StateType getType() {
return type;
}
/**
* @return The SQL expression to invoke.
*/
public String getSql() {
assert type == StateType.SQL;
return sql;
}
/**
* @return The metadata operation to invoke.
*/
public MetaDataOperation getMetaDataOperation() {
assert type == StateType.METADATA;
return metaDataOperation;
}
/**
* @return The Arguments for the given metadata operation.
*/
public Object[] getOperationArgs() {
assert type == StateType.METADATA;
return operationArgs;
}
public ResultSet invoke(Connection conn, Statement statement) throws SQLException {
switch (type) {
case SQL:
boolean ret = Objects.requireNonNull(statement).execute(sql);
ResultSet results = statement.getResultSet();
// Either execute(sql) returned true or the resultSet was null
assert ret || null == results;
return results;
case METADATA:
DatabaseMetaData metadata = Objects.requireNonNull(conn).getMetaData();
switch (metaDataOperation) {
case GET_ATTRIBUTES:
verifyOpArgs(4);
return metadata.getAttributes((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3]);
case GET_BEST_ROW_IDENTIFIER:
verifyOpArgs(5);
return metadata.getBestRowIdentifier((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(int) operationArgs[3],
(boolean) operationArgs[4]);
case GET_CATALOGS:
verifyOpArgs(0);
return metadata.getCatalogs();
case GET_COLUMNS:
verifyOpArgs(4);
return metadata.getColumns((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3]);
case GET_COLUMN_PRIVILEGES:
verifyOpArgs(4);
return metadata.getColumnPrivileges((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3]);
case GET_CROSS_REFERENCE:
verifyOpArgs(6);
return metadata.getCrossReference((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3],
(String) operationArgs[4],
(String) operationArgs[5]);
case GET_EXPORTED_KEYS:
verifyOpArgs(3);
return metadata.getExportedKeys((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_FUNCTIONS:
verifyOpArgs(3);
return metadata.getFunctions((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_FUNCTION_COLUMNS:
verifyOpArgs(4);
return metadata.getFunctionColumns((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3]);
case GET_IMPORTED_KEYS:
verifyOpArgs(3);
return metadata.getImportedKeys((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_INDEX_INFO:
verifyOpArgs(5);
return metadata.getIndexInfo((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(boolean) operationArgs[3],
(boolean) operationArgs[4]);
case GET_PRIMARY_KEYS:
verifyOpArgs(3);
return metadata.getPrimaryKeys((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_PROCEDURES:
verifyOpArgs(3);
return metadata.getProcedures((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_PROCEDURE_COLUMNS:
verifyOpArgs(4);
return metadata.getProcedureColumns((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3]);
case GET_PSEUDO_COLUMNS:
verifyOpArgs(4);
return metadata.getPseudoColumns((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String) operationArgs[3]);
case GET_SCHEMAS:
verifyOpArgs(0);
return metadata.getSchemas();
case GET_SCHEMAS_WITH_ARGS:
verifyOpArgs(2);
return metadata.getSchemas((String) operationArgs[0],
(String) operationArgs[1]);
case GET_SUPER_TABLES:
verifyOpArgs(3);
return metadata.getSuperTables((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_SUPER_TYPES:
verifyOpArgs(3);
return metadata.getSuperTypes((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_TABLES:
verifyOpArgs(4);
return metadata.getTables((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(String[]) operationArgs[3]);
case GET_TABLE_PRIVILEGES:
verifyOpArgs(3);
return metadata.getTablePrivileges((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
case GET_TABLE_TYPES:
verifyOpArgs(0);
return metadata.getTableTypes();
case GET_TYPE_INFO:
verifyOpArgs(0);
return metadata.getTypeInfo();
case GET_UDTS:
verifyOpArgs(4);
return metadata.getUDTs((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2],
(int[]) operationArgs[3]);
case GET_VERSION_COLUMNS:
verifyOpArgs(3);
return metadata.getVersionColumns((String) operationArgs[0],
(String) operationArgs[1],
(String) operationArgs[2]);
default:
throw new IllegalArgumentException("Unhandled Metadata operation: " + metaDataOperation);
}
default:
throw new IllegalArgumentException("Unable to process QueryState of type " + type);
}
}
private void verifyOpArgs(int expectedArgs) {
if (expectedArgs != operationArgs.length) {
throw new RuntimeException("Expected " + expectedArgs + " arguments, but got "
+ Arrays.toString(operationArgs));
}
}
public Common.QueryState toProto() {
Common.QueryState.Builder builder = Common.QueryState.newBuilder();
// Required
switch (type) {
case SQL:
builder.setType(Common.StateType.SQL);
break;
case METADATA:
builder.setType(Common.StateType.METADATA);
break;
default:
throw new IllegalStateException("Unhandled type: " + type);
}
// Optional SQL
if (null != sql) {
builder.setSql(sql).setHasSql(true);
}
// Optional metaDataOperation
if (null != metaDataOperation) {
builder.setOp(metaDataOperation.toProto()).setHasOp(true);
}
// Optional operationArgs
if (null != operationArgs) {
builder.setHasArgs(true);
for (Object arg : operationArgs) {
MetaDataOperationArgument.Builder argBuilder = MetaDataOperationArgument.newBuilder();
if (null == arg) {
builder.addArgs(argBuilder.setType(ArgumentType.NULL).build());
} else if (arg instanceof String) {
builder.addArgs(argBuilder.setType(ArgumentType.STRING)
.setStringValue((String) arg).build());
} else if (arg instanceof Integer) {
builder.addArgs(argBuilder.setType(ArgumentType.INT).setIntValue((int) arg).build());
} else if (arg instanceof Boolean) {
builder.addArgs(
argBuilder.setType(ArgumentType.BOOL).setBoolValue((boolean) arg).build());
} else if (arg instanceof String[]) {
argBuilder.setType(ArgumentType.REPEATED_STRING);
for (String strArg : (String[]) arg) {
argBuilder.addStringArrayValues(strArg);
}
builder.addArgs(argBuilder.build());
} else if (arg instanceof int[]) {
argBuilder.setType(ArgumentType.REPEATED_INT);
for (int intArg : (int[]) arg) {
argBuilder.addIntArrayValues(intArg);
}
builder.addArgs(argBuilder.build());
} else {
throw new RuntimeException("Unexpected operation argument: " + arg.getClass());
}
}
} else {
builder.setHasArgs(false);
}
return builder.build();
}
public static QueryState fromProto(Common.QueryState protoState) {
StateType type = StateType.fromProto(protoState.getType());
String sql = protoState.getHasSql() ? protoState.getSql() : null;
MetaDataOperation op = protoState.getHasOp()
? MetaDataOperation.fromProto(protoState.getOp()) : null;
Object[] opArgs = null;
if (protoState.getHasArgs()) {
opArgs = new Object[protoState.getArgsCount()];
int i = 0;
for (Common.MetaDataOperationArgument arg : protoState.getArgsList()) {
switch (arg.getType()) {
case STRING:
opArgs[i] = arg.getStringValue();
break;
case BOOL:
opArgs[i] = arg.getBoolValue();
break;
case INT:
opArgs[i] = arg.getIntValue();
break;
case REPEATED_STRING:
opArgs[i] = arg.getStringArrayValuesList().toArray(
new String[arg.getStringArrayValuesCount()]);
break;
case REPEATED_INT:
int[] arr = new int[arg.getIntArrayValuesCount()];
int offset = 0;
for (Integer val : arg.getIntArrayValuesList()) {
arr[offset] = val;
offset++;
}
opArgs[i] = arr;
break;
case NULL:
opArgs[i] = null;
break;
default:
throw new RuntimeException("Could not interpret " + arg.getType());
}
i++;
}
}
return new QueryState(type, sql, op, opArgs);
}
@Override public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + ((metaDataOperation == null) ? 0 : metaDataOperation.hashCode());
result = prime * result + Arrays.hashCode(operationArgs);
result = prime * result + ((sql == null) ? 0 : sql.hashCode());
return result;
}
@Override public boolean equals(Object obj) {
if (this == obj) {
return true;
}
if (!(obj instanceof QueryState)) {
return false;
}
QueryState other = (QueryState) obj;
if (metaDataOperation != other.metaDataOperation) {
return false;
}
if (!Arrays.deepEquals(operationArgs, other.operationArgs)) {
return false;
}
if (sql == null) {
if (other.sql != null) {
return false;
}
} else if (!sql.equals(other.sql)) {
return false;
}
return true;
}
}
// End QueryState.java