Examples with SelectedRole com.facebook.presto.spi.security.SelectedRole used on opensource projects

Example 1 with SelectedRole

use of com.facebook.presto.spi.security.SelectedRole in project presto by prestodb.

the class TestHiveRoles method testSetRole.

@Test
public void testSetRole() throws Exception {
    executeFromAdmin("CREATE ROLE set_role_1");
    executeFromAdmin("CREATE ROLE set_role_2");
    executeFromAdmin("CREATE ROLE set_role_3");
    executeFromAdmin("CREATE ROLE set_role_4");
    executeFromAdmin("GRANT set_role_1 TO USER set_user_1");
    executeFromAdmin("GRANT set_role_2 TO ROLE set_role_1");
    executeFromAdmin("GRANT set_role_3 TO ROLE set_role_2");
    Session unsetRole = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty())).build();
    Session setRoleAll = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(SelectedRole.Type.ALL, Optional.empty())), ImmutableMap.of(), ImmutableMap.of())).build();
    Session setRoleNone = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(SelectedRole.Type.NONE, Optional.empty())), ImmutableMap.of(), ImmutableMap.of())).build();
    Session setRole1 = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(SelectedRole.Type.ROLE, Optional.of("set_role_1"))), ImmutableMap.of(), ImmutableMap.of())).build();
    Session setRole2 = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(SelectedRole.Type.ROLE, Optional.of("set_role_2"))), ImmutableMap.of(), ImmutableMap.of())).build();
    Session setRole3 = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(SelectedRole.Type.ROLE, Optional.of("set_role_3"))), ImmutableMap.of(), ImmutableMap.of())).build();
    Session setRole4 = Session.builder(getQueryRunner().getDefaultSession()).setIdentity(new Identity("set_user_1", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(SelectedRole.Type.ROLE, Optional.of("set_role_4"))), ImmutableMap.of(), ImmutableMap.of())).build();
    MaterializedResult actual = getQueryRunner().execute(unsetRole, "SELECT * FROM hive.information_schema.applicable_roles");
    MaterializedResult expected = MaterializedResult.resultBuilder(unsetRole, createUnboundedVarcharType(), createUnboundedVarcharType(), createUnboundedVarcharType(), createUnboundedVarcharType()).row("set_user_1", "USER", "public", "NO").row("set_user_1", "USER", "set_role_1", "NO").row("set_role_1", "ROLE", "set_role_2", "NO").row("set_role_2", "ROLE", "set_role_3", "NO").build();
    assertEqualsIgnoreOrder(actual, expected);
    actual = getQueryRunner().execute(unsetRole, "SELECT * FROM hive.information_schema.enabled_roles");
    expected = MaterializedResult.resultBuilder(unsetRole, createUnboundedVarcharType()).row("public").row("set_role_1").row("set_role_2").row("set_role_3").build();
    assertEqualsIgnoreOrder(actual, expected);
    actual = getQueryRunner().execute(setRoleAll, "SELECT * FROM hive.information_schema.enabled_roles");
    expected = MaterializedResult.resultBuilder(setRoleAll, createUnboundedVarcharType()).row("public").row("set_role_1").row("set_role_2").row("set_role_3").build();
    assertEqualsIgnoreOrder(actual, expected);
    actual = getQueryRunner().execute(setRoleNone, "SELECT * FROM hive.information_schema.enabled_roles");
    expected = MaterializedResult.resultBuilder(setRoleNone, createUnboundedVarcharType()).row("public").build();
    assertEqualsIgnoreOrder(actual, expected);
    actual = getQueryRunner().execute(setRole1, "SELECT * FROM hive.information_schema.enabled_roles");
    expected = MaterializedResult.resultBuilder(setRole1, createUnboundedVarcharType()).row("public").row("set_role_1").row("set_role_2").row("set_role_3").build();
    assertEqualsIgnoreOrder(actual, expected);
    actual = getQueryRunner().execute(setRole2, "SELECT * FROM hive.information_schema.enabled_roles");
    expected = MaterializedResult.resultBuilder(setRole2, createUnboundedVarcharType()).row("public").row("set_role_2").row("set_role_3").build();
    assertEqualsIgnoreOrder(actual, expected);
    actual = getQueryRunner().execute(setRole3, "SELECT * FROM hive.information_schema.enabled_roles");
    expected = MaterializedResult.resultBuilder(setRole3, createUnboundedVarcharType()).row("public").row("set_role_3").build();
    assertEqualsIgnoreOrder(actual, expected);
    assertQueryFails(setRole4, "SELECT * FROM hive.information_schema.enabled_roles", ".*?Cannot set role set_role_4");
    executeFromAdmin("DROP ROLE set_role_1");
    executeFromAdmin("DROP ROLE set_role_2");
    executeFromAdmin("DROP ROLE set_role_3");
    executeFromAdmin("DROP ROLE set_role_4");
}

Example 2 with SelectedRole

use of com.facebook.presto.spi.security.SelectedRole in project presto by prestodb.

the class TestHttpRequestSessionContext method testSessionContext.

@Test
public void testSessionContext() {
    HttpServletRequest request = new MockHttpServletRequest(ImmutableListMultimap.<String, String>builder().put(PRESTO_USER, "testUser").put(PRESTO_SOURCE, "testSource").put(PRESTO_CATALOG, "testCatalog").put(PRESTO_SCHEMA, "testSchema").put(PRESTO_LANGUAGE, "zh-TW").put(PRESTO_TIME_ZONE, "Asia/Taipei").put(PRESTO_CLIENT_INFO, "client-info").put(PRESTO_SESSION, QUERY_MAX_MEMORY + "=1GB").put(PRESTO_SESSION, JOIN_DISTRIBUTION_TYPE + "=partitioned," + HASH_PARTITION_COUNT + " = 43").put(PRESTO_PREPARED_STATEMENT, "query1=select * from foo,query2=select * from bar").put(PRESTO_ROLE, "foo_connector=ALL").put(PRESTO_ROLE, "bar_connector=NONE").put(PRESTO_ROLE, "foobar_connector=ROLE{role}").put(PRESTO_EXTRA_CREDENTIAL, "test.token.foo=bar").put(PRESTO_EXTRA_CREDENTIAL, "test.token.abc=xyz").put(PRESTO_SESSION_FUNCTION, format("%s=%s,%s=%s", urlEncode(SERIALIZED_SQL_FUNCTION_ID_ADD), urlEncode(SERIALIZED_SQL_FUNCTION_ADD), urlEncode(SERIALIZED_SQL_FUNCTION_ID_ADD_1_TO_INT_ARRAY), urlEncode(SERIALIZED_SQL_FUNCTION_ADD_1_to_INT_ARRAY))).build(), "testRemote");
    HttpRequestSessionContext context = new HttpRequestSessionContext(request, new SqlParserOptions());
    assertEquals(context.getSource(), "testSource");
    assertEquals(context.getCatalog(), "testCatalog");
    assertEquals(context.getSchema(), "testSchema");
    assertEquals(context.getIdentity(), new Identity("testUser", Optional.empty()));
    assertEquals(context.getClientInfo(), "client-info");
    assertEquals(context.getLanguage(), "zh-TW");
    assertEquals(context.getTimeZoneId(), "Asia/Taipei");
    assertEquals(context.getSystemProperties(), ImmutableMap.of(QUERY_MAX_MEMORY, "1GB", JOIN_DISTRIBUTION_TYPE, "partitioned", HASH_PARTITION_COUNT, "43"));
    assertEquals(context.getPreparedStatements(), ImmutableMap.of("query1", "select * from foo", "query2", "select * from bar"));
    assertEquals(context.getIdentity().getRoles(), ImmutableMap.of("foo_connector", new SelectedRole(SelectedRole.Type.ALL, Optional.empty()), "bar_connector", new SelectedRole(SelectedRole.Type.NONE, Optional.empty()), "foobar_connector", new SelectedRole(SelectedRole.Type.ROLE, Optional.of("role"))));
    assertEquals(context.getIdentity().getExtraCredentials(), ImmutableMap.of("test.token.foo", "bar", "test.token.abc", "xyz"));
    assertEquals(context.getSessionFunctions(), ImmutableMap.of(SQL_FUNCTION_ID_ADD, SQL_FUNCTION_ADD, SQL_FUNCTION_ID_ADD1_TO_INT_ARRAY, SQL_FUNCTION_ADD_1_TO_INT_ARRAY));
}

Example 3 with SelectedRole

use of com.facebook.presto.spi.security.SelectedRole in project presto by prestodb.

the class PrestoStatement method internalExecute.

final boolean internalExecute(String sql) throws SQLException {
    clearCurrentResults();
    checkOpen();
    StatementClient client = null;
    PrestoResultSet resultSet = null;
    boolean intercepted = false;
    try {
        WarningsManager warningsManager = new WarningsManager();
        currentWarningsManager.set(Optional.of(warningsManager));
        int statementDepth = this.statementDepth.incrementAndGet();
        boolean shouldIntercept = !connection().getQueryInterceptorInstances().isEmpty() && statementDepth == 1;
        if (shouldIntercept) {
            Optional<PrestoResultSet> newResultSet = connection().invokeQueryInterceptorsPre(sql, this);
            if (newResultSet.isPresent()) {
                resultSet = newResultSet.get();
            }
        }
        // Check if no resultSet is returned from an interceptor
        if (resultSet != null) {
            currentResult.set(resultSet);
            intercepted = true;
        } else {
            client = connection().startQuery(sql, getStatementSessionProperties());
            if (client.isFinished()) {
                QueryStatusInfo finalStatusInfo = client.finalStatusInfo();
                if (finalStatusInfo.getError() != null) {
                    throw resultsException(finalStatusInfo);
                }
            }
            executingClient.set(client);
            resultSet = new PrestoResultSet(this, client, maxRows.get(), progressConsumer, warningsManager);
            for (Map.Entry<String, SelectedRole> entry : client.getSetRoles().entrySet()) {
                connection.get().setRole(entry.getKey(), entry.getValue());
            }
        }
        // check if this is a query
        if (intercepted || client.currentStatusInfo().getUpdateType() == null) {
            currentResult.set(resultSet);
            if (shouldIntercept) {
                resultSet = connection().invokeQueryInterceptorsPost(sql, this, resultSet);
                verifyNotNull(resultSet, "invokeQueryInterceptorsPost should never return a null ResultSet");
                currentResult.set(resultSet);
            }
            return true;
        }
        // this is an update, not a query
        while (resultSet.next()) {
        // ignore rows
        }
        connection().updateSession(client);
        Long updateCount = client.finalStatusInfo().getUpdateCount();
        currentUpdateCount.set((updateCount != null) ? updateCount : 0);
        currentUpdateType.set(client.finalStatusInfo().getUpdateType());
        warningsManager.addWarnings(client.finalStatusInfo().getWarnings());
        return false;
    } catch (ClientException e) {
        throw new SQLException(e.getMessage(), e);
    } catch (RuntimeException e) {
        throw new SQLException("Error executing query", e);
    } finally {
        this.statementDepth.decrementAndGet();
        executingClient.set(null);
        if (currentResult.get() == null) {
            if (resultSet != null) {
                resultSet.close();
            }
            if (client != null) {
                client.close();
            }
        }
    }
}

Example 4 with SelectedRole

use of com.facebook.presto.spi.security.SelectedRole in project presto by prestodb.

the class TestPrestoDriver method testSetRole.

@Test
public void testSetRole() throws Exception {
    try (PrestoConnection connection = createConnection(TEST_CATALOG, "tiny").unwrap(PrestoConnection.class)) {
        try (Statement statement = connection.createStatement()) {
            statement.executeUpdate("SET ROLE ALL");
        }
        assertEquals(connection.getRoles(), ImmutableMap.of(TEST_CATALOG, new SelectedRole(SelectedRole.Type.ALL, Optional.empty())));
        try (Statement statement = connection.createStatement()) {
            statement.executeUpdate("SET ROLE NONE");
        }
        assertEquals(connection.getRoles(), ImmutableMap.of(TEST_CATALOG, new SelectedRole(SelectedRole.Type.NONE, Optional.empty())));
        try (Statement statement = connection.createStatement()) {
            statement.executeUpdate("SET ROLE bar");
        }
        assertEquals(connection.getRoles(), ImmutableMap.of(TEST_CATALOG, new SelectedRole(SelectedRole.Type.ROLE, Optional.of("bar"))));
    }
}

Example 5 with SelectedRole

use of com.facebook.presto.spi.security.SelectedRole in project presto by prestodb.

the class TestIcebergMetadataListing method createQueryRunner.

@Override
protected QueryRunner createQueryRunner() throws Exception {
    Session session = testSessionBuilder().setIdentity(new Identity("hive", Optional.empty(), ImmutableMap.of("hive", new SelectedRole(ROLE, Optional.of("admin"))), ImmutableMap.of(), ImmutableMap.of())).build();
    DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session).build();
    Path catalogDir = queryRunner.getCoordinator().getBaseDataDir().resolve("iceberg_data").resolve("catalog");
    queryRunner.installPlugin(new IcebergPlugin());
    Map<String, String> icebergProperties = ImmutableMap.<String, String>builder().put("hive.metastore", "file").put("hive.metastore.catalog.dir", catalogDir.toFile().toURI().toString()).build();
    queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg", icebergProperties);
    queryRunner.installPlugin(new HivePlugin("hive"));
    Map<String, String> hiveProperties = ImmutableMap.<String, String>builder().put("hive.metastore", "file").put("hive.metastore.catalog.dir", catalogDir.toFile().toURI().toString()).put("hive.security", "sql-standard").build();
    queryRunner.createCatalog("hive", "hive", hiveProperties);
    return queryRunner;
}