DataProvider.java

package it.inaf.oats.ia2.tapschemamanager.datalayer;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import javax.persistence.Persistence;
import javax.sql.DataSource;

/**
 *
 * @author Sonia Zorba <zorba at oats.inaf.it>
 */
public class DataProvider {

    public static List<String> getAllSchemasNames(DataSource dataSource, DatabaseType dbType) throws SQLException {

        String query;
        if (dbType == DatabaseType.MYSQL) {
            query = "SHOW DATABASES";
        } else if (dbType == DatabaseType.POSTGRES) {
            query = "SELECT schema_name FROM information_schema.schemata";
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        List<String> allSchemas = new ArrayList<>();

        try (Connection connection = dataSource.getConnection();
                Statement statement = connection.createStatement();
                ResultSet resultSet = statement.executeQuery(query)) {
            while (resultSet.next()) {
                allSchemas.add(resultSet.getString(1));
            }
        }

        return allSchemas;
    }

    public static List<String> getAllTAPSchemasNames(DBWrapper dbs, List<String> allSchemas) throws SQLException {

        List<String> allTAPSchemas = new ArrayList<>();

        for (String schemaName : allSchemas) {

            boolean schemas = false,
                    tables = false,
                    columns = false,
                    keys = false,
                    keyColumns = false;

            DatabaseType dbType = dbs.getTapSchemaDatabaseType();

            String query;
            if (dbType == DatabaseType.MYSQL) {
                query = "SHOW TABLES FROM `" + schemaName + "`";
            } else if (dbType == DatabaseType.POSTGRES) {
                query = "SELECT tablename FROM pg_catalog.pg_tables where schemaname = '" + schemaName + "'";
            } else {
                throw new UnsupportedOperationException("Database type " + dbType + " not supported");
            }

            try (Connection connection = dbs.getTapSchemaConnection();
                    Statement statement = connection.createStatement();
                    ResultSet resultSet = statement.executeQuery(query)) {
                while (resultSet.next()) {
                    String shortTableName = resultSet.getString(1);

                    if (null != shortTableName) {
                        switch (shortTableName) {
                            case "schemas":
                                schemas = true;
                                break;
                            case "tables":
                                tables = true;
                                break;
                            case "columns":
                                columns = true;
                                break;
                            case "keys":
                                keys = true;
                                break;
                            case "key_columns":
                                keyColumns = true;
                                break;
                        }
                    }
                }
            }

            if (schemas && tables && columns && keys && keyColumns) {
                // the schema is a TAP_SCHEMA
                allTAPSchemas.add(schemaName);
            }
        }

        return allTAPSchemas;
    }

    public static List<String> getExposedSchemas(DBWrapper dbs, String tapSchemaName) throws SQLException {

        final List<String> exposedSchemas = new ArrayList<>();

        DatabaseType dbType = dbs.getTapSchemaDatabaseType();

        String query;
        if (dbType == DatabaseType.MYSQL) {
            query = "SELECT schema_name FROM `" + tapSchemaName + "`.`schemas`";
        } else if (dbType == DatabaseType.POSTGRES) {
            query = "SELECT schema_name FROM \"" + tapSchemaName + "\".\"schemas\"";
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        try (Connection connection = dbs.getTapSchemaConnection();
                Statement statement = connection.createStatement();
                ResultSet resultSet = statement.executeQuery(query)) {
            while (resultSet.next()) {
                exposedSchemas.add(resultSet.getString(1));
            }
        }

        return exposedSchemas;
    }

    /**
     * @param tapSchemaCredentials
     * @param tapSchemaName
     * @return JPA entity manager for given TAP_SCHEMA name and credentials
     */
    protected static EntityManager getEntityManager(Credentials tapSchemaCredentials, String tapSchemaName) {
        Map<String, String> persistenceMap = new HashMap<>();

        DatabaseType dbType = tapSchemaCredentials.getDatabaseType();

        String jdbcUrl;
        if (dbType == DatabaseType.MYSQL) {
            jdbcUrl = String.format("jdbc:mysql://%s:%s/%s", tapSchemaCredentials.getHostname(), tapSchemaCredentials.getPort(), tapSchemaName);
        } else if (dbType == DatabaseType.POSTGRES) {
            jdbcUrl = String.format("jdbc:postgresql://%s:%s/%s?currentSchema=%s", tapSchemaCredentials.getHostname(), tapSchemaCredentials.getPort(), tapSchemaCredentials.getDatabase(), tapSchemaName);
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        persistenceMap.put("javax.persistence.jdbc.url", jdbcUrl);
        persistenceMap.put("javax.persistence.jdbc.user", tapSchemaCredentials.getUsername());
        persistenceMap.put("javax.persistence.jdbc.password", tapSchemaCredentials.getPassword());
        persistenceMap.put("javax.persistence.jdbc.driver", tapSchemaCredentials.getDatabaseDriverClass());

        EntityManagerFactory managerFactory = Persistence.createEntityManagerFactory("pu", persistenceMap);

        EntityManager em = managerFactory.createEntityManager();

//        if (dbType == DatabaseType.POSTGRES) {
//            ((EntityManagerImpl) em).getSession().getLogin().setTableQualifier(tapSchemaName);
//        }
        return em;
    }

    protected static void createTapSchemaSchema(DBWrapper dbs, String tapSchemaName) throws SQLException {

        DatabaseType dbType = dbs.getTapSchemaDatabaseType();

        try (Connection connection = dbs.getTapSchemaConnection();
                Statement statement = connection.createStatement()) {

            if (dbType == DatabaseType.MYSQL) {
                statement.executeUpdate("CREATE DATABASE IF NOT EXISTS `" + tapSchemaName + "`");
            } else if (dbType == DatabaseType.POSTGRES) {
                statement.executeUpdate("CREATE SCHEMA IF NOT EXISTS \"" + tapSchemaName + "\"");
            } else {
                throw new UnsupportedOperationException("Database type " + dbType + " not supported");
            }
        }
    }

    protected static ArrayList<String> getAllTablesNames(DataSource dataSource, DatabaseType dbType, String schemaName) throws SQLException {

        String query;
        if (dbType == DatabaseType.MYSQL) {
            query = "SHOW TABLES FROM `" + schemaName + "`";
        } else if (dbType == DatabaseType.POSTGRES) {
            query = "SELECT tablename FROM pg_catalog.pg_tables where schemaname = '" + schemaName + "'";
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        ArrayList<String> allTables = new ArrayList<>();
        try (Connection connection = dataSource.getConnection();
                Statement statement = connection.createStatement();
                ResultSet resultSet = statement.executeQuery(query)) {
            while (resultSet.next()) {
                allTables.add(resultSet.getString(1));
            }
        }

        return allTables;
    }

    protected static Map<String, String> getTablesTypes(DataSource dataSource, DatabaseType dbType, String schemaName) throws SQLException {
        final Map<String, String> tablesTypes = new HashMap<>();

        String query;
        if (dbType == DatabaseType.MYSQL) {
            query = "SELECT table_name, table_type FROM information_schema.tables WHERE table_schema = '" + schemaName + "'";
        } else if (dbType == DatabaseType.POSTGRES) {
            query = "SELECT tablename AS table_name, 'table' AS table_type\n"
                    + "FROM pg_catalog.pg_tables WHERE schemaname = '" + schemaName + "'\n"
                    + "UNION\n"
                    + "SELECT table_name AS table_name, 'view' AS table_type\n"
                    + "FROM INFORMATION_SCHEMA.views\n"
                    + "WHERE table_schema = '" + schemaName + "'";
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        try (Connection connection = dataSource.getConnection();
                Statement statement = connection.createStatement();
                ResultSet resultSet = statement.executeQuery(query)) {
            while (resultSet.next()) {
                String tableName = resultSet.getString("table_name");
                String tableType = resultSet.getString("table_type").equalsIgnoreCase("VIEW") ? "view" : "table";
                tablesTypes.put(tableName, tableType);
            }
        }

        return tablesTypes;
    }

    /**
     * Generate list of KeyEntity for a given schema. This keys are without id!
     * It has to be set when a table is added to a schema.
     *
     * @param sourceConnection
     * @param schemaName
     * @return
     * @throws SQLException
     */
    protected static List<KeyEntity> getSchemaKeys(DataSource dataSource, DatabaseType dbType, String schemaName) throws SQLException {

        final Map<String, KeyEntity> schemaKeysMap = new HashMap<>();

        String query;

        if (dbType == DatabaseType.MYSQL) {
            query = "SELECT k.TABLE_NAME, k.COLUMN_NAME, "
                    + "k.REFERENCED_TABLE_NAME, k.REFERENCED_COLUMN_NAME "
                    + "FROM information_schema.TABLE_CONSTRAINTS i "
                    + "LEFT JOIN information_schema.KEY_COLUMN_USAGE k "
                    + "ON i.CONSTRAINT_NAME = k.CONSTRAINT_NAME "
                    + "WHERE i.CONSTRAINT_TYPE = 'FOREIGN KEY' "
                    + "AND i.TABLE_SCHEMA = '" + schemaName + "' "
                    + "AND k.TABLE_SCHEMA = '" + schemaName + "'";
        } else if (dbType == DatabaseType.POSTGRES) {
            query = "SELECT \n"
                    + "conrelid::regclass AS from_table,\n"
                    + "confrelid::regclass AS target_table,\n"
                    + "c.column_name AS from_column,\n"
                    + "a.attname AS target_column\n"
                    + "FROM pg_catalog.pg_constraint r\n"
                    + "JOIN pg_catalog.pg_tables t ON r.conrelid = (t.schemaname || '.' || t.tablename)::regclass::oid\n"
                    + "JOIN information_schema.columns c ON c.table_schema = t.schemaname AND c.table_name = t.tablename AND c.ordinal_position = ANY(r.conkey)\n"
                    + "JOIN pg_catalog.pg_attribute a ON r.conindid = a.attrelid\n"
                    + "WHERE contype = 'f' AND t.schemaname = '" + schemaName + "'";
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        try (Connection connection = dataSource.getConnection();
                Statement statement = connection.createStatement();
                ResultSet resultSet = statement.executeQuery(query)) {
            while (resultSet.next()) {
                String fromTable, targetTable, fromColumn, targetColumn;

                if (dbType == DatabaseType.MYSQL) {
                    fromTable = schemaName + "." + resultSet.getString("k.TABLE_NAME");
                    targetTable = schemaName + "." + resultSet.getString("k.REFERENCED_TABLE_NAME");
                    fromColumn = resultSet.getString("k.COLUMN_NAME");
                    targetColumn = resultSet.getString("k.REFERENCED_COLUMN_NAME");
                } else if (dbType == DatabaseType.POSTGRES) {
                    fromTable = resultSet.getString("from_table");
                    targetTable = resultSet.getString("target_table");
                    fromColumn = resultSet.getString("from_column").replace(schemaName + ".", "");
                    targetColumn = resultSet.getString("target_column").replace(schemaName + ".", "");
                } else {
                    throw new UnsupportedOperationException("Database type " + dbType + " not supported");
                }

                String mapKey = fromTable + "_" + targetTable;

                KeyEntity keyEntity = schemaKeysMap.get(mapKey);
                if (keyEntity == null) {
                    keyEntity = new KeyEntity();
                    keyEntity.setFromTableFullName(fromTable);
                    keyEntity.setTargetTableFullName(targetTable);
                    schemaKeysMap.put(mapKey, keyEntity);
                }

                KeyColumnEntity keyColumnEntity = new KeyColumnEntity();
                keyColumnEntity.setFromColumn(fromColumn);
                keyColumnEntity.setTargetColumn(targetColumn);
                keyEntity.getKeyColumns().add(keyColumnEntity);
            }
        }

        return new ArrayList(schemaKeysMap.values());
    }

    private static boolean equalsOneOf(String string, String... values) {
        for (String value : values) {
            if (string.equals(value)) {
                return true;
            }
        }
        return false;
    }

    protected static List<ColumnInfo> getAllColumns(DataSource dataSource, DatabaseType dbType, TableEntity table) throws SQLException {
        final List<ColumnInfo> allColumns = new ArrayList<>();

        String query;
        if (dbType == DatabaseType.MYSQL) {
            query = String.format("SHOW COLUMNS FROM `%s`.`%s`", table.getSchemaName(), table.getShortTableName());
        } else if (dbType == DatabaseType.POSTGRES) {
            query = "SELECT c.column_name, c.data_type, r.contype AS column_type, c.character_maximum_length, c.numeric_precision\n" //, c.numeric_precision_radix
                    + "FROM information_schema.columns c\n"
                    + "JOIN pg_catalog.pg_tables t ON c.table_schema = t.schemaname AND c.table_name = t.tablename\n"
                    + "LEFT JOIN pg_catalog.pg_constraint r ON c.ordinal_position = ANY(r.conkey) AND r.conrelid = (t.schemaname || '.' || t.tablename)::regclass::oid\n"
                    + "WHERE t.schemaname = '" + table.getSchemaName() + "' AND t.tablename = '" + table.getShortTableName() + "'";
        } else {
            throw new UnsupportedOperationException("Database type " + dbType + " not supported");
        }

        try (Connection connection = dataSource.getConnection();
                Statement statement = connection.createStatement();
                ResultSet resultSet = statement.executeQuery(query)) {

            while (resultSet.next()) {

                String columnName;
                if (dbType == DatabaseType.MYSQL) {
                    columnName = resultSet.getString("Field");
                } else if (dbType == DatabaseType.POSTGRES) {
                    columnName = resultSet.getString("column_name");
                } else {
                    throw new UnsupportedOperationException("Database type " + dbType + " not supported");
                }

                ColumnEntity columnEntity = DLUtil.getEntityByName(table.getColumns(), columnName);
                boolean alreadyLoaded = columnEntity != null;

                if (columnEntity == null) {
                    columnEntity = new ColumnEntity(columnName);
                }

                columnEntity.setFullTableName(table.getName());
                columnEntity.setTable(table);

                ColumnInfo columnInfo = new ColumnInfo(columnEntity);

                // Key type
                if (dbType == DatabaseType.MYSQL) {
                    String key = resultSet.getString("Key");
                    columnInfo.setPrimaryKey(key.equals("PRI"));
                    columnEntity.setIndexed(equalsOneOf(key, "PRI", "UNI", "MUL") ? 1 : 0);
                    if (!alreadyLoaded) {
                        columnEntity.setPrincipal(key.equals("PRI") ? 1 : 0);
                    }
                } else if (dbType == DatabaseType.POSTGRES) {
                    String columnType = resultSet.getString("column_type");
                    if (columnType != null) {
                        columnInfo.setPrimaryKey("p".equals(columnType));
                        columnEntity.setIndexed(equalsOneOf(columnType, "p", "f", "u") ? 1 : 0);
                    }
                    if (!alreadyLoaded) {
                        columnEntity.setPrincipal("p".equals(columnType) ? 1 : 0);
                    }
                } else {
                    throw new UnsupportedOperationException("Database type " + dbType + " not supported");
                }

                // Datatype and Size
                int size = 0;
                String datatype;

                if (dbType == DatabaseType.MYSQL) {
                    String type = resultSet.getString("Type").toLowerCase();

                    if (type.startsWith("int")) {
                        datatype = "adql:INTEGER";
                    } else if (type.startsWith("smallint")) {
                        datatype = "adql:SMALLINT";
                    } else if (type.startsWith("bigint")) {
                        datatype = "adql:BIGINT";
                    } else if (type.startsWith("float")) {
                        datatype = "adql:REAL";
                    } else if (type.startsWith("char")) {
                        int beginIndex = type.indexOf('(');
                        int endIndex = type.indexOf(')');
                        size = Integer.parseInt(type.substring(beginIndex + 1, endIndex));
                        datatype = "adql:CHAR";
                    } else if (type.startsWith("varchar")) {
                        int beginIndex = type.indexOf('(');
                        int endIndex = type.indexOf(')');
                        size = Integer.parseInt(type.substring(beginIndex + 1, endIndex));
                        datatype = "adql:VARCHAR";
                    } else if (type.contains("timestamp")) {
                        datatype = "adql:TIMESTAMP";
                    } else {
                        datatype = "adql:" + type.toUpperCase();
                    }
                } else if (dbType == DatabaseType.POSTGRES) {
                    String type = resultSet.getString("data_type");

                    if (type.startsWith("int")) {
                        datatype = "adql:INTEGER";
                    } else if (type.startsWith("smallint")) {
                        datatype = "adql:SMALLINT";
                    } else if (type.startsWith("bigint")) {
                        datatype = "adql:BIGINT";
                    } else if (type.startsWith("double") || type.startsWith("real")) {
                        datatype = "adql:REAL";
                    } else if (type.startsWith("character varying")) {
                        datatype = "adql:VARCHAR";
                        size = resultSet.getInt("character_maximum_length");
                    } else if (type.startsWith("char")) {
                        datatype = "adql:CHAR";
                        size = resultSet.getInt("character_maximum_length");
                    } else if (type.contains("timestamp")) {
                        datatype = "adql:TIMESTAMP";
                    } else {
                        datatype = "adql:" + type.toUpperCase();
                    }
                } else {
                    throw new UnsupportedOperationException("Database type " + dbType + " not supported");
                }

                columnEntity.setDatatype(datatype);
                columnEntity.setSize(size);

                if (!alreadyLoaded) {
                    columnEntity.setStd(0);
                }

                allColumns.add(columnInfo);
            }
        }

        return allColumns;
    }
}