diff --git a/phoenix-core/src/main/java/org/apache/phoenix/compile/QueryCompiler.java b/phoenix-core/src/main/java/org/apache/phoenix/compile/QueryCompiler.java index a4b6749f613..be8428358f2 100644 --- a/phoenix-core/src/main/java/org/apache/phoenix/compile/QueryCompiler.java +++ b/phoenix-core/src/main/java/org/apache/phoenix/compile/QueryCompiler.java @@ -125,6 +125,9 @@ public QueryCompiler(PhoenixStatement statement, SelectStatement select, ColumnR this.select = select; this.resolver = resolver; this.scan = new Scan(); + if (statement.getQueryId() != null) { + this.scan.setId(statement.getQueryId()); + } this.targetColumns = targetColumns; this.parallelIteratorFactory = parallelIteratorFactory; this.sequenceManager = sequenceManager; diff --git a/phoenix-core/src/main/java/org/apache/phoenix/compile/UpsertCompiler.java b/phoenix-core/src/main/java/org/apache/phoenix/compile/UpsertCompiler.java index 3bd6d4eaec5..0a2c7bf8789 100644 --- a/phoenix-core/src/main/java/org/apache/phoenix/compile/UpsertCompiler.java +++ b/phoenix-core/src/main/java/org/apache/phoenix/compile/UpsertCompiler.java @@ -182,8 +182,14 @@ private static void setValues(byte[][] values, int[] pkSlotIndex, int[] columnIn ptr.set(ScanRanges.prefixKey(ptr.get(), 0, ptr.getLength(), regionPrefix, regionPrefix.length)); } - } - mutation.put(ptr, new RowMutationState(columnValues, columnValueSize, statement.getConnection().getStatementExecutionCounter(), rowTsColInfo, onDupKeyBytes)); + } + RowMutationState + rowMutationState = + new RowMutationState(columnValues, columnValueSize, + statement.getConnection().getStatementExecutionCounter(), rowTsColInfo, + onDupKeyBytes); + rowMutationState.setQueryId(statement.getQueryId()); + mutation.put(ptr, rowMutationState); } public static String getExceedMaxHBaseClientKeyValueAllowanceRowkeyAndColumnInfo( diff --git a/phoenix-core/src/main/java/org/apache/phoenix/execute/MutationState.java b/phoenix-core/src/main/java/org/apache/phoenix/execute/MutationState.java index 91545d0329a..3b2d9f2e1da 100644 --- a/phoenix-core/src/main/java/org/apache/phoenix/execute/MutationState.java +++ b/phoenix-core/src/main/java/org/apache/phoenix/execute/MutationState.java @@ -681,11 +681,24 @@ private void generateMutations(final TableRef tableRef, final long mutationTimes } annotateMutationsWithMetadata(table, rowMutations); mutationList.addAll(rowMutations); - if (mutationsPertainingToIndex != null) mutationsPertainingToIndex.addAll(rowMutationsPertainingToIndex); + if (mutationsPertainingToIndex != null){ + mutationsPertainingToIndex.addAll(rowMutationsPertainingToIndex); + } + addQueryIdToRowMutation(state.getQueryId(), rowMutations); + addQueryIdToRowMutation(state.getQueryId(), mutationsPertainingToIndex); } values.putAll(modifiedValues); } + private void addQueryIdToRowMutation(String queryId, List mutations) { + if (mutations == null || queryId == null) { + return; + } + for (Mutation mutation : mutations) { + mutation.setId(queryId); + } + } + private void annotateMutationsWithMetadata(PTable table, List rowMutations) { //only annotate if the change detection flag is on the table and HBase supports // preWALAppend coprocs server-side @@ -1708,6 +1721,7 @@ public void clear() { public Collection values() { return rowKeyToRowMutationState.values(); } + } public static class RowMutationState { @@ -1718,6 +1732,7 @@ public static class RowMutationState { private final RowTimestampColInfo rowTsColInfo; private byte[] onDupKeyBytes; private long colValuesSize; + private String queryId; public RowMutationState(@Nonnull Map columnValues, long colValuesSize, int statementIndex, @Nonnull RowTimestampColInfo rowTsColInfo, byte[] onDupKeyBytes) { @@ -1773,6 +1788,14 @@ RowTimestampColInfo getRowTimestampColInfo() { return rowTsColInfo; } + public String getQueryId() { + return queryId; + } + + public void setQueryId(String queryId) { + this.queryId = queryId; + } + } public ReadMetricQueue getReadMetricQueue() { diff --git a/phoenix-core/src/main/java/org/apache/phoenix/jdbc/PhoenixStatement.java b/phoenix-core/src/main/java/org/apache/phoenix/jdbc/PhoenixStatement.java index 10c379bb24c..7e20d859f11 100644 --- a/phoenix-core/src/main/java/org/apache/phoenix/jdbc/PhoenixStatement.java +++ b/phoenix-core/src/main/java/org/apache/phoenix/jdbc/PhoenixStatement.java @@ -257,6 +257,7 @@ public String toString() { private int maxRows; private int fetchSize = -1; private int queryTimeoutMillis; + private String queryId; public PhoenixStatement(PhoenixConnection connection) { this.connection = connection; @@ -2262,4 +2263,12 @@ private void checkIfDDLStatementandMutationState(final CompilableStatement stmt, } } } + + public String getQueryId() { + return queryId; + } + + public void setQueryId(String queryId) { + this.queryId = queryId; + } } diff --git a/phoenix-core/src/test/java/org/apache/phoenix/compile/QueryCompilerTest.java b/phoenix-core/src/test/java/org/apache/phoenix/compile/QueryCompilerTest.java index d7c8fb94dec..1a098faa740 100644 --- a/phoenix-core/src/test/java/org/apache/phoenix/compile/QueryCompilerTest.java +++ b/phoenix-core/src/test/java/org/apache/phoenix/compile/QueryCompilerTest.java @@ -38,39 +38,26 @@ import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; import java.sql.Statement; -import java.util.ArrayList; -import java.util.Collections; -import java.util.Iterator; -import java.util.List; -import java.util.Properties; -import java.util.Set; +import java.util.*; +import com.google.inject.internal.util.$AbstractMapEntry; import org.apache.hadoop.hbase.HRegionLocation; +import org.apache.hadoop.hbase.client.Mutation; import org.apache.hadoop.hbase.client.Scan; import org.apache.hadoop.hbase.filter.Filter; import org.apache.hadoop.hbase.filter.FirstKeyOnlyFilter; import org.apache.hadoop.hbase.filter.PageFilter; import org.apache.hadoop.hbase.util.Bytes; +import org.apache.hadoop.hbase.util.Pair; import org.apache.phoenix.compile.JoinCompiler.JoinTable; import org.apache.phoenix.compile.JoinCompiler.Table; import org.apache.phoenix.compile.OrderByCompiler.OrderBy; import org.apache.phoenix.coprocessor.BaseScannerRegionObserver; import org.apache.phoenix.exception.SQLExceptionCode; -import org.apache.phoenix.execute.AggregatePlan; -import org.apache.phoenix.execute.ClientAggregatePlan; -import org.apache.phoenix.execute.ClientScanPlan; -import org.apache.phoenix.execute.CursorFetchPlan; -import org.apache.phoenix.execute.HashJoinPlan; +import org.apache.phoenix.execute.*; import org.apache.phoenix.execute.HashJoinPlan.HashSubPlan; import org.apache.phoenix.execute.HashJoinPlan.SubPlan; import org.apache.phoenix.execute.HashJoinPlan.WhereClauseSubPlan; -import org.apache.phoenix.execute.LiteralResultIterationPlan; -import org.apache.phoenix.execute.ScanPlan; -import org.apache.phoenix.execute.SortMergeJoinPlan; -import org.apache.phoenix.execute.TupleProjectionPlan; -import org.apache.phoenix.execute.TupleProjector; -import org.apache.phoenix.execute.UnionPlan; -import org.apache.phoenix.execute.UnnestArrayPlan; import org.apache.phoenix.execute.visitor.QueryPlanVisitor; import org.apache.phoenix.expression.Expression; import org.apache.phoenix.expression.LiteralExpression; @@ -6781,4 +6768,53 @@ public void testHashJoinBug6232() throws Exception { conn.close(); } } + + @Test public void testQueryIdForScan() throws SQLException { + String query = "select * from atable"; + Properties props = PropertiesUtil.deepCopy(TEST_PROPERTIES); + Connection conn = DriverManager.getConnection(getUrl(), props); + String queryId = "test-query"; + try { + PhoenixPreparedStatement + statement = + conn.prepareStatement(query).unwrap(PhoenixPreparedStatement.class); + statement.setQueryId(queryId); + QueryPlan plan = statement.compileQuery(query); + plan.iterator(); + Scan scan = plan.getContext().getScan(); + assertEquals(queryId, scan.getId()); + } finally { + conn.close(); + } + } + + @Test + public void testQueryIdForPut() throws SQLException { + Properties connectionProperties = new Properties(); + Connection connection = DriverManager.getConnection(getUrl(), connectionProperties); + String queryId = "test-upsert-queryId"; + PhoenixPreparedStatement + stmt = + (PhoenixPreparedStatement) connection.prepareStatement( + "UPSERT INTO " + ATABLE + " (organization_id, entity_id, a_integer) " + + "VALUES (?,?,?)"); + stmt.setQueryId(queryId); + stmt.setString(1, "AAA"); + stmt.setString(2, "BBB"); + stmt.setInt(3, 1); + try { + MutationPlan mutationPlan = stmt.compileMutation(); + MutationState mutationState = mutationPlan.execute(); + Iterator>> iterator = mutationState.toMutations(); + while (iterator.hasNext()) { + Pair> entry = iterator.next(); + List mutations = entry.getSecond(); + for (Mutation mutation : mutations) { + assertEquals(queryId, mutation.getId()); + } + } + } finally { + connection.close(); + } + } }