@JsonCreator public static StageId valueOf(String stageId) { List<String> ids = QueryId.parseDottedId(stageId, 2, "stageId"); return new StageId(new QueryId(ids.get(0)), ids.get(1)); }
@JsonCreator public static QueryId valueOf(String queryId) { List<String> ids = parseDottedId(queryId, 1, "queryId"); return new QueryId(ids.get(0)); }
public StageId getStageId() { List<String> ids = QueryId.parseDottedId(fullId, 3, "taskId"); return new StageId(new QueryId(ids.get(0)), ids.get(1)); }
public StageId(String queryId, String id) { this.queryId = new QueryId(queryId); this.id = validateId(id); }
public QueryId getQueryId() { return new QueryId(QueryId.parseDottedId(fullId, 3, "taskId").get(0)); }
@UsedByGeneratedCode public void killQuery(String queryId) { queryManager.cancelQuery(new QueryId(queryId)); }
public Session toSession(SessionPropertyManager sessionPropertyManager) { return new Session( new QueryId(queryId), transactionId, clientTransactionSupport, new Identity(user, Optional.empty()), source, catalog, schema, timeZoneKey, locale, remoteUserAddress, userAgent, startTime, systemProperties, catalogProperties, sessionPropertyManager); } }
public static TaskContext createTaskContext(Executor executor, Session session, DataSize maxMemory, DataSize preallocated) { MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); MemoryPool systemMemoryPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)); QueryContext queryContext = new QueryContext(new QueryId("test_query"), maxMemory, memoryPool, systemMemoryPool, executor); return createTaskContext(queryContext, executor, session, preallocated); }
@Test public void test() throws Exception { Session session = TEST_SESSION.withSystemProperty("foo", "bar").withCatalogProperty("catalog", "baz", "blah"); TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "reset foo", session, URI.create("fake://uri"), false, transactionManager, executor); new ResetSessionTask().execute(new ResetSession(QualifiedName.of("catalog", "baz")), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); Set<String> sessionProperties = stateMachine.getResetSessionProperties(); assertEquals(sessionProperties, ImmutableSet.of("catalog.baz")); } }
private void testSetSession(Expression expression, String expectedValue) throws Exception { TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "set foo.bar = 'baz'", TEST_SESSION, URI.create("fake://uri"), false, transactionManager, executor); new SetSessionTask().execute(new SetSession(QualifiedName.of("foo", "bar"), expression), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); Map<String, String> sessionProperties = stateMachine.getSetSessionProperties(); assertEquals(sessionProperties, ImmutableMap.of("foo.bar", expectedValue)); } }
@Test public void testCommit() throws Exception { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .setTransactionId(transactionManager.beginTransaction(false)) .build(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "COMMIT", session, URI.create("fake://uri"), true, transactionManager, executor); assertTrue(stateMachine.getSession().getTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); new CommitTask().execute(new Commit(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertTrue(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); assertFalse(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }
@Test public void testUnknownTransactionRollback() throws Exception { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .setTransactionId(TransactionId.create()) // Use a random transaction ID that is unknown to the system .build(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "ROLLBACK", session, URI.create("fake://uri"), true, transactionManager, executor); new RollbackTask().execute(new Rollback(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertTrue(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); // Still issue clear signal assertFalse(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }
@Test public void testRollback() throws Exception { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .setTransactionId(transactionManager.beginTransaction(false)) .build(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "ROLLBACK", session, URI.create("fake://uri"), true, transactionManager, executor); assertTrue(stateMachine.getSession().getTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); new RollbackTask().execute(new Rollback(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertTrue(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); assertFalse(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }
@Test public void testStartTransaction() throws Exception { Session session = sessionBuilder() .setClientTransactionSupport() .build(); TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "START TRANSACTION", session, URI.create("fake://uri"), true, transactionManager, executor); assertFalse(stateMachine.getSession().getTransactionId().isPresent()); new StartTransactionTask().execute(new StartTransaction(ImmutableList.of()), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertFalse(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); assertTrue(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); TransactionInfo transactionInfo = transactionManager.getTransactionInfo(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().get()); assertFalse(transactionInfo.isAutoCommitContext()); }
@Test public void testNoTransactionRollback() throws Exception { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .build(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "ROLLBACK", session, URI.create("fake://uri"), true, transactionManager, executor); try { try { new RollbackTask().execute(new Rollback(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); fail(); } catch (CompletionException e) { throw Throwables.propagate(e.getCause()); } } catch (PrestoException e) { assertEquals(e.getErrorCode(), NOT_IN_TRANSACTION.toErrorCode()); } assertFalse(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); assertFalse(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }
@Test public void testStartTransactionExplicitModes() throws Exception { Session session = sessionBuilder() .setClientTransactionSupport() .build(); TransactionManager transactionManager = createTestTransactionManager(); QueryStateMachine stateMachine = QueryStateMachine.begin(new QueryId("query"), "START TRANSACTION", session, URI.create("fake://uri"), true, transactionManager, executor); assertFalse(stateMachine.getSession().getTransactionId().isPresent()); new StartTransactionTask().execute(new StartTransaction(ImmutableList.of(new Isolation(Isolation.Level.SERIALIZABLE), new TransactionAccessMode(true))), transactionManager, metadata, new AllowAllAccessControl(), stateMachine).join(); assertFalse(stateMachine.getQueryInfoWithoutDetails().isClearTransactionId()); assertTrue(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); TransactionInfo transactionInfo = transactionManager.getTransactionInfo(stateMachine.getQueryInfoWithoutDetails().getStartedTransactionId().get()); assertEquals(transactionInfo.getIsolationLevel(), IsolationLevel.SERIALIZABLE); assertTrue(transactionInfo.isReadOnly()); assertFalse(transactionInfo.isAutoCommitContext()); }
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), "stage"); SqlStageExecution stage = new SqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(executor), TEST_SESSION, nodeTaskMap, executor); stage.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
public static Session toSession(TransactionId transactionId, ConnectorSession session) { return Session.builder(new SessionPropertyManager(SYSTEM_SESSION_PROPERTIES)) .setQueryId(new QueryId(session.getQueryId())) .setTransactionId(transactionId) .setCatalog("catalog") .setSchema("schema") .setIdentity(session.getIdentity()) .setTimeZoneKey(session.getTimeZoneKey()) .setLocale(session.getLocale()) .setStartTime(session.getStartTime()) .build(); }
public SqlTask createInitialTask() { TaskId taskId = new TaskId("query", "stage", "task" + nextTaskId.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); return new SqlTask( taskId, location, new QueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)), taskNotificationExecutor), sqlTaskExecutionFactory, taskNotificationExecutor, Functions.<SqlTask>identity(), new DataSize(32, MEGABYTE)); } }
private InternalTable getInternalTable(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, ConnectorSplit connectorSplit, List<ColumnHandle> columns) { InformationSchemaTransactionHandle transaction = checkType(transactionHandle, InformationSchemaTransactionHandle.class, "transaction"); InformationSchemaSplit split = checkType(connectorSplit, InformationSchemaSplit.class, "split"); requireNonNull(columns, "columns is null"); InformationSchemaTableHandle handle = split.getTableHandle(); Map<String, NullableValue> filters = split.getFilters(); Session session = Session.builder(metadata.getSessionPropertyManager()) .setTransactionId(transaction.getTransactionId()) .setQueryId(new QueryId(connectorSession.getQueryId())) .setIdentity(connectorSession.getIdentity()) .setSource("information_schema") .setCatalog("") // default catalog is not be used .setSchema("") // default schema is not be used .setTimeZoneKey(connectorSession.getTimeZoneKey()) .setLocale(connectorSession.getLocale()) .setStartTime(connectorSession.getStartTime()) .build(); return getInformationSchemaTable(session, handle.getCatalogName(), handle.getSchemaTableName(), filters); }