Examples with SessionBuilder io.trino.Session.SessionBuilder used on opensource projects

Example 1 with SessionBuilder

use of io.trino.Session.SessionBuilder in project trino by trinodb.

the class QuerySessionSupplier method createSession.

@Override
public Session createSession(QueryId queryId, SessionContext context) {
    Identity identity = context.getIdentity();
    accessControl.checkCanSetUser(identity.getPrincipal(), identity.getUser());
    // authenticated identity is not present for HTTP or if authentication is not setup
    if (context.getAuthenticatedIdentity().isPresent()) {
        Identity authenticatedIdentity = context.getAuthenticatedIdentity().get();
        // only check impersonation if authenticated user is not the same as the explicitly set user
        if (!authenticatedIdentity.getUser().equals(identity.getUser())) {
            // add enabled roles for authenticated identity, so impersonation permissions can be assigned to roles
            authenticatedIdentity = addEnabledRoles(authenticatedIdentity, context.getSelectedRole(), metadata);
            accessControl.checkCanImpersonateUser(authenticatedIdentity, identity.getUser());
        }
    }
    // add the enabled roles
    identity = addEnabledRoles(identity, context.getSelectedRole(), metadata);
    SessionBuilder sessionBuilder = Session.builder(sessionPropertyManager).setQueryId(queryId).setIdentity(identity).setPath(context.getPath().or(() -> defaultPath).map(SqlPath::new)).setSource(context.getSource()).setRemoteUserAddress(context.getRemoteUserAddress()).setUserAgent(context.getUserAgent()).setClientInfo(context.getClientInfo()).setClientTags(context.getClientTags()).setClientCapabilities(context.getClientCapabilities()).setTraceToken(context.getTraceToken()).setResourceEstimates(context.getResourceEstimates()).setProtocolHeaders(context.getProtocolHeaders());
    if (context.getCatalog().isPresent()) {
        sessionBuilder.setCatalog(context.getCatalog());
        sessionBuilder.setSchema(context.getSchema());
    } else {
        defaultCatalog.ifPresent(sessionBuilder::setCatalog);
        defaultSchema.ifPresent(sessionBuilder::setSchema);
    }
    if (forcedSessionTimeZone.isPresent()) {
        sessionBuilder.setTimeZoneKey(forcedSessionTimeZone.get());
    } else {
        String sessionTimeZoneId = context.getSystemProperties().get(TIME_ZONE_ID);
        if (sessionTimeZoneId != null) {
            sessionBuilder.setTimeZoneKey(getTimeZoneKey(sessionTimeZoneId));
        } else {
            sessionBuilder.setTimeZoneKey(context.getTimeZoneId().map(TimeZoneKey::getTimeZoneKey));
        }
    }
    context.getLanguage().ifPresent(s -> sessionBuilder.setLocale(Locale.forLanguageTag(s)));
    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();
    }
    return sessionBuilder.build();
}

Example 2 with SessionBuilder

use of io.trino.Session.SessionBuilder in project trino by trinodb.

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 = LocalQueryRunner.create(session);
    // add tpch
    localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of());
    return localQueryRunner;
}

Example 3 with SessionBuilder

use of io.trino.Session.SessionBuilder in project trino by trinodb.

the class TestCompletedEventWarnings method setUp.

@BeforeMethod
public void setUp() throws Exception {
    SessionBuilder sessionBuilder = testSessionBuilder();
    generatedEvents = new EventsCollector();
    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.reset(EXPECTED_EVENTS);
}

Example 4 with SessionBuilder

use of io.trino.Session.SessionBuilder in project trino by trinodb.

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(TrinoWarning::getWarningCode).collect(toImmutableSet());
    for (WarningCode warningCode : expectedWarnings) {
        if (!warnings.contains(warningCode)) {
            fail("Expected warning: " + warningCode);
        }
    }
}

Example 5 with SessionBuilder

use of io.trino.Session.SessionBuilder in project trino by trinodb.

the class TestTpchCostBasedPlan method createLocalQueryRunner.

/*
     * CAUTION: The expected plans here are not necessarily optimal yet. Their role is to prevent
     * inadvertent regressions. A conscious improvement to the planner may require changing some
     * of the expected plans, but any such change should be verified on an actual cluster with
     * large amount of data.
     */
@Override
protected LocalQueryRunner createLocalQueryRunner() {
    String catalog = "local";
    SessionBuilder sessionBuilder = testSessionBuilder().setCatalog(catalog).setSchema("sf3000.0").setSystemProperty("task_concurrency", // these tests don't handle exchanges from local parallel
    "1").setSystemProperty(JOIN_REORDERING_STRATEGY, JoinReorderingStrategy.AUTOMATIC.name()).setSystemProperty(JOIN_DISTRIBUTION_TYPE, JoinDistributionType.AUTOMATIC.name());
    LocalQueryRunner queryRunner = LocalQueryRunner.builder(sessionBuilder.build()).withNodeCountForStats(8).build();
    queryRunner.createCatalog(catalog, new TpchConnectorFactory(1, false, false), ImmutableMap.of(TPCH_COLUMN_NAMING_PROPERTY, ColumnNaming.SIMPLIFIED.name()));
    return queryRunner;
}