src/main/java/com/googlesource/gerrit/modules/cache/pg/PgSqlHandle.java

// Copyright (C) 2018 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.googlesource.gerrit.modules.cache.pg;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectStreamClass;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

class PgSqlHandle {
  private static final Logger log = LoggerFactory.getLogger(PgSqlHandle.class);

  Connection conn;
  PreparedStatement get;
  PreparedStatement put;
  PreparedStatement touch;
  PreparedStatement invalidate;

  private final String name;

  PgSqlHandle(PgSqlSource store, String name, KeyType<?> type) throws SQLException {
    this.name = name;
    this.conn = store.getConnection();
    try (Statement stmt = conn.createStatement()) {
      stmt.addBatch(
          "CREATE TABLE IF NOT EXISTS \"data_"
              + this.name
              + "\" (k "
              + type.columnType()
              + " NOT NULL PRIMARY KEY" // H2 uses HASH however it is discouraged in Postgres doc
              + ",v BYTEA NOT NULL"
              + ",created TIMESTAMP NOT NULL"
              + ",accessed TIMESTAMP NOT NULL"
              + ",space BIGINT"
              + ")");
      // Postgres has no 'computed columns' concept, however one can achieve that
      // with trigger
      stmt.addBatch(
          "CREATE OR REPLACE FUNCTION compute_space()\n"
              + "RETURNS trigger\n"
              + "LANGUAGE plpgsql\n"
              + "SECURITY DEFINER\n"
              + "AS $BODY$\n"
              + "BEGIN\n"
              + "    NEW.space = OCTET_LENGTH(NEW.k) + OCTET_LENGTH(NEW.v);\n"
              + "    RETURN NEW;\n"
              + "END\n"
              + "$BODY$;");
      stmt.addBatch("DROP TRIGGER IF EXISTS computed_space ON \"data_" + this.name + "\"");
      stmt.addBatch(
          "CREATE TRIGGER computed_space "
              + "BEFORE INSERT OR UPDATE "
              + "ON \"data_"
              + this.name
              + "\" "
              + "FOR EACH ROW "
              + "EXECUTE PROCEDURE compute_space()");
      stmt.executeBatch();
    }
  }

  void close() {
    get = closeStatement(get);
    put = closeStatement(put);
    touch = closeStatement(touch);
    invalidate = closeStatement(invalidate);

    if (conn != null) {
      try {
        conn.close();
      } catch (SQLException e) {
        log.warn("Cannot close connection to " + name, e);
      } finally {
        conn = null;
      }
    }
  }

  private PreparedStatement closeStatement(PreparedStatement ps) {
    if (ps != null) {
      try {
        ps.close();
      } catch (SQLException e) {
        log.warn("Cannot close statement for " + name, e);
      }
    }
    return null;
  }

  static Object deserialize(byte[] d) {
    try {
      ByteArrayInputStream in = new ByteArrayInputStream(d);
      ClassLoader loader = Thread.currentThread().getContextClassLoader();
      ObjectInputStream is =
          new ObjectInputStream(in) {
            @Override
            protected Class<?> resolveClass(ObjectStreamClass desc)
                throws IOException, ClassNotFoundException {
              try {
                return Class.forName(desc.getName(), true, loader);
              } catch (ClassNotFoundException e) {
                return super.resolveClass(desc);
              }
            }
          };
      return is.readObject();
    } catch (Throwable e) {
      log.error("Object deserialization failed", e);
      throw new RuntimeException(e);
    }
  }

  static byte[] serialize(Object o) {
    try (ByteArrayOutputStream out = new ByteArrayOutputStream();
        ObjectOutputStream os = new ObjectOutputStream(out)) {
      os.writeObject(o);
      return out.toByteArray();
    } catch (Throwable e) {
      log.error("Object serialization failed {}", o, e);
      throw new RuntimeException(e);
    }
  }
}