public Map<RoleType, List<String>> getAvailableTables(RdsQuery rdsQuery) throws SQLException{
    String address = rdsQuery.getAddress();
    Map<RoleType, List<String>> results = new HashMap<>();
    PGPoolingDataSource dataSource = connect(address, gkUserCredentialsProvider.getGatekeeperSecret(rdsQuery));
    JdbcTemplate conn = new JdbcTemplate(dataSource);

    logger.info("Getting available schema information for " + address);
    for(RoleType roleType : RoleType.values()) {
        try {
            List<String> schemas = conn.queryForList(getSchemas, new Object[]{roleType.getDbRole()}, String.class);
            results.put(roleType, !schemas.isEmpty() ? schemas : Collections.singletonList("No Schemas are available for role " + roleType.getDbRole() + " at this time."));
            logger.info("Retrieved available schema information for database " + address + " for role " + roleType.getDbRole());
        } catch (Exception ex) {
            logger.error("Could not retrieve available role information for database " + address + " for role " +  roleType.getDbRole(), ex);
            results.put(roleType, Collections.singletonList("Unable to get available schemas for role " + roleType.getDbRole()));
        }
    }
    dataSource.close();
    return results;
}
 

public List<DbUser> getUsers(RdsQuery rdsQuery) throws SQLException{
    String address = rdsQuery.getAddress();
    PGPoolingDataSource dataSource = connect(address, gkUserCredentialsProvider.getGatekeeperSecret(rdsQuery));
    JdbcTemplate conn = new JdbcTemplate(dataSource);
    List<DbUser> results;
    logger.info("Getting available schema information for " + address);
    try {
        results = conn.query(getUsers, new PostgresDbUserMapper());
        logger.info("Retrieved users for database " + address);
    } catch (Exception ex) {
        logger.error("Could not retrieve list of users for database " + address, ex);
        results = Collections.emptyList();
    }
    dataSource.close();
    return results;
}
 

public List<String> getAvailableRoles(RdsQuery rdsQuery) throws SQLException{
    String address = rdsQuery.getAddress();

    PGPoolingDataSource dataSource = connect(address, gkUserCredentialsProvider.getGatekeeperSecret(rdsQuery));
    JdbcTemplate conn = new JdbcTemplate(dataSource);
    List<String> results;
    logger.info("Getting available roles for " + address);
    try {
        results = conn.queryForList("select rolname from pg_catalog.pg_roles where rolname like 'gk_%' and rolcanlogin = false", String.class);
    } catch (Exception ex) {
        logger.error("Could not retrieve list of roles for database " + address, ex);
        throw ex;
    }finally {
        dataSource.close();
    }
    return results;
}
 

private PGPoolingDataSource connectHelper(String address, String gkUserPassword) {
    PGPoolingDataSource dataSource = new PGPoolingDataSource();
    String dbUrl = "jdbc:postgresql://" + address;

    dataSource.setDataSourceName(address);
    dataSource.setUrl(dbUrl);
    dataSource.setUser(gkUserName);
    dataSource.setPassword(gkUserPassword);
    dataSource.setConnectTimeout(connectTimeout);
    dataSource.setSsl(ssl);
    dataSource.setSslMode(sslMode);
    dataSource.setSslRootCert(sslCert);
    //Do not want to keep the connection after execution

    try {
        new JdbcTemplate(dataSource).queryForList("select 1"); // Tests the connection
    } catch (Exception e) {
        logger.error("Failed to connect to " + address);
        dataSource.close(); // close the datasource
        throw e;
    }
    logger.info("Using the following properties with the connection: " + ssl);
    return dataSource;

}
 

public boolean revokeAccess(RdsRevokeAccessQuery rdsRevokeAccessQuery) throws SQLException{
    String address = rdsRevokeAccessQuery.getAddress();
    String user = rdsRevokeAccessQuery.getUser();
    RoleType role = rdsRevokeAccessQuery.getRole();

    PGPoolingDataSource dataSource = null;
    try {
        dataSource = connect(address, gkUserCredentialsProvider.getGatekeeperSecret(rdsRevokeAccessQuery));
        JdbcTemplate conn = new JdbcTemplate(dataSource);
        logger.info("Removing " + user + " from " + address + " if they exist.");
        if(role != null) {
            //if roles is provided revoke the user with the suffix (from activiti)
            revokeUser(conn, user + "_" + role.getShortSuffix());
        }else{
            //if roles is not provided just revoke the user (forced removal)
            revokeUser(conn, user);
        }
        return true;

    }catch(Exception ex){
        String username = role == null ? user : user + "_" + role.getShortSuffix();
        logger.error("An exception was thrown while trying to revoke user " + username + " from address " + address, ex);
        return false;
    } finally {
        if(dataSource != null) {
            dataSource.close();
        }
    }
}
 

/**
 * Check to see if this user is the owner of any tables on the DB
 *
 * @param rdsCheckUsersTableQuery - the query details for the db
 * @return List of String - List of users that still own tables
 *
 * @throws SQLException - if there's an issue executing the query on the database
 */
public List<String> checkIfUsersHasTables(RdsCheckUsersTableQuery rdsCheckUsersTableQuery) throws SQLException{
    String address = rdsCheckUsersTableQuery.getAddress();
    List<String> users = rdsCheckUsersTableQuery.getUsers();
    PGPoolingDataSource dataSource = null;
    try {
        dataSource = connect(address, gkUserCredentialsProvider.getGatekeeperSecret(rdsCheckUsersTableQuery));
        JdbcTemplate conn = new JdbcTemplate(dataSource);
        StringBuilder sb = new StringBuilder();
        users.forEach(user -> {
            sb.append("?,");
        });

        String query = "SELECT distinct tableowner FROM pg_tables t where t.tableowner in ("+ sb.deleteCharAt(sb.length() - 1).toString() + ")";

        List<String> outcome = conn.queryForList(query, users.toArray(), String.class);

        return outcome;

    }catch(SQLException ex){
        logger.error("An Error occured while checking to see if the user owns any tables on the database", ex);
        return users;
    }finally {
        if(dataSource != null) {
            dataSource.close();
        }
    }
}
 

private PGPoolingDataSource connect(String url, String gkUserPassword) throws SQLException {
    String dbUrl = url.split("/")[0];
    logger.info("Getting connection for " + dbUrl);
    logger.info("Creating Datasource connection for " + dbUrl);
    String pgUrl = dbUrl + "/postgres"; // url with postgres instead of whatever was on the AWS console
    try {
        return connectHelper(pgUrl, gkUserPassword); // Try postgres first since it is a default db.
    } catch (Exception e){
        logger.info("postgres database not present for " + dbUrl.split("/")[0] + " Attempting connection to " + url + " as fallback.");
        return connectHelper(url, gkUserPassword); // Fall-back if postgres isn't there
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setInitialConnections(3);
    source.setMaxConnections(10);
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setInitialConnections(3);
    source.setMaxConnections(10);
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

public static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

public static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setInitialConnections(3);
    source.setMaxConnections(10);
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setInitialConnections(3);
    source.setMaxConnections(10);
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private static Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setInitialConnections(3);
    source.setMaxConnections(10);
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private Connection getDbConnection(){
    PGPoolingDataSource source = new PGPoolingDataSource();
    source.setServerName("localhost");
    source.setDatabaseName("cookbook");
    source.setInitialConnections(3);
    source.setMaxConnections(10);
    source.setLoginTimeout(10);
    try {
        return source.getConnection();
    }
    catch(Exception ex) {
        ex.printStackTrace();
        return null;
    }
}
 

private DataSource setDataSource(EmbeddedPostgres server) throws Exception {
  PGPoolingDataSource dataSource = new PGPoolingDataSource();
  dataSource.setUser("kloopzcm");
  dataSource.setPassword("testpwd");
  dataSource.setPortNumber(server.getPort());
  dataSource.setDatabaseName("kloopzdb");

  Flyway flyway = new Flyway();
  flyway.setPlaceholderReplacement(false);
  flyway.setLocations("classpath:deployer");
  flyway.setDataSource(dataSource);
  flyway.migrate();

  return dataSource;
}
 

public List<String> checkDb(RdsQuery rdsQuery) throws GKUnsupportedDBException {
    String address = rdsQuery.getAddress();

    String gkUserCreateRoleCheck = "select rolcreaterole from pg_roles where rolname = 'gatekeeper'";
    String gkRoleCheck = "select rolname from pg_roles where rolname in ('gk_datafix','gk_dba','gk_readonly')";

    List<String> issues = new ArrayList<>();
    List<String> gkRoles = new ArrayList<>();
    gkRoles.addAll(Arrays.asList("gk_datafix", "gk_readonly", "gk_dba"));
    PGPoolingDataSource dataSource = null;

    try{
        logger.info("Checking the gatekeeper setup for " + address);
        dataSource = connect(address, gkUserCredentialsProvider.getGatekeeperSecret(rdsQuery));
        JdbcTemplate conn = new JdbcTemplate(dataSource);
        Boolean createRolePermCheckResult = conn.queryForObject(gkUserCreateRoleCheck, Boolean.class);
        List<String> roleCheckResult = conn.queryForList(gkRoleCheck, String.class);

        if(!createRolePermCheckResult){
            issues.add("gatekeeper user missing createrole");
        }
        gkRoles.removeAll(roleCheckResult);
        if(!gkRoles.isEmpty()) {
            issues.add("missing the following roles: " + gkRoles);
        }
    }catch(SQLException e){
        logger.error("Error running check query", e);
    } catch(CannotGetJdbcConnectionException ex){
        logger.error("Failed to connect to DB", ex);
        if(ex.getMessage().contains("password")) {
            issues.add("Password authentication failed for gatekeeper user");
        }else{
            issues.add("Unable to connect to DB (" + ex.getCause().getMessage() + ")");
        }
    }finally{
        if(dataSource != null) {
            dataSource.close();
        }
    }

    return issues;

}