content: /*
   * 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.arrow.driver.jdbc;
  
  import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.replaceSemiColons;
  
  import java.io.BufferedReader;
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.InputStreamReader;
  import java.io.Reader;
  import java.net.URI;
  import java.nio.charset.StandardCharsets;
  import java.sql.SQLException;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Properties;
  import java.util.logging.Logger;
  import org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty;
  import org.apache.arrow.driver.jdbc.utils.UrlParser;
  import org.apache.arrow.flight.FlightRuntimeException;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.util.Preconditions;
  import org.apache.arrow.util.VisibleForTesting;
  import org.apache.calcite.avatica.AvaticaConnection;
  import org.apache.calcite.avatica.DriverVersion;
  import org.apache.calcite.avatica.Meta;
  import org.apache.calcite.avatica.UnregisteredDriver;
  
  /** JDBC driver for querying data from an Apache Arrow Flight server. */
  public class ArrowFlightJdbcDriver extends UnregisteredDriver {
    private static final String CONNECT_STRING_PREFIX = "jdbc:arrow-flight-sql://";
    private static final String CONNECT_STRING_PREFIX_DEPRECATED = "jdbc:arrow-flight://";
    private static final String CONNECTION_STRING_EXPECTED =
        "jdbc:arrow-flight-sql://[host][:port][?param1=value&...]";
    private static DriverVersion version;
  
    static {
      // Special code for supporting Java9 and higher.
      // Netty requires some extra properties to unlock some native memory management api
      // Setting this property if not already set externally
      // This has to be done before any netty class is being loaded
      final String key = "io.netty.tryReflectionSetAccessible";
      final String tryReflectionSetAccessible = System.getProperty(key);
      if (tryReflectionSetAccessible == null) {
        System.setProperty(key, Boolean.TRUE.toString());
      }
  
      new ArrowFlightJdbcDriver().register();
    }
  
    @Override
    public Logger getParentLogger() {
      // Return the logger associated with the driver package ('org.apache.arrow.driver.jdbc')
      // When packaged in flight-sql-jdbc-driver, it will also apply to all shaded dependencies
      return Logger.getLogger(getClass().getPackage().getName());
    }
  
    @Override
    public ArrowFlightConnection connect(final String url, final Properties info)
        throws SQLException {
      final Properties properties = new Properties(info);
      properties.putAll(info);
  
      if (url != null) {
        final Optional<Map<Object, Object>> maybeProperties = getUrlsArgs(url);
        if (!maybeProperties.isPresent()) {
          return null;
        }
        final Map<Object, Object> propertiesFromUrl = maybeProperties.get();
        properties.putAll(propertiesFromUrl);
      }
  
      try {
        return ArrowFlightConnection.createNewConnection(
            this, factory, url, lowerCasePropertyKeys(properties), new RootAllocator(Long.MAX_VALUE));
      } catch (final FlightRuntimeException e) {
        throw new SQLException("Failed to connect.", e);
      }
    }
  
    @Override
    protected String getFactoryClassName(final JdbcVersion jdbcVersion) {
      return ArrowFlightJdbcFactory.class.getName();
    }
  
    @Override
    @SuppressWarnings("StringSplitter")
    protected DriverVersion createDriverVersion() {
      if (version == null) {
        final InputStream flightProperties =
            this.getClass().getResourceAsStream("/properties/flight.properties");
        if (flightProperties == null) {
          throw new RuntimeException(
              "Flight Properties not found. Ensure the JAR was built properly.");
        }
        try (final Reader reader =
            new BufferedReader(new InputStreamReader(flightProperties, StandardCharsets.UTF_8))) {
          final Properties properties = new Properties();
          properties.load(reader);
  
          final String parentName = properties.getProperty("org.apache.arrow.flight.name");
          final String parentVersion = properties.getProperty("org.apache.arrow.flight.version");
          final String[] pVersion = parentVersion.split("\\.");
  
          final int parentMajorVersion = Integer.parseInt(pVersion[0]);
          final int parentMinorVersion = Integer.parseInt(pVersion[1]);
  
          final String childName = properties.getProperty("org.apache.arrow.flight.jdbc-driver.name");
          final String childVersion =
              properties.getProperty("org.apache.arrow.flight.jdbc-driver.version");
          final String[] cVersion = childVersion.split("\\.");
  
          final int childMajorVersion = Integer.parseInt(cVersion[0]);
          final int childMinorVersion = Integer.parseInt(cVersion[1]);
  
          version =
              new DriverVersion(
                  childName,
                  childVersion,
                  parentName,
                  parentVersion,
                  true,
                  childMajorVersion,
                  childMinorVersion,
                  parentMajorVersion,
                  parentMinorVersion);
        } catch (final IOException e) {
          throw new RuntimeException("Failed to load driver version.", e);
        }
      }
  
      return version;
    }
  
    @Override
    public Meta createMeta(final AvaticaConnection connection) {
      return new ArrowFlightMetaImpl(connection);
    }
  
    @Override
    protected String getConnectStringPrefix() {
      return CONNECT_STRING_PREFIX;
    }
  
    @Override
    public boolean acceptsURL(final String url) {
      Preconditions.checkNotNull(url);
      return url.startsWith(CONNECT_STRING_PREFIX)
          || url.startsWith(CONNECT_STRING_PREFIX_DEPRECATED);
    }
  
    /**
     * Parses the provided url based on the format this driver accepts, retrieving arguments after the
     * {@link #CONNECT_STRING_PREFIX}.
     *
     * <p>This method gets the args if the provided URL follows this pattern: {@code
     * jdbc:arrow-flight-sql://<host>:<port>[/?key1=val1&key2=val2&(...)]}
     *
     * <table border="1">
     *    <tr>
     *        <td>Group</td>
     *        <td>Definition</td>
     *        <td>Value</td>
     *    </tr>
     *    <tr>
     *        <td>? — inaccessible</td>
     *        <td>{@link #getConnectStringPrefix}</td>
     *        <td>
     *            the URL prefix accepted by this driver, i.e.,
     *            {@code "jdbc:arrow-flight-sql://"}
     *        </td>
     *    </tr>
     *    <tr>
     *        <td>1</td>
     *        <td>IPv4 host name</td>
     *        <td>
     *            first word after previous group and before "{@code :}"
     *        </td>
     *    </tr>
     *    <tr>
     *        <td>2</td>
     *        <td>IPv4 port number</td>
     *        <td>
     *            first number after previous group and before "{@code /?}"
     *        </td>
     *    </tr>
     *    <tr>
     *        <td>3</td>
     *        <td>custom call parameters</td>
     *        <td>
     *            all parameters provided after "{@code /?}" — must follow the
     *            pattern: "{@code key=value}" with "{@code &}" separating a
     *            parameter from another
     *        </td>
     *    </tr>
     * </table>
     *
     * @param url The url to parse.
     * @return the parsed arguments, or an empty optional if the driver does not handle this URL.
     * @throws SQLException If an error occurs while trying to parse the URL.
     */
    @VisibleForTesting // ArrowFlightJdbcDriverTest
    Optional<Map<Object, Object>> getUrlsArgs(String url) throws SQLException {
  
      /*
       *
       * Perhaps this logic should be inside a utility class, separated from this
       * one, so as to better delegate responsibilities and concerns throughout
       * the code and increase maintainability.
       *
       * =====
       *
       * Keep in mind that the URL must ALWAYS follow the pattern:
       * "jdbc:arrow-flight-sql://<host>:<port>[/?param1=value1&param2=value2&(...)]."
       *
       */
  
      final Properties resultMap = new Properties();
      url = replaceSemiColons(url);
  
      if (!url.startsWith("jdbc:")) {
        throw new SQLException(
            "Connection string must start with 'jdbc:'. Expected format: "
                + CONNECTION_STRING_EXPECTED);
      }
  
      // It's necessary to use a string without "jdbc:" at the beginning to be parsed as a valid URL.
      url = url.substring(5);
  
      final URI uri;
  
      try {
        uri = URI.create(url);
      } catch (final IllegalArgumentException e) {
        throw new SQLException("Malformed/invalid URL!", e);
      }
  
      if (!Objects.equals(uri.getScheme(), "arrow-flight")
          && !Objects.equals(uri.getScheme(), "arrow-flight-sql")) {
        return Optional.empty();
      }
  
      if (uri.getHost() == null) {
        throw new SQLException(
            "URL must have a host. Expected format: " + CONNECTION_STRING_EXPECTED);
      } else if (uri.getPort() < 0) {
        throw new SQLException(
            "URL must have a port. Expected format: " + CONNECTION_STRING_EXPECTED);
      }
      resultMap.put(ArrowFlightConnectionProperty.HOST.camelName(), uri.getHost()); // host
      resultMap.put(ArrowFlightConnectionProperty.PORT.camelName(), uri.getPort()); // port
  
      final String extraParams = uri.getRawQuery(); // optional params
      if (extraParams != null) {
        final Map<String, String> keyValuePairs = UrlParser.parse(extraParams, "&");
        resultMap.putAll(keyValuePairs);
      }
  
      return Optional.of(resultMap);
    }
  
    static Properties lowerCasePropertyKeys(final Properties properties) {
      final Properties resultProperty = new Properties();
      properties.forEach((k, v) -> resultProperty.put(k.toString().toLowerCase(), v));
      return resultProperty;
    }
  }
  