package datability;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

public class Postgresql implements Database {

    private static final Logger LOG = LoggerFactory.getLogger(Postgresql.class);
    private final Connection connection;

    public Postgresql(Connection connection) {
        this.connection = Assert.notNull(connection, "A connection is required");
        assertValidConnection(connection);
    }

    @Override
    public Database dropNotNulls(String... tables) {
        if (tables != null) {
            try {
                for (String table : tables) {
                    for (String column : findNotNullColumns(table)) {
                        LOG.info("Disabling not-null ON column '{}' of table '{}'", column, table);
                        executeSql("ALTER TABLE " + table + " ALTER " + column + " DROP NOT NULL");
                    }
                }
            } catch (SQLException e) {
                throw new DatabaseException(e);
            }
        }

        return this;
    }

    @Override
    public Database dropPrimaryKeys(String... tables) {
        if (tables != null) {
            try {
                for (String table : tables) {
                    for (String primaryKey : findIndexes(table)) {
                        LOG.info("Disabling primary key '{}' of table '{}'", primaryKey, table);
                        executeSql("ALTER TABLE " + table + " DROP CONSTRAINT " + primaryKey);
                    }
                }
            } catch (SQLException e) {
                throw new DatabaseException(e);
            }
        }

        return this;
    }

    @Override
    public Database dropForeignKeys(String... tables) {
        if (tables != null) {
            try {
                for (String table : tables) {
                    for (String foreignKey : findForeignKeys(table)) {
                        LOG.info("Disabling foreign key '{}' of table '{}'", foreignKey, table);
                        executeSql("ALTER TABLE " + table + " DROP CONSTRAINT " + foreignKey);
                    }
                }
            } catch (SQLException e) {
                throw new DatabaseException(e);
            }
        }

        return this;
    }

    @Override
    public Database dropAll(String... tables) {
        if (tables != null) {
            dropForeignKeys(tables);
            dropPrimaryKeys(tables);
            dropNotNulls(tables);
        }
        return this;
    }

    private List<String> findNotNullColumns(String table) throws SQLException {
        List<String> notNullsColumns = new ArrayList<String>();

        String sql = "SELECT DISTINCT column_name FROM INFORMATION_SCHEMA.COLUMNS" +
                " WHERE column_name IS NOT NULL" +
                " AND columns.table_name = '" + table + "'" +
                " AND columns.is_nullable='NO'";

        try {
            notNullsColumns.addAll(getFirstColumnValues(sql));
        } catch (SQLException e) {
            throw new DatabaseException(sql, e);
        }

        // Cannot drop-non-null from columns part of a primary key
        notNullsColumns.removeAll(findPrimaryKeys(table));

        return notNullsColumns;
    }

    private List<String> findPrimaryKeys(String table) throws SQLException {
        return getFirstColumnValues("SELECT a.attname" +
                " FROM pg_index i JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY (i.indkey)" +
                " WHERE i.indrelid = '" + table + "'::REGCLASS AND i.indisprimary");

    }

    private List<String> findIndexes(String table) throws SQLException {
        return getFirstColumnValues("SELECT i.relname as index_name" +
                " FROM pg_class t JOIN pg_attribute a ON a.attrelid = t.oid" +
                "  JOIN pg_index ix ON t.oid = ix.indrelid AND a.attnum = ANY(ix.indkey)" +
                "  JOIN pg_class i ON i.oid = ix.indexrelid" +
                " WHERE t.relkind = 'r' AND t.relname in ('" + table + "')");
    }

    private List<String> findForeignKeys(String table) throws SQLException {
        return getFirstColumnValues("SELECT constraint_name" +
                " FROM information_schema.table_constraints" +
                " WHERE constraint_type = 'FOREIGN KEY'" +
                " AND table_name='" + table + "'");
    }

    private List<String> getFirstColumnValues(String sql) throws SQLException {
        List<String> values = new ArrayList<String>();
        ResultSet rs = connection.createStatement().executeQuery(sql);
        while (rs.next()) {
            values.add(rs.getString(1));
        }
        return values;
    }

    private void executeSql(String sql) {
        try {
            connection.createStatement().execute(sql);
        } catch (SQLException e) {
            throw new DatabaseException(sql, e);
        }
    }

    private void assertValidConnection(Connection connection) {
        try {
            if (connection.isClosed()) {
                throw new DatabaseException("Connection is already closed");
            }
        } catch (SQLException e) {
            throw new DatabaseException(e);
        }
    }
}

@@ -1382,22 +1382,36 @@
 y : find
-Indexe
+PrimaryKeyConstraint
 s(table)
@@ -2005,24 +2005,34 @@
 ndForeignKey
+Constraint
 s(table)) %7B%0A
@@ -2770,30 +2770,24 @@
 w ArrayList%3C
-String
 %3E();%0A%0A      
@@ -3311,24 +3311,30 @@
 ndPrimaryKey
+Column
 s(table));%0A%0A
@@ -3403,32 +3403,38 @@
 %3E findPrimaryKey
+Column
 s(String table) 
@@ -3752,14 +3752,28 @@
 find
-Indexe
+PrimaryKeyConstraint
 s(St
@@ -3854,26 +3854,18 @@
 ECT 
-i.relname as index
+constraint
 _nam
@@ -3896,277 +3896,140 @@
 ROM 
-pg_class t JOIN pg_attribute a ON a.attrelid = t.oid%22 +%0A                %22  JOIN pg_index ix ON t.oid = ix.indrelid AND a.attnum = ANY(ix.indkey)%22 +%0A                %22  JOIN pg_class i ON i.oid = ix.indexrelid%22 +%0A                %22 WHERE t.relkind = 'r' AND t.relname in (
+information_schema.table_constraints%22 +%0A                %22 WHERE constraint_type = 'PRIMARY KEY'%22 +%0A                %22 AND table_name=
 '%22 +
@@ -4039,17 +4039,16 @@
 ble + %22'
-)
 %22);%0A    
@@ -4081,32 +4081,42 @@
 %3E findForeignKey
+Constraint
 s(String table) 
@@ -4505,22 +4505,16 @@
 rayList%3C
-String
 %3E();%0A   
@@ -4738,32 +4738,81 @@
 %7B%0A        try %7B%0A
+            LOG.debug(%22Executing sql: %7B%7D%22, sql);%0A
             conn

Drop primary keys by reading `table_constraints` instead of indexes