Refine search
@Test public void testRollback() { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .setTransactionId(transactionManager.beginTransaction(false)) .build(); QueryStateMachine stateMachine = createQueryStateMachine("ROLLBACK", session, transactionManager); assertTrue(stateMachine.getSession().getTransactionId().isPresent()); assertEquals(transactionManager.getAllTransactionInfos().size(), 1); getFutureValue(new RollbackTask().execute(new Rollback(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine, emptyList())); assertTrue(stateMachine.getQueryInfo(Optional.empty()).isClearTransactionId()); assertFalse(stateMachine.getQueryInfo(Optional.empty()).getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }
if (owner.isPresent() && !owner.get().equals(session.getIdentity().getUser())) { identity = new Identity(owner.get(), Optional.empty()); viewAccessControl = new ViewAccessControl(accessControl); identity = session.getIdentity(); viewAccessControl = accessControl; Session viewSession = Session.builder(metadata.getSessionPropertyManager()) .setQueryId(session.getQueryId()) .setTransactionId(session.getTransactionId().orElse(null)) .setIdentity(identity) .setSource(session.getSource().orElse(null)) .setCatalog(catalog.orElse(null)) .setSchema(schema.orElse(null)) .setPath(session.getPath()) .setTimeZoneKey(session.getTimeZoneKey()) .setLocale(session.getLocale()) .setRemoteUserAddress(session.getRemoteUserAddress().orElse(null)) .setUserAgent(session.getUserAgent().orElse(null)) .setClientInfo(session.getClientInfo().orElse(null)) .setStartTime(session.getStartTime()) .build();
private static LocalQueryRunner createQueryRunner(Map<String, String> sessionProperties) { Session.SessionBuilder sessionBuilder = testSessionBuilder() .setCatalog("local") .setSchema("tiny") .setSystemProperty("task_concurrency", "1"); // these tests don't handle exchanges from local parallel sessionProperties.entrySet().forEach(entry -> sessionBuilder.setSystemProperty(entry.getKey(), entry.getValue())); LocalQueryRunner queryRunner = new LocalQueryRunner(sessionBuilder.build()); queryRunner.createCatalog(queryRunner.getDefaultSession().getCatalog().get(), new TpchConnectorFactory(1), ImmutableMap.of()); return queryRunner; }
requireNonNull(query, "query is null"); checkArgument(!query.isEmpty(), "query must not be empty string"); checkArgument(!queryTracker.tryGetQuery(queryId).isPresent(), "query %s already exists", queryId); Optional<String> queryType = getQueryType(preparedQuery.getStatement().getClass()).map(Enum::name); selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria( sessionContext.getIdentity().getPrincipal().isPresent(), sessionContext.getIdentity().getUser(), Optional.ofNullable(sessionContext.getSource()), queryExecution.getSession().getTransactionId().ifPresent(transactionManager::trySetInactive); session = Session.builder(new SessionPropertyManager()) .setQueryId(queryId) .setIdentity(sessionContext.getIdentity()) .setPath(new SqlPath(Optional.empty())) .build(); QUERY_STATE_LOG.debug(e, "Query %s failed", session.getQueryId());
@Test public void testQueryTypeBasedSelection() throws InterruptedException { Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("sf100000") .build(); QueryId queryId = createQuery(queryRunner, session, "EXPLAIN " + LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED)); Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId(); assertTrue(resourceGroupId.isPresent(), "Query should have a resource group"); assertEquals(resourceGroupId.get(), createResourceGroupId("explain")); }
SessionBuilder sessionBuilder = Session.builder(sessionPropertyManager) .setQueryId(queryId) .setIdentity(identity) .setSource(context.getSource()) .setCatalog(context.getCatalog()) .setSchema(context.getSchema()) .setPath(new SqlPath(path)) .setRemoteUserAddress(context.getRemoteUserAddress()) .setUserAgent(context.getUserAgent()) .setClientInfo(context.getClientInfo()) .setClientTags(context.getClientTags()) .setClientCapabilities(context.getClientCapabilities()) .setTraceToken(context.getTraceToken()) .setResourceEstimates(context.getResourceEstimates()); sessionBuilder.setPath(new SqlPath(Optional.of(context.getPath()))); sessionBuilder.setTimeZoneKey(getTimeZoneKey(context.getTimeZoneId())); sessionBuilder.setLocale(Locale.forLanguageTag(context.getLanguage())); sessionBuilder.setSystemProperty(entry.getKey(), entry.getValue()); sessionBuilder.setCatalogSessionProperty(catalog, entry.getKey(), entry.getValue()); sessionBuilder.addPreparedStatement(preparedStatement.getKey(), preparedStatement.getValue()); sessionBuilder.setClientTransactionSupport();
private Session getSession(String user) { return testSessionBuilder() .setCatalog(queryRunner.getDefaultSession().getCatalog().get()) .setSchema(queryRunner.getDefaultSession().getSchema().get()) .setIdentity(new Identity(user, Optional.empty())).build(); } }
@Test public void testCurrentUserInView() { checkState(getSession().getCatalog().isPresent(), "catalog is not set"); checkState(getSession().getSchema().isPresent(), "schema is not set"); String testAccountsUnqualifiedName = "test_accounts"; String testAccountsViewUnqualifiedName = "test_accounts_view"; String testAccountsViewFullyQualifiedName = format("%s.%s.%s", getSession().getCatalog().get(), getSession().getSchema().get(), testAccountsViewUnqualifiedName); assertUpdate(format("CREATE TABLE %s AS SELECT user_name, account_name" + " FROM (VALUES ('user1', 'account1'), ('user2', 'account2'))" + " t (user_name, account_name)", testAccountsUnqualifiedName), 2); assertUpdate(format("CREATE VIEW %s AS SELECT account_name FROM test_accounts WHERE user_name = CURRENT_USER", testAccountsViewUnqualifiedName)); assertUpdate(format("GRANT SELECT ON %s TO user1", testAccountsViewFullyQualifiedName)); assertUpdate(format("GRANT SELECT ON %s TO user2", testAccountsViewFullyQualifiedName)); Session user1 = testSessionBuilder() .setCatalog(getSession().getCatalog().get()) .setSchema(getSession().getSchema().get()) .setIdentity(new Identity("user1", getSession().getIdentity().getPrincipal())) .build(); Session user2 = testSessionBuilder() .setCatalog(getSession().getCatalog().get()) .setSchema(getSession().getSchema().get()) .setIdentity(new Identity("user2", getSession().getIdentity().getPrincipal())) .build(); assertQuery(user1, "SELECT account_name FROM test_accounts_view", "VALUES 'account1'"); assertQuery(user2, "SELECT account_name FROM test_accounts_view", "VALUES 'account2'"); assertUpdate("DROP VIEW test_accounts_view"); assertUpdate("DROP TABLE test_accounts"); }
@Test public void testUnknownTransactionRollback() { TransactionManager transactionManager = createTestTransactionManager(); Session session = sessionBuilder() .setTransactionId(TransactionId.create()) // Use a random transaction ID that is unknown to the system .build(); QueryStateMachine stateMachine = createQueryStateMachine("ROLLBACK", session, transactionManager); getFutureValue(new RollbackTask().execute(new Rollback(), transactionManager, metadata, new AllowAllAccessControl(), stateMachine, emptyList())); assertTrue(stateMachine.getQueryInfo(Optional.empty()).isClearTransactionId()); // Still issue clear signal assertFalse(stateMachine.getQueryInfo(Optional.empty()).getStartedTransactionId().isPresent()); assertTrue(transactionManager.getAllTransactionInfos().isEmpty()); }
@Test public void testViewAccessControl() Session viewOwnerSession = TestingSession.testSessionBuilder() .setIdentity(new Identity("test_view_access_owner", Optional.empty())) .setCatalog(getSession().getCatalog().get()) .setSchema(getSession().getSchema().get()) .build(); "SELECT * FROM test_view_access", "View owner 'test_view_access_owner' cannot create view that selects from .*.orders.*", privilege(viewOwnerSession.getUser(), "orders", CREATE_VIEW_WITH_SELECT_COLUMNS)); privilege(getSession().getUser(), "orders", SELECT_COLUMN)); Session nestedViewOwnerSession = TestingSession.testSessionBuilder() .setIdentity(new Identity("test_nested_view_access_owner", Optional.empty())) .setCatalog(getSession().getCatalog().get()) .setSchema(getSession().getSchema().get()) .build();
@Test public void testApplyDefaultProperties() Session session = Session.builder(new SessionPropertyManager()) .setQueryId(new QueryId("test_query_id")) .setIdentity(new Identity("testUser", Optional.empty())) .setSystemProperty(QUERY_MAX_MEMORY, "1GB") .setSystemProperty(JOIN_DISTRIBUTION_TYPE, "partitioned") .setSystemProperty(HASH_PARTITION_COUNT, "43") .setCatalogSessionProperty("testCatalog", "explicit_set", "explicit_set") .build(); assertEquals(session.getSystemProperties(), ImmutableMap.<String, String>builder() .put(QUERY_MAX_MEMORY, "1GB") .put(JOIN_DISTRIBUTION_TYPE, "partitioned") .put(HASH_PARTITION_COUNT, "43") .build()); assertEquals( session.getUnprocessedCatalogProperties(), .build())); session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, Optional.empty(), TEST_RESOURCE_GROUP_ID); assertEquals(session.getSystemProperties(), ImmutableMap.<String, String>builder() .put(QUERY_MAX_MEMORY, "1GB") .put(JOIN_DISTRIBUTION_TYPE, "partitioned")
@Test public void testDescribeOutput() { Session session = Session.builder(getSession()) .addPreparedStatement("my_query", "SELECT * FROM nation") .build(); MaterializedResult actual = computeActual(session, "DESCRIBE OUTPUT my_query"); MaterializedResult expected = resultBuilder(session, VARCHAR, VARCHAR, VARCHAR, VARCHAR, VARCHAR, BIGINT, BOOLEAN) .row("nationkey", session.getCatalog().get(), session.getSchema().get(), "nation", "bigint", 8, false) .row("name", session.getCatalog().get(), session.getSchema().get(), "nation", "varchar(25)", 0, false) .row("regionkey", session.getCatalog().get(), session.getSchema().get(), "nation", "bigint", 8, false) .row("comment", session.getCatalog().get(), session.getSchema().get(), "nation", "varchar(152)", 0, false) .build(); assertEqualsIgnoreOrder(actual, expected); }
public static SessionBuilder testSessionBuilder(SessionPropertyManager sessionPropertyManager) { return Session.builder(sessionPropertyManager) .setQueryId(queryIdGenerator.createNextQueryId()) .setIdentity(new Identity("user", Optional.empty())) .setSource("test") .setCatalog("catalog") .setSchema("schema") .setPath(new SqlPath(Optional.of("path"))) .setTimeZoneKey(DEFAULT_TIME_ZONE_KEY) .setLocale(ENGLISH) .setRemoteUserAddress("address") .setUserAgent("agent"); }
@BeforeClass private void setUp() throws Exception { session = testSessionBuilder() .setSystemProperty("task_concurrency", "1") .setCatalog("tpch") .setSchema("tiny") .setClientInfo("{\"clientVersion\":\"testVersion\"}") .build(); queryRunner = new DistributedQueryRunner(session, 1); queryRunner.installPlugin(new TpchPlugin()); queryRunner.installPlugin(new TestingEventListenerPlugin(generatedEvents)); queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of("tpch.splits-per-node", Integer.toString(SPLITS_PER_NODE))); queryRunner.getCoordinator().getResourceGroupManager().get() .setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_simple.json"))); }
@Test public void testNonLeafGroup() throws Exception { Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("sf100000") .setSource("non-leaf") .build(); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); InternalResourceGroupManager manager = queryRunner.getCoordinator().getResourceGroupManager().get(); DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager(); int originalSize = getSelectors(queryRunner).size(); // Add a selector for a non leaf group dao.insertSelector(3, 100, "user.*", "(?i).*non-leaf.*", null, null, null); dbConfigurationManager.load(); while (getSelectors(queryRunner).size() != originalSize + 1) { MILLISECONDS.sleep(500); } // Submit query with side effect of creating resource groups QueryId firstDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, firstDashboardQuery, RUNNING); cancelQuery(queryRunner, firstDashboardQuery); waitForQueryState(queryRunner, firstDashboardQuery, FAILED); // Submit a query to a non-leaf resource group QueryId invalidResourceGroupQuery = createQuery(queryRunner, session, LONG_LASTING_QUERY); waitForQueryState(queryRunner, invalidResourceGroupQuery, FAILED); assertEquals(queryRunner.getQueryInfo(invalidResourceGroupQuery).getErrorCode(), INVALID_RESOURCE_GROUP.toErrorCode()); }
@Test public void testCurrentUser() { Session session = testSessionBuilder().setIdentity(new Identity("test_current_user", Optional.empty())).build(); try (QueryAssertions queryAssertions = new QueryAssertions(session)) { queryAssertions.assertQuery("SELECT CURRENT_USER", "SELECT CAST('" + session.getUser() + "' AS VARCHAR)"); } }
@Test public void testPrepareNameExists() { Session session = testSessionBuilder() .addPreparedStatement("my_query", "SELECT bar, baz from foo") .build(); Query query = simpleQuery(selectList(new AllColumns()), table(QualifiedName.of("foo"))); String sqlString = "PREPARE my_query FROM SELECT * FROM foo"; Map<String, String> statements = executePrepare("my_query", query, sqlString, session); assertEquals(statements, ImmutableMap.of("my_query", "SELECT *\nFROM\n foo\n")); }
@Test public void testCurrentPath() { Session session = testSessionBuilder() .setPath(new SqlPath(Optional.of("testPath"))) .build(); try (QueryAssertions queryAssertions = new QueryAssertions(session)) { queryAssertions.assertQuery("SELECT CURRENT_PATH", "SELECT CAST('" + session.getPath().toString() + "' AS VARCHAR)"); } Session emptyPathSession = testSessionBuilder() .setPath(new SqlPath(Optional.of("\"\""))) .build(); try (QueryAssertions queryAssertions = new QueryAssertions(emptyPathSession)) { queryAssertions.assertQuery("SELECT CURRENT_PATH", "SELECT CAST('" + emptyPathSession.getPath().toString() + "' AS VARCHAR)"); } } }
@Test public void testNameExpansion() { Session session = TestingSession.testSessionBuilder() .setIdentity(new Identity("bob", Optional.empty())) .setSource("the-internet") .build(); QueryQueueDefinition definition = new QueryQueueDefinition("user.${USER}", 1, 1); assertEquals(definition.getExpandedTemplate(session), "user.bob"); definition = new QueryQueueDefinition("source.${SOURCE}", 1, 1); assertEquals(definition.getExpandedTemplate(session), "source.the-internet"); definition = new QueryQueueDefinition("${USER}.${SOURCE}", 1, 1); assertEquals(definition.getExpandedTemplate(session), "bob.the-internet"); definition = new QueryQueueDefinition("global", 1, 1); assertEquals(definition.getExpandedTemplate(session), "global"); }
private TaskInfo createTask(SqlTaskManager sqlTaskManager, TaskId taskId, OutputBuffers outputBuffers) { sqlTaskManager.getQueryContext(taskId.getQueryId()) .addTaskContext(new TaskStateMachine(taskId, directExecutor()), testSessionBuilder().build(), false, false, OptionalInt.empty()); return sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(), outputBuffers, OptionalInt.empty()); }