software.amazon.jdbc.Driver Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of aws-advanced-jdbc-wrapper Show documentation
Amazon Web Services (AWS) Advanced JDBC Wrapper
The newest version!
/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * 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 software.amazon.jdbc;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.DriverPropertyInfo;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Properties;
import java.util.logging.ConsoleHandler;
import java.util.logging.Handler;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import software.amazon.jdbc.profile.ConfigurationProfile;
import software.amazon.jdbc.profile.DriverConfigurationProfiles;
import software.amazon.jdbc.states.ResetSessionStateOnCloseCallable;
import software.amazon.jdbc.states.TransferSessionStateOnSwitchCallable;
import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
import software.amazon.jdbc.targetdriverdialect.TargetDriverDialectManager;
import software.amazon.jdbc.util.ConnectionUrlParser;
import software.amazon.jdbc.util.DriverInfo;
import software.amazon.jdbc.util.Messages;
import software.amazon.jdbc.util.PropertyUtils;
import software.amazon.jdbc.util.StringUtils;
import software.amazon.jdbc.util.telemetry.DefaultTelemetryFactory;
import software.amazon.jdbc.util.telemetry.TelemetryContext;
import software.amazon.jdbc.util.telemetry.TelemetryFactory;
import software.amazon.jdbc.util.telemetry.TelemetryTraceLevel;
import software.amazon.jdbc.wrapper.ConnectionWrapper;

public class Driver implements java.sql.Driver {

  private static final String PROTOCOL_PREFIX = "jdbc:aws-wrapper:";
  private static final Logger PARENT_LOGGER = Logger.getLogger("software.amazon.jdbc");
  private static final Logger LOGGER = Logger.getLogger("software.amazon.jdbc.Driver");
  private static @Nullable Driver registeredDriver;

  private static ResetSessionStateOnCloseCallable resetSessionStateOnCloseCallable = null;
  private static TransferSessionStateOnSwitchCallable transferSessionStateOnSwitchCallable = null;

  static {
    try {
      register();
    } catch (final SQLException e) {
      throw new ExceptionInInitializerError(e);
    }
  }

  public static void register() throws SQLException {
    if (isRegistered()) {
      throw new IllegalStateException(
          Messages.get("Driver.alreadyRegistered"));
    }
    final Driver driver = new Driver();
    DriverManager.registerDriver(driver);
    registeredDriver = driver;
  }

  public static void deregister() throws SQLException {
    if (registeredDriver == null) {
      throw new IllegalStateException(
          Messages.get("Driver.notRegistered"));
    }
    DriverManager.deregisterDriver(registeredDriver);
    registeredDriver = null;
  }

  public static boolean isRegistered() {
    if (registeredDriver != null) {
      final List registeredDrivers = Collections.list(DriverManager.getDrivers());
      for (final java.sql.Driver d : registeredDrivers) {
        if (d == registeredDriver) {
          return true;
        }
      }

      // Driver isn't registered with DriverManager.
      registeredDriver = null;
    }
    return false;
  }

  @Override
  public Connection connect(final String url, final Properties info) throws SQLException {
    if (!acceptsURL(url)) {
      return null;
    }

    LOGGER.finest("Opening connection to " + url);

    ConnectionUrlParser.parsePropertiesFromUrl(url, info);
    final Properties props = PropertyUtils.copyProperties(info);

    final String databaseName = ConnectionUrlParser.parseDatabaseFromUrl(url);
    if (!StringUtils.isNullOrEmpty(databaseName)) {
      PropertyDefinition.DATABASE.set(props, databaseName);
    }

    LOGGER.finest(() -> PropertyUtils.logProperties(
        PropertyUtils.maskProperties(props), "Connecting with properties: \n"));

    final String profileName = PropertyDefinition.PROFILE_NAME.getString(props);
    ConfigurationProfile configurationProfile = null;
    if (!StringUtils.isNullOrEmpty(profileName)) {
      configurationProfile = DriverConfigurationProfiles.getProfileConfiguration(profileName);
      if (configurationProfile != null) {
        PropertyUtils.addProperties(props, configurationProfile.getProperties());
      } else {
        throw new SQLException(
            Messages.get(
                "Driver.configurationProfileNotFound",
                new Object[] {profileName}));
      }
    }

    TelemetryFactory telemetryFactory = new DefaultTelemetryFactory(props);
    TelemetryContext context = telemetryFactory.openTelemetryContext(
        "software.amazon.jdbc.Driver.connect", TelemetryTraceLevel.TOP_LEVEL);

    try {
      final String driverUrl = url.replaceFirst(PROTOCOL_PREFIX, "jdbc:");

      TargetDriverHelper helper = new TargetDriverHelper();
      java.sql.Driver driver = helper.getTargetDriver(driverUrl, props);

      final String logLevelStr = PropertyDefinition.LOGGER_LEVEL.getString(props);
      if (!StringUtils.isNullOrEmpty(logLevelStr)) {
        final Level logLevel = Level.parse(logLevelStr.toUpperCase());
        final Logger rootLogger = Logger.getLogger("");
        for (final Handler handler : rootLogger.getHandlers()) {
          if (handler instanceof ConsoleHandler) {
            if (handler.getLevel().intValue() > logLevel.intValue()) {
              // Set higher (more detailed) level as requested
              handler.setLevel(logLevel);
            }
          }
        }
        PARENT_LOGGER.setLevel(logLevel);
      }

      TargetDriverDialect targetDriverDialect = configurationProfile == null
          ? null
          : configurationProfile.getTargetDriverDialect();

      if (targetDriverDialect == null) {
        final TargetDriverDialectManager targetDriverDialectManager = new TargetDriverDialectManager();
        targetDriverDialect = targetDriverDialectManager.getDialect(driver, props);
      }

      final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver);

      ConnectionProvider effectiveConnectionProvider = null;
      if (configurationProfile != null) {
        effectiveConnectionProvider = configurationProfile.getConnectionProvider();
      }

      return new ConnectionWrapper(
          props,
          driverUrl,
          defaultConnectionProvider,
          effectiveConnectionProvider,
          targetDriverDialect,
          configurationProfile,
          telemetryFactory);

    } catch (Exception ex) {
      context.setException(ex);
      context.setSuccess(false);
      throw ex;
    } finally {
      context.closeContext();
    }
  }

  @Override
  public boolean acceptsURL(final String url) throws SQLException {
    if (url == null) {
      throw new SQLException(Messages.get("Driver.nullUrl"));
    }

    return url.startsWith(PROTOCOL_PREFIX);
  }

  @Override
  public DriverPropertyInfo[] getPropertyInfo(final String url, final Properties info) throws SQLException {
    final Properties copy = new Properties(info);
    final String databaseName = ConnectionUrlParser.parseDatabaseFromUrl(url);
    if (!StringUtils.isNullOrEmpty(databaseName)) {
      PropertyDefinition.DATABASE.set(copy, databaseName);
    }
    ConnectionUrlParser.parsePropertiesFromUrl(url, copy);

    final Collection knownProperties = PropertyDefinition.allProperties();
    final DriverPropertyInfo[] props = new DriverPropertyInfo[knownProperties.size()];
    int i = 0;
    for (final AwsWrapperProperty prop : knownProperties) {
      props[i++] = prop.toDriverPropertyInfo(copy);
    }

    return props;
  }

  @Override
  public int getMajorVersion() {
    return DriverInfo.MAJOR_VERSION;
  }

  @Override
  public int getMinorVersion() {
    return DriverInfo.MINOR_VERSION;
  }

  @Override
  public boolean jdbcCompliant() {
    return false;
  }

  @Override
  public Logger getParentLogger() throws SQLFeatureNotSupportedException {
    return PARENT_LOGGER;
  }

  public static void setResetSessionStateOnCloseFunc(final @NonNull ResetSessionStateOnCloseCallable func) {
    resetSessionStateOnCloseCallable = func;
  }

  public static void resetResetSessionStateOnCloseFunc() {
    resetSessionStateOnCloseCallable = null;
  }

  public static ResetSessionStateOnCloseCallable getResetSessionStateOnCloseFunc() {
    return resetSessionStateOnCloseCallable;
  }

  public static void setTransferSessionStateOnSwitchFunc(final @NonNull TransferSessionStateOnSwitchCallable func) {
    transferSessionStateOnSwitchCallable = func;
  }

  public static void resetTransferSessionStateOnSwitchFunc() {
    transferSessionStateOnSwitchCallable = null;
  }

  public static TransferSessionStateOnSwitchCallable getTransferSessionStateOnSwitchFunc() {
    return transferSessionStateOnSwitchCallable;
  }
}