JDBCConnection.java

package gov.usgs.earthquake.util;

import gov.usgs.earthquake.aws.AwsProductSender;
import gov.usgs.util.Config;
import gov.usgs.util.DefaultConfigurable;

import java.io.StringReader;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.json.Json;
import javax.json.JsonObject;
import javax.json.JsonReader;
import javax.naming.ConfigurationException;
import gov.usgs.earthquake.aws.SecretResolver;

/**
 * Utility class for JDBC Connection.
 *
 * Sub-classes must implement the connect method, and extend startup and
 * shutdown methods. The {@link #verifyConnection()} method tests whether the
 * connection is active, and will shutdown() and startup() to reinitialize if it
 * is not active.
 *
 * @author jmfee
 */
public class JDBCConnection extends DefaultConfigurable implements AutoCloseable {
  public static final String DATABASE_SECRET_ARN_PROPERTY = "databaseSecretArn";

  private static final Logger LOGGER = Logger.getLogger(JDBCConnection.class.getName());

  /** Connection object. */
  private Connection connection;

  /** JDBC driver class. */
  private String driver;

  /** JDBC connect url. */
  private String url;

  /**
   * Create a new JDBCConnection object.
   */
  public JDBCConnection() {
    this.connection = null;
  }

  /**
   * Create a new JDBCConnection object with specific driver and URL
   *
   * @param driver String of driver
   * @param url    String of URL
   */
  public JDBCConnection(final String driver, final String url) {
    this.driver = driver;
    this.url = url;
  }

  /**
   * Implement autocloseable.
   *
   * Calls {@link #shutdown()}.
   *
   * @throws Exception Exception
   */
  @Override
  public void close() throws Exception {
    shutdown();
  }

  /**
   * Implement Configurable
   *
   * @param config Config to set driver and URL in
   * @throws Exception Exception
   */
  @Override
  public void configure(final Config config) throws Exception {
    setDriver(config.getProperty("driver"));
    String databaseSecretArn = config.getProperty(DATABASE_SECRET_ARN_PROPERTY);
    String secretResolverProperty = AwsProductSender.SECRET_RESOLVER_PROPERTY;

    final String secretResolverConfigSection = config.getProperty(secretResolverProperty);
    SecretResolver secretResolver = null;
    if (secretResolverConfigSection != null) {
      secretResolver = (SecretResolver) Config.getConfig().getObject(secretResolverConfigSection);
    }
    if (databaseSecretArn != null) {
      if (secretResolver == null) {
        throw new ConfigurationException(
            DATABASE_SECRET_ARN_PROPERTY + " requires " + secretResolverProperty);
      }
      String secretValue = secretResolver.getPlaintextSecret(databaseSecretArn);
      try (final JsonReader reader = Json.createReader(new StringReader(secretValue))) {
        // parse message
        final JsonObject secret = reader.readObject();
        this.setUrl(this.buildUrl(secret));
      } catch (Exception e) {
        LOGGER.log(Level.SEVERE, e.getMessage());
        throw new ConfigurationException("Error configuring JDBC connection using secret");
      }
    } else {
      setUrl(config.getProperty("url"));
    }
  }

  /**
   * Connect to the database.
   *
   * Sub-classes determine how connection is made.
   *
   * @return the connection.
   * @throws Exception if unable to connect.
   */
  protected Connection connect() throws Exception {
    // load driver if needed
    Class.forName(driver);
    final Connection conn = DriverManager.getConnection(url);
    return conn;
  }

  /**
   * Formats a URL encodes appended values and replaces plus signs with %20. this
   * is required because java encoding replaces spaces with +
   *
   * @param secret Json object
   * @return formatted url string
   * @throws UnsupportedEncodingException
   */
  public String buildUrl(JsonObject secret) throws UnsupportedEncodingException {
    String formattedUrl = new StringBuffer()
        .append("jdbc:")
        .append(encode(secret.getString("engine")))
        .append("://")
        .append(encode(secret.getString("host")))
        .append("/")
        .append(encode(secret.getString("dbname")))
        .append("?user=")
        .append(encode(secret.getString("username")))
        .append("&password=")
        .append(encode(secret.getString("password")))
        .toString();
    return formattedUrl;
  }

  private String encode(String value) {
    return URLEncoder.encode(value, StandardCharsets.UTF_8).replace("+", "%20");
  }

  /**
   * Initialize the database connection.
   *
   * Sub-classes should call super.startup(), before preparing any statements.
   *
   * @throws Exception if error occurs
   */
  @Override
  public void startup() throws Exception {
    this.connection = connect();
  }

  /**
   * Shutdown the database connection.
   *
   * Sub-classes should close any prepared statements (catching any exceptions),
   * and then call super.shutdown() to close the database connection.
   *
   * @throws Exception if error occurs
   */
  @Override
  public void shutdown() throws Exception {
    try {
      if (connection != null) {
        connection.close();
      }
    } catch (Exception e) {
      // log
      e.printStackTrace();
    } finally {
      connection = null;
    }
  }

  /**
   * Open a transaction on the database connection
   *
   * @throws Exception if error occurs
   */
  public synchronized void beginTransaction() throws Exception {
    Connection conn = this.verifyConnection();
    conn.setAutoCommit(false);
  }

  /**
   * Finalize the transaction by committing all the changes and closing the
   * transaction.
   *
   * @throws Exception if error occurs
   */
  public synchronized void commitTransaction() throws Exception {
    getConnection().setAutoCommit(true);
  }

  /**
   * Undo all of the changes made during the current transaction
   *
   * @throws Exception if error occurs
   */
  public synchronized void rollbackTransaction() throws Exception {
    getConnection().rollback();
  }

  /**
   * @return current connection object, or null if not connected.
   */
  public Connection getConnection() {
    return this.connection;
  }

  /**
   * Check whether database connection is closed, and reconnect if needed.
   *
   * Executes the query "select 1" using the current database connection. If this
   * doesn't succeed, reinitializes the database connection by calling shutdown()
   * then startup().
   *
   * @return Valid connection object.
   * @throws Exception if unable to (re)connect.
   */
  public synchronized Connection verifyConnection() throws Exception {
    try {
      // usually throws an exception when connection is closed
      if (connection.isClosed()) {
        shutdown();
      }
    } catch (Exception e) {
      shutdown();
    }

    if (connection == null) {
      // connection is null after shutdown()
      startup();
    }

    // isClosed() doesn't check if we can still communicate with the server.
    // current mysql driver doesn't implement isValid(), so check manually.
    String query_text = "SELECT 1";

    try {
      Statement statement = null;
      ResultSet results = null;
      try {
        statement = connection.createStatement();
        results = statement.executeQuery(query_text);
        while (results.next()) {
          if (results.getInt(1) != 1) {
            throw new Exception("[" + getName() + "] Problem checking database connection");
          }
        }
      } finally {
        // close result and statement no matter what
        try {
          results.close();
        } catch (Exception e2) {
          // ignore
        }
        try {
          statement.close();
        } catch (Exception e2) {
          // ignore
        }
      }
    } catch (Exception e) {
      // The connection was dead, so lets try to restart it
      LOGGER.log(Level.FINE, "[" + getName() + "] Restarting database connection");
      shutdown();
      startup();
    }

    return this.connection;
  }

  /** @return driver */
  public String getDriver() {
    return this.driver;
  }

  /** @param driver Driver to set */
  public void setDriver(final String driver) {
    this.driver = driver;
  }

  /** @return URL */
  public String getUrl() {
    return this.url;
  }

  /**
   * @param url URL to set
   */
  public void setUrl(final String url) {
    this.url = url;
  }

}