@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")); }
defaultSession = defaultSession.toSessionRepresentation().toSession(coordinator.getMetadata().getSessionPropertyManager()); this.prestoClient = closer.register(new TestingPrestoClient(coordinator, defaultSession)); Catalog bogusTestingCatalog = createBogusTestingCatalog(TESTING_CATALOG); server.getCatalogManager().registerCatalog(bogusTestingCatalog);
public static Catalog createBogusTestingCatalog(String catalogName) { ConnectorId connectorId = new ConnectorId(catalogName); return new Catalog( catalogName, connectorId, createTestSessionConnector(), createInformationSchemaConnectorId(connectorId), createTestSessionConnector(), createSystemTablesConnectorId(connectorId), createTestSessionConnector()); }
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; }
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 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())
@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"))); }
private LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig, NodeSpillConfig nodeSpillConfig, boolean withInitialTransaction, boolean alwaysRevokeMemory, int nodeCountForStats) checkArgument(!defaultSession.getTransactionId().isPresent() || !withInitialTransaction, "Already in transaction"); connectorManager.createConnection(GlobalSystemConnector.NAME, GlobalSystemConnector.NAME, ImmutableMap.of()); catalogManager.registerCatalog(createBogusTestingCatalog(TESTING_CATALOG)); this.defaultSession = new Session( defaultSession.getQueryId(), withInitialTransaction ? Optional.of(transactionManager.beginTransaction(false)) : defaultSession.getTransactionId(), defaultSession.isClientTransactionSupport(), defaultSession.getIdentity(), defaultSession.getPreparedStatements()); dataDefinitionTask = ImmutableMap.<Class<? extends Statement>, DataDefinitionTask<?>>builder() .put(CreateTable.class, new CreateTableTask()) .put(CreateView.class, new CreateViewTask(jsonCodec(ViewDefinition.class), sqlParser, new FeaturesConfig()))
public RuleTester(List<Plugin> plugins, Map<String, String> sessionProperties, Optional<Integer> nodeCountForStats) { Session.SessionBuilder sessionBuilder = testSessionBuilder() .setCatalog(CATALOG_ID) .setSchema("tiny") .setSystemProperty("task_concurrency", "1"); // these tests don't handle exchanges from local parallel for (Map.Entry<String, String> entry : sessionProperties.entrySet()) { sessionBuilder.setSystemProperty(entry.getKey(), entry.getValue()); } session = sessionBuilder.build(); queryRunner = nodeCountForStats .map(nodeCount -> LocalQueryRunner.queryRunnerWithFakeNodeCountForStats(session, nodeCount)) .orElseGet(() -> new LocalQueryRunner(session)); queryRunner.createCatalog(session.getCatalog().get(), new TpchConnectorFactory(1), ImmutableMap.of()); plugins.stream().forEach(queryRunner::installPlugin); this.metadata = queryRunner.getMetadata(); this.transactionManager = queryRunner.getTransactionManager(); this.splitManager = queryRunner.getSplitManager(); this.pageSourceManager = queryRunner.getPageSourceManager(); this.accessControl = queryRunner.getAccessControl(); this.sqlParser = queryRunner.getSqlParser(); }
public static void assertPlannerWarnings(LocalQueryRunner queryRunner, @Language("SQL") String sql, Map<String, String> sessionProperties, List<WarningCode> expectedWarnings, Optional<List<Rule<?>>> rules) { Session.SessionBuilder sessionBuilder = testSessionBuilder() .setCatalog(queryRunner.getDefaultSession().getCatalog().get()) .setSchema(queryRunner.getDefaultSession().getSchema().get()); sessionProperties.forEach(sessionBuilder::setSystemProperty); WarningCollector warningCollector = new DefaultWarningCollector(new WarningCollectorConfig()); try { queryRunner.inTransaction(sessionBuilder.build(), transactionSession -> { if (rules.isPresent()) { createPlan(queryRunner, transactionSession, sql, warningCollector, rules.get()); } else { queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.CREATED, false, warningCollector); } return null; }); } catch (SemanticException e) { // ignore } Set<WarningCode> warnings = warningCollector.getWarnings().stream() .map(PrestoWarning::getWarningCode) .collect(toImmutableSet()); for (WarningCode expectedWarning : expectedWarnings) { if (!warnings.contains(expectedWarning)) { fail("Expected warning: " + expectedWarning); } } }
public static LocalQueryRunner createLocalQueryRunner(File tempDir) Session session = testSessionBuilder() .setCatalog("hive") .setSchema("tpch") .build(); localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of()); "hive", HiveBenchmarkQueryRunner.class.getClassLoader(), Optional.of(metastore)); Map<String, String> hiveCatalogConfig = ImmutableMap.<String, String>builder() .put("hive.max-split-size", "10GB") .build();
@Test public void testStatisticsAggregation() throws Exception ColumnStatisticMetadata statisticMetadata = new ColumnStatisticMetadata("column", MAX_VALUE); StatisticAggregationsDescriptor<Integer> descriptor = new StatisticAggregationsDescriptor<>( ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(statisticMetadata, 0)); Session session = testSessionBuilder() .setSystemProperty("statistics_cpu_timer_enabled", "true") .build(); TableFinishOperatorFactory operatorFactory = new TableFinishOperatorFactory( 0, new PlanNodeId("test"), AggregationNode.Step.SINGLE, ImmutableList.of(LONG_MAX.bind(ImmutableList.of(2), Optional.empty())), true), descriptor, assertEquals(driverContext.getMemoryUsage(), 0); assertTrue(operator.isBlocked().isDone()); assertTrue(operator.needsInput());
@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")); }
public static LocalQueryRunner createLocalQueryRunner(Map<String, String> extraSessionProperties) { SessionBuilder sessionBuilder = testSessionBuilder() .setCatalog("tpch") .setSchema(TINY_SCHEMA_NAME); extraSessionProperties.forEach(sessionBuilder::setSystemProperty); Session session = sessionBuilder.build(); LocalQueryRunner localQueryRunner = new LocalQueryRunner(session); // add tpch localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of()); return localQueryRunner; } }
@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 testStatisticsAggregation() throws Exception pageSinkManager.addConnectorPageSinkProvider(CONNECTOR_ID, new ConstantPageSinkProvider(new TableWriteInfoTestPageSink())); ImmutableList<Type> outputTypes = ImmutableList.of(BIGINT, VARBINARY, BIGINT); Session session = testSessionBuilder() .setSystemProperty("statistics_cpu_timer_enabled", "true") .build(); DriverContext driverContext = createTaskContext(executor, scheduledExecutor, session) .addPipelineContext(0, true, true, false) new PlanNodeId("test"), AggregationNode.Step.SINGLE, ImmutableList.of(LONG_MAX.bind(ImmutableList.of(0), Optional.empty())), true), outputTypes, operator.addInput(rowPagesBuilder(BIGINT).row(43).build().get(0)); assertTrue(operator.isBlocked().isDone()); assertTrue(operator.needsInput()); assertEquals(driverContext.getMemoryUsage(), 0);
@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 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()); }