org.apache.kyuubi.jdbc.hive.KyuubiStatement 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.kyuubi.jdbc.hive;
import org.apache.kyuubi.shade.com.google.common.annotations.VisibleForTesting;
import java.sql.*;
import java.util.*;
import org.apache.kyuubi.shade.org.apache.commons.lang3.StringUtils;
import org.apache.kyuubi.jdbc.hive.adapter.SQLStatement;
import org.apache.kyuubi.jdbc.hive.cli.FetchType;
import org.apache.kyuubi.jdbc.hive.cli.RowSet;
import org.apache.kyuubi.jdbc.hive.cli.RowSetFactory;
import org.apache.kyuubi.jdbc.hive.logs.InPlaceUpdateStream;
import org.apache.kyuubi.jdbc.hive.logs.KyuubiLoggable;
import org.apache.kyuubi.shaded.hive.service.rpc.thrift.*;
import org.apache.kyuubi.shaded.thrift.TException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/** KyuubiStatement. */
public class KyuubiStatement implements SQLStatement, KyuubiLoggable {
public static final Logger LOG = LoggerFactory.getLogger(KyuubiStatement.class.getName());
public static final int DEFAULT_FETCH_SIZE = 1000;
public static final String DEFAULT_RESULT_FORMAT = "thrift";
public static final String DEFAULT_ARROW_TIMESTAMP_AS_STRING = "false";
private final KyuubiConnection connection;
private TCLIService.Iface client;
private TOperationHandle stmtHandle = null;
private final TSessionHandle sessHandle;
Map sessConf = new HashMap<>();
private int fetchSize = DEFAULT_FETCH_SIZE;
private boolean isScrollableResultset = false;
private boolean isOperationComplete = false;
private Map properties = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
/**
* We need to keep a reference to the result set to support the following:
* statement.execute(String sql);
* statement.getResultSet();
* .
*/
private ResultSet resultSet = null;
/**
* Sets the limit for the maximum number of rows that any ResultSet object produced by this
* Statement can contain to the given number. If the limit is exceeded, the excess rows are
* silently dropped. The value must be >= 0, and 0 means there is not limit.
*/
private int maxRows = 0;
/** Add SQLWarnings to the warningChain if needed. */
private SQLWarning warningChain = null;
/** Keep state so we can fail certain calls made after close(). */
private boolean isClosed = false;
/** Keep state so we can fail certain calls made after cancel(). */
private boolean isCancelled = false;
/** Keep this state so we can know whether the query in this statement is closed. */
private boolean isQueryClosed = false;
/** Keep this state so we can know whether the query logs are being generated in HS2. */
private boolean isLogBeingGenerated = true;
/**
* Keep this state so we can know whether the statement is submitted to HS2 and start execution
* successfully.
*/
private boolean isExecuteStatementFailed = false;
private int queryTimeout = 0;
private InPlaceUpdateStream inPlaceUpdateStream = InPlaceUpdateStream.NO_OP;
public KyuubiStatement(
KyuubiConnection connection, TCLIService.Iface client, TSessionHandle sessHandle) {
this(connection, client, sessHandle, false, DEFAULT_FETCH_SIZE);
}
public KyuubiStatement(
KyuubiConnection connection,
TCLIService.Iface client,
TSessionHandle sessHandle,
int fetchSize) {
this(connection, client, sessHandle, false, fetchSize);
}
public KyuubiStatement(
KyuubiConnection connection,
TCLIService.Iface client,
TSessionHandle sessHandle,
boolean isScrollableResultset) {
this(connection, client, sessHandle, isScrollableResultset, DEFAULT_FETCH_SIZE);
}
public KyuubiStatement(
KyuubiConnection connection,
TCLIService.Iface client,
TSessionHandle sessHandle,
boolean isScrollableResultset,
int fetchSize) {
this.connection = connection;
this.client = client;
this.sessHandle = sessHandle;
this.isScrollableResultset = isScrollableResultset;
this.fetchSize = fetchSize;
}
@Override
public void cancel() throws SQLException {
checkConnection("cancel");
if (isCancelled) {
return;
}
try {
if (stmtHandle != null) {
TCancelOperationReq cancelReq = new TCancelOperationReq(stmtHandle);
TCancelOperationResp cancelResp = client.CancelOperation(cancelReq);
Utils.verifySuccessWithInfo(cancelResp.getStatus());
}
} catch (SQLException e) {
throw e;
} catch (Exception e) {
throw new KyuubiSQLException(e.toString(), "08S01", e);
}
isCancelled = true;
}
@Override
public void clearWarnings() throws SQLException {
warningChain = null;
}
/**
* Closes the statement if there is one running. Do not change the the flags.
*
* @throws SQLException If there is an error closing the statement
*/
private void closeStatementIfNeeded() throws SQLException {
try {
if (stmtHandle != null) {
TCloseOperationReq closeReq = new TCloseOperationReq(stmtHandle);
TCloseOperationResp closeResp = client.CloseOperation(closeReq);
Utils.verifySuccessWithInfo(closeResp.getStatus());
stmtHandle = null;
}
} catch (SQLException e) {
throw e;
} catch (Exception e) {
throw new KyuubiSQLException(e.toString(), "08S01", e);
}
}
void closeClientOperation() throws SQLException {
closeStatementIfNeeded();
isQueryClosed = true;
isExecuteStatementFailed = false;
stmtHandle = null;
}
@Override
public void close() throws SQLException {
if (isClosed) {
return;
}
closeClientOperation();
client = null;
closeResultSet();
isClosed = true;
}
@Override
public boolean execute(String sql) throws SQLException {
return executeWithConfOverlay(sql, null);
}
private boolean executeWithConfOverlay(String sql, Map confOverlay)
throws SQLException {
runAsyncOnServer(sql, confOverlay);
TGetOperationStatusResp status = waitForOperationToComplete();
// The query should be completed by now
if (!status.isHasResultSet() && !stmtHandle.isHasResultSet()) {
return false;
}
TGetResultSetMetadataResp metadata = getResultSetMetadata();
List columnNames = new ArrayList<>();
List columnTypes = new ArrayList<>();
List columnAttributes = new ArrayList<>();
parseMetadata(metadata, columnNames, columnTypes, columnAttributes);
String resultFormat =
properties.getOrDefault("__kyuubi_operation_result_format__", DEFAULT_RESULT_FORMAT);
LOG.debug("kyuubi.operation.result.format: {}", resultFormat);
switch (resultFormat) {
case "arrow":
boolean timestampAsString =
Boolean.parseBoolean(
properties.getOrDefault(
"__kyuubi_operation_result_arrow_timestampAsString__",
DEFAULT_ARROW_TIMESTAMP_AS_STRING));
resultSet =
new KyuubiArrowQueryResultSet.Builder(this)
.setClient(client)
.setSessionHandle(sessHandle)
.setStmtHandle(stmtHandle)
.setMaxRows(maxRows)
.setFetchSize(fetchSize)
.setScrollable(isScrollableResultset)
.setSchema(columnNames, columnTypes, columnAttributes)
.setTimestampAsString(timestampAsString)
.build();
break;
default:
resultSet =
new KyuubiQueryResultSet.Builder(this)
.setClient(client)
.setSessionHandle(sessHandle)
.setStmtHandle(stmtHandle)
.setMaxRows(maxRows)
.setFetchSize(fetchSize)
.setScrollable(isScrollableResultset)
.setSchema(columnNames, columnTypes, columnAttributes)
.build();
}
return true;
}
/**
* Starts the query execution asynchronously on the server, and immediately returns to the client.
* The client subsequently blocks on ResultSet#next or Statement#getUpdateCount, depending on the
* query type. Users should call ResultSet.next or Statement#getUpdateCount (depending on whether
* query returns results) to ensure that query completes successfully. Calling another execute*
* method, or close before query completion would result in the async query getting killed if it
* is not already finished. Note: This method is an API for limited usage outside of Hive by
* applications like Apache Ambari, although it is not part of the interface java.sql.Statement.
*
* @param sql
* @return true if the first result is a ResultSet object; false if it is an update count or there
* are no results
* @throws SQLException
*/
public boolean executeAsync(String sql) throws SQLException {
runAsyncOnServer(sql);
TGetOperationStatusResp status = waitForResultSetStatus();
if (!status.isHasResultSet()) {
return false;
}
TGetResultSetMetadataResp metadata = getResultSetMetadata();
List columnNames = new ArrayList<>();
List columnTypes = new ArrayList<>();
List columnAttributes = new ArrayList<>();
parseMetadata(metadata, columnNames, columnTypes, columnAttributes);
String resultFormat =
properties.getOrDefault("__kyuubi_operation_result_format__", DEFAULT_RESULT_FORMAT);
LOG.debug("kyuubi.operation.result.format: {}", resultFormat);
switch (resultFormat) {
case "arrow":
boolean timestampAsString =
Boolean.parseBoolean(
properties.getOrDefault(
"__kyuubi_operation_result_arrow_timestampAsString__",
DEFAULT_ARROW_TIMESTAMP_AS_STRING));
resultSet =
new KyuubiArrowQueryResultSet.Builder(this)
.setClient(client)
.setSessionHandle(sessHandle)
.setStmtHandle(stmtHandle)
.setMaxRows(maxRows)
.setFetchSize(fetchSize)
.setScrollable(isScrollableResultset)
.setSchema(columnNames, columnTypes, columnAttributes)
.setTimestampAsString(timestampAsString)
.build();
break;
default:
resultSet =
new KyuubiQueryResultSet.Builder(this)
.setClient(client)
.setSessionHandle(sessHandle)
.setStmtHandle(stmtHandle)
.setMaxRows(maxRows)
.setFetchSize(fetchSize)
.setScrollable(isScrollableResultset)
.setSchema(columnNames, columnTypes, columnAttributes)
.build();
}
return true;
}
private void runAsyncOnServer(String sql) throws SQLException {
runAsyncOnServer(sql, null);
}
private void runAsyncOnServer(String sql, Map confOneTime) throws SQLException {
checkConnection("execute");
reInitState();
TExecuteStatementReq execReq = new TExecuteStatementReq(sessHandle, sql);
/**
* Run asynchronously whenever possible Currently only a SQLOperation can be run asynchronously,
* in a background operation thread Compilation can run asynchronously or synchronously and
* execution run asynchronously
*/
execReq.setRunAsync(true);
if (confOneTime != null) {
Map confOverlay = new HashMap(sessConf);
confOverlay.putAll(confOneTime);
execReq.setConfOverlay(confOverlay);
} else {
execReq.setConfOverlay(sessConf);
}
execReq.setQueryTimeout(queryTimeout);
try {
TExecuteStatementResp execResp = client.ExecuteStatement(execReq);
Utils.verifySuccessWithInfo(execResp.getStatus());
stmtHandle = execResp.getOperationHandle();
isExecuteStatementFailed = false;
} catch (SQLException eS) {
isExecuteStatementFailed = true;
isLogBeingGenerated = false;
throw eS;
} catch (Exception ex) {
isExecuteStatementFailed = true;
isLogBeingGenerated = false;
throw new KyuubiSQLException(ex.toString(), "08S01", ex);
}
}
/**
* Poll the result set status by checking if isSetHasResultSet is set
*
* @return
* @throws SQLException
*/
private TGetOperationStatusResp waitForResultSetStatus() throws SQLException {
TGetOperationStatusReq statusReq = new TGetOperationStatusReq(stmtHandle);
TGetOperationStatusResp statusResp = null;
while (statusResp == null || !statusResp.isSetHasResultSet()) {
try {
statusResp = client.GetOperationStatus(statusReq);
} catch (TException e) {
isLogBeingGenerated = false;
throw new KyuubiSQLException(e.toString(), "08S01", e);
}
}
return statusResp;
}
TGetOperationStatusResp waitForOperationToComplete() throws SQLException {
TGetOperationStatusReq statusReq = new TGetOperationStatusReq(stmtHandle);
boolean shouldGetProgressUpdate = inPlaceUpdateStream != InPlaceUpdateStream.NO_OP;
statusReq.setGetProgressUpdate(shouldGetProgressUpdate);
if (!shouldGetProgressUpdate) {
/** progress bar is completed if there is nothing we want to request in the first place. */
inPlaceUpdateStream.getEventNotifier().progressBarCompleted();
}
TGetOperationStatusResp statusResp = null;
// Poll on the operation status, till the operation is complete
do {
try {
/**
* For an async SQLOperation, GetOperationStatus will use the long polling approach It will
* essentially return after the HIVE_SERVER2_LONG_POLLING_TIMEOUT (a server config) expires
*/
statusResp = client.GetOperationStatus(statusReq);
inPlaceUpdateStream.update(statusResp.getProgressUpdateResponse());
Utils.verifySuccessWithInfo(statusResp.getStatus());
if (statusResp.isSetOperationState()) {
switch (statusResp.getOperationState()) {
case CLOSED_STATE:
case FINISHED_STATE:
isOperationComplete = true;
isLogBeingGenerated = false;
break;
case CANCELED_STATE:
// 01000 -> warning
String errMsg = statusResp.getErrorMessage();
if (errMsg != null && !errMsg.isEmpty()) {
throw new KyuubiSQLException("Query was cancelled. " + errMsg, "01000");
} else {
throw new KyuubiSQLException("Query was cancelled", "01000");
}
case TIMEDOUT_STATE:
throw new SQLTimeoutException("Query timed out after " + queryTimeout + " seconds");
case ERROR_STATE:
// Get the error details from the underlying exception
throw new KyuubiSQLException(
statusResp.getErrorMessage(),
statusResp.getSqlState(),
statusResp.getErrorCode());
case UKNOWN_STATE:
throw new KyuubiSQLException("Unknown query", "HY000");
case INITIALIZED_STATE:
case PENDING_STATE:
case RUNNING_STATE:
break;
}
}
} catch (SQLException e) {
isLogBeingGenerated = false;
throw e;
} catch (Exception e) {
isLogBeingGenerated = false;
throw new KyuubiSQLException(e.toString(), "08S01", e);
}
} while (!isOperationComplete);
/*
we set progress bar to be completed when hive query execution has completed
*/
inPlaceUpdateStream.getEventNotifier().progressBarCompleted();
return statusResp;
}
private void checkConnection(String action) throws SQLException {
if (isClosed) {
throw new KyuubiSQLException("Can't " + action + " after statement has been closed");
}
}
/**
* Close statement if needed, and reset the flags.
*
* @throws SQLException
*/
private void reInitState() throws SQLException {
closeStatementIfNeeded();
isCancelled = false;
isQueryClosed = false;
isLogBeingGenerated = true;
isExecuteStatementFailed = false;
isOperationComplete = false;
}
@Override
public ResultSet executeQuery(String sql) throws SQLException {
if (!execute(sql)) {
throw new KyuubiSQLException("The query did not generate a result set!");
}
return resultSet;
}
public ResultSet executeScala(String code) throws SQLException {
if (!executeWithConfOverlay(
code, Collections.singletonMap("kyuubi.operation.language", "SCALA"))) {
throw new KyuubiSQLException("The query did not generate a result set!");
}
return resultSet;
}
public ResultSet executePython(String code) throws SQLException {
if (!executeWithConfOverlay(
code, Collections.singletonMap("kyuubi.operation.language", "PYTHON"))) {
throw new KyuubiSQLException("The query did not generate a result set!");
}
return resultSet;
}
public void executeSetCurrentCatalog(String sql, String catalog) throws SQLException {
if (executeWithConfOverlay(
sql, Collections.singletonMap("kyuubi.operation.set.current.catalog", catalog))) {
closeResultSet();
}
}
public ResultSet executeGetCurrentCatalog(String sql) throws SQLException {
if (!executeWithConfOverlay(
sql, Collections.singletonMap("kyuubi.operation.get.current.catalog", ""))) {
throw new KyuubiSQLException("The query did not generate a result set!");
}
return resultSet;
}
public void executeSetCurrentDatabase(String sql, String database) throws SQLException {
if (executeWithConfOverlay(
sql, Collections.singletonMap("kyuubi.operation.set.current.database", database))) {
closeResultSet();
}
}
public ResultSet executeGetCurrentDatabase(String sql) throws SQLException {
if (!executeWithConfOverlay(
sql, Collections.singletonMap("kyuubi.operation.get.current.database", ""))) {
throw new KyuubiSQLException("The query did not generate a result set!");
}
return resultSet;
}
@Override
public int executeUpdate(String sql) throws SQLException {
execute(sql);
return getUpdateCount();
}
@Override
public Connection getConnection() throws SQLException {
checkConnection("getConnection");
return this.connection;
}
@Override
public int getFetchDirection() throws SQLException {
checkConnection("getFetchDirection");
return ResultSet.FETCH_FORWARD;
}
@Override
public int getFetchSize() throws SQLException {
checkConnection("getFetchSize");
return fetchSize;
}
@Override
public int getMaxRows() throws SQLException {
checkConnection("getMaxRows");
return maxRows;
}
@Override
public boolean getMoreResults(int current) throws SQLException {
if (current == Statement.CLOSE_CURRENT_RESULT) {
closeResultSet();
return false;
}
if (current != KEEP_CURRENT_RESULT && current != CLOSE_ALL_RESULTS) {
throw new SQLException("Invalid argument: " + current);
}
throw new SQLFeatureNotSupportedException("Multiple open results not supported");
}
@Override
public boolean getMoreResults() throws SQLException {
// The javadoc of this method says, that it should implicitly close any current
// ResultSet object, i.e. is similar to calling the getMoreResults(int)
// method with Statement.CLOSE_CURRENT_RESULT argument.
// this is an additional enhancement on top of HIVE-7680
return getMoreResults(Statement.CLOSE_CURRENT_RESULT);
}
@Override
public int getQueryTimeout() throws SQLException {
checkConnection("getQueryTimeout");
return 0;
}
@Override
public ResultSet getResultSet() throws SQLException {
checkConnection("getResultSet");
return resultSet;
}
@Override
public int getResultSetType() throws SQLException {
checkConnection("getResultSetType");
return ResultSet.TYPE_FORWARD_ONLY;
}
@Override
public int getUpdateCount() throws SQLException {
checkConnection("getUpdateCount");
/**
* Poll on the operation status, till the operation is complete. We want to ensure that since a
* client might end up using executeAsync and then call this to check if the query run is
* finished.
*/
long numModifiedRows = -1L;
TGetOperationStatusResp resp = waitForOperationToComplete();
if (resp != null && resp.isSetNumModifiedRows()) {
numModifiedRows = resp.getNumModifiedRows();
}
if (numModifiedRows == -1L || numModifiedRows > Integer.MAX_VALUE) {
LOG.warn("Invalid number of updated rows: {}", numModifiedRows);
return -1;
}
return (int) numModifiedRows;
}
@Override
public SQLWarning getWarnings() throws SQLException {
checkConnection("getWarnings");
return warningChain;
}
@Override
public boolean isClosed() throws SQLException {
return isClosed;
}
@Override
public boolean isCloseOnCompletion() throws SQLException {
return false;
}
@Override
public boolean isPoolable() throws SQLException {
return false;
}
@Override
public void setFetchDirection(int direction) throws SQLException {
checkConnection("setFetchDirection");
if (direction != ResultSet.FETCH_FORWARD) {
throw new KyuubiSQLException("Not supported direction " + direction);
}
}
@Override
public void setFetchSize(int rows) throws SQLException {
checkConnection("setFetchSize");
if (rows > 0) {
fetchSize = rows;
} else if (rows == 0) {
// Javadoc for Statement interface states that if the value is zero
// then "fetch size" hint is ignored.
// In this case it means reverting it to the default value.
fetchSize = DEFAULT_FETCH_SIZE;
} else {
throw new KyuubiSQLException("Fetch size must be greater or equal to 0");
}
}
@Override
public void setMaxRows(int max) throws SQLException {
checkConnection("setMaxRows");
if (max < 0) {
throw new KyuubiSQLException("max must be >= 0");
}
maxRows = max;
}
@Override
public void setQueryTimeout(int seconds) throws SQLException {
this.queryTimeout = seconds;
}
@Override
public boolean isWrapperFor(Class iface) throws SQLException {
return false;
}
@Override
public T unwrap(Class iface) throws SQLException {
throw new KyuubiSQLException("Cannot unwrap to " + iface);
}
/**
* Check whether query execution might be producing more logs to be fetched. This method is a
* public API for usage outside of Hive, although it is not part of the interface
* java.sql.Statement.
*
* @return true if query execution might be producing more logs. It does not indicate if last log
* lines have been fetched by getQueryLog.
*/
@Override
public boolean hasMoreLogs() {
return isLogBeingGenerated;
}
/**
* Get the execution logs of the given SQL statement. This method is a public API for usage
* outside of Hive, although it is not part of the interface java.sql.Statement. This method gets
* the incremental logs during SQL execution, and uses fetchSize holden by HiveStatement object.
*
* @return a list of logs. It can be empty if there are no new logs to be retrieved at that time.
* @throws SQLException
* @throws ClosedOrCancelledException if statement has been cancelled or closed
*/
@Override
public List getExecLog() throws SQLException, ClosedOrCancelledException {
return getQueryLog(true, fetchSize);
}
/**
* Get the execution logs of the given SQL statement. This method is a public API for usage
* outside of Hive, although it is not part of the interface java.sql.Statement.
*
* @param incremental indicate getting logs either incrementally or from the beginning, when it is
* true or false.
* @param fetchSize the number of lines to fetch
* @return a list of logs. It can be empty if there are no new logs to be retrieved at that time.
* @throws SQLException
* @throws ClosedOrCancelledException if statement has been cancelled or closed
*/
public List getQueryLog(boolean incremental, int fetchSize)
throws SQLException, ClosedOrCancelledException {
checkConnection("getQueryLog");
if (isCancelled) {
throw new ClosedOrCancelledException(
"Method getQueryLog() failed. The " + "statement has been closed or cancelled.");
}
List logs = new ArrayList();
TFetchResultsResp tFetchResultsResp = null;
try {
if (stmtHandle != null) {
TFetchResultsReq tFetchResultsReq =
new TFetchResultsReq(stmtHandle, getFetchOrientation(incremental), fetchSize);
tFetchResultsReq.setFetchType(FetchType.LOG.toTFetchType());
tFetchResultsResp = client.FetchResults(tFetchResultsReq);
Utils.verifySuccessWithInfo(tFetchResultsResp.getStatus());
} else {
if (isQueryClosed) {
throw new ClosedOrCancelledException(
"Method getQueryLog() failed. The " + "statement has been closed or cancelled.");
} else {
return logs;
}
}
} catch (SQLException e) {
throw e;
} catch (TException e) {
throw new KyuubiSQLException("Error when getting query log: " + e, e);
} catch (Exception e) {
throw new KyuubiSQLException("Error when getting query log: " + e, e);
}
try {
RowSet rowSet;
rowSet = RowSetFactory.create(tFetchResultsResp.getResults(), connection.getProtocol());
for (Object[] row : rowSet) {
logs.add(String.valueOf(row[0]));
}
} catch (TException e) {
throw new KyuubiSQLException("Error building result set for query log: " + e, e);
}
return logs;
}
private TFetchOrientation getFetchOrientation(boolean incremental) {
if (incremental) {
return TFetchOrientation.FETCH_NEXT;
} else {
return TFetchOrientation.FETCH_FIRST;
}
}
/**
* Returns the Yarn ATS GUID. This method is a public API for usage outside of Hive, although it
* is not part of the interface java.sql.Statement.
*
* @return Yarn ATS GUID or null if it hasn't been created yet.
*/
public String getYarnATSGuid() {
if (stmtHandle != null) {
// Set on the server side.
// @see org.apache.hive.service.cli.operation.SQLOperation#prepare
String guid64 =
Base64.getUrlEncoder().encodeToString(stmtHandle.getOperationId().getGuid()).trim();
return guid64;
}
return null;
}
/**
* Returns the Query ID if it is running. This method is a public API for usage outside of Hive,
* although it is not part of the interface java.sql.Statement.
*
* @return Valid query ID if it is running else returns NULL.
* @throws SQLException If any internal failures.
*/
@VisibleForTesting
public String getQueryId() throws SQLException {
if (stmtHandle == null) {
// If query is not running or already closed.
return null;
}
try {
final String queryId = client.GetQueryId(new TGetQueryIdReq(stmtHandle)).getQueryId();
// queryId can be empty string if query was already closed. Need to return null in such case.
return StringUtils.isBlank(queryId) ? null : queryId;
} catch (TException e) {
throw new KyuubiSQLException(e);
}
}
/**
* This is only used by the beeline client to set the stream on which in place progress updates
* are to be shown
*/
public void setInPlaceUpdateStream(InPlaceUpdateStream stream) {
this.inPlaceUpdateStream = stream;
}
private TGetResultSetMetadataResp getResultSetMetadata() throws SQLException {
try {
TGetResultSetMetadataReq metadataReq = new TGetResultSetMetadataReq(stmtHandle);
// TODO need session handle
TGetResultSetMetadataResp metadataResp;
metadataResp = client.GetResultSetMetadata(metadataReq);
Utils.verifySuccess(metadataResp.getStatus());
return metadataResp;
} catch (SQLException eS) {
throw eS; // rethrow the SQLException as is
} catch (Exception ex) {
throw new KyuubiSQLException("Could not create ResultSet: " + ex.getMessage(), ex);
}
}
private void parseMetadata(
TGetResultSetMetadataResp metadataResp,
List columnNames,
List columnTypes,
List columnAttributes)
throws KyuubiSQLException {
TTableSchema schema = metadataResp.getSchema();
if (schema == null || !schema.isSetColumns()) {
throw new KyuubiSQLException("the result set schema is null");
}
// parse kyuubi hint
List infoMessages = metadataResp.getStatus().getInfoMessages();
if (infoMessages != null) {
metadataResp.getStatus().getInfoMessages().stream()
.filter(hint -> Utils.isKyuubiOperationHint(hint))
.forEach(
line -> {
String[] keyValue = line.toLowerCase(Locale.ROOT).split("=");
assert keyValue.length == 2 : "Illegal Kyuubi operation hint found!";
String key = keyValue[0];
String value = keyValue[1];
properties.put(key, value);
});
}
// parse metadata
List columns = schema.getColumns();
for (int pos = 0; pos < schema.getColumnsSize(); pos++) {
String columnName = columns.get(pos).getColumnName();
columnNames.add(columnName);
TPrimitiveTypeEntry primitiveTypeEntry =
columns.get(pos).getTypeDesc().getTypes().get(0).getPrimitiveEntry();
columnTypes.add(primitiveTypeEntry.getType());
columnAttributes.add(KyuubiArrowQueryResultSet.getColumnAttributes(primitiveTypeEntry));
}
}
private void closeResultSet() throws SQLException {
if (resultSet != null) {
resultSet.close();
resultSet = null;
}
}
}