diff --git a/itests/hive-unit/src/test/java/org/apache/hive/jdbc/TestJdbcDriver2.java b/itests/hive-unit/src/test/java/org/apache/hive/jdbc/TestJdbcDriver2.java index 4eaff10..b427dc1 100644 --- a/itests/hive-unit/src/test/java/org/apache/hive/jdbc/TestJdbcDriver2.java +++ b/itests/hive-unit/src/test/java/org/apache/hive/jdbc/TestJdbcDriver2.java @@ -2343,6 +2343,15 @@ public void testShowRoleGrant() throws SQLException { } /** + * Useful for modifying outer class context from anonymous inner class + */ + public static interface Holder { + public void set(T obj); + + public T get(); + } + + /** * Test the cancellation of a query that is running. * We spawn 2 threads - one running the query and * the other attempting to cancel. @@ -2394,6 +2403,75 @@ public void run() { stmt.close(); } + /** + * Test the non-null value of the Yarn ATS GUID. + * We spawn 2 threads - one running the query and + * the other attempting to read the ATS GUID. + * We're using a dummy udf to simulate a query, + * that runs for a sufficiently long time. + * @throws Exception + */ + @Test + public void testYarnATSGuid() throws Exception { + String udfName = SleepUDF.class.getName(); + Statement stmt1 = con.createStatement(); + stmt1.execute("create temporary function sleepUDF as '" + udfName + "'"); + stmt1.close(); + final Statement stmt = con.createStatement(); + final Holder yarnATSGuidSet = new Holder() { + public Boolean b = false; + + public void set(Boolean b) { + this.b = b; + } + + public Boolean get() { + return this.b; + } + }; + + // Thread executing the query + Thread tExecute = new Thread(new Runnable() { + @Override + public void run() { + try { + System.out.println("Executing query: "); + stmt.executeQuery("select sleepUDF(t1.under_col) as u0, t1.under_col as u1, " + + "t2.under_col as u2 from " + tableName + " t1 join " + tableName + + " t2 on t1.under_col = t2.under_col"); + } catch (SQLException e) { + // No op + } + } + }); + // Thread reading the ATS GUID + Thread tGuid = new Thread(new Runnable() { + @Override + public void run() { + try { + Thread.sleep(10000); + } catch (InterruptedException e) { + e.printStackTrace(); + } + String atsGuid = ((HiveStatement) stmt).getYarnATSGuid(); + if (atsGuid != null) { + yarnATSGuidSet.set(true); + System.out.println("Yarn ATS GUID: " + atsGuid); + } else { + yarnATSGuidSet.set(false); + } + } + }); + tExecute.start(); + tGuid.start(); + tExecute.join(); + tGuid.join(); + if (!yarnATSGuidSet.get()) { + fail("Failed to set the YARN ATS Guid"); + } + stmt.close(); + } + // A udf which sleeps for 100ms to simulate a long running query public static class SleepUDF extends UDF { public Integer evaluate(final Integer value) { diff --git a/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java b/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java index b4dba44..0bbd0e3 100644 --- a/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java +++ b/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java @@ -18,6 +18,7 @@ package org.apache.hive.jdbc; +import org.apache.commons.codec.binary.Base64; import org.apache.hive.service.cli.RowSet; import org.apache.hive.service.cli.RowSetFactory; import org.apache.hive.service.rpc.thrift.TCLIService; @@ -856,4 +857,21 @@ private TFetchOrientation getFetchOrientation(boolean incremental) { 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.encodeBase64URLSafeString(stmtHandle.getOperationId().getGuid()).trim(); + return guid64; + } + return null; + } }