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.calcite.avatica.remote;
  
  import org.apache.calcite.avatica.ConnectionSpec;
  import org.apache.calcite.avatica.Meta;
  import org.apache.calcite.avatica.NoSuchStatementException;
  import org.apache.calcite.avatica.jdbc.JdbcMeta;
  import org.apache.calcite.avatica.metrics.MetricsSystemConfiguration;
  import org.apache.calcite.avatica.remote.Driver.Serialization;
  import org.apache.calcite.avatica.server.AvaticaHandler;
  import org.apache.calcite.avatica.server.AvaticaJsonHandler;
  import org.apache.calcite.avatica.server.AvaticaProtobufHandler;
  import org.apache.calcite.avatica.server.AvaticaServerConfiguration;
  import org.apache.calcite.avatica.server.HandlerFactory;
  import org.apache.calcite.avatica.server.HttpServer;
  import org.apache.calcite.avatica.server.Main;
  
  import org.mockito.Mockito;
  import org.mockito.internal.stubbing.answers.CallsRealMethods;
  import org.mockito.invocation.InvocationOnMock;
  
  import java.sql.SQLException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Properties;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.matches;
  /**
   * Utility class which encapsulates the setup required to write Avatica tests that run against
   * servers using each serialization approach.
   */
  public class AvaticaServersForTest {
    private static final ConnectionSpec CONNECTION_SPEC = ConnectionSpec.HSQLDB;
    private static final String[] SERVER_ARGS = { FullyRemoteJdbcMetaFactory.class.getName() };
  
    private final Map<Serialization, HttpServer> serversBySerialization;
  
    public AvaticaServersForTest() {
      serversBySerialization = new HashMap<>();
    }
  
    /**
     * Starts an Avatica server for each serialization type.
     */
    public void startServers() throws Exception {
      // Bind to '0' to pluck an ephemeral port instead of expecting a certain one to be free
      final HttpServer jsonServer = Main.start(SERVER_ARGS, 0, new Main.HandlerFactory() {
        @Override public AvaticaJsonHandler createHandler(Service service) {
          return new AvaticaJsonHandler(service);
        }
      });
      serversBySerialization.put(Serialization.JSON, jsonServer);
  
      final HttpServer protobufServer = Main.start(SERVER_ARGS, 0, new Main.HandlerFactory() {
        @Override public AvaticaProtobufHandler createHandler(Service service) {
          return new AvaticaProtobufHandler(service);
        }
      });
      serversBySerialization.put(Serialization.PROTOBUF, protobufServer);
    }
  
    /**
     * Starts Avatica servers for each serialization type with the properties.
     */
    public void startServers(Properties properties) {
      final HandlerFactory factory = new HandlerFactory();
  
      // Construct the JSON server
      Service jsonService =
              new LocalService(PropertyRemoteJdbcMetaFactory.getInstance(properties));
      startServer(factory, jsonService, Serialization.JSON, null, null);
  
      // Construct the Protobuf server
      Service protobufService =
              new LocalService(PropertyRemoteJdbcMetaFactory.getInstance(properties));
      startServer(factory, protobufService, Serialization.PROTOBUF, null, null);
    }
  
    /**
     * Starts Avatica servers for each serialization type with the provided {@code serverConfig}.
     */
    public void startServers(AvaticaServerConfiguration serverConfig) {
      final HandlerFactory factory = new HandlerFactory();
  
      // Construct the JSON server
      Service jsonService = new LocalService(FullyRemoteJdbcMetaFactory.getInstance());
      startServer(factory, jsonService, Serialization.JSON, null, serverConfig);
  
      // Construct the Protobuf server
      Service protobufService = new LocalService(FullyRemoteJdbcMetaFactory.getInstance());
      startServer(factory, protobufService, Serialization.PROTOBUF, null, serverConfig);
    }
  
    /**
     * Starts Avatica server and cache.
     */
    public void startServer(HandlerFactory factory, Service service, Serialization serialization,
                            MetricsSystemConfiguration metricsConfig,
                            AvaticaServerConfiguration serverConfig) {
      AvaticaHandler handler = factory.getHandler(service, serialization,
              metricsConfig, serverConfig);
      final HttpServer server = new HttpServer.Builder().withHandler(handler)
              .withPort(0).build();
      server.start();
      serversBySerialization.put(serialization, server);
    }
  
    /**
     * Stops the servers currently running.
     *
     * @throws Exception If there is an error stopping a server
     */
    public void stopServers() throws Exception {
      Iterator<Entry<Serialization, HttpServer>> servers =
          serversBySerialization.entrySet().iterator();
      while (servers.hasNext()) {
        try {
          servers.next().getValue().stop();
        } finally {
          // Still remove it if we failed to stop it
          servers.remove();
        }
      }
    }
  
    /**
     * Computes an array of parameters to support JUnit's parameterized tests. The Object array
     * actually contains a {@link Serialization} and the {@link HttpServer} instance in that order.
     *
     * @return A list of arrays of Serialization and HttpServer pairs.
     */
    public List<Object[]> getJUnitParameters() {
      List<Object[]> params = new ArrayList<>(serversBySerialization.size());
  
      for (Entry<Serialization, HttpServer> servers : serversBySerialization.entrySet()) {
        params.add(new Object[]{ servers.getKey(), servers.getValue() });
      }
  
      return params;
    }
  
    /**
     * Computes the JDBC url for the Avatica server running on localhost, bound to the given port,
     * and using the given serialization.
     *
     * @param port The port the Avatica server is listening on.
     * @param serialization The serialization the Avatica server is using.
     * @return A JDBC URL to the local Avatica server.
     */
    public String getJdbcUrl(int port, Serialization serialization) {
      return getJdbcUrl(port, serialization, "");
    }
  
    /**
     * Computes the JDBC url for the Avatica server running on localhost, bound to the given port,
     * using the given serialization, with the user-provided suffix for the HTTP URL to the server.
     *
     * @param port The port the Avatica server is listening on.
     * @param serialization The serialization the Avatica server is using.
     * @param urlSuffix Extra information to apend to the HTTP URL to the Avatica server (optional).
     * @return A JDBC URL to the local Avatica server.
     */
    public String getJdbcUrl(int port, Serialization serialization, String urlSuffix) {
      return "jdbc:avatica:remote:url=http://localhost:" + port + urlSuffix + ";serialization="
          + serialization.name();
    }
  
    /** Factory that provides a {@link JdbcMeta}. */
    public static class FullyRemoteJdbcMetaFactory implements Meta.Factory {
  
      private static JdbcMeta instance = null;
  
      static JdbcMeta getInstance() {
        if (instance == null) {
          try {
            JdbcMeta realInstance = new JdbcMeta(CONNECTION_SPEC.url, CONNECTION_SPEC.username,
                    CONNECTION_SPEC.password);
            // Add a configurable delay to the Statement.execute() method based on a DELAY SQL comment
            instance = Mockito.spy(realInstance);
            Mockito.doAnswer(new CallsRealMethods() {
              @Override
              public Object answer(InvocationOnMock invocation) throws Throwable {
                Pattern pattern = Pattern.compile("/\\* DELAY=(\\d+) \\*/");
                Matcher matcher = pattern.matcher((String) invocation.getArgument(1));
                if (matcher.find()) {
                  long delay =  Long.parseLong(matcher.group(1));
                  Thread.sleep(delay);
                }
                return super.answer(invocation);
              }
            })
                .when(instance).prepareAndExecute(any(Meta.StatementHandle.class),
                matches("/\\* DELAY=(\\d+) \\*/"),
                any(long.class), any(int.class), any(Meta.PrepareCallback.class));
          } catch (SQLException | NoSuchStatementException e) {
            throw new RuntimeException("Error when applying test DELAY via Mockito", e);
          }
        }
        return instance;
      }
  
      @Override public Meta create(List<String> args) {
        return getInstance();
      }
    }
  
    /** Factory that provides a {@link JdbcMeta} with properties. */
    public static class PropertyRemoteJdbcMetaFactory implements Meta.Factory {
  
      private static Map<Properties, JdbcMeta> instances = new HashMap<>();
  
      static JdbcMeta getInstance(Properties properties) {
        try {
          if (properties == null) {
            return new JdbcMeta(CONNECTION_SPEC.url, CONNECTION_SPEC.username,
                      CONNECTION_SPEC.password);
          }
          if (instances.get(properties) == null) {
            properties.put("user", CONNECTION_SPEC.username);
            properties.put("password", CONNECTION_SPEC.password);
            JdbcMeta instance = new JdbcMeta(CONNECTION_SPEC.url, properties);
            instances.put(properties, instance);
          }
        } catch (SQLException e) {
          throw new RuntimeException(e);
        }
        return instances.get(properties);
      }
  
      @Override public Meta create(List<String> args) {
        return getInstance(new Properties());
      }
    }
  }
  
  // End AvaticaServersForTest.java
  