org.apache.calcite.avatica.remote.Service Java Examples

@Override public boolean syncResults(final StatementHandle h, final QueryState state,
    final long offset) throws NoSuchStatementException {
  try {
    return connection.invokeWithRetries(
        new CallableWithoutException<Boolean>() {
          public Boolean call() {
            final Service.SyncResultsResponse response =
                service.apply(
                    new Service.SyncResultsRequest(h.connectionId, h.id, state, offset));
            if (response.missingStatement) {
              throw new RuntimeException(new NoSuchStatementException(h));
            }
            return response.moreResults;
          }
        });
  } catch (RuntimeException e) {
    Throwable cause = e.getCause();
    if (cause instanceof NoSuchStatementException) {
      throw (NoSuchStatementException) cause;
    }
    throw e;
  }
}
 

/**
 * Instantiates the Handler for use by the Avatica (Jetty) server.
 *
 * @param service The Avatica Service implementation
 * @param handlerFactory Factory used for creating a Handler
 * @return The Handler to use.
 */
Handler getHandler(Service service, HandlerFactory handlerFactory) {
  String serializationName = "PROTOBUF";
  Driver.Serialization serialization;
  try {
    serialization = Driver.Serialization.valueOf(serializationName);
  } catch (Exception e) {
    LOG.error("Unknown message serialization type for " + serializationName);
    throw e;
  }

  Handler handler = handlerFactory.getHandler(service, serialization);
  LOG.info("Instantiated " + handler.getClass() + " for Quark Server");

  return handler;
}
 

private MetaResultSet toResultSet(Class clazz,
    Service.ResultSetResponse response) {
  if (response.updateCount != -1) {
    return MetaResultSet.count(response.connectionId, response.statementId,
        response.updateCount);
  }
  Signature signature0 = response.signature;
  if (signature0 == null) {
    final List<ColumnMetaData> columns =
        clazz == null
            ? Collections.<ColumnMetaData>emptyList()
            : fieldMetaData(clazz).columns;
    signature0 = Signature.create(columns,
        "?", Collections.<AvaticaParameter>emptyList(), CursorFactory.ARRAY,
        StatementType.SELECT);
  }
  return MetaResultSet.create(response.connectionId, response.statementId,
      response.ownStatement, signature0, response.firstFrame);
}
 

public AvaticaJsonHandler(Service service, MetricsSystem metrics,
    AvaticaServerConfiguration serverConfig) {
  this.service = Objects.requireNonNull(service);
  this.metrics = Objects.requireNonNull(metrics);
  // Avatica doesn't have a Guava dependency
  this.jsonHandler = new JsonHandler(service, this.metrics);

  // Metrics
  this.requestTimer = this.metrics.getTimer(
      concat(AvaticaJsonHandler.class, MetricsAwareAvaticaHandler.REQUEST_TIMER_NAME));

  this.threadLocalBuffer = new ThreadLocal<UnsynchronizedBuffer>() {
    @Override public UnsynchronizedBuffer initialValue() {
      return new UnsynchronizedBuffer();
    }
  };

  this.serverConfig = serverConfig;
}
 

@Test public void testExecuteRequestWithNumberParameter() {
  final List<TypedValue> expectedParameterValues = new ArrayList<>();
  final Service service = new ParameterValuesCheckingService(expectedParameterValues);
  final JsonService jsonService = new LocalJsonService(service);
  final JsonHandler jsonHandler = new JsonHandler(jsonService, NoopMetricsSystem.getInstance());

  final List<TypedValue> parameterValues = Arrays.asList(
      TypedValue.create("NUMBER", new BigDecimal("123")),
      TypedValue.create("STRING", "calcite"));

  jsonHandler.apply(
      "{'request':'execute',"
      + "'parameterValues':[{'type':'NUMBER','value':123},"
      + "{'type':'STRING','value':'calcite'}]}");
  assertThat(expectedParameterValues.size(), is(2));
  assertThat(expectedParameterValues.get(0), is(parameterValues.get(0)));
  assertThat(expectedParameterValues.get(1), is(parameterValues.get(1)));
}
 

@Override public ExecuteResponse apply(ExecuteRequest request) {
  expectedParameterValues.addAll(request.parameterValues);

  final Meta.Signature signature =
      new Meta.Signature(Collections.<ColumnMetaData>emptyList(),
          "SELECT 1 FROM VALUE()",
          Collections.<AvaticaParameter>emptyList(),
          Collections.<String, Object>emptyMap(),
          CursorFactory.LIST, Meta.StatementType.SELECT);

  final Service.ResultSetResponse resultSetResponse =
      new Service.ResultSetResponse(UUID.randomUUID().toString(),
          RANDOM.nextInt(), false, signature, Meta.Frame.EMPTY, -1L, null);

  return new Service.ExecuteResponse(
      Collections.singletonList(resultSetResponse), false, null);
}
 

/**
 * Creates a {@link Service} with the given {@link AvaticaConnection} and configuration.
 *
 * @param connection The {@link AvaticaConnection} to use.
 * @param config Configuration properties
 * @return A Service implementation.
 */
Service createService(AvaticaConnection connection, ConnectionConfig config) {
  final Service.Factory metaFactory = config.factory();
  final Service service;
  if (metaFactory != null) {
    service = metaFactory.create(connection);
  } else if (config.url() != null) {
    final AvaticaHttpClient httpClient = getHttpClient(connection, config);
    final Serialization serializationType = getSerialization(config);

    LOG.debug("Instantiating {} service", serializationType);
    switch (serializationType) {
    case JSON:
      service = new RemoteService(httpClient);
      break;
    case PROTOBUF:
      service = new RemoteProtobufService(httpClient, new ProtobufTranslationImpl());
      break;
    default:
      throw new IllegalArgumentException("Unhandled serialization type: " + serializationType);
    }
  } else {
    service = new MockJsonService(Collections.<String, String>emptyMap());
  }
  return service;
}
 

public static void main(String[] args) throws Exception {
    OptionsParser parser = new OptionsParser(args);
    final int port = Integer.parseInt(parser.getOptionValue(SubmitOption.PORT));
    final String[] mainArgs = new String[]{FullyRemoteJdbcMetaFactory.class.getName()};

    HttpServer server = Main.start(mainArgs, port, new HandlerFactory() {
        @Override
        public AbstractHandler createHandler(Service service) {
            return new AvaticaJsonHandler(service);
        }
    });
    InetAddress address = InetAddress.getLocalHost();
    String hostName = "localhost";
    if (address != null) {
         hostName = StringUtils.isNotBlank(address.getHostName()) ? address.getHostName() : address
            .getHostAddress();
    }
    String url = Driver.CONNECT_STRING_PREFIX + "url=http://" + hostName + ":" + server.getPort();
    System.out.println("Quicksql server started, Please connect : " + url);
    server.join();
}
 

@SuppressWarnings("unchecked") // needed for the mocked customizers, not the builder
@Test public void serverCustomizersInvoked() {
  ServerCustomizer<Server> mockCustomizer1 =
      (ServerCustomizer<Server>) mock(ServerCustomizer.class);
  ServerCustomizer<Server> mockCustomizer2 =
      (ServerCustomizer<Server>) mock(ServerCustomizer.class);
  Service service = new LocalService(mockMeta);
  HttpServer server =
      HttpServer.Builder.<Server>newBuilder().withHandler(service, Driver.Serialization.PROTOBUF)
          .withServerCustomizers(Arrays.asList(mockCustomizer1, mockCustomizer2), Server.class)
          .withPort(0).build();
  try {
    server.start();
    verify(mockCustomizer2).customize(any(Server.class));
    verify(mockCustomizer1).customize(any(Server.class));
  } finally {
    server.stop();
  }
}
 

/**
 * Constructs the desired implementation for the given serialization method and server
 * configuration with metrics.
 *
 * @param service The underlying {@link Service}
 * @param serialization The serializatio mechanism to use
 * @param metricsConfig Configuration for the {@link MetricsSystem}.
 * @param serverConfig Avatica server configuration or null
 * @return An {@link AvaticaHandler}
 */
public AvaticaHandler getHandler(Service service, Driver.Serialization serialization,
    MetricsSystemConfiguration<?> metricsConfig, AvaticaServerConfiguration serverConfig) {
  if (null == metricsConfig) {
    metricsConfig = NoopMetricsSystemConfiguration.getInstance();
  }
  MetricsSystem metrics = MetricsSystemLoader.load(metricsConfig);

  switch (serialization) {
  case JSON:
    return new AvaticaJsonHandler(service, metrics, serverConfig);
  case PROTOBUF:
    return new AvaticaProtobufHandler(service, metrics, serverConfig);
  default:
    throw new IllegalArgumentException("Unknown Avatica handler for " + serialization.name());
  }
}
 

public AvaticaProtobufHandler(Service service, MetricsSystem metrics,
    AvaticaServerConfiguration serverConfig) {
  this.service = Objects.requireNonNull(service);
  this.metrics = Objects.requireNonNull(metrics);

  this.requestTimer = this.metrics.getTimer(
      MetricsHelper.concat(AvaticaProtobufHandler.class,
          MetricsAwareAvaticaHandler.REQUEST_TIMER_NAME));

  this.protobufTranslation = new ProtobufTranslationImpl();
  this.pbHandler = new ProtobufHandler(service, protobufTranslation, metrics);

  this.threadLocalBuffer = new ThreadLocal<UnsynchronizedBuffer>() {
    @Override public UnsynchronizedBuffer initialValue() {
      return new UnsynchronizedBuffer();
    }
  };

  this.serverConfig = serverConfig;
}
 

@Test public void lotsOfExtraRoles() {
  final String[] extraRoles = new String[] {"BAR.COM", "BAZ.COM", "FOO.COM"};
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertArrayEquals(extraRoles, server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM", "BAR.COM", "BAZ.COM", "FOO.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Test public void extraAllowedRolesConfiguredWithExplitRealm() {
  final String[] extraRoles = new String[] {"BAR.COM"};
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", "EXAMPLE.COM", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertArrayEquals(extraRoles, server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM", "BAR.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Test public void extraAllowedRolesConfigured() {
  final String[] extraRoles = new String[] {"BAR.COM"};
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertArrayEquals(extraRoles, server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM", "BAR.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Override public MetaResultSet getSchemas(final ConnectionHandle ch, final String catalog,
    final Pat schemaPattern) {
  return connection.invokeWithRetries(
      new CallableWithoutException<MetaResultSet>() {
        public MetaResultSet call() {
          final Service.ResultSetResponse response =
              service.apply(
                  new Service.SchemasRequest(ch.id, catalog, schemaPattern.s));
          return toResultSet(MetaSchema.class, response);
        }
      });
}
 

@Override public Meta createMeta(AvaticaConnection connection) {
  final ConnectionConfig config = connection.config();

  // Perform the login and launch the renewal thread if necessary
  final KerberosConnection kerberosUtil = createKerberosUtility(config);
  if (null != kerberosUtil) {
    kerberosUtil.login();
    connection.setKerberosConnection(kerberosUtil);
  }

  // Create a single Service and set it on the Connection instance
  final Service service = createService(connection, config);
  connection.setService(service);
  return new QuicksqlRemoteMeta(connection, service);
}
 

@Override public MetaResultSet getTables(final ConnectionHandle ch, final String catalog,
    final Pat schemaPattern, final Pat tableNamePattern, final List<String> typeList) {
  return connection.invokeWithRetries(
      new CallableWithoutException<MetaResultSet>() {
        public MetaResultSet call() {
          final Service.ResultSetResponse response =
              service.apply(
                  new Service.TablesRequest(ch.id, catalog, schemaPattern.s,
                      tableNamePattern.s, typeList));
          return toResultSet(MetaTable.class, response);
        }
      });
}
 

@Override public MetaResultSet getTypeInfo(final ConnectionHandle ch) {
  return connection.invokeWithRetries(
      new CallableWithoutException<MetaResultSet>() {
        public MetaResultSet call() {
          final Service.ResultSetResponse response =
              service.apply(new Service.TypeInfoRequest(ch.id));
          return toResultSet(MetaTypeInfo.class, response);
        }
      });
}
 

@Test public void nullExtraRoles() {
  final String[] extraRoles = null;
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertNull(server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Test public void emptyExtraRoles() {
  final String[] extraRoles = new String[0];
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertArrayEquals(extraRoles, server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Override public Service create(AvaticaConnection connection) {
  try {
    Connection conn = makeConnection();
    final CalciteMetaImpl meta =
        new CalciteMetaImpl(conn.unwrap(CalciteConnectionImpl.class));
    return new LocalService(meta);
  } catch (Exception e) {
    throw TestUtil.rethrow(e);
  }
}
 

@Test public void lotsOfExtraRolesWithExplitRealm() {
  final String[] extraRoles = new String[] {"BAR.COM", "BAZ.COM", "FOO.COM"};
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", "EXAMPLE.COM", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertArrayEquals(extraRoles, server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM", "BAR.COM", "BAZ.COM", "FOO.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Test public void nullExtraRolesWithExplitRealm() {
  final String[] extraRoles = null;
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", "EXAMPLE.COM", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertNull(server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Test public void emptyExtraRolesWithExplitRealm() {
  final String[] extraRoles = new String[0];
  final Service mockService = Mockito.mock(Service.class);
  HttpServer server = new HttpServer.Builder()
      .withSpnego("HTTP/[email protected]", "EXAMPLE.COM", extraRoles)
      .withHandler(mockService, Serialization.JSON)
      .build();

  assertArrayEquals(extraRoles, server.getConfig().getAllowedRoles());

  assertArrayEquals(new String[] {"EXAMPLE.COM"},
      server.getAllowedRealms("EXAMPLE.COM", server.getConfig()));
}
 

@Test public void onlyJettyCustomizersAllowed() {
  Service service = new LocalService(mockMeta);
  List<ServerCustomizer<UnsupportedServer>> unsupportedCustomizers = new ArrayList<>();
  unsupportedCustomizers.add(new ServerCustomizer<UnsupportedServer>() {
    @Override public void customize(UnsupportedServer server) {
    }
  });
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("Only Jetty Server customizers are supported");
  HttpServer.Builder.<UnsupportedServer>newBuilder()
      .withHandler(service, Driver.Serialization.PROTOBUF)
      .withServerCustomizers(unsupportedCustomizers, UnsupportedServer.class).withPort(0).build();
}
 

@Override public Service create(AvaticaConnection connection) {
  try {
    return new LocalService(
        new JdbcMeta(CONNECTION_SPEC.url, CONNECTION_SPEC.username,
            CONNECTION_SPEC.password));
  } catch (SQLException e) {
    throw new RuntimeException(e);
  }
}
 

public void run(String[] args) throws Exception {
  logJVMInfo();
  try {

    Class<? extends Meta.Factory> factoryClass = QuarkMetaFactoryImpl.class;

    Meta.Factory factory =
        factoryClass.getDeclaredConstructor().newInstance();
    Meta meta = factory.create(Arrays.asList(args));

    int port = 8765;
    if (QuarkMetaFactoryImpl.serverConfig.port != 0) {
      port = QuarkMetaFactoryImpl.serverConfig.port;
    }
    LOG.debug("Listening on port " + port);

    final HandlerFactory handlerFactory = new HandlerFactory();
    Service service = new LocalService(meta);
    server = new HttpServer(port, getHandler(service, handlerFactory));
    Class.forName("com.qubole.quark.fatjdbc.QuarkDriver");
    server.start();
    runningLatch.countDown();
    server.join();
  } catch (Throwable t) {
    LOG.fatal("Unrecoverable service error. Shutting down.", t);
    this.t = t;
  }
}
 

@Override public Service create(AvaticaConnection connection) {
  try {
    Connection conn = makeConnection();
    final CalciteMetaImpl meta =
        new CalciteMetaImpl(conn.unwrap(CalciteConnectionImpl.class));
    return new LocalService(meta);
  } catch (Exception e) {
    throw TestUtil.rethrow(e);
  }
}
 

public Service create(AvaticaConnection connection) {
  try {
    Connection localConnection = CalciteAssert.hr().connect();
    final Meta meta = CalciteConnectionImpl.TROJAN
        .getMeta((CalciteConnectionImpl) localConnection);
    return new LocalJsonService(new LocalService(meta));
  } catch (Exception e) {
    throw TestUtil.rethrow(e);
  }
}
 

@Override public Service create(AvaticaConnection connection) {
  try {
    Connection conn = JdbcFrontLinqBackTest.makeConnection();
    final CalciteMetaImpl meta =
        new CalciteMetaImpl(
            conn.unwrap(CalciteConnectionImpl.class));
    return new LocalService(meta);
  } catch (Exception e) {
    throw TestUtil.rethrow(e);
  }
}