Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you 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 org.apache.druid.metadata.storage.postgresql; import com.google.common.base.Supplier; import com.google.inject.Inject; import org.apache.commons.dbcp2.BasicDataSource; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.metadata.MetadataStorageConnectorConfig; import org.apache.druid.metadata.MetadataStorageTablesConfig; import org.apache.druid.metadata.SQLMetadataConnector; import org.postgresql.PGProperty; import org.skife.jdbi.v2.DBI; import org.skife.jdbi.v2.Handle; import org.skife.jdbi.v2.tweak.HandleCallback; import org.skife.jdbi.v2.util.StringMapper; import java.sql.DatabaseMetaData; import java.sql.SQLException; public class PostgreSQLConnector extends SQLMetadataConnector { private static final Logger log = new Logger(PostgreSQLConnector.class); private static final String PAYLOAD_TYPE = "BYTEA"; private static final String SERIAL_TYPE = "BIGSERIAL"; private static final String QUOTE_STRING = "\\\""; public static final int DEFAULT_STREAMING_RESULT_SIZE = 100; private final DBI dbi; private volatile Boolean canUpsert; @Inject public PostgreSQLConnector(Supplier<MetadataStorageConnectorConfig> config, Supplier<MetadataStorageTablesConfig> dbTables, PostgreSQLConnectorConfig connectorConfig) { super(config, dbTables); final BasicDataSource datasource = getDatasource(); // PostgreSQL driver is classloader isolated as part of the extension // so we need to help JDBC find the driver datasource.setDriverClassLoader(getClass().getClassLoader()); datasource.setDriverClassName("org.postgresql.Driver"); // SSL Configuration if (connectorConfig.isUseSSL()) { log.info("SSL is enabled on this PostgreSQL connection."); datasource.addConnectionProperty(PGProperty.SSL.getName(), String.valueOf(connectorConfig.isUseSSL())); if (connectorConfig.getPassword() != null) { datasource.addConnectionProperty(PGProperty.SSL_PASSWORD.getName(), connectorConfig.getPassword()); } if (connectorConfig.getSslFactory() != null) { datasource.addConnectionProperty(PGProperty.SSL_FACTORY.getName(), connectorConfig.getSslFactory()); } if (connectorConfig.getSslFactoryArg() != null) { datasource.addConnectionProperty(PGProperty.SSL_FACTORY_ARG.getName(), connectorConfig.getSslFactoryArg()); } if (connectorConfig.getSslMode() != null) { datasource.addConnectionProperty(PGProperty.SSL_MODE.getName(), connectorConfig.getSslMode()); } if (connectorConfig.getSslCert() != null) { datasource.addConnectionProperty(PGProperty.SSL_CERT.getName(), connectorConfig.getSslCert()); } if (connectorConfig.getSslKey() != null) { datasource.addConnectionProperty(PGProperty.SSL_KEY.getName(), connectorConfig.getSslKey()); } if (connectorConfig.getSslRootCert() != null) { datasource.addConnectionProperty(PGProperty.SSL_ROOT_CERT.getName(), connectorConfig.getSslRootCert()); } if (connectorConfig.getSslHostNameVerifier() != null) { datasource.addConnectionProperty(PGProperty.SSL_HOSTNAME_VERIFIER.getName(), connectorConfig.getSslHostNameVerifier()); } if (connectorConfig.getSslPasswordCallback() != null) { datasource.addConnectionProperty(PGProperty.SSL_PASSWORD_CALLBACK.getName(), connectorConfig.getSslPasswordCallback()); } } this.dbi = new DBI(datasource); log.info("Configured PostgreSQL as metadata storage"); } @Override protected String getPayloadType() { return PAYLOAD_TYPE; } @Override protected String getSerialType() { return SERIAL_TYPE; } @Override public String getQuoteString() { return QUOTE_STRING; } @Override protected int getStreamingFetchSize() { return DEFAULT_STREAMING_RESULT_SIZE; } protected boolean canUpsert(Handle handle) throws SQLException { if (canUpsert == null) { DatabaseMetaData metaData = handle.getConnection().getMetaData(); canUpsert = metaData.getDatabaseMajorVersion() > 9 || (metaData.getDatabaseMajorVersion() == 9 && metaData.getDatabaseMinorVersion() >= 5); } return canUpsert; } @Override public boolean tableExists(final Handle handle, final String tableName) { return !handle.createQuery( "SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = 'public' AND tablename ILIKE :tableName") .bind("tableName", tableName).map(StringMapper.FIRST).list().isEmpty(); } @Override public Void insertOrUpdate(final String tableName, final String keyColumn, final String valueColumn, final String key, final byte[] value) { return getDBI().withHandle(new HandleCallback<Void>() { @Override public Void withHandle(Handle handle) throws Exception { if (canUpsert(handle)) { handle.createStatement(StringUtils.format( "INSERT INTO %1$s (%2$s, %3$s) VALUES (:key, :value) ON CONFLICT (%2$s) DO UPDATE SET %3$s = EXCLUDED.%3$s", tableName, keyColumn, valueColumn)).bind("key", key).bind("value", value).execute(); } else { handle.createStatement(StringUtils.format("BEGIN;\n" + "LOCK TABLE %1$s IN SHARE ROW EXCLUSIVE MODE;\n" + "WITH upsert AS (UPDATE %1$s SET %3$s=:value WHERE %2$s=:key RETURNING *)\n" + " INSERT INTO %1$s (%2$s, %3$s) SELECT :key, :value WHERE NOT EXISTS (SELECT * FROM upsert)\n;" + "COMMIT;", tableName, keyColumn, valueColumn)).bind("key", key).bind("value", value) .execute(); } return null; } }); } @Override public DBI getDBI() { return dbi; } @Override protected boolean connectorIsTransientException(Throwable e) { if (e instanceof SQLException) { final String sqlState = ((SQLException) e).getSQLState(); // limited to errors that are likely to be resolved within a few retries // retry on connection errors and insufficient resources // see http://www.postgresql.org/docs/current/static/errcodes-appendix.html for details return sqlState != null && (sqlState.startsWith("08") || sqlState.startsWith("53")); } return false; } }