Search in sources :

Example 1 with SessionBuilder

use of io.prestosql.Session.SessionBuilder in project hetu-core by openlookeng.

the class BenchmarkQueryRunner method createLocalQueryRunner.

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;
}
Also used : TpchConnectorFactory(io.prestosql.plugin.tpch.TpchConnectorFactory) SessionBuilder(io.prestosql.Session.SessionBuilder) TestingSession.testSessionBuilder(io.prestosql.testing.TestingSession.testSessionBuilder) LocalQueryRunner(io.prestosql.testing.LocalQueryRunner) Session(io.prestosql.Session)

Example 2 with SessionBuilder

use of io.prestosql.Session.SessionBuilder in project hetu-core by openlookeng.

the class QuerySessionSupplier method createSession.

@Override
public Session createSession(QueryId queryId, SessionContext context) {
    Identity identity = context.getIdentity();
    accessControl.checkCanSetUser(identity.getPrincipal(), identity.getUser());
    AccessControlUtil.checkCanImpersonateUser(accessControl, context);
    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());
    if (context.getPath() != null) {
        sessionBuilder.setPath(new SqlPath(Optional.of(context.getPath())));
    }
    if (context.getTimeZoneId() != null) {
        sessionBuilder.setTimeZoneKey(getTimeZoneKey(context.getTimeZoneId()));
    }
    if (context.getLanguage() != null) {
        sessionBuilder.setLocale(Locale.forLanguageTag(context.getLanguage()));
    }
    for (Entry<String, String> entry : context.getSystemProperties().entrySet()) {
        sessionBuilder.setSystemProperty(entry.getKey(), entry.getValue());
    }
    for (Entry<String, Map<String, String>> catalogProperties : context.getCatalogSessionProperties().entrySet()) {
        String catalog = catalogProperties.getKey();
        for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
            sessionBuilder.setCatalogSessionProperty(catalog, entry.getKey(), entry.getValue());
        }
    }
    for (Entry<String, String> preparedStatement : context.getPreparedStatements().entrySet()) {
        sessionBuilder.addPreparedStatement(preparedStatement.getKey(), preparedStatement.getValue());
    }
    if (context.supportClientTransaction()) {
        sessionBuilder.setClientTransactionSupport();
    }
    Session session = sessionBuilder.build();
    if (context.getTransactionId().isPresent()) {
        session = session.beginTransactionId(context.getTransactionId().get(), transactionManager, accessControl);
    }
    return session;
}
Also used : SqlPath(io.prestosql.sql.SqlPath) SessionBuilder(io.prestosql.Session.SessionBuilder) Identity(io.prestosql.spi.security.Identity) Map(java.util.Map) Session(io.prestosql.Session)

Example 3 with SessionBuilder

use of io.prestosql.Session.SessionBuilder in project hetu-core by openlookeng.

the class TestCompletedEventWarnings method setUp.

@BeforeMethod
public void setUp() throws Exception {
    SessionBuilder sessionBuilder = testSessionBuilder();
    generatedEvents = new EventsBuilder();
    queryRunner = DistributedQueryRunner.builder(sessionBuilder.build()).setExtraProperties(ImmutableMap.of("testing-warning-collector.preloaded-warnings", String.valueOf(TEST_WARNINGS))).setNodeCount(1).build();
    queryRunner.installPlugin(new TestingEventListenerPlugin(generatedEvents));
    generatedEvents.initialize(EXPECTED_EVENTS);
}
Also used : EventsBuilder(io.prestosql.execution.TestEventListener.EventsBuilder) TestingEventListenerPlugin(io.prestosql.execution.TestEventListenerPlugin.TestingEventListenerPlugin) SessionBuilder(io.prestosql.Session.SessionBuilder) TestingSession.testSessionBuilder(io.prestosql.testing.TestingSession.testSessionBuilder) BeforeMethod(org.testng.annotations.BeforeMethod)

Example 4 with SessionBuilder

use of io.prestosql.Session.SessionBuilder in project hetu-core by openlookeng.

the class TestCompletedEventWarnings method assertWarnings.

private void assertWarnings(@Language("SQL") String sql, Map<String, String> sessionProperties, List<WarningCode> expectedWarnings) throws InterruptedException {
    // Task concurrency must be 1 otherwise these tests fail due to change in the number of EXPECTED_EVENTS
    SessionBuilder sessionBuilder = testSessionBuilder().setSystemProperty("task_concurrency", "1");
    sessionProperties.forEach(sessionBuilder::setSystemProperty);
    queryRunner.execute(sessionBuilder.build(), sql);
    generatedEvents.waitForEvents(10);
    Set<WarningCode> warnings = getOnlyElement(generatedEvents.getQueryCompletedEvents()).getWarnings().stream().map(PrestoWarning::getWarningCode).collect(toImmutableSet());
    for (WarningCode warningCode : expectedWarnings) {
        if (!warnings.contains(warningCode)) {
            fail("Expected warning: " + warningCode);
        }
    }
}
Also used : SessionBuilder(io.prestosql.Session.SessionBuilder) TestingSession.testSessionBuilder(io.prestosql.testing.TestingSession.testSessionBuilder) WarningCode(io.prestosql.spi.WarningCode)

Aggregations

SessionBuilder (io.prestosql.Session.SessionBuilder)4 TestingSession.testSessionBuilder (io.prestosql.testing.TestingSession.testSessionBuilder)3 Session (io.prestosql.Session)2 EventsBuilder (io.prestosql.execution.TestEventListener.EventsBuilder)1 TestingEventListenerPlugin (io.prestosql.execution.TestEventListenerPlugin.TestingEventListenerPlugin)1 TpchConnectorFactory (io.prestosql.plugin.tpch.TpchConnectorFactory)1 WarningCode (io.prestosql.spi.WarningCode)1 Identity (io.prestosql.spi.security.Identity)1 SqlPath (io.prestosql.sql.SqlPath)1 LocalQueryRunner (io.prestosql.testing.LocalQueryRunner)1 Map (java.util.Map)1 BeforeMethod (org.testng.annotations.BeforeMethod)1