com.eventsourcing.postgresql.PostgreSQLSerialization Maven / Gradle / Ivy

Go to download
/**
 * Copyright (c) 2016, All Contributors (see CONTRIBUTORS file)
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */
package com.eventsourcing.postgresql;

import com.eventsourcing.layout.Layout;
import com.eventsourcing.layout.Property;
import com.eventsourcing.layout.TypeHandler;
import com.eventsourcing.layout.types.*;
import com.google.common.io.BaseEncoding;
import com.impossibl.postgres.api.jdbc.PGConnection;
import com.impossibl.postgres.jdbc.PGConnectionImpl;
import com.impossibl.postgres.jdbc.PGStruct;
import com.impossibl.postgres.system.Version;
import com.impossibl.postgres.system.tables.PgAttribute;
import com.impossibl.postgres.system.tables.PgProc;
import com.impossibl.postgres.system.tables.PgType;
import io.netty.buffer.ByteBufInputStream;
import lombok.SneakyThrows;

import java.io.DataInput;
import java.math.BigDecimal;
import java.security.MessageDigest;
import java.sql.*;
import java.time.Instant;
import java.util.*;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;

public class PostgreSQLSerialization {

    private static List pgTypes = new ArrayList<>();
    private static List pgAttrs = new ArrayList<>();
    private static List pgProcs = new ArrayList<>();

    @SneakyThrows
    public static void refreshTypes(Connection connection) {
        PGConnection conn = getPgConnection(connection);

        PGConnectionImpl connImpl = (PGConnectionImpl) conn;
        Version serverVersion = connImpl.getServerVersion();

        //Load types
        String typeSQL = PgType.INSTANCE.getSQL(serverVersion);
        pgTypes = connImpl.queryResults(typeSQL, PgType.Row.class);

        //Load attributes
        String attrsSQL = PgAttribute.INSTANCE.getSQL(serverVersion);
        pgAttrs = connImpl.queryResults(attrsSQL, PgAttribute.Row.class);

        //Load procs
        String procsSQL = PgProc.INSTANCE.getSQL(serverVersion);
        pgProcs = connImpl.queryResults(procsSQL, PgProc.Row.class);
    }

    protected static PGConnection getPgConnection(Connection connection) {
        PGConnection conn;
        if (connection instanceof Wrapper) {
            try {
                conn = connection.unwrap(PGConnection.class);
            } catch (SQLException e) {
                throw new RuntimeException("pgjdbc-ng is required");
            }
        } else if (connection instanceof PGConnection) {
            conn = (PGConnection) connection;
        } else {
            throw new RuntimeException("pgjdbc-ng is required");
        }
        return conn;
    }

    public static void refreshConnectionRegistry(Connection connection) {
        PGConnection conn = getPgConnection(connection);

        PGConnectionImpl connImpl = (PGConnectionImpl) conn;
        //Update the registry with known types
        connImpl.getRegistry().update(pgTypes, pgAttrs, pgProcs);
    }



    @SneakyThrows
    public static String getMappedType(Connection connection, TypeHandler typeHandler) {
        if (typeHandler instanceof BigDecimalTypeHandler) {
            return "NUMERIC";
        }
        if (typeHandler instanceof BooleanTypeHandler) {
            return "BOOLEAN";
        }
        if (typeHandler instanceof ByteArrayTypeHandler) {
            return "BYTEA";
        }
        if (typeHandler instanceof ByteTypeHandler) {
            return "SMALLINT";
        }
        if (typeHandler instanceof DateTypeHandler) {
            return "TIMESTAMP";
        }
        if (typeHandler instanceof DoubleTypeHandler) {
            return "DOUBLE PRECISION";
        }
        if (typeHandler instanceof EnumTypeHandler) {
            return "INTEGER";
        }
        if (typeHandler instanceof FloatTypeHandler) {
            return "REAL";
        }
        if (typeHandler instanceof IntegerTypeHandler) {
            return "INTEGER";
        }
        if (typeHandler instanceof ListTypeHandler) {
            TypeHandler wrappedHandler = ((ListTypeHandler) typeHandler).getWrappedHandler();
            String mappedType = getMappedType(connection, wrappedHandler);
            if (wrappedHandler instanceof ListTypeHandler) {
                mappedType = mappedType.replaceAll("\\[\\]$", "");
                String typname = "layout_v1_arr_" + mappedType;

                PreparedStatement check = connection
                        .prepareStatement("SELECT * FROM pg_catalog.pg_type WHERE lower(typname) = lower(?)");

                check.setString(1, typname);

                boolean shouldCreateType;
                try (ResultSet resultSet = check.executeQuery()) {
                    shouldCreateType = !resultSet.next();
                }
                check.close();

                if (shouldCreateType) {
                    String createType = "CREATE TYPE " + typname + " AS (" +
                            "\"value\" " + mappedType + "[])";
                    PreparedStatement s = connection.prepareStatement(createType);
                    s.execute();
                    s.close();
                    s = connection.prepareStatement("COMMENT ON TYPE " + typname + " IS '" + mappedType + " array'");
                    s.execute();
                    s.close();
                    refreshTypes(connection);
                }

                return typname + "[]";
            } else {
                return mappedType + "[]";
            }
        }
        if (typeHandler instanceof LongTypeHandler) {
            return "BIGINT";
        }
        if (typeHandler instanceof ObjectTypeHandler) {
            Layout layout = ((ObjectTypeHandler) typeHandler).getLayout();
            byte[] fingerprint = layout.getHash();
            String encoded = BaseEncoding.base16().encode(fingerprint);
            String typname = "layout_v1_" + encoded;

            PreparedStatement check = connection
                    .prepareStatement("SELECT * FROM pg_catalog.pg_type WHERE lower(typname) = lower(?)");

            check.setString(1, typname);

            boolean shouldCreateType;
            try (ResultSet resultSet = check.executeQuery()) {
                shouldCreateType = !resultSet.next();
            }
            check.close();

            if (shouldCreateType) {
                String columns = PostgreSQLJournal.defineColumns(connection, layout);
                String createType = "CREATE TYPE " + typname + " AS (" +
                        columns +
                        ")";
                PreparedStatement s = connection.prepareStatement(createType);
                s.execute();
                s.close();
                s = connection.prepareStatement("COMMENT ON TYPE " + typname + " IS '" + layout.getName() + "'");
                s.execute();
                s.close();
                refreshTypes(connection);
            }

            return typname.toLowerCase();
        }
        if (typeHandler instanceof OptionalTypeHandler) {
            return getMappedType(connection, ((OptionalTypeHandler) typeHandler).getWrappedHandler());
        }
        if (typeHandler instanceof ShortTypeHandler) {
            return "SMALLINT";
        }
        if (typeHandler instanceof StringTypeHandler) {
            return "TEXT";
        }
        if (typeHandler instanceof UUIDTypeHandler) {
            return "UUID";
        }
        if (typeHandler instanceof MapTypeHandler) {
            String keyType = getMappedType(connection, ((MapTypeHandler) typeHandler).getWrappedKeyHandler());
            String valueType = getMappedType(connection, ((MapTypeHandler) typeHandler).getWrappedValueHandler());

            MessageDigest digest = MessageDigest.getInstance("SHA-1");
            digest.update(keyType.getBytes());
            digest.update(valueType.getBytes());
            String encodedHash = BaseEncoding.base16().encode(digest.digest());

            String typname = "map_v2_" + encodedHash;

            // Upgrade v1 to v2. Luckily the actual type hasn't changed,
            // just the naming.
            String v1typname = ("map_v1_" + keyType + "_" + valueType).replaceAll("\\[\\]", "__");
            try (PreparedStatement check = connection
                    .prepareStatement("SELECT lower(typname) FROM pg_catalog.pg_type WHERE lower(typname) = lower(?)" +
                                              "::name")) {
                check.setString(1, v1typname);
                try (ResultSet resultSet = check.executeQuery()) {
                    if (resultSet.next()) {
                        try (PreparedStatement s = connection
                                .prepareStatement("ALTER TABLE " + resultSet.getString(1) + " RENAME TO " + typname)) {
                            s.executeUpdate();
                        }
                    }
                }
            }
            //

            boolean shouldCreateType;

            try (PreparedStatement check = connection
                    .prepareStatement("SELECT * FROM pg_catalog.pg_type WHERE lower(typname) = lower(?)")) {
                check.setString(1, typname);

                try (ResultSet resultSet = check.executeQuery()) {
                    shouldCreateType = !resultSet.next();
                }
            }

            if (shouldCreateType) {
                String columns = "key " + keyType + ", value " + valueType;
                String createType = "CREATE TYPE " + typname + " AS (" +
                        columns +
                        ")";
                PreparedStatement s = connection.prepareStatement(createType);
                s.execute();
                s.close();
                s = connection.prepareStatement("COMMENT ON TYPE " + typname + " IS 'Map[" + keyType + "][" + valueType + "]'");
                s.execute();
                s.close();
                refreshTypes(connection);
            }

            return typname.toLowerCase() + "[]";
        }
        throw new RuntimeException("Unsupported type handler " + typeHandler.getClass());
    }

    public static int getMappedSqlType(TypeHandler typeHandler) {
        if (typeHandler instanceof BigDecimalTypeHandler) {
            return Types.DECIMAL;
        }
        if (typeHandler instanceof BooleanTypeHandler) {
            return Types.BOOLEAN;
        }
        if (typeHandler instanceof ByteArrayTypeHandler) {
            return Types.BINARY;
        }
        if (typeHandler instanceof ByteTypeHandler) {
            return Types.SMALLINT;
        }
        if (typeHandler instanceof DateTypeHandler) {
            return Types.TIMESTAMP;
        }
        if (typeHandler instanceof DoubleTypeHandler) {
            return Types.DOUBLE;
        }
        if (typeHandler instanceof EnumTypeHandler) {
            return Types.INTEGER;
        }
        if (typeHandler instanceof FloatTypeHandler) {
            return Types.FLOAT;
        }
        if (typeHandler instanceof IntegerTypeHandler) {
            return Types.INTEGER;
        }
        if (typeHandler instanceof ListTypeHandler) {
            return Types.ARRAY;
        }
        if (typeHandler instanceof MapTypeHandler) {
            return Types.ARRAY;
        }
        if (typeHandler instanceof LongTypeHandler) {
            return Types.BIGINT;
        }
        if (typeHandler instanceof ObjectTypeHandler) {
            return Types.VARCHAR;
        }
        if (typeHandler instanceof OptionalTypeHandler) {
            return getMappedSqlType(((OptionalTypeHandler) typeHandler).getWrappedHandler());
        }
        if (typeHandler instanceof ShortTypeHandler) {
            return Types.SMALLINT;
        }
        if (typeHandler instanceof StringTypeHandler) {
            return Types.VARCHAR;
        }
        if (typeHandler instanceof UUIDTypeHandler) {
            return Types.VARCHAR;
        }
        throw new RuntimeException("Unsupported type handler " + typeHandler.getClass());
    }

    @SneakyThrows
    private static Object prepareValue(Connection connection, TypeHandler typeHandler, Object value) {
        if (typeHandler instanceof BigDecimalTypeHandler) {
            return value == null ? BigDecimal.ZERO : (BigDecimal) value;
        } else if (typeHandler instanceof BooleanTypeHandler) {
            return value == null ? false: (Boolean) value;
        } else if (typeHandler instanceof ByteArrayTypeHandler) {
            if (((ByteArrayTypeHandler) typeHandler).isPrimitive()) {
                return value == null ? new byte[]{} : (byte[]) value;
            } else {
                return value == null ?  new byte[]{} :
                        (byte[]) ((ByteArrayTypeHandler) typeHandler).toPrimitive(value);
            }
        } else if (typeHandler instanceof ByteTypeHandler) {
            return value == null ? 0 : (Byte) value;
        } else if (typeHandler instanceof DateTypeHandler) {
            return value == null ? Timestamp.from(Instant.EPOCH) : Timestamp.from(((java.util.Date)value).toInstant());
        } else if (typeHandler instanceof DoubleTypeHandler) {
            return value == null ? 0 : (Double) value;
        } else if (typeHandler instanceof EnumTypeHandler) {
            return value == null ? 0 : ((Enum)value).ordinal();
        } else if (typeHandler instanceof FloatTypeHandler) {
            return value == null ? 0 : (Float) value;
        } else if (typeHandler instanceof IntegerTypeHandler) {
            return value == null ? 0 : (Integer) value;
        } else if (typeHandler instanceof ListTypeHandler) {
            TypeHandler handler = ((ListTypeHandler)typeHandler).getWrappedHandler();
            List val = value == null ? Collections.emptyList() : (List) value;
            String typeName = getMappedType(connection, handler).toLowerCase();
            if (handler instanceof ListTypeHandler) {
                String arrTypeName = getMappedType(connection, typeHandler).toLowerCase().replaceAll("\\[\\]$", "");
                Object[] arr = val.stream().map(new Function() {
                    @SneakyThrows
                    @Override public Object apply(Object v) {
                        return connection.createStruct(arrTypeName,
                                                       new Object[]{prepareValue(connection,
                                                                                 handler, v)});
                    }
                }).toArray();
                return connection.createArrayOf(arrTypeName, arr);
            } else {
                Object[] arr = val.stream().map(v -> prepareValue(connection, handler, v)).toArray();
                return connection.createArrayOf(typeName, arr);
            }
        } else if (typeHandler instanceof MapTypeHandler) {
            if (value == null) {
                value = new HashMap<>();
            }
            String typeName = getMappedType(connection, typeHandler).replaceAll("\\[\\]", "");
            Object[] arr = ((Map)value).entrySet().stream()
                                       .map(new Function() {
                                           @SneakyThrows
                                           @Override public Object apply(Object e) {
                                               Map.Entry entry = (Map.Entry) e;
                                               return connection.createStruct(typeName,
                                                                              new Object[]{
                                                                                      prepareValue(connection,
                                                                                                   ((MapTypeHandler) typeHandler)
                                                                                                           .getWrappedKeyHandler(),
                                                                                                   entry.getKey()),
                                                                                      prepareValue(connection,
                                                                                                   ((MapTypeHandler) typeHandler)
                                                                                                           .getWrappedValueHandler(),
                                                                                                   entry.getValue())});
                                           }
                                       }).toArray();
            return connection.createArrayOf(typeName, arr);
        } else if (typeHandler instanceof LongTypeHandler) {
            return value == null ? 0 : (Long) value;
        } else if (typeHandler instanceof ObjectTypeHandler) {
            Layout