diff --git a/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java b/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java index 54516007d32a..e82794f39f82 100644 --- a/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java +++ b/common/src/java/org/apache/hadoop/hive/conf/HiveConf.java @@ -3886,8 +3886,26 @@ public static enum ConfVars { "hs2ActivePassiveHA", "When HiveServer2 Active/Passive High Availability is enabled, uses this namespace for registering HS2\n" + "instances with zookeeper"), - HIVE_SERVER2_ACTIVE_PASSIVE_HA_HEALTHCHECK_PORT("hive.server2.active.passive.ha.healthcheck.port", 11002, + HIVE_SERVER2_ACTIVE_PASSIVE_HA_HEALTHCHECK_PORT("hive.server2.active.passive.ha.healthcheck.port", 11002, "The port the HiveServer2 ha-healthcheck web app will listen on"), + + // Persistable session state store configs + HIVE_SERVER2_SESSION_STATE_STORE_CLASS("hive.server2.session.state.store.class", + "", + "Implementation class for the session state store. Empty means disabled. Options:\n" + + " org.apache.hive.service.cli.session.store.ZooKeeperSessionStateStore\n" + + " org.apache.hive.service.cli.session.store.RedisSessionStateStore"), + HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY("hive.server2.session.state.store.fetch.strategy", + "NEVER", + new StringSet("NEVER", "ALWAYS", "FETCH_WHEN_MISSING"), + "Session fetch strategy from shared store:\n" + + " NEVER - only use local session state\n" + + " ALWAYS - on every access, compare local lastAccessTime with remote; if remote is newer, re-hydrate\n" + + " FETCH_WHEN_MISSING - fetch from store only when session not found locally"), + HIVE_SERVER2_SESSION_STATE_STORE_TTL("hive.server2.session.state.store.ttl", "24h", + new TimeValidator(TimeUnit.SECONDS), + "TTL for session snapshots in the state store. Abandoned sessions auto-expire after this duration."), + HIVE_SERVER2_TEZ_INTERACTIVE_QUEUE("hive.server2.tez.interactive.queue", "", "A single YARN queues to use for Hive Interactive sessions. When this is specified,\n" + "workload management is enabled and used for these sessions."), diff --git a/itests/hive-unit/pom.xml b/itests/hive-unit/pom.xml index 978c1788a4a9..357a32376609 100644 --- a/itests/hive-unit/pom.xml +++ b/itests/hive-unit/pom.xml @@ -54,6 +54,18 @@ org.apache.hive hive-service + + org.apache.hive + hive-service-session-store + ${project.version} + test + + + redis.clients + jedis + ${jedis.version} + test + org.apache.hive hive-llap-server diff --git a/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionBase.java b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionBase.java new file mode 100644 index 000000000000..48c3cb8ad755 --- /dev/null +++ b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionBase.java @@ -0,0 +1,667 @@ +/* + * 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.hive.service.cli.session; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.apache.hive.jdbc.miniHS2.MiniHS2; +import org.apache.hive.service.cli.CLIServiceClient; +import org.apache.hive.service.cli.HiveSQLException; +import org.apache.hive.service.cli.OperationHandle; +import org.apache.hive.service.cli.OperationState; +import org.apache.hive.service.cli.RowSet; +import org.apache.hive.service.cli.SessionHandle; +import org.apache.hive.service.cli.session.store.HiveSessionSnapshot; +import org.apache.hive.service.cli.session.store.SessionStateStore; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Abstract integration test for Persistable Sessions feature. + * Tests that session state (configs, temp tables, data) is persisted to a shared + * store and recovered by another HS2 instance on failover. + * + * Subclasses provide the store implementation class and infrastructure setup + * (e.g. ZooKeeper TestingServer, Redis container) via {@link #getStoreClassName()} + * and {@link #configureStore(HiveConf)}. + */ +public abstract class TestPersistableSessionBase { + + protected MiniHS2 miniHs2First; + protected MiniHS2 miniHs2Second; + protected HiveConf hiveConf1; + protected HiveConf hiveConf2; + + protected abstract String getStoreClassName(); + + protected abstract void configureStore(HiveConf conf); + + protected abstract SessionStateStore createVerifyStore() throws Exception; + + @Before + public void setUp() throws Exception { + hiveConf1 = new HiveConf(); + hiveConf1.setBoolVar(ConfVars.HIVE_SUPPORT_CONCURRENCY, false); + configurePersistableSession(hiveConf1); + miniHs2First = new MiniHS2.Builder().withConf(hiveConf1).withHTTPTransport() + .cleanupLocalDirOnStartup(false).build(); + + hiveConf2 = new HiveConf(); + hiveConf2.setBoolVar(ConfVars.HIVE_SUPPORT_CONCURRENCY, false); + configurePersistableSession(hiveConf2); + miniHs2Second = new MiniHS2.Builder().withConf(hiveConf2).withHTTPTransport() + .cleanupLocalDirOnStartup(false).build(); + } + + @After + public void tearDown() { + if (miniHs2First != null && miniHs2First.isStarted()) { + miniHs2First.stop(); + } + if (miniHs2Second != null && miniHs2Second.isStarted()) { + miniHs2Second.stop(); + } + } + + private void executeStatementAndWait(CLIServiceClient client, SessionHandle sessHandle, + String statement, Map confOverlay) throws Exception { + OperationHandle opHandle = client.executeStatementAsync(sessHandle, statement, confOverlay); + long timeout = System.currentTimeMillis() + 60000; + while (true) { + OperationState state = client.getOperationStatus(opHandle, false).getState(); + if (state == OperationState.FINISHED) { + break; + } + if (state == OperationState.ERROR || state == OperationState.CANCELED) { + fail("Operation failed with state: " + state + " for statement: " + statement); + } + if (System.currentTimeMillis() > timeout) { + fail("Timed out waiting for: " + statement); + } + Thread.sleep(200); + } + } + + /** + * Polls the session state store until the snapshot satisfies the given condition or times out. + * Needed because the snapshot save runs asynchronously on the background thread after the + * operation state becomes FINISHED (there is a small window between state visibility and + * save completion). + */ + private HiveSessionSnapshot waitForSnapshotCondition(SessionStateStore store, + String storeKey, java.util.function.Predicate condition, + long timeoutMs) throws Exception { + long deadline = System.currentTimeMillis() + timeoutMs; + HiveSessionSnapshot snapshot = null; + while (System.currentTimeMillis() < deadline) { + snapshot = store.getSnapshot(storeKey); + if (snapshot != null && condition.test(snapshot)) { + return snapshot; + } + Thread.sleep(100); + } + return snapshot; + } + + private void configurePersistableSession(HiveConf conf) { + conf.setVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_CLASS, getStoreClassName()); + conf.setVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY, "FETCH_WHEN_MISSING"); + conf.setVar(ConfVars.HIVE_EXECUTION_ENGINE, "tez"); + conf.setBoolean("tez.local.mode", true); + conf.setBoolean("tez.local.mode.without.network", true); + conf.setVar(ConfVars.HIVE_JAR_DIRECTORY, System.getProperty("java.io.tmpdir")); + configureStore(conf); + } + + @Test(timeout = 120000) + public void testSessionRecoveryOnFailover() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition=true", confOverlay); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition.mode=nonstrict", confOverlay); + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY TABLE tmp_failover_test (id INT, name STRING)", confOverlay); + executeStatementAndWait(client1, sessHandle, + "INSERT INTO tmp_failover_test VALUES (1, 'alice'), (2, 'bob'), (3, 'charlie')", confOverlay); + + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Present the SAME session handle — triggers recovery from shared store + OperationHandle opHandle = client2.executeStatement(sessHandle, "SELECT 1", confOverlay); + RowSet rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + + // Verify configs recovered + opHandle = client2.executeStatement(sessHandle, "SET hive.exec.dynamic.partition", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertTrue(rowSet.iterator().next()[0].toString().contains("true")); + + opHandle = client2.executeStatement(sessHandle, + "SET hive.exec.dynamic.partition.mode", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertTrue(rowSet.iterator().next()[0].toString().contains("nonstrict")); + + // Verify temp table data is recovered (shared filesystem, LOCATION preserved) + opHandle = client2.executeStatement(sessHandle, + "SELECT id FROM tmp_failover_test ORDER BY id", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals(3, rowSet.numRows()); + + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testFetchStrategyNeverNoRecovery() throws Exception { + hiveConf1.setVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY, "NEVER"); + miniHs2First = new MiniHS2.Builder().withConf(hiveConf1).withHTTPTransport() + .cleanupLocalDirOnStartup(false).build(); + + hiveConf2.setVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY, "NEVER"); + miniHs2Second = new MiniHS2.Builder().withConf(hiveConf2).withHTTPTransport() + .cleanupLocalDirOnStartup(false).build(); + + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition=true", confOverlay); + + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + try { + client2.executeStatement(sessHandle, "SELECT 1", confOverlay); + fail("Expected HiveSQLException for invalid session handle"); + } catch (HiveSQLException e) { + assertTrue("Expected 'Invalid SessionHandle' error, got: " + e.getMessage(), + e.getMessage().contains("Invalid SessionHandle")); + } + } + + @Test(timeout = 120000) + public void testConfigsRecoveredAfterFailover() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition=true", confOverlay); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition.mode=nonstrict", confOverlay); + executeStatementAndWait(client1, sessHandle, "SET hive.mapred.mode=strict", confOverlay); + + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + OperationHandle opHandle = client2.executeStatement(sessHandle, + "SET hive.exec.dynamic.partition", confOverlay); + RowSet rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertTrue(rowSet.iterator().next()[0].toString().contains("true")); + + opHandle = client2.executeStatement(sessHandle, + "SET hive.exec.dynamic.partition.mode", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertTrue(rowSet.iterator().next()[0].toString().contains("nonstrict")); + + opHandle = client2.executeStatement(sessHandle, "SET hive.mapred.mode", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertTrue(rowSet.iterator().next()[0].toString().contains("strict")); + + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testTempTableWithDataRecoveredAfterFailover() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY TABLE tmp_data_test (id INT, name STRING)", confOverlay); + executeStatementAndWait(client1, sessHandle, + "INSERT INTO tmp_data_test VALUES (10, 'hive'), (20, 'hadoop')", confOverlay); + + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Verify temp table schema is recovered (table exists after failover) + OperationHandle opHandle = client2.executeStatement(sessHandle, + "DESCRIBE tmp_data_test", confOverlay); + RowSet rowSet = client2.fetchResults(opHandle); + assertTrue("Temp table schema should be recovered", rowSet.numRows() > 0); + + // Verify the data is recovered (shared filesystem, LOCATION preserved in DDL) + opHandle = client2.executeStatement(sessHandle, + "SELECT id, name FROM tmp_data_test ORDER BY id", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals(2, rowSet.numRows()); + Iterator it = rowSet.iterator(); + Object[] row1 = it.next(); + assertEquals("10", row1[0].toString()); + assertEquals("hive", row1[1].toString()); + Object[] row2 = it.next(); + assertEquals("20", row2[0].toString()); + assertEquals("hadoop", row2[1].toString()); + + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testQueryRetrySucceedsOnSecondHS2AfterCrash() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + + // Create a regular table and insert data on HS2-1 + client1.executeStatement(sessHandle, + "CREATE TABLE IF NOT EXISTS retry_test (id INT, val STRING)", confOverlay); + client1.executeStatement(sessHandle, + "INSERT INTO retry_test VALUES (1, 'one'), (2, 'two'), (3, 'three')", confOverlay); + + // Verify query works on HS2-1 + OperationHandle opHandle = client1.executeStatement(sessHandle, + "SELECT count(*) FROM retry_test", confOverlay); + RowSet rowSet = client1.fetchResults(opHandle); + assertEquals(3L, Long.parseLong(rowSet.iterator().next()[0].toString())); + + // HS2-1 crashes + miniHs2First.stop(); + + // HS2-2 comes up — same shared store + miniHs2Second.start(confOverlay); + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Retry the query with the SAME session handle on HS2-2 + // Without persistable sessions this would throw "Invalid SessionHandle" + opHandle = client2.executeStatement(sessHandle, + "SELECT count(*) FROM retry_test", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertEquals(3L, Long.parseLong(rowSet.iterator().next()[0].toString())); + + // Further queries on the recovered session also work + opHandle = client2.executeStatement(sessHandle, + "SELECT val FROM retry_test WHERE id = 2", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertEquals("two", rowSet.iterator().next()[0].toString()); + + // Cleanup + client2.executeStatement(sessHandle, "DROP TABLE retry_test", confOverlay); + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testSnapshotWipedOnSessionClose() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition=true", confOverlay); + + SessionStateStore verifyStore = createVerifyStore(); + String storeKey = sessHandle.getHandleIdentifier().getPublicId().toString() + ":" + + sessHandle.getHandleIdentifier().getSecretId().toString(); + assertNotNull("Snapshot should exist before close", verifyStore.getSnapshot(storeKey)); + + client1.closeSession(sessHandle); + + assertNull("Snapshot should be wiped after session close", verifyStore.getSnapshot(storeKey)); + verifyStore.close(); + } + + @Test(timeout = 120000) + public void testTempTablesAcrossDifferentDatabasesRecovered() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + + // Create two databases and a temp table in each + executeStatementAndWait(client1, sessHandle, "CREATE DATABASE IF NOT EXISTS db_alpha", confOverlay); + executeStatementAndWait(client1, sessHandle, "CREATE DATABASE IF NOT EXISTS db_beta", confOverlay); + + executeStatementAndWait(client1, sessHandle, "USE db_alpha", confOverlay); + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY TABLE tmp_cross_db (id INT, label STRING)", confOverlay); + executeStatementAndWait(client1, sessHandle, + "INSERT INTO tmp_cross_db VALUES (1, 'alpha_row')", confOverlay); + + executeStatementAndWait(client1, sessHandle, "USE db_beta", confOverlay); + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY TABLE tmp_cross_db (id INT, label STRING)", confOverlay); + executeStatementAndWait(client1, sessHandle, + "INSERT INTO tmp_cross_db VALUES (2, 'beta_row')", confOverlay); + + // Failover + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Verify the temp table in db_alpha has the correct data + client2.executeStatement(sessHandle, "USE db_alpha", confOverlay); + OperationHandle opHandle = client2.executeStatement(sessHandle, + "SELECT id, label FROM tmp_cross_db", confOverlay); + RowSet rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + Object[] row = rowSet.iterator().next(); + assertEquals("1", row[0].toString()); + assertEquals("alpha_row", row[1].toString()); + + // Verify the temp table in db_beta has the correct data + client2.executeStatement(sessHandle, "USE db_beta", confOverlay); + opHandle = client2.executeStatement(sessHandle, + "SELECT id, label FROM tmp_cross_db", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + row = rowSet.iterator().next(); + assertEquals("2", row[0].toString()); + assertEquals("beta_row", row[1].toString()); + + // Verify the current database is restored to what it was before failover (db_beta) + opHandle = client2.executeStatement(sessHandle, "SELECT current_database()", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals("db_beta", rowSet.iterator().next()[0].toString()); + + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testTempTableWithPartitionsAndComplexTypesRecovered() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY TABLE tmp_complex (" + + "id INT, " + + "info STRUCT, " + + "tags ARRAY, " + + "metadata MAP" + + ") PARTITIONED BY (dt STRING) " + + "TBLPROPERTIES ('custom.key'='custom.value', 'transient_lastDdlTime'='0')", + confOverlay); + executeStatementAndWait(client1, sessHandle, + "INSERT INTO tmp_complex PARTITION(dt='2024-01-01') " + + "VALUES (1, named_struct('name','alice','age',30), " + + "array('tag1','tag2'), map('k1','v1'))", + confOverlay); + + // Verify partition metadata is captured in the snapshot after INSERT + SessionStateStore verifyStore = createVerifyStore(); + String storeKey = sessHandle.getHandleIdentifier().getPublicId().toString() + ":" + + sessHandle.getHandleIdentifier().getSecretId().toString(); + HiveSessionSnapshot snapshot = waitForSnapshotCondition(verifyStore, storeKey, + s -> s.getTempTablePartitionDefinitions() != null + && !s.getTempTablePartitionDefinitions().isEmpty(), + 10000); + assertNotNull("Snapshot should exist after INSERT into partitioned temp table", snapshot); + assertTrue("tempTablePartitionDefinitions should contain partition metadata for tmp_complex", + snapshot.getTempTablePartitionDefinitions().values().stream() + .flatMap(List::stream) + .anyMatch(p -> p.getValues().contains("2024-01-01") + && p.getLocation() != null && !p.getLocation().isEmpty())); + verifyStore.close(); + + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Verify table schema is recovered with complex types and partitions + OperationHandle opHandle = client2.executeStatement(sessHandle, + "DESCRIBE tmp_complex", confOverlay); + RowSet rowSet = client2.fetchResults(opHandle); + assertTrue("Partitioned temp table with complex types should be recovered", + rowSet.numRows() > 0); + + // Verify pre-failover data is queryable (partition metadata + LOCATION restored) + opHandle = client2.executeStatement(sessHandle, + "SELECT id, info.name, tags[0], metadata['k1'], dt FROM tmp_complex " + + "WHERE dt='2024-01-01'", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + Object[] row = rowSet.iterator().next(); + assertEquals("1", row[0].toString()); + assertEquals("alice", row[1].toString()); + assertEquals("tag1", row[2].toString()); + assertEquals("v1", row[3].toString()); + assertEquals("2024-01-01", row[4].toString()); + + // Verify new inserts still work after recovery + executeStatementAndWait(client2, sessHandle, + "INSERT INTO tmp_complex PARTITION(dt='2024-02-01') " + + "VALUES (2, named_struct('name','bob','age',25), " + + "array('x'), map('k2','v2'))", + confOverlay); + opHandle = client2.executeStatement(sessHandle, + "SELECT id, info.name, tags[0], metadata['k2'], dt FROM tmp_complex " + + "WHERE dt='2024-02-01'", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + row = rowSet.iterator().next(); + assertEquals("2", row[0].toString()); + assertEquals("bob", row[1].toString()); + assertEquals("x", row[2].toString()); + assertEquals("v2", row[3].toString()); + assertEquals("2024-02-01", row[4].toString()); + + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testAddedFilesRecoveredAfterFailover() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + + // Create the file before ADD FILE + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY TABLE tmp_file_helper (line STRING)", confOverlay); + executeStatementAndWait(client1, sessHandle, + "INSERT INTO tmp_file_helper VALUES ('hello')", confOverlay); + + executeStatementAndWait(client1, sessHandle, + "INSERT OVERWRITE LOCAL DIRECTORY '/tmp/test_persistable_file_dir' " + + "SELECT 'test_content' FROM tmp_file_helper LIMIT 1", confOverlay); + executeStatementAndWait(client1, sessHandle, + "ADD FILE /tmp/test_persistable_file_dir/000000_0", confOverlay); + + // Verify the snapshot contains the file — poll briefly because the snapshot save + // runs on the background thread after the operation state becomes FINISHED + SessionStateStore verifyStore = createVerifyStore(); + String storeKey = sessHandle.getHandleIdentifier().getPublicId().toString() + ":" + + sessHandle.getHandleIdentifier().getSecretId().toString(); + HiveSessionSnapshot snapshot = waitForSnapshotCondition(verifyStore, storeKey, + s -> !s.getAddedFiles().isEmpty(), 5000); + assertNotNull("Snapshot should exist", snapshot); + assertTrue("addedFiles should contain the file", + snapshot.getAddedFiles().stream() + .anyMatch(f -> f.contains("000000_0"))); + + // Failover to second HS2 + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Verify session is recovered and file resource is available + OperationHandle opHandle = client2.executeStatement(sessHandle, + "SELECT 1", confOverlay); + RowSet rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + + // Verify the recovered snapshot on the second HS2 also has the file + HiveSessionSnapshot recoveredSnapshot = verifyStore.getSnapshot(storeKey); + assertNotNull("Snapshot should exist after recovery", recoveredSnapshot); + assertTrue("addedFiles should be preserved after recovery", + recoveredSnapshot.getAddedFiles().stream() + .anyMatch(f -> f.contains("000000_0"))); + + client2.closeSession(sessHandle); + verifyStore.close(); + } + + @Test(timeout = 120000) + public void testTempFunctionRecoveredAfterFailover() throws Exception { + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + + // Register a temporary function using GenericUDFUpper (always on HS2 classpath) + executeStatementAndWait(client1, sessHandle, + "CREATE TEMPORARY FUNCTION tmp_my_upper AS " + + "'org.apache.hadoop.hive.ql.udf.generic.GenericUDFUpper'", + confOverlay); + + // Verify snapshot has the function captured — poll briefly because the snapshot save + // runs on the background thread after the operation state becomes FINISHED + SessionStateStore verifyStore = createVerifyStore(); + String storeKey = sessHandle.getHandleIdentifier().getPublicId().toString() + ":" + + sessHandle.getHandleIdentifier().getSecretId().toString(); + HiveSessionSnapshot snapshot = waitForSnapshotCondition(verifyStore, storeKey, + s -> !s.getTempFunctionDefinitions().isEmpty(), 5000); + assertNotNull("Snapshot should exist after CREATE TEMPORARY FUNCTION", snapshot); + assertTrue("tempFunctionDefinitions should contain tmp_my_upper, but got: " + + snapshot.getTempFunctionDefinitions(), + snapshot.getTempFunctionDefinitions().stream() + .anyMatch(d -> d.contains("tmp_my_upper"))); + verifyStore.close(); + + // Verify it works before failover + OperationHandle opHandle = client1.executeStatement(sessHandle, + "SELECT tmp_my_upper('hello')", confOverlay); + RowSet rowSet = client1.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + assertEquals("HELLO", rowSet.iterator().next()[0].toString()); + + // Failover + miniHs2First.stop(); + miniHs2Second.start(confOverlay); + + CLIServiceClient client2 = miniHs2Second.getServiceClient(); + + // Verify the temp function is usable after recovery + opHandle = client2.executeStatement(sessHandle, + "SELECT tmp_my_upper('recovered')", confOverlay); + rowSet = client2.fetchResults(opHandle); + assertEquals(1, rowSet.numRows()); + assertEquals("RECOVERED", rowSet.iterator().next()[0].toString()); + + client2.closeSession(sessHandle); + } + + @Test(timeout = 120000) + public void testAlwaysStrategySyncsFromRemoteWhenStale() throws Exception { + hiveConf1.setVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY, "ALWAYS"); + miniHs2First = new MiniHS2.Builder().withConf(hiveConf1).withHTTPTransport() + .cleanupLocalDirOnStartup(false).build(); + + Map confOverlay = new HashMap<>(); + miniHs2First.start(confOverlay); + + CLIServiceClient client1 = miniHs2First.getServiceClient(); + SessionHandle sessHandle = client1.openSession("foo", "bar"); + executeStatementAndWait(client1, sessHandle, "SET hive.exec.dynamic.partition=true", confOverlay); + + // Now simulate another HS2 updating the store with a newer snapshot + // that has an additional config and a later lastAccessTime + SessionStateStore verifyStore = createVerifyStore(); + String storeKey = sessHandle.getHandleIdentifier().getPublicId().toString() + ":" + + sessHandle.getHandleIdentifier().getSecretId().toString(); + HiveSessionSnapshot current = verifyStore.getSnapshot(storeKey); + assertNotNull(current); + + Map updatedConfigs = new HashMap<>(); + if (current.getOverriddenConfigurations() != null) { + updatedConfigs.putAll(current.getOverriddenConfigurations()); + } + updatedConfigs.put("hive.mapred.mode", "strict"); + + HiveSessionSnapshot newerSnapshot = HiveSessionSnapshot.builder() + .sessionHandleId(current.getSessionHandleId()) + .username(current.getUsername()) + .ipAddress(current.getIpAddress()) + .currentDatabase(current.getCurrentDatabase()) + .overriddenConfigurations(updatedConfigs) + .addedJars(current.getAddedJars() != null ? current.getAddedJars() : new ArrayList<>()) + .addedFiles(current.getAddedFiles() != null ? current.getAddedFiles() : new ArrayList<>()) + .tempTableDefinitions(current.getTempTableDefinitions()) + .tempTablePartitionDefinitions(current.getTempTablePartitionDefinitions()) + .tempFunctionDefinitions(current.getTempFunctionDefinitions() != null + ? current.getTempFunctionDefinitions() : new ArrayList<>()) + .protocolVersion(current.getProtocolVersion()) + .creationTime(current.getCreationTime()) + .lastAccessTime(System.currentTimeMillis() + 60000) + .build(); + verifyStore.saveSnapshot(storeKey, newerSnapshot); + + // Access the session again — ALWAYS strategy should detect the remote is newer and re-hydrate + OperationHandle opHandle = client1.executeStatement(sessHandle, + "SET hive.mapred.mode", confOverlay); + RowSet rowSet = client1.fetchResults(opHandle); + assertTrue(rowSet.numRows() > 0); + assertTrue("ALWAYS strategy should have synced hive.mapred.mode=strict from remote", + rowSet.iterator().next()[0].toString().contains("strict")); + + client1.closeSession(sessHandle); + verifyStore.close(); + } +} diff --git a/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionWithRedis.java b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionWithRedis.java new file mode 100644 index 000000000000..c92296c583d4 --- /dev/null +++ b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionWithRedis.java @@ -0,0 +1,82 @@ +/* + * 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.hive.service.cli.session; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.apache.hive.jdbc.miniHS2.MiniHS2; +import org.apache.hive.service.cli.session.store.RedisSessionStateStore; +import org.apache.hive.service.cli.session.store.SessionStateStore; +import org.junit.AfterClass; +import org.junit.Assume; +import org.junit.BeforeClass; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.utility.DockerImageName; + +public class TestPersistableSessionWithRedis extends TestPersistableSessionBase { + + private static GenericContainer redisContainer; + + @BeforeClass + public static void beforeTest() throws Exception { + MiniHS2.cleanupLocalDir(); + try { + redisContainer = new GenericContainer<>(DockerImageName.parse("redis:7-alpine")) + .withExposedPorts(6379); + redisContainer.start(); + } catch (Exception e) { + Assume.assumeTrue("Docker not available, skipping Redis integration tests", false); + } + } + + @AfterClass + public static void afterTest() throws Exception { + if (redisContainer != null) { + redisContainer.stop(); + } + MiniHS2.cleanupLocalDir(); + } + + @Override + protected String getStoreClassName() { + return "org.apache.hive.service.cli.session.store.RedisSessionStateStore"; + } + + @Override + protected void configureStore(HiveConf conf) { + Assume.assumeTrue("Redis container not running", + redisContainer != null && redisContainer.isRunning()); + conf.set(RedisSessionStateStore.CONF_REDIS_HOST, redisContainer.getHost()); + conf.set(RedisSessionStateStore.CONF_REDIS_PORT, + String.valueOf(redisContainer.getMappedPort(6379))); + conf.set(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_TTL.varname, "3600s"); + } + + @Override + protected SessionStateStore createVerifyStore() throws Exception { + HiveConf verifyConf = new HiveConf(); + verifyConf.set(RedisSessionStateStore.CONF_REDIS_HOST, redisContainer.getHost()); + verifyConf.set(RedisSessionStateStore.CONF_REDIS_PORT, + String.valueOf(redisContainer.getMappedPort(6379))); + verifyConf.set(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_TTL.varname, "3600s"); + RedisSessionStateStore store = new RedisSessionStateStore(); + store.init(verifyConf); + return store; + } +} diff --git a/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionWithZooKeeper.java b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionWithZooKeeper.java new file mode 100644 index 000000000000..ef6cd3b6d29c --- /dev/null +++ b/itests/hive-unit/src/test/java/org/apache/hive/service/cli/session/TestPersistableSessionWithZooKeeper.java @@ -0,0 +1,84 @@ +/* + * 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.hive.service.cli.session; + +import java.util.concurrent.TimeUnit; + +import org.apache.curator.test.TestingServer; +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.apache.hive.jdbc.miniHS2.MiniHS2; +import org.apache.hive.service.cli.session.store.SessionStateStore; +import org.apache.hive.service.cli.session.store.ZooKeeperSessionStateStore; +import org.junit.AfterClass; +import org.junit.BeforeClass; + +public class TestPersistableSessionWithZooKeeper extends TestPersistableSessionBase { + + private static TestingServer zkServer; + private static final String ZK_SESSION_PATH = "/test_persistable_sessions"; + + @BeforeClass + public static void beforeTest() throws Exception { + System.setProperty("zookeeper.extendedTypesEnabled", "true"); + MiniHS2.cleanupLocalDir(); + zkServer = new TestingServer(); + zkServer.start(); + } + + @AfterClass + public static void afterTest() throws Exception { + if (zkServer != null) { + zkServer.close(); + zkServer = null; + } + MiniHS2.cleanupLocalDir(); + } + + @Override + protected String getStoreClassName() { + return "org.apache.hive.service.cli.session.store.ZooKeeperSessionStateStore"; + } + + @Override + protected void configureStore(HiveConf conf) { + conf.setVar(ConfVars.HIVE_ZOOKEEPER_QUORUM, zkServer.getConnectString()); + conf.set(ZooKeeperSessionStateStore.CONF_ZK_PATH, ZK_SESSION_PATH); + conf.setTimeDuration(ConfVars.HIVE_ZOOKEEPER_SESSION_TIMEOUT.varname, 2, TimeUnit.SECONDS); + conf.setTimeDuration(ConfVars.HIVE_ZOOKEEPER_CONNECTION_TIMEOUT.varname, 2, TimeUnit.SECONDS); + conf.setTimeDuration(ConfVars.HIVE_ZOOKEEPER_CONNECTION_BASESLEEPTIME.varname, + 100, TimeUnit.MILLISECONDS); + conf.setInt(ConfVars.HIVE_ZOOKEEPER_CONNECTION_MAX_RETRIES.varname, 1); + } + + @Override + protected SessionStateStore createVerifyStore() throws Exception { + HiveConf verifyConf = new HiveConf(); + verifyConf.setVar(ConfVars.HIVE_ZOOKEEPER_QUORUM, zkServer.getConnectString()); + verifyConf.set(ZooKeeperSessionStateStore.CONF_ZK_PATH, ZK_SESSION_PATH); + verifyConf.setTimeDuration(ConfVars.HIVE_ZOOKEEPER_SESSION_TIMEOUT.varname, 2, TimeUnit.SECONDS); + verifyConf.setTimeDuration(ConfVars.HIVE_ZOOKEEPER_CONNECTION_TIMEOUT.varname, 2, TimeUnit.SECONDS); + verifyConf.setTimeDuration(ConfVars.HIVE_ZOOKEEPER_CONNECTION_BASESLEEPTIME.varname, + 100, TimeUnit.MILLISECONDS); + verifyConf.setInt(ConfVars.HIVE_ZOOKEEPER_CONNECTION_MAX_RETRIES.varname, 1); + ZooKeeperSessionStateStore store = new ZooKeeperSessionStateStore(); + store.init(verifyConf); + return store; + } +} diff --git a/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java b/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java index 7e75cd1b3906..eb26338bdacb 100644 --- a/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java +++ b/jdbc/src/java/org/apache/hive/jdbc/HiveConnection.java @@ -311,6 +311,12 @@ protected int getNumRetries() { return maxRetries; } + boolean isPersistableSession() { + String fetchStrategy = connParams.getHiveConfs().get( + Utils.JdbcConnectionParams.HIVE_CONF_PREFIX + ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY.varname); + return fetchStrategy != null && !"NEVER".equalsIgnoreCase(fetchStrategy); + } + @VisibleForTesting protected HiveConnection(String uri, Properties info, IJdbcBrowserClientFactory browserClientFactory) throws SQLException { @@ -560,6 +566,18 @@ private void openTransport() throws Exception { logZkDiscoveryMessage("Connected to " + connParams.getHost() + ":" + connParams.getPort()); } + void reconnect() throws SQLException { + try { + if (transport != null && transport.isOpen()) { + transport.close(); + } + openTransport(); + client = newSynchronizedClient(new TCLIService.Client(new TBinaryProtocol(transport))); + } catch (Exception e) { + throw new SQLException("Failed to reconnect transport", "08S01", e); + } + } + public String getConnectedUrl() { return jdbcUriString; } diff --git a/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java b/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java index 8531780e3f7d..b978593db5fa 100644 --- a/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java +++ b/jdbc/src/java/org/apache/hive/jdbc/HiveStatement.java @@ -341,6 +341,8 @@ public boolean executeAsync(String sql) throws SQLException { return true; } + private static final String DECOMMISSIONED_ERROR = "HiveServer2 is decommissioned or inactive"; + private void runAsyncOnServer(String sql) throws SQLException { checkConnection("execute"); @@ -356,27 +358,57 @@ private void runAsyncOnServer(String sql) throws SQLException { execReq.setRunAsync(true); execReq.setConfOverlay(sessConf); execReq.setQueryTimeout(queryTimeout); - try { - LOG.debug("Submitting statement [{}]: {}", sessHandle, sql); - TExecuteStatementResp execResp = client.ExecuteStatement(execReq); - Utils.verifySuccessWithInfo(execResp.getStatus()); - List infoMessages = execResp.getStatus().getInfoMessages(); - if (infoMessages != null) { - for (String message : infoMessages) { - LOG.info(message); + + int maxRetries = connection.getNumRetries(); + for (int attempt = 0; ; attempt++) { + try { + LOG.debug("Submitting statement [{}]: {}", sessHandle, sql); + TExecuteStatementResp execResp = client.ExecuteStatement(execReq); + Utils.verifySuccessWithInfo(execResp.getStatus()); + List infoMessages = execResp.getStatus().getInfoMessages(); + if (infoMessages != null) { + for (String message : infoMessages) { + LOG.info(message); + } + } + stmtHandle = Optional.of(execResp.getOperationHandle()); + LOG.debug("Running with statement handle: {}", stmtHandle.get()); + return; + } catch (SQLException eS) { + if (isDecommissionedError(eS) && attempt < maxRetries && connection.isPersistableSession()) { + LOG.warn("HiveServer2 is decommissioned. Reconnecting and retrying attempt {} of {}.", + attempt + 1, maxRetries); + try { + Thread.sleep(1000L); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + isLogBeingGenerated = false; + throw eS; + } + try { + connection.reconnect(); + client = connection.getClient(); + } catch (SQLException reconnectEx) { + LOG.error("Failed to reconnect after decommissioning error", reconnectEx); + isLogBeingGenerated = false; + throw eS; + } + continue; } + isLogBeingGenerated = false; + throw eS; + } catch (Exception ex) { + isLogBeingGenerated = false; + throw new SQLException("Failed to run async statement", "08S01", ex); } - stmtHandle = Optional.of(execResp.getOperationHandle()); - LOG.debug("Running with statement handle: {}", stmtHandle.get()); - } catch (SQLException eS) { - isLogBeingGenerated = false; - throw eS; - } catch (Exception ex) { - isLogBeingGenerated = false; - throw new SQLException("Failed to run async statement", "08S01", ex); } } + private static boolean isDecommissionedError(SQLException e) { + String msg = e.getMessage(); + return msg != null && msg.contains(DECOMMISSIONED_ERROR); + } + /** * Poll the result set status by checking if isSetHasResultSet is set * @return diff --git a/jdbc/src/test/org/apache/hive/jdbc/TestHiveStatement.java b/jdbc/src/test/org/apache/hive/jdbc/TestHiveStatement.java index ee88d1c492cc..eb6b74406ce2 100644 --- a/jdbc/src/test/org/apache/hive/jdbc/TestHiveStatement.java +++ b/jdbc/src/test/org/apache/hive/jdbc/TestHiveStatement.java @@ -18,14 +18,23 @@ package org.apache.hive.jdbc; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; import java.sql.Statement; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.hive.service.rpc.thrift.TCLIService.Iface; +import org.apache.hive.service.rpc.thrift.TExecuteStatementReq; import org.apache.hadoop.hive.conf.HiveConf; import org.apache.hive.service.rpc.thrift.TSessionHandle; import org.junit.Test; @@ -139,4 +148,49 @@ public void testaddBatch() throws SQLException { stmt.addBatch(null); } } + + @Test + public void testDecommissionRetryWhenPersistableSessionEnabled() throws Exception { + verifyDecommissionRetryBehavior(true); + } + + @Test + public void testDecommissionNoRetryWhenPersistableSessionDisabled() throws Exception { + verifyDecommissionRetryBehavior(false); + } + + private void verifyDecommissionRetryBehavior(boolean persistableEnabled) throws Exception { + final HiveConnection connection = mock(HiveConnection.class); + final Iface client = mock(Iface.class); + final TSessionHandle handle = mock(TSessionHandle.class); + connection.fetchSize = 100; + + when(connection.getNumRetries()).thenReturn(1); + when(connection.isPersistableSession()).thenReturn(persistableEnabled); + when(connection.getClient()).thenReturn(client); + + AtomicInteger callCount = new AtomicInteger(0); + when(client.ExecuteStatement(any(TExecuteStatementReq.class))).thenAnswer(invocation -> { + if (callCount.getAndIncrement() == 0) { + throw new SQLException( + "Unable to run new queries as HiveServer2 is decommissioned or inactive"); + } + throw new SQLException("Some other error after reconnect"); + }); + + try (HiveStatement stmt = new HiveStatement(connection, client, handle, false, 100)) { + stmt.executeAsync("SELECT 1"); + fail("Expected SQLException"); + } catch (SQLException e) { + if (persistableEnabled) { + assertEquals("Some other error after reconnect", e.getMessage()); + } else { + assertTrue(e.getMessage().contains("decommissioned or inactive")); + } + } + + int expectedCalls = persistableEnabled ? 2 : 1; + verify(connection, times(persistableEnabled ? 1 : 0)).reconnect(); + verify(client, times(expectedCalls)).ExecuteStatement(any(TExecuteStatementReq.class)); + } } diff --git a/packaging/src/main/assembly/src.xml b/packaging/src/main/assembly/src.xml index fac8a76515f7..a7eb633f8e07 100644 --- a/packaging/src/main/assembly/src.xml +++ b/packaging/src/main/assembly/src.xml @@ -100,6 +100,7 @@ ql/**/* serde/**/* service-rpc/**/* + service-session-store/**/* service/**/* shims/**/* storage-api/**/* diff --git a/pom.xml b/pom.xml index d9f299351b3f..235cb9f27023 100644 --- a/pom.xml +++ b/pom.xml @@ -47,6 +47,7 @@ serde service-rpc service + service-session-store streaming llap-common llap-client @@ -243,6 +244,7 @@ 2025-01-01T00:00:00Z 26.0.6 11.28 + 5.1.0 diff --git a/ql/src/java/org/apache/hadoop/hive/ql/session/SessionState.java b/ql/src/java/org/apache/hadoop/hive/ql/session/SessionState.java index 53ee29743449..4aa578404820 100644 --- a/ql/src/java/org/apache/hadoop/hive/ql/session/SessionState.java +++ b/ql/src/java/org/apache/hadoop/hive/ql/session/SessionState.java @@ -832,7 +832,9 @@ private void createSessionDirs(String userName) throws IOException { String sessionId = getSessionId(); // 4. HDFS session path hdfsSessionPath = new Path(hdfsScratchDirURIString, sessionId); - createPath(conf, hdfsSessionPath, scratchDirPermission, false, true); + boolean deleteOnExit = StringUtils.isBlank( + conf.getVar(HiveConf.ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_CLASS)); + createPath(conf, hdfsSessionPath, scratchDirPermission, false, deleteOnExit); conf.set(HDFS_SESSION_PATH_KEY, hdfsSessionPath.toUri().toString()); // 5. hold a lock file in HDFS session dir to indicate the it is in use if (conf.getBoolVar(HiveConf.ConfVars.HIVE_SCRATCH_DIR_LOCK)) { @@ -1136,6 +1138,10 @@ public static Registry getRegistry() { return session != null ? session.registry : null; } + public Registry getSessionRegistry() { + return registry; + } + public static Registry getRegistryForWrite() { Registry registry = getRegistry(); if (registry == null) { diff --git a/service-session-store/pom.xml b/service-session-store/pom.xml new file mode 100644 index 000000000000..3ad6a9ec701c --- /dev/null +++ b/service-session-store/pom.xml @@ -0,0 +1,91 @@ + + + + 4.0.0 + + org.apache.hive + hive + 4.3.0-SNAPSHOT + ../pom.xml + + hive-service-session-store + jar + Hive Service Session Store + + .. + + + + + org.apache.hive + hive-common + ${project.version} + + + + com.fasterxml.jackson.core + jackson-databind + + + + org.apache.curator + curator-client + provided + + + org.apache.curator + curator-framework + provided + + + org.apache.curator + curator-recipes + provided + + + + redis.clients + jedis + ${jedis.version} + provided + + + + org.slf4j + slf4j-api + + + + org.apache.hadoop + hadoop-mapreduce-client-core + test + + + junit + junit + test + + + org.apache.curator + curator-test + test + + + org.testcontainers + testcontainers + test + + + diff --git a/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/HiveSessionSnapshot.java b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/HiveSessionSnapshot.java new file mode 100644 index 000000000000..8c10772cd787 --- /dev/null +++ b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/HiveSessionSnapshot.java @@ -0,0 +1,237 @@ +/* + * 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.hive.service.cli.session.store; + +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.ArrayList; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +public class HiveSessionSnapshot { + + private final String sessionHandleId; + private final String username; + private final String ipAddress; + private final String currentDatabase; + private final Map overriddenConfigurations; + private final List addedJars; + private final List addedFiles; + private final Map tempTableDefinitions; + private final Map> tempTablePartitionDefinitions; + private final List tempFunctionDefinitions; + private final int protocolVersion; + private final long creationTime; + private final long lastAccessTime; + + @JsonCreator + public HiveSessionSnapshot( + @JsonProperty("sessionHandleId") String sessionHandleId, + @JsonProperty("username") String username, + @JsonProperty("ipAddress") String ipAddress, + @JsonProperty("currentDatabase") String currentDatabase, + @JsonProperty("overriddenConfigurations") Map overriddenConfigurations, + @JsonProperty("addedJars") List addedJars, + @JsonProperty("addedFiles") List addedFiles, + @JsonProperty("tempTableDefinitions") Map tempTableDefinitions, + @JsonProperty("tempTablePartitionDefinitions") Map> tempTablePartitionDefinitions, + @JsonProperty("tempFunctionDefinitions") List tempFunctionDefinitions, + @JsonProperty("protocolVersion") int protocolVersion, + @JsonProperty("creationTime") long creationTime, + @JsonProperty("lastAccessTime") long lastAccessTime) { + this.sessionHandleId = sessionHandleId; + this.username = username; + this.ipAddress = ipAddress; + this.currentDatabase = currentDatabase; + this.overriddenConfigurations = overriddenConfigurations != null + ? new HashMap<>(overriddenConfigurations) : Collections.emptyMap(); + this.addedJars = addedJars != null ? new ArrayList<>(addedJars) : Collections.emptyList(); + this.addedFiles = addedFiles != null ? new ArrayList<>(addedFiles) : Collections.emptyList(); + this.tempTableDefinitions = tempTableDefinitions != null + ? new HashMap<>(tempTableDefinitions) : Collections.emptyMap(); + this.tempTablePartitionDefinitions = tempTablePartitionDefinitions != null + ? new HashMap<>(tempTablePartitionDefinitions) : Collections.emptyMap(); + this.tempFunctionDefinitions = tempFunctionDefinitions != null + ? new ArrayList<>(tempFunctionDefinitions) : Collections.emptyList(); + this.protocolVersion = protocolVersion; + this.creationTime = creationTime; + this.lastAccessTime = lastAccessTime; + } + + @JsonProperty("sessionHandleId") + public String getSessionHandleId() { + return sessionHandleId; + } + + @JsonProperty("username") + public String getUsername() { + return username; + } + + @JsonProperty("ipAddress") + public String getIpAddress() { + return ipAddress; + } + + @JsonProperty("currentDatabase") + public String getCurrentDatabase() { + return currentDatabase; + } + + @JsonProperty("overriddenConfigurations") + public Map getOverriddenConfigurations() { + return overriddenConfigurations; + } + + @JsonProperty("addedJars") + public List getAddedJars() { + return addedJars; + } + + @JsonProperty("addedFiles") + public List getAddedFiles() { + return addedFiles; + } + + @JsonProperty("tempTableDefinitions") + public Map getTempTableDefinitions() { + return tempTableDefinitions; + } + + @JsonProperty("tempTablePartitionDefinitions") + public Map> getTempTablePartitionDefinitions() { + return tempTablePartitionDefinitions; + } + + @JsonProperty("tempFunctionDefinitions") + public List getTempFunctionDefinitions() { + return tempFunctionDefinitions; + } + + @JsonProperty("protocolVersion") + public int getProtocolVersion() { + return protocolVersion; + } + + @JsonProperty("creationTime") + public long getCreationTime() { + return creationTime; + } + + @JsonProperty("lastAccessTime") + public long getLastAccessTime() { + return lastAccessTime; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private String sessionHandleId; + private String username; + private String ipAddress; + private String currentDatabase; + private Map overriddenConfigurations; + private List addedJars; + private List addedFiles; + private Map tempTableDefinitions; + private Map> tempTablePartitionDefinitions; + private List tempFunctionDefinitions; + private int protocolVersion; + private long creationTime; + private long lastAccessTime; + + public Builder sessionHandleId(String sessionHandleId) { + this.sessionHandleId = sessionHandleId; + return this; + } + + public Builder username(String username) { + this.username = username; + return this; + } + + public Builder ipAddress(String ipAddress) { + this.ipAddress = ipAddress; + return this; + } + + public Builder currentDatabase(String currentDatabase) { + this.currentDatabase = currentDatabase; + return this; + } + + public Builder overriddenConfigurations(Map overriddenConfigurations) { + this.overriddenConfigurations = overriddenConfigurations; + return this; + } + + public Builder addedJars(List addedJars) { + this.addedJars = addedJars; + return this; + } + + public Builder addedFiles(List addedFiles) { + this.addedFiles = addedFiles; + return this; + } + + public Builder tempTableDefinitions(Map tempTableDefinitions) { + this.tempTableDefinitions = tempTableDefinitions; + return this; + } + + public Builder tempTablePartitionDefinitions( + Map> tempTablePartitionDefinitions) { + this.tempTablePartitionDefinitions = tempTablePartitionDefinitions; + return this; + } + + public Builder tempFunctionDefinitions(List tempFunctionDefinitions) { + this.tempFunctionDefinitions = tempFunctionDefinitions; + return this; + } + + public Builder protocolVersion(int protocolVersion) { + this.protocolVersion = protocolVersion; + return this; + } + + public Builder creationTime(long creationTime) { + this.creationTime = creationTime; + return this; + } + + public Builder lastAccessTime(long lastAccessTime) { + this.lastAccessTime = lastAccessTime; + return this; + } + + public HiveSessionSnapshot build() { + return new HiveSessionSnapshot(sessionHandleId, username, ipAddress, currentDatabase, + overriddenConfigurations, addedJars, addedFiles, tempTableDefinitions, + tempTablePartitionDefinitions, tempFunctionDefinitions, protocolVersion, creationTime, + lastAccessTime); + } + } +} diff --git a/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/RedisSessionStateStore.java b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/RedisSessionStateStore.java new file mode 100644 index 000000000000..08bda6abfa56 --- /dev/null +++ b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/RedisSessionStateStore.java @@ -0,0 +1,125 @@ +/* + * 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.hive.service.cli.session.store; + +import java.util.concurrent.TimeUnit; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import redis.clients.jedis.DefaultJedisClientConfig; +import redis.clients.jedis.HostAndPort; +import redis.clients.jedis.JedisPool; +import redis.clients.jedis.JedisPoolConfig; +import redis.clients.jedis.Jedis; + +public class RedisSessionStateStore implements SessionStateStore { + + private static final Logger LOG = LoggerFactory.getLogger(RedisSessionStateStore.class); + private static final String KEY_PREFIX = "hive:session:"; + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + public static final String CONF_REDIS_HOST = "hive.server2.session.state.store.redis.host"; + public static final String CONF_REDIS_HOST_DEFAULT = "localhost"; + public static final String CONF_REDIS_PORT = "hive.server2.session.state.store.redis.port"; + public static final int CONF_REDIS_PORT_DEFAULT = 6379; + public static final String CONF_REDIS_PASSWORD = "hive.server2.session.state.store.redis.password"; + public static final String CONF_REDIS_SSL = "hive.server2.session.state.store.redis.ssl"; + + private JedisPool jedisPool; + private long ttlSeconds; + + @Override + public void init(HiveConf conf) { + String host = conf.get(CONF_REDIS_HOST, CONF_REDIS_HOST_DEFAULT); + int port = Integer.parseInt(conf.get(CONF_REDIS_PORT, + String.valueOf(CONF_REDIS_PORT_DEFAULT))); + String password = conf.get(CONF_REDIS_PASSWORD); + boolean useSsl = Boolean.parseBoolean(conf.get(CONF_REDIS_SSL, "false")); + this.ttlSeconds = conf.getTimeVar( + ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_TTL, TimeUnit.SECONDS); + + JedisPoolConfig poolConfig = new JedisPoolConfig(); + poolConfig.setMaxTotal(16); + poolConfig.setMaxIdle(8); + poolConfig.setMinIdle(2); + + DefaultJedisClientConfig.Builder clientConfigBuilder = DefaultJedisClientConfig.builder(); + if (password != null && !password.isEmpty()) { + clientConfigBuilder.password(password); + } + if (useSsl) { + clientConfigBuilder.ssl(true); + } + jedisPool = new JedisPool(poolConfig, new HostAndPort(host, port), + clientConfigBuilder.build()); + LOG.info("Initialized RedisSessionStateStore with host={}:{}, ssl={}, ttl={}s", + host, port, useSsl, ttlSeconds); + } + + @Override + public void saveSnapshot(String sessionHandleId, HiveSessionSnapshot snapshot) { + String key = KEY_PREFIX + sessionHandleId; + try (Jedis jedis = jedisPool.getResource()) { + String json = OBJECT_MAPPER.writeValueAsString(snapshot); + jedis.setex(key, ttlSeconds, json); + LOG.debug("Saved session snapshot to Redis: {}", key); + } catch (Exception e) { + LOG.error("Failed to save session snapshot to Redis: {}", key, e); + throw new RuntimeException("Failed to save session snapshot", e); + } + } + + @Override + public HiveSessionSnapshot getSnapshot(String sessionHandleId) { + String key = KEY_PREFIX + sessionHandleId; + try (Jedis jedis = jedisPool.getResource()) { + String json = jedis.get(key); + if (json == null) { + return null; + } + return OBJECT_MAPPER.readValue(json, HiveSessionSnapshot.class); + } catch (Exception e) { + LOG.error("Failed to get session snapshot from Redis: {}", key, e); + throw new RuntimeException("Failed to get session snapshot", e); + } + } + + @Override + public void deleteSnapshot(String sessionHandleId) { + String key = KEY_PREFIX + sessionHandleId; + try (Jedis jedis = jedisPool.getResource()) { + jedis.del(key); + LOG.debug("Deleted session snapshot from Redis: {}", key); + } catch (Exception e) { + LOG.error("Failed to delete session snapshot from Redis: {}", key, e); + throw new RuntimeException("Failed to delete session snapshot", e); + } + } + + @Override + public void close() { + if (jedisPool != null) { + jedisPool.close(); + jedisPool = null; + } + } +} diff --git a/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/SessionStateStore.java b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/SessionStateStore.java new file mode 100644 index 000000000000..ce6d96b16cb3 --- /dev/null +++ b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/SessionStateStore.java @@ -0,0 +1,34 @@ +/* + * 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.hive.service.cli.session.store; + +import org.apache.hadoop.hive.conf.HiveConf; + +public interface SessionStateStore { + + void init(HiveConf conf); + + void saveSnapshot(String sessionHandleId, HiveSessionSnapshot snapshot); + + HiveSessionSnapshot getSnapshot(String sessionHandleId); + + void deleteSnapshot(String sessionHandleId); + + void close(); +} diff --git a/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/TempTablePartitionSnapshot.java b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/TempTablePartitionSnapshot.java new file mode 100644 index 000000000000..0668b168a31a --- /dev/null +++ b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/TempTablePartitionSnapshot.java @@ -0,0 +1,54 @@ +/* + * 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.hive.service.cli.session.store; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** + * Session-local partition metadata for a partitioned temporary table. + * Temp table partitions are not stored in HMS and must be persisted separately. + */ +public class TempTablePartitionSnapshot { + + private final List values; + private final String location; + + @JsonCreator + public TempTablePartitionSnapshot( + @JsonProperty("values") List values, + @JsonProperty("location") String location) { + this.values = values != null ? new ArrayList<>(values) : Collections.emptyList(); + this.location = location; + } + + @JsonProperty("values") + public List getValues() { + return values; + } + + @JsonProperty("location") + public String getLocation() { + return location; + } +} diff --git a/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/ZooKeeperSessionStateStore.java b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/ZooKeeperSessionStateStore.java new file mode 100644 index 000000000000..2dcdca19985d --- /dev/null +++ b/service-session-store/src/main/java/org/apache/hive/service/cli/session/store/ZooKeeperSessionStateStore.java @@ -0,0 +1,144 @@ +/* + * 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.hive.service.cli.session.store; + +import java.util.concurrent.TimeUnit; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.curator.framework.CuratorFramework; +import org.apache.curator.framework.CuratorFrameworkFactory; +import org.apache.curator.retry.ExponentialBackoffRetry; +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.apache.zookeeper.CreateMode; +import org.apache.zookeeper.KeeperException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ZooKeeperSessionStateStore implements SessionStateStore { + + private static final Logger LOG = LoggerFactory.getLogger(ZooKeeperSessionStateStore.class); + + public static final String CONF_ZK_PATH = "hive.server2.session.state.store.zk.path"; + public static final String CONF_ZK_PATH_DEFAULT = "/hive_sessions"; + + private CuratorFramework zkClient; + private String zkBasePath; + private long ttlMillis; + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + @Override + public void init(HiveConf conf) { + String quorum = conf.getVar(ConfVars.HIVE_ZOOKEEPER_QUORUM); + int sessionTimeout = (int) conf.getTimeVar( + ConfVars.HIVE_ZOOKEEPER_SESSION_TIMEOUT, TimeUnit.MILLISECONDS); + int connectionTimeout = (int) conf.getTimeVar( + ConfVars.HIVE_ZOOKEEPER_CONNECTION_TIMEOUT, TimeUnit.MILLISECONDS); + int baseSleepTime = (int) conf.getTimeVar( + ConfVars.HIVE_ZOOKEEPER_CONNECTION_BASESLEEPTIME, TimeUnit.MILLISECONDS); + int maxRetries = conf.getIntVar(ConfVars.HIVE_ZOOKEEPER_CONNECTION_MAX_RETRIES); + + this.zkBasePath = conf.get(CONF_ZK_PATH, CONF_ZK_PATH_DEFAULT); + this.ttlMillis = conf.getTimeVar( + ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_TTL, TimeUnit.MILLISECONDS); + + zkClient = CuratorFrameworkFactory.builder() + .connectString(quorum) + .sessionTimeoutMs(sessionTimeout) + .connectionTimeoutMs(connectionTimeout) + .retryPolicy(new ExponentialBackoffRetry(baseSleepTime, maxRetries)) + .build(); + zkClient.start(); + + try { + if (zkClient.checkExists().forPath(zkBasePath) == null) { + zkClient.create().creatingParentsIfNeeded().forPath(zkBasePath); + } + } catch (Exception e) { + LOG.error("Failed to create ZooKeeper base path: {}", zkBasePath, e); + throw new RuntimeException("Failed to initialize ZooKeeperSessionStateStore", e); + } + + LOG.info("Initialized ZooKeeperSessionStateStore with quorum={}, basePath={}, ttl={}ms", + quorum, zkBasePath, ttlMillis); + } + + @Override + public void saveSnapshot(String sessionHandleId, HiveSessionSnapshot snapshot) { + String path = getNodePath(sessionHandleId); + try { + byte[] data = OBJECT_MAPPER.writeValueAsBytes(snapshot); + if (zkClient.checkExists().forPath(path) != null) { + zkClient.delete().forPath(path); + } + zkClient.create().withTtl(ttlMillis).creatingParentsIfNeeded() + .withMode(CreateMode.PERSISTENT_WITH_TTL) + .forPath(path, data); + LOG.debug("Saved session snapshot to ZooKeeper: {}", path); + } catch (Exception e) { + LOG.error("Failed to save session snapshot to ZooKeeper: {}", path, e); + throw new RuntimeException("Failed to save session snapshot", e); + } + } + + @Override + public HiveSessionSnapshot getSnapshot(String sessionHandleId) { + String path = getNodePath(sessionHandleId); + try { + if (zkClient.checkExists().forPath(path) == null) { + return null; + } + byte[] data = zkClient.getData().forPath(path); + return OBJECT_MAPPER.readValue(data, HiveSessionSnapshot.class); + } catch (KeeperException.NoNodeException e) { + return null; + } catch (Exception e) { + LOG.error("Failed to get session snapshot from ZooKeeper: {}", path, e); + throw new RuntimeException("Failed to get session snapshot", e); + } + } + + @Override + public void deleteSnapshot(String sessionHandleId) { + String path = getNodePath(sessionHandleId); + try { + if (zkClient.checkExists().forPath(path) != null) { + zkClient.delete().forPath(path); + } + LOG.debug("Deleted session snapshot from ZooKeeper: {}", path); + } catch (KeeperException.NoNodeException e) { + // Already gone, ignore + } catch (Exception e) { + LOG.error("Failed to delete session snapshot from ZooKeeper: {}", path, e); + throw new RuntimeException("Failed to delete session snapshot", e); + } + } + + @Override + public void close() { + if (zkClient != null) { + zkClient.close(); + zkClient = null; + } + } + + private String getNodePath(String sessionHandleId) { + return zkBasePath + "/" + sessionHandleId; + } +} diff --git a/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestRedisSessionStateStore.java b/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestRedisSessionStateStore.java new file mode 100644 index 000000000000..bd3393a3bfa1 --- /dev/null +++ b/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestRedisSessionStateStore.java @@ -0,0 +1,66 @@ +/* + * 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.hive.service.cli.session.store; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.junit.AfterClass; +import org.junit.Assume; +import org.junit.BeforeClass; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.utility.DockerImageName; + +import static org.apache.hive.service.cli.session.store.RedisSessionStateStore.CONF_REDIS_HOST; +import static org.apache.hive.service.cli.session.store.RedisSessionStateStore.CONF_REDIS_PORT; + +public class TestRedisSessionStateStore extends TestSessionStateStoreBase { + + private static GenericContainer redisContainer; + + @BeforeClass + public static void startRedis() { + try { + redisContainer = new GenericContainer<>(DockerImageName.parse("redis:7-alpine")) + .withExposedPorts(6379); + redisContainer.start(); + } catch (Exception e) { + // Docker not available, skip tests + Assume.assumeTrue("Docker not available, skipping Redis tests", false); + } + } + + @AfterClass + public static void stopRedis() { + if (redisContainer != null) { + redisContainer.stop(); + } + } + + @Override + protected SessionStateStore createStore() { + Assume.assumeTrue("Redis container not running", redisContainer != null && redisContainer.isRunning()); + HiveConf conf = new HiveConf(); + conf.set(CONF_REDIS_HOST, redisContainer.getHost()); + conf.set(CONF_REDIS_PORT, String.valueOf(redisContainer.getMappedPort(6379))); + conf.set(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_TTL.varname, "3600s"); + RedisSessionStateStore store = new RedisSessionStateStore(); + store.init(conf); + return store; + } +} diff --git a/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestSessionStateStoreBase.java b/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestSessionStateStoreBase.java new file mode 100644 index 000000000000..4268d770cbcd --- /dev/null +++ b/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestSessionStateStoreBase.java @@ -0,0 +1,180 @@ +/* + * 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.hive.service.cli.session.store; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +public abstract class TestSessionStateStoreBase { + + protected SessionStateStore store; + + protected abstract SessionStateStore createStore() throws Exception; + + @Before + public void setUp() throws Exception { + store = createStore(); + } + + @After + public void tearDown() { + if (store != null) { + store.close(); + } + } + + @Test + public void testSaveAndGetSnapshot() { + String sessionId = UUID.randomUUID().toString(); + HiveSessionSnapshot snapshot = createTestSnapshot(sessionId, "testuser", "test_db"); + + store.saveSnapshot(sessionId, snapshot); + HiveSessionSnapshot retrieved = store.getSnapshot(sessionId); + + assertNotNull(retrieved); + assertEquals(sessionId, retrieved.getSessionHandleId()); + assertEquals("testuser", retrieved.getUsername()); + assertEquals("127.0.0.1", retrieved.getIpAddress()); + assertEquals("test_db", retrieved.getCurrentDatabase()); + assertEquals(2, retrieved.getOverriddenConfigurations().size()); + assertEquals("true", retrieved.getOverriddenConfigurations().get("hive.exec.dynamic.partition")); + assertEquals("nonstrict", retrieved.getOverriddenConfigurations().get("hive.exec.dynamic.partition.mode")); + assertEquals(2, retrieved.getAddedJars().size()); + assertEquals("hdfs:///user/hive/udfs/my-udf.jar", retrieved.getAddedJars().get(0)); + assertEquals(2, retrieved.getAddedFiles().size()); + assertEquals("hdfs:///user/hive/files/data.csv", retrieved.getAddedFiles().get(0)); + assertEquals("/tmp/local_file.txt", retrieved.getAddedFiles().get(1)); + assertEquals(1, retrieved.getTempTableDefinitions().size()); + assertEquals("CREATE TEMPORARY TABLE tmp_t (col1 STRING, col2 INT)", + retrieved.getTempTableDefinitions().get("tmp_t")); + assertEquals(10, retrieved.getProtocolVersion()); + assertEquals(1000L, retrieved.getCreationTime()); + assertEquals(2000L, retrieved.getLastAccessTime()); + } + + @Test + public void testDeleteSnapshot() { + String sessionId = UUID.randomUUID().toString(); + HiveSessionSnapshot snapshot = createTestSnapshot(sessionId, "testuser", "default"); + + store.saveSnapshot(sessionId, snapshot); + assertNotNull(store.getSnapshot(sessionId)); + + store.deleteSnapshot(sessionId); + assertNull(store.getSnapshot(sessionId)); + } + + @Test + public void testOverwriteSnapshot() { + String sessionId = UUID.randomUUID().toString(); + + HiveSessionSnapshot snapshot1 = createTestSnapshot(sessionId, "user1", "db1"); + store.saveSnapshot(sessionId, snapshot1); + + HiveSessionSnapshot snapshot2 = createTestSnapshot(sessionId, "user1", "db2"); + store.saveSnapshot(sessionId, snapshot2); + + HiveSessionSnapshot retrieved = store.getSnapshot(sessionId); + assertNotNull(retrieved); + assertEquals("db2", retrieved.getCurrentDatabase()); + } + + @Test + public void testGetNonExistent() { + String sessionId = UUID.randomUUID().toString(); + assertNull(store.getSnapshot(sessionId)); + } + + @Test + public void testMultipleSessions() { + String sessionId1 = UUID.randomUUID().toString(); + String sessionId2 = UUID.randomUUID().toString(); + String sessionId3 = UUID.randomUUID().toString(); + + store.saveSnapshot(sessionId1, createTestSnapshot(sessionId1, "user1", "db1")); + store.saveSnapshot(sessionId2, createTestSnapshot(sessionId2, "user2", "db2")); + store.saveSnapshot(sessionId3, createTestSnapshot(sessionId3, "user3", "db3")); + + assertEquals("db1", store.getSnapshot(sessionId1).getCurrentDatabase()); + assertEquals("db2", store.getSnapshot(sessionId2).getCurrentDatabase()); + assertEquals("db3", store.getSnapshot(sessionId3).getCurrentDatabase()); + + store.deleteSnapshot(sessionId2); + assertNotNull(store.getSnapshot(sessionId1)); + assertNull(store.getSnapshot(sessionId2)); + assertNotNull(store.getSnapshot(sessionId3)); + } + + @Test + public void testDeleteNonExistent() { + String sessionId = UUID.randomUUID().toString(); + store.deleteSnapshot(sessionId); + assertNull(store.getSnapshot(sessionId)); + } + + @Test + public void testSnapshotWipedOnSessionClose() { + String sessionId1 = UUID.randomUUID().toString(); + String sessionId2 = UUID.randomUUID().toString(); + + store.saveSnapshot(sessionId1, createTestSnapshot(sessionId1, "user1", "db1")); + store.saveSnapshot(sessionId2, createTestSnapshot(sessionId2, "user2", "db2")); + + // Simulate session close — snapshot should be completely removed + store.deleteSnapshot(sessionId1); + + assertNull("Snapshot should be wiped after session close", store.getSnapshot(sessionId1)); + // Other sessions remain unaffected + assertNotNull("Other session should still exist", store.getSnapshot(sessionId2)); + assertEquals("db2", store.getSnapshot(sessionId2).getCurrentDatabase()); + } + + protected HiveSessionSnapshot createTestSnapshot(String sessionId, String username, String database) { + Map configs = new HashMap<>(); + configs.put("hive.exec.dynamic.partition", "true"); + configs.put("hive.exec.dynamic.partition.mode", "nonstrict"); + + Map tempTables = new HashMap<>(); + tempTables.put("tmp_t", "CREATE TEMPORARY TABLE tmp_t (col1 STRING, col2 INT)"); + + return HiveSessionSnapshot.builder() + .sessionHandleId(sessionId) + .username(username) + .ipAddress("127.0.0.1") + .currentDatabase(database) + .overriddenConfigurations(configs) + .addedJars(Arrays.asList("hdfs:///user/hive/udfs/my-udf.jar", "hdfs:///user/hive/udfs/other.jar")) + .addedFiles(Arrays.asList("hdfs:///user/hive/files/data.csv", "/tmp/local_file.txt")) + .tempTableDefinitions(tempTables) + .protocolVersion(10) + .creationTime(1000L) + .lastAccessTime(2000L) + .build(); + } +} diff --git a/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestZooKeeperSessionStateStore.java b/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestZooKeeperSessionStateStore.java new file mode 100644 index 000000000000..24d19aa5c843 --- /dev/null +++ b/service-session-store/src/test/java/org/apache/hive/service/cli/session/store/TestZooKeeperSessionStateStore.java @@ -0,0 +1,56 @@ +/* + * 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.hive.service.cli.session.store; + +import org.apache.curator.test.TestingServer; +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.conf.HiveConf.ConfVars; +import org.junit.AfterClass; +import org.junit.BeforeClass; + +import static org.apache.hive.service.cli.session.store.ZooKeeperSessionStateStore.CONF_ZK_PATH; + +public class TestZooKeeperSessionStateStore extends TestSessionStateStoreBase { + + private static TestingServer zkServer; + + @BeforeClass + public static void startZk() throws Exception { + System.setProperty("zookeeper.extendedTypesEnabled", "true"); + zkServer = new TestingServer(); + zkServer.start(); + } + + @AfterClass + public static void stopZk() throws Exception { + if (zkServer != null) { + zkServer.close(); + } + } + + @Override + protected SessionStateStore createStore() { + HiveConf conf = new HiveConf(); + conf.setVar(ConfVars.HIVE_ZOOKEEPER_QUORUM, zkServer.getConnectString()); + conf.set(CONF_ZK_PATH, "/test_hive_sessions"); + ZooKeeperSessionStateStore store = new ZooKeeperSessionStateStore(); + store.init(conf); + return store; + } +} diff --git a/service/pom.xml b/service/pom.xml index 3f224f4e3310..74ec4137a76f 100644 --- a/service/pom.xml +++ b/service/pom.xml @@ -56,6 +56,11 @@ hive-service-rpc ${project.version} + + org.apache.hive + hive-service-session-store + ${project.version} + org.apache.hive hive-llap-server diff --git a/service/src/java/org/apache/hive/service/cli/operation/HiveCommandOperation.java b/service/src/java/org/apache/hive/service/cli/operation/HiveCommandOperation.java index c216851dc808..d3f3a2a5aecb 100644 --- a/service/src/java/org/apache/hive/service/cli/operation/HiveCommandOperation.java +++ b/service/src/java/org/apache/hive/service/cli/operation/HiveCommandOperation.java @@ -44,6 +44,8 @@ import org.apache.hive.service.cli.RowSetFactory; import org.apache.hive.service.cli.TableSchema; import org.apache.hive.service.cli.session.HiveSession; +import org.apache.hive.service.cli.session.HiveSessionImpl; +import org.apache.hive.service.cli.session.PersistableSessionUtils; /** * Executes a HiveCommand @@ -132,6 +134,17 @@ public void runInternal() throws HiveSQLException { setState(OperationState.FINISHED); } + @Override + protected void onNewState(OperationState state, OperationState prevState) { + super.onNewState(state, prevState); + if (state == OperationState.FINISHED) { + HiveSessionImpl impl = PersistableSessionUtils.unwrapSession(parentSession); + if (impl != null) { + impl.onOperationFinished(statement); + } + } + } + /* (non-Javadoc) * @see org.apache.hive.service.cli.operation.Operation#close() */ diff --git a/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java b/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java index 4caa963e2a6f..3ff59f2489ac 100644 --- a/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java +++ b/service/src/java/org/apache/hive/service/cli/operation/SQLOperation.java @@ -78,6 +78,8 @@ import org.apache.hive.service.cli.RowSetFactory; import org.apache.hive.service.cli.TableSchema; import org.apache.hive.service.cli.session.HiveSession; +import org.apache.hive.service.cli.session.HiveSessionImpl; +import org.apache.hive.service.cli.session.PersistableSessionUtils; import org.apache.hive.service.server.ThreadWithGarbageCleanup; import static org.apache.hadoop.hive.shims.HadoopShims.USER_ID; @@ -648,6 +650,10 @@ protected void onNewState(final OperationState state, final OperationState prevS } markQueryMetric(MetricsFactory.getInstance(), MetricsConstant.HS2_SUCCEEDED_QUERIES); queryInfo.updateState(state.toString()); + HiveSessionImpl impl = PersistableSessionUtils.unwrapSession(parentSession); + if (impl != null) { + impl.onOperationFinished(statement); + } break; case INITIALIZED: /* fall through */ diff --git a/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java b/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java index 77c3d321f416..8df76d349628 100644 --- a/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java +++ b/service/src/java/org/apache/hive/service/cli/session/HiveSessionImpl.java @@ -64,6 +64,7 @@ import org.apache.hive.service.cli.RowSet; import org.apache.hive.service.cli.SessionHandle; import org.apache.hive.service.cli.TableSchema; +import org.apache.hive.service.cli.session.store.HiveSessionSnapshot; import org.apache.hive.service.cli.operation.ExecuteStatementOperation; import org.apache.hive.service.cli.operation.GetCatalogsOperation; import org.apache.hive.service.cli.operation.GetColumnsOperation; @@ -819,6 +820,25 @@ private void cleanupSessionLogDir() { } } + public void onOperationFinished(String statement) { + if (sessionManager == null || !sessionManager.isPersistableSessionsEnabled()) { + return; + } + notifyIfStateChanging(statement); + } + + private void notifyIfStateChanging(String statement) { + if (PersistableSessionUtils.shouldPersistSnapshot(statement, sessionState) + && sessionManager != null) { + sessionManager.notifySessionStateChanged(sessionHandle); + } + } + + public HiveSessionSnapshot captureSnapshot() { + return PersistableSessionUtils.captureSnapshot(sessionHandle, username, ipAddress, + sessionState, sessionConf, getProtocolVersion(), creationTime, lastAccessTime); + } + @Override public SessionState getSessionState() { return sessionState; diff --git a/service/src/java/org/apache/hive/service/cli/session/HiveSessionProxy.java b/service/src/java/org/apache/hive/service/cli/session/HiveSessionProxy.java index a01bef71ca62..51b5bfbd5901 100644 --- a/service/src/java/org/apache/hive/service/cli/session/HiveSessionProxy.java +++ b/service/src/java/org/apache/hive/service/cli/session/HiveSessionProxy.java @@ -42,6 +42,10 @@ public HiveSessionProxy(HiveSession hiveSession, UserGroupInformation ugi) { this.ugi = ugi; } + public HiveSession getBaseSession() { + return base; + } + public static HiveSession getProxy(HiveSession hiveSession, UserGroupInformation ugi) throws IllegalArgumentException, HiveSQLException { return (HiveSession)Proxy.newProxyInstance(HiveSession.class.getClassLoader(), diff --git a/service/src/java/org/apache/hive/service/cli/session/PersistableSessionUtils.java b/service/src/java/org/apache/hive/service/cli/session/PersistableSessionUtils.java new file mode 100644 index 000000000000..ab335d2782ce --- /dev/null +++ b/service/src/java/org/apache/hive/service/cli/session/PersistableSessionUtils.java @@ -0,0 +1,471 @@ +/* + * 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.hive.service.cli.session; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.regex.Pattern; + +import java.util.LinkedHashMap; + +import org.apache.hadoop.fs.Path; +import org.apache.hadoop.hive.common.TableName; +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.metastore.Warehouse; +import org.apache.hadoop.hive.metastore.api.AlreadyExistsException; +import org.apache.hadoop.hive.metastore.api.FieldSchema; +import org.apache.hadoop.hive.metastore.api.MetaException; +import org.apache.hadoop.hive.ql.exec.DDLPlanUtils; +import org.apache.hadoop.hive.ql.exec.FunctionInfo; +import org.apache.hadoop.hive.ql.exec.FunctionInfo.FunctionResource; +import org.apache.hadoop.hive.ql.exec.Registry; +import org.apache.hadoop.hive.ql.metadata.HiveException; +import org.apache.hadoop.hive.ql.metadata.Partition; +import org.apache.hadoop.hive.ql.metadata.Table; +import org.apache.hadoop.hive.ql.metadata.TempTable; +import org.apache.hadoop.hive.ql.session.SessionState; +import org.apache.hive.service.cli.HiveSQLException; +import org.apache.hive.service.cli.OperationHandle; +import org.apache.hive.service.cli.SessionHandle; +import org.apache.hive.service.cli.session.store.HiveSessionSnapshot; +import org.apache.hive.service.cli.session.store.SessionStateStore; +import org.apache.hive.service.cli.session.store.TempTablePartitionSnapshot; +import org.apache.hive.service.rpc.thrift.TProtocolVersion; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Utility methods for the Persistable Sessions feature. + * Extracted from SessionManager and HiveSessionImpl to keep + * the feature's logic isolated from existing core classes. + */ +public final class PersistableSessionUtils { + + public enum FetchStrategy { + NEVER, + ALWAYS, + FETCH_WHEN_MISSING + } + + private static final Logger LOG = LoggerFactory.getLogger(PersistableSessionUtils.class); + + /** + * Returns a store key that incorporates both the public and secret UUIDs + * of the session handle, preventing hijacking via the public ID alone. + */ + public static String storeKey(SessionHandle sessionHandle) { + return sessionHandle.getHandleIdentifier().getPublicId().toString() + ":" + + sessionHandle.getHandleIdentifier().getSecretId().toString(); + } + + private static final Pattern STATE_CHANGING_PATTERN = Pattern.compile( + "(?i)^\\s*(USE\\b|SET\\b|ADD\\s+(JAR|FILE)\\b|DELETE\\s+(JAR|FILE)\\b" + + "|(CREATE|DROP)\\s+TEMPORARY\\s+(TABLE|FUNCTION)\\b).*"); + + private static final Pattern INSERT_TARGET_PATTERN = Pattern.compile( + "(?i)^\\s*INSERT\\s+(?:INTO|OVERWRITE)\\s+(?:TABLE\\s+)?([^\\s(]+)"); + private static final Pattern LOAD_TARGET_PATTERN = Pattern.compile( + "(?i)^\\s*LOAD\\s+DATA\\s+(?:LOCAL\\s+)?INPATH\\s+.+\\s+INTO\\s+TABLE\\s+([^\\s(]+)"); + + private PersistableSessionUtils() { + } + + /** + * Determines whether a SQL statement changes session state that should + * be persisted (database, configs, JARs, temp tables, temp functions). + */ + public static boolean isStateChangingCommand(String statement) { + if (statement == null) { + return false; + } + return STATE_CHANGING_PATTERN.matcher(statement).matches(); + } + + /** + * Returns true when a finished statement may have changed persisted session state, + * including DML that adds data or partition metadata to temporary tables. + */ + public static boolean shouldPersistSnapshot(String statement, SessionState sessionState) { + if (statement == null) { + return false; + } + if (isStateChangingCommand(statement)) { + return true; + } + String targetTable = extractTempTableDmlTarget(statement); + return targetTable != null && isSessionTempTable(sessionState, targetTable); + } + + /** + * Extracts the target table from INSERT or LOAD DATA statements, or returns null. + */ + static String extractTempTableDmlTarget(String statement) { + if (statement == null) { + return null; + } + java.util.regex.Matcher insertMatcher = INSERT_TARGET_PATTERN.matcher(statement); + if (insertMatcher.find()) { + return normalizeTableReference(insertMatcher.group(1)); + } + java.util.regex.Matcher loadMatcher = LOAD_TARGET_PATTERN.matcher(statement); + if (loadMatcher.find()) { + return normalizeTableReference(loadMatcher.group(1)); + } + return null; + } + + private static String normalizeTableReference(String tableRef) { + return tableRef.replace("`", ""); + } + + /** + * Returns true if the given table reference resolves to a session-local temp table. + */ + static boolean isSessionTempTable(SessionState sessionState, String tableReference) { + if (sessionState == null || tableReference == null) { + return false; + } + try { + TableName tableName = TableName.fromString(tableReference, null, + sessionState.getCurrentDatabase()); + Map> tempTables = sessionState.getTempTables(); + if (tempTables == null || tempTables.isEmpty()) { + return false; + } + Map dbTables = tempTables.get(tableName.getDb().toLowerCase()); + return dbTables != null && dbTables.containsKey(tableName.getTable().toLowerCase()); + } catch (IllegalArgumentException e) { + return false; + } + } + + /** + * Captures the current session state into a snapshot DTO. + */ + public static HiveSessionSnapshot captureSnapshot(SessionHandle sessionHandle, + String username, String ipAddress, SessionState sessionState, + HiveConf sessionConf, TProtocolVersion protocol, + long creationTime, long lastAccessTime) { + List jars = new ArrayList<>(); + List files = new ArrayList<>(); + if (sessionState != null) { + Set jarSet = sessionState.list_resource(SessionState.ResourceType.JAR, null); + if (jarSet != null) { + jars.addAll(jarSet); + } + Set fileSet = sessionState.list_resource(SessionState.ResourceType.FILE, null); + if (fileSet != null) { + files.addAll(fileSet); + } + } + + Map tempTableDefs = new HashMap<>(); + Map> tempTablePartitionDefs = new HashMap<>(); + if (sessionState != null && sessionState.getTempTables() != null) { + DDLPlanUtils ddlPlanUtils = new DDLPlanUtils(); + for (Map.Entry> dbEntry : + sessionState.getTempTables().entrySet()) { + String dbName = dbEntry.getKey(); + for (Map.Entry tableEntry : dbEntry.getValue().entrySet()) { + String tableName = tableEntry.getKey(); + Table table = tableEntry.getValue(); + String tableKey = TableName.getDbTable(dbName, tableName); + String ddl = generateTempTableDDL(ddlPlanUtils, table); + if (ddl != null) { + tempTableDefs.put(tableKey, ddl); + } + List partitionSnapshots = captureTempTablePartitions( + sessionState, dbName, tableName, table); + if (!partitionSnapshots.isEmpty()) { + tempTablePartitionDefs.put(tableKey, partitionSnapshots); + } + } + } + } + + List tempFuncDefs = captureTempFunctions(sessionState); + + return HiveSessionSnapshot.builder() + .sessionHandleId(storeKey(sessionHandle)) + .username(username) + .ipAddress(ipAddress) + .currentDatabase(sessionState != null ? sessionState.getCurrentDatabase() : null) + .overriddenConfigurations(sessionState != null + ? new HashMap<>(sessionState.getOverriddenConfigurations()) : null) + .addedJars(jars) + .addedFiles(files) + .tempTableDefinitions(tempTableDefs) + .tempTablePartitionDefinitions(tempTablePartitionDefs) + .tempFunctionDefinitions(tempFuncDefs) + .protocolVersion(protocol.getValue()) + .creationTime(creationTime) + .lastAccessTime(lastAccessTime) + .build(); + } + + /** + * Generates the CREATE TEMPORARY TABLE DDL for a temp table using DDLPlanUtils, + * which handles partitions, table properties, complex types, bucket specs, etc. + */ + static String generateTempTableDDL(DDLPlanUtils ddlPlanUtils, Table table) { + try { + return ddlPlanUtils.getCreateTableCommand(table, true); + } catch (Exception e) { + LOG.warn("Failed to generate DDL for temp table: {}", table.getTableName(), e); + return null; + } + } + + /** + * Captures session-local partition metadata for a partitioned temp table. + */ + static List captureTempTablePartitions(SessionState sessionState, + String dbName, String tableName, Table table) { + List partitions = new ArrayList<>(); + if (sessionState == null || !table.isPartitioned()) { + return partitions; + } + Map tempPartitions = sessionState.getTempPartitions(); + if (tempPartitions == null || tempPartitions.isEmpty()) { + return partitions; + } + String qualifiedKey = Warehouse.getQualifiedName(dbName.toLowerCase(), tableName.toLowerCase()); + TempTable tempTable = tempPartitions.get(qualifiedKey); + if (tempTable == null) { + return partitions; + } + for (org.apache.hadoop.hive.metastore.api.Partition apiPartition : tempTable.listPartitions()) { + String location = apiPartition.getSd() != null ? apiPartition.getSd().getLocation() : null; + partitions.add(new TempTablePartitionSnapshot( + new ArrayList<>(apiPartition.getValues()), location)); + } + return partitions; + } + + /** + * Captures temporary function definitions as CREATE TEMPORARY FUNCTION DDL statements. + * Uses the passed sessionState's registry directly rather than the thread-local, + * since the snapshot may be captured on a thread different from the session's own. + */ + static List captureTempFunctions(SessionState sessionState) { + List funcDefs = new ArrayList<>(); + if (sessionState == null) { + return funcDefs; + } + Registry registry = sessionState.getSessionRegistry(); + if (registry == null) { + return funcDefs; + } + for (String funcName : registry.getCurrentFunctionNames()) { + try { + FunctionInfo info = registry.getFunctionInfo(funcName); + if (info == null || info.getFunctionType() != FunctionInfo.FunctionType.TEMPORARY) { + continue; + } + String className = info.getClassName(); + if (className == null) { + Class funcClass = info.getFunctionClass(); + if (funcClass != null) { + className = funcClass.getName(); + } + } + if (className == null) { + continue; + } + StringBuilder ddl = new StringBuilder("CREATE TEMPORARY FUNCTION "); + ddl.append(funcName).append(" AS '").append(className).append("'"); + FunctionResource[] resources = info.getResources(); + if (resources != null && resources.length > 0) { + ddl.append(" USING"); + for (int i = 0; i < resources.length; i++) { + if (i > 0) { + ddl.append(","); + } + ddl.append(" ").append(resources[i].getResourceType().name()) + .append(" '").append(resources[i].getResourceURI()).append("'"); + } + } + funcDefs.add(ddl.toString()); + } catch (Exception e) { + LOG.warn("Failed to capture temp function: {}", funcName, e); + } + } + return funcDefs; + } + + /** + * Hydrates a recovered session from a snapshot: restores database, configs, + * JARs, files, temp functions, and temp tables. + */ + public static void hydrateSession(HiveSession session, HiveSessionSnapshot snapshot) + throws HiveSQLException { + try { + SessionState sessionState = session.getSessionState(); + if (snapshot.getCurrentDatabase() != null) { + sessionState.setCurrentDatabase(snapshot.getCurrentDatabase()); + } + if (snapshot.getOverriddenConfigurations() != null) { + for (Map.Entry entry : snapshot.getOverriddenConfigurations().entrySet()) { + session.getHiveConf().set(entry.getKey(), entry.getValue()); + sessionState.getOverriddenConfigurations().put(entry.getKey(), entry.getValue()); + } + } + if (snapshot.getAddedJars() != null) { + for (String jar : snapshot.getAddedJars()) { + sessionState.add_resource(SessionState.ResourceType.JAR, jar); + } + } + if (snapshot.getAddedFiles() != null) { + for (String file : snapshot.getAddedFiles()) { + sessionState.add_resource(SessionState.ResourceType.FILE, file); + } + } + if (snapshot.getTempFunctionDefinitions() != null) { + restoreTempFunctions(session, snapshot.getTempFunctionDefinitions()); + } + if (snapshot.getTempTableDefinitions() != null) { + restoreTempTables(session, sessionState, snapshot.getTempTableDefinitions(), + snapshot.getTempTablePartitionDefinitions()); + } + } catch (Exception e) { + LOG.error("Failed to hydrate session: {}", session.getSessionHandle(), e); + throw new HiveSQLException("Failed to hydrate recovered session", e); + } + } + + private static void restoreTempFunctions(HiveSession session, List tempFuncDefs) { + for (String ddl : tempFuncDefs) { + try { + OperationHandle opHandle = session.executeStatement(ddl, null); + session.closeOperation(opHandle); + } catch (Exception e) { + LOG.warn("Failed to restore temporary function: {}", ddl, e); + } + } + } + + private static void restoreTempTables(HiveSession session, SessionState sessionState, + Map tempTableDefs, Map> tempTablePartitionDefs) + throws HiveSQLException { + String currentDb = sessionState.getCurrentDatabase(); + for (Map.Entry entry : tempTableDefs.entrySet()) { + try { + TableName tn = TableName.fromString(entry.getKey(), null, currentDb); + String db = tn.getDb(); + if (!db.equals(sessionState.getCurrentDatabase())) { + sessionState.setCurrentDatabase(db); + } + OperationHandle opHandle = session.executeStatement(entry.getValue(), null); + session.closeOperation(opHandle); + List partitions = tempTablePartitionDefs != null + ? tempTablePartitionDefs.get(entry.getKey()) : null; + if (partitions != null && !partitions.isEmpty()) { + Table table = sessionState.getTempTables().get(db).get(tn.getTable()); + restoreTempTablePartitions(sessionState, table, partitions); + } + } catch (Exception e) { + LOG.warn("Failed to restore temporary table {}", entry.getKey(), e); + throw new HiveSQLException("Failed to restore temporary table " + entry.getKey(), e); + } + } + sessionState.setCurrentDatabase(currentDb); + } + + private static void restoreTempTablePartitions(SessionState sessionState, Table table, + List partitions) + throws HiveException, MetaException, AlreadyExistsException { + String qualifiedKey = Warehouse.getQualifiedName( + table.getDbName().toLowerCase(), table.getTableName().toLowerCase()); + TempTable tempTable = sessionState.getTempPartitions().get(qualifiedKey); + if (tempTable == null) { + throw new HiveException("TempTable partition metadata missing for " + qualifiedKey); + } + List toAdd = new ArrayList<>(partitions.size()); + List partCols = table.getPartitionKeys(); + for (TempTablePartitionSnapshot snapshot : partitions) { + Map partSpec = new LinkedHashMap<>(); + List values = snapshot.getValues(); + for (int i = 0; i < partCols.size(); i++) { + partSpec.put(partCols.get(i).getName(), values.get(i)); + } + Path location = snapshot.getLocation() != null ? new Path(snapshot.getLocation()) : null; + Partition qlPart = new Partition(table, partSpec, location); + toAdd.add(qlPart.getTPartition()); + } + tempTable.addPartitions(toAdd, true); + } + + /** + * Unwraps a HiveSession proxy to get the underlying HiveSessionImpl. + * Returns null if the session cannot be unwrapped. + */ + public static HiveSessionImpl unwrapSession(HiveSession session) { + if (session instanceof HiveSessionImpl impl) { + return impl; + } + if (Proxy.isProxyClass(session.getClass())) { + InvocationHandler handler = Proxy.getInvocationHandler(session); + if (handler instanceof HiveSessionProxy proxy) { + HiveSession base = proxy.getBaseSession(); + if (base instanceof HiveSessionImpl impl) { + return impl; + } + } + } + return null; + } + + /** + * Saves the session snapshot to the state store. + */ + public static void saveSnapshot(SessionStateStore store, HiveSession session) { + if (store == null) { + return; + } + try { + HiveSessionImpl impl = unwrapSession(session); + if (impl == null) { + return; + } + HiveSessionSnapshot snapshot = impl.captureSnapshot(); + store.saveSnapshot(storeKey(session.getSessionHandle()), snapshot); + } catch (Exception e) { + LOG.warn("Failed to save session snapshot for: {}", session.getSessionHandle(), e); + } + } + + /** + * Deletes the session snapshot from the state store. + */ + public static void deleteSnapshot(SessionStateStore store, SessionHandle sessionHandle) { + if (store == null) { + return; + } + try { + store.deleteSnapshot(storeKey(sessionHandle)); + } catch (Exception e) { + LOG.warn("Failed to delete session snapshot for: {}", sessionHandle, e); + } + } +} diff --git a/service/src/java/org/apache/hive/service/cli/session/SessionManager.java b/service/src/java/org/apache/hive/service/cli/session/SessionManager.java index c792eb6bbd92..210d94b0b7ae 100644 --- a/service/src/java/org/apache/hive/service/cli/session/SessionManager.java +++ b/service/src/java/org/apache/hive/service/cli/session/SessionManager.java @@ -28,6 +28,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.BlockingQueue; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Future; @@ -58,6 +59,9 @@ import org.apache.hive.service.cli.operation.Operation; import org.apache.hive.service.cli.operation.OperationLogManager; import org.apache.hive.service.cli.operation.OperationManager; +import org.apache.hive.service.cli.session.store.HiveSessionSnapshot; +import org.apache.hive.service.cli.session.store.SessionStateStore; +import org.apache.hive.service.cli.session.PersistableSessionUtils.FetchStrategy; import org.apache.hive.service.rpc.thrift.TOpenSessionReq; import org.apache.hive.service.rpc.thrift.TProtocolVersion; import org.apache.hive.service.server.HiveServer2; @@ -110,7 +114,12 @@ public class SessionManager extends CompositeService { private String sessionImplclassName; private CleanupService cleanupService; // Tracks which LLAP target gauges have been lazily registered. - private final java.util.Set registeredLlapTargetGauges = ConcurrentHashMap.newKeySet(); + private final Set registeredLlapTargetGauges = ConcurrentHashMap.newKeySet(); + // Persistable session state store + private SessionStateStore sessionStateStore; + private FetchStrategy fetchStrategy; + private final ThreadLocal recoveringSession = ThreadLocal.withInitial(() -> false); + private final ConcurrentHashMap recoverySessions = new ConcurrentHashMap<>(); public SessionManager(HiveServer2 hiveServer2, boolean allowSessions) { super(SessionManager.class.getSimpleName()); @@ -136,8 +145,8 @@ public synchronized void init(HiveConf hiveConf) { initSessionImplClassName(); Metrics metrics = MetricsFactory.getInstance(); if(metrics != null){ - registerOpenSesssionMetrics(metrics); - registerActiveSesssionMetrics(metrics); + registerOpenSessionMetrics(metrics); + registerActiveSessionMetrics(metrics); } userLimit = hiveConf.getIntVar(ConfVars.HIVE_SERVER2_LIMIT_CONNECTIONS_PER_USER); @@ -154,10 +163,36 @@ public synchronized void init(HiveConf hiveConf) { cleanupService = SyncCleanupService.INSTANCE; } cleanupService.start(); + initSessionStateStore(); super.init(hiveConf); } - private void registerOpenSesssionMetrics(Metrics metrics) { + private void initSessionStateStore() { + String storeClassName = hiveConf.getVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_CLASS); + if (storeClassName == null || storeClassName.isEmpty()) { + LOG.info("Session state store not configured. Persistable sessions disabled."); + this.sessionStateStore = null; + this.fetchStrategy = FetchStrategy.NEVER; + return; + } + this.fetchStrategy = FetchStrategy.valueOf( + hiveConf.getVar(ConfVars.HIVE_SERVER2_SESSION_STATE_STORE_FETCH_STRATEGY)); + try { + Class storeClass = Class.forName(storeClassName); + this.sessionStateStore = (SessionStateStore) storeClass.getDeclaredConstructor().newInstance(); + this.sessionStateStore.init(hiveConf); + LOG.info("Initialized session state store: {}, fetch strategy: {}", storeClassName, fetchStrategy); + } catch (ClassNotFoundException e) { + LOG.warn("Session state store class not found: {}. Persistable sessions disabled.", storeClassName); + this.sessionStateStore = null; + this.fetchStrategy = FetchStrategy.NEVER; + } catch (Exception e) { + LOG.error("Failed to initialize session state store: {}", storeClassName, e); + throw new RuntimeException("Failed to initialize session state store", e); + } + } + + private void registerOpenSessionMetrics(Metrics metrics) { MetricsVariable openSessionCnt = new MetricsVariable() { @Override public Integer getValue() { @@ -180,7 +215,7 @@ public Integer getValue() { metrics.addRatio(MetricsConstant.HS2_AVG_OPEN_SESSION_TIME, openSessionTime, openSessionCnt); } - private void registerActiveSesssionMetrics(Metrics metrics) { + private void registerActiveSessionMetrics(Metrics metrics) { MetricsVariable activeSessionCnt = new MetricsVariable() { @Override public Integer getValue() { @@ -432,6 +467,13 @@ public synchronized void stop() { } cleanupLoggingRootDir(); logManager.ifPresent(lm -> lm.stop()); + if (sessionStateStore != null) { + try { + sessionStateStore.close(); + } catch (Exception e) { + LOG.warn("Error closing session state store", e); + } + } } private void cleanupLoggingRootDir() { @@ -577,6 +619,9 @@ public HiveSession createSession(SessionHandle sessionHandle, TProtocolVersion p throw new HiveSQLException(FAIL_CLOSE_ERROR_MESSAGE); } registerLlapTargetGaugeIfNeeded(session); + if (!recoveringSession.get()) { + saveSessionSnapshot(session); + } LOG.info("Session opened, " + session.getSessionHandle() + ", current sessions:" + getOpenSessionCount()); return session; @@ -684,6 +729,7 @@ public void closeSession(SessionHandle sessionHandle) throws HiveSQLException { } LOG.info("Session closed, " + sessionHandle + ", current sessions:" + getOpenSessionCount()); } + deleteSessionSnapshot(sessionHandle); closeSessionInternal(session); } @@ -715,10 +761,80 @@ public void run() { public HiveSession getSession(SessionHandle sessionHandle) throws HiveSQLException { HiveSession session = handleToSession.get(sessionHandle); - if (session == null) { + if (session != null) { + if (fetchStrategy == FetchStrategy.ALWAYS && sessionStateStore != null) { + syncFromRemoteIfStale(session); + } + return session; + } + if (fetchStrategy == FetchStrategy.NEVER) { throw new HiveSQLException("Invalid SessionHandle: " + sessionHandle); } - return session; + try { + return recoverySessions.computeIfAbsent(sessionHandle, handle -> { + try { + return recoverSession(handle); + } catch (HiveSQLException e) { + throw new RuntimeException(e); + } + }); + } catch (RuntimeException e) { + if (e.getCause() instanceof HiveSQLException) { + throw (HiveSQLException) e.getCause(); + } + throw e; + } finally { + recoverySessions.remove(sessionHandle); + } + } + + private void syncFromRemoteIfStale(HiveSession session) { + try { + String handleId = PersistableSessionUtils.storeKey(session.getSessionHandle()); + HiveSessionSnapshot remoteSnapshot = sessionStateStore.getSnapshot(handleId); + if (remoteSnapshot == null) { + return; + } + if (remoteSnapshot.getLastAccessTime() > session.getLastAccessTime()) { + LOG.info("Remote snapshot is newer for session {}, re-hydrating", session.getSessionHandle()); + hydrateSession(session, remoteSnapshot); + } + } catch (Exception e) { + LOG.warn("Failed to sync session from remote store: {}", session.getSessionHandle(), e); + } + } + + private HiveSession recoverSession(SessionHandle sessionHandle) throws HiveSQLException { + String handleId = PersistableSessionUtils.storeKey(sessionHandle); + HiveSessionSnapshot snapshot = sessionStateStore.getSnapshot(handleId); + if (snapshot == null) { + throw new HiveSQLException("Invalid SessionHandle: " + sessionHandle); + } + LOG.info("Recovering session from state store: {}", sessionHandle); + TProtocolVersion protocol = TProtocolVersion.findByValue(snapshot.getProtocolVersion()); + if (protocol == null) { + protocol = TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V1; + } + SessionHandle recoveredHandle = new SessionHandle( + sessionHandle.getHandleIdentifier(), protocol); + boolean withImpersonation = hiveConf.getBoolVar(ConfVars.HIVE_SERVER2_ENABLE_DOAS) + && snapshot.getUsername() != null; + recoveringSession.set(true); + try { + HiveSession recovered = createSession(recoveredHandle, protocol, + snapshot.getUsername(), null, snapshot.getIpAddress(), + null, withImpersonation, null); + hydrateSession(recovered, snapshot); + saveSessionSnapshot(recovered); + LOG.info("Successfully recovered session: {}", sessionHandle); + return recovered; + } finally { + recoveringSession.remove(); + } + } + + private void hydrateSession(HiveSession session, HiveSessionSnapshot snapshot) throws HiveSQLException { + PersistableSessionUtils.hydrateSession(session, snapshot); } public OperationManager getOperationManager() { @@ -838,5 +954,35 @@ public void allowSessions(boolean b) { this.allowSessions = b; } } + + public boolean isPersistableSessionsEnabled() { + return sessionStateStore != null; + } + + + public void notifySessionStateChanged(SessionHandle sessionHandle) { + HiveSession session = handleToSession.get(sessionHandle); + if (session != null) { + PersistableSessionUtils.saveSnapshot(sessionStateStore, session); + } + } + + private void saveSessionSnapshot(HiveSession session) { + PersistableSessionUtils.saveSnapshot(sessionStateStore, session); + } + + private void deleteSessionSnapshot(SessionHandle sessionHandle) { + PersistableSessionUtils.deleteSnapshot(sessionStateStore, sessionHandle); + } + + @VisibleForTesting + public SessionStateStore getSessionStateStore() { + return sessionStateStore; + } + + @VisibleForTesting + public FetchStrategy getFetchStrategy() { + return fetchStrategy; + } } diff --git a/service/src/test/org/apache/hive/service/cli/session/TestPersistableSessionUtils.java b/service/src/test/org/apache/hive/service/cli/session/TestPersistableSessionUtils.java new file mode 100644 index 000000000000..4d13ab8d038c --- /dev/null +++ b/service/src/test/org/apache/hive/service/cli/session/TestPersistableSessionUtils.java @@ -0,0 +1,74 @@ +/* + * 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.hive.service.cli.session; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.util.HashMap; +import java.util.Map; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.ql.metadata.Table; +import org.apache.hadoop.hive.ql.session.SessionState; +import org.junit.Before; +import org.junit.Test; + +public class TestPersistableSessionUtils { + + private SessionState sessionState; + + @Before + public void setUp() { + sessionState = new SessionState(new HiveConf()); + sessionState.setCurrentDatabase("default"); + Map tempTables = new HashMap<>(); + tempTables.put("tmp", new Table("default", "tmp")); + sessionState.getTempTables().put("default", tempTables); + } + + @Test + public void testShouldPersistSnapshotForInsertIntoTempTable() { + assertTrue(PersistableSessionUtils.shouldPersistSnapshot( + "INSERT INTO tmp PARTITION(dt='2024-01-01') VALUES (1)", sessionState)); + assertTrue(PersistableSessionUtils.shouldPersistSnapshot( + "insert overwrite table tmp partition (dt='x') select 1", sessionState)); + assertTrue(PersistableSessionUtils.shouldPersistSnapshot( + "LOAD DATA INPATH '/tmp/data' INTO TABLE tmp PARTITION (dt='x')", sessionState)); + } + + @Test + public void testShouldNotPersistSnapshotForDmlOnPersistentTable() { + assertFalse(PersistableSessionUtils.shouldPersistSnapshot( + "INSERT INTO permanent_table VALUES (1)", sessionState)); + assertFalse(PersistableSessionUtils.shouldPersistSnapshot( + "INSERT OVERWRITE TABLE permanent_table SELECT 1", sessionState)); + assertFalse(PersistableSessionUtils.shouldPersistSnapshot( + "LOAD DATA INPATH '/tmp/data' INTO TABLE permanent_table", sessionState)); + } + + @Test + public void testShouldPersistSnapshotForStateChangingCommands() { + assertTrue(PersistableSessionUtils.shouldPersistSnapshot( + "SET hive.exec.mode=strict", sessionState)); + assertTrue(PersistableSessionUtils.shouldPersistSnapshot( + "CREATE TEMPORARY TABLE tmp (id INT)", sessionState)); + assertFalse(PersistableSessionUtils.shouldPersistSnapshot("SELECT 1", sessionState)); + } +}