pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/RealtimeKinesisIntegrationTest.java - pinot - Git at Google

 /**
  * 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.pinot.integration.tests;

 import cloud.localstack.Localstack;
 import cloud.localstack.ServiceName;
 import cloud.localstack.docker.annotation.LocalstackDockerAnnotationProcessor;
 import cloud.localstack.docker.annotation.LocalstackDockerConfiguration;
 import cloud.localstack.docker.annotation.LocalstackDockerProperties;
 import cloud.localstack.docker.command.Command;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.JsonNodeType;
 import com.google.common.base.Function;
 import java.io.BufferedReader;
 import java.io.File;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.URI;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 import java.sql.DriverManager;
 import java.sql.PreparedStatement;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import javax.activation.UnsupportedDataTypeException;
 import javax.annotation.Nullable;
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.pinot.client.ResultSet;
 import org.apache.pinot.plugin.stream.kinesis.KinesisConfig;
 import org.apache.pinot.plugin.stream.kinesis.KinesisConsumerFactory;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TableType;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.stream.StreamConfig;
 import org.apache.pinot.spi.stream.StreamConfigProperties;
 import org.apache.pinot.spi.utils.JsonUtils;
 import org.apache.pinot.spi.utils.StringUtil;
 import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
 import org.apache.pinot.util.TestUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
 import org.testng.SkipException;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.http.SdkHttpConfigurationOption;
 import software.amazon.awssdk.http.apache.ApacheSdkHttpService;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.kinesis.KinesisClient;
 import software.amazon.awssdk.services.kinesis.model.CreateStreamRequest;
 import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest;
 import software.amazon.awssdk.services.kinesis.model.PutRecordRequest;
 import software.amazon.awssdk.services.kinesis.model.PutRecordResponse;
 import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
 import software.amazon.awssdk.utils.AttributeMap;


 @LocalstackDockerProperties(services = {ServiceName.KINESIS}, imageTag = "0.12.15")
 public class RealtimeKinesisIntegrationTest extends BaseClusterIntegrationTestSet {
   private static final Logger LOGGER = LoggerFactory.getLogger(RealtimeKinesisIntegrationTest.class);

   private static final LocalstackDockerAnnotationProcessor PROCESSOR = new LocalstackDockerAnnotationProcessor();
   private static final String STREAM_NAME = "kinesis-test";
   private static final String STREAM_TYPE = "kinesis";
   public static final int MAX_RECORDS_TO_FETCH = Integer.MAX_VALUE;

   public static final String REGION = "us-east-1";
   public static final String LOCALSTACK_KINESIS_ENDPOINT = "http://localhost:4566";
   public static final int NUM_SHARDS = 10;

   // Localstack Kinesis doesn't support large rows.
   // So, this airlineStats data file consists of only few fields and rows from the original data
   public static final String SCHEMA_FILE_PATH = "kinesis/airlineStats_data_reduced.schema";
   public static final String DATA_FILE_PATH = "kinesis/airlineStats_data_reduced.json";

   private final Localstack _localstackDocker = Localstack.INSTANCE;

   private static KinesisClient _kinesisClient = null;

   private long _totalRecordsPushedInStream = 0;

   List<String> _h2FieldNameAndTypes = new ArrayList<>();

   private boolean _skipTestNoDockerInstalled = false;

   @BeforeClass(enabled = false)
   public void setUp()
       throws Exception {
     try {
       DockerInfoCommand dockerInfoCommand = new DockerInfoCommand();
       dockerInfoCommand.execute();
     } catch (IllegalStateException e) {
       _skipTestNoDockerInstalled = true;
       LOGGER.warn("Skipping test! Docker is not found running", e);
       throw new SkipException(e.getMessage());
     }

     TestUtils.ensureDirectoriesExistAndEmpty(_tempDir);

     // Start the Pinot cluster
     startZk();
     startController();
     startBroker();
     startServer();

     // Start Kinesis
     startKinesis();

     // Create and upload the schema and table config
     addSchema(createKinesisSchema());
     addTableConfig(createKinesisTableConfig());

     createH2ConnectionAndTable();

     // Push data into Kinesis
     publishRecordsToKinesis();

     // Wait for all documents loaded
     waitForAllDocsLoadedKinesis(120_000L);
   }

   public Schema createKinesisSchema()
       throws Exception {
     URL resourceUrl = BaseClusterIntegrationTest.class.getClassLoader().getResource(SCHEMA_FILE_PATH);
     Assert.assertNotNull(resourceUrl);
     return Schema.fromFile(new File(resourceUrl.getFile()));
   }

   protected void waitForAllDocsLoadedKinesis(long timeoutMs)
       throws Exception {
     waitForAllDocsLoadedKinesis(timeoutMs, true);
   }

   protected void waitForAllDocsLoadedKinesis(long timeoutMs, boolean raiseError) {
     TestUtils.waitForCondition(new Function<Void, Boolean>() {
       @Nullable
       @Override
       public Boolean apply(@Nullable Void aVoid) {
         try {
           return getCurrentCountStarResult() >= _totalRecordsPushedInStream;
         } catch (Exception e) {
           LOGGER.warn("Could not fetch current number of rows in pinot table " + getTableName(), e);
           return null;
         }
       }
     }, 1000L, timeoutMs, "Failed to load " + _totalRecordsPushedInStream + " documents", raiseError);
   }

   public TableConfig createKinesisTableConfig() {
     return new TableConfigBuilder(TableType.REALTIME).setTableName(getTableName()).setSchemaName(getTableName())
         .setTimeColumnName("DaysSinceEpoch").setFieldConfigList(getFieldConfigs()).setNumReplicas(getNumReplicas())
         .setSegmentVersion(getSegmentVersion()).setLoadMode(getLoadMode()).setTaskConfig(getTaskConfig())
         .setBrokerTenant(getBrokerTenant()).setServerTenant(getServerTenant()).setIngestionConfig(getIngestionConfig())
         .setLLC(true).setStreamConfigs(createKinesisStreamConfig()).setNullHandlingEnabled(getNullHandlingEnabled())
         .build();
   }

   public Map<String, String> createKinesisStreamConfig() {
     Map<String, String> streamConfigMap = new HashMap<>();
     String streamType = "kinesis";
     streamConfigMap.put(StreamConfigProperties.STREAM_TYPE, streamType);

     streamConfigMap.put(
         StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_TOPIC_NAME),
         STREAM_NAME);

     streamConfigMap.put(
         StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_FETCH_TIMEOUT_MILLIS),
         "30000");
     streamConfigMap.put(
         StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_CONSUMER_TYPES),
         StreamConfig.ConsumerType.LOWLEVEL.toString());
     streamConfigMap.put(StreamConfigProperties.constructStreamProperty(STREAM_TYPE,
         StreamConfigProperties.STREAM_CONSUMER_FACTORY_CLASS), KinesisConsumerFactory.class.getName());
     streamConfigMap.put(
         StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_DECODER_CLASS),
         "org.apache.pinot.plugin.inputformat.json.JSONMessageDecoder");
     streamConfigMap.put(KinesisConfig.REGION, REGION);
     streamConfigMap.put(KinesisConfig.MAX_RECORDS_TO_FETCH, String.valueOf(MAX_RECORDS_TO_FETCH));
     streamConfigMap.put(KinesisConfig.SHARD_ITERATOR_TYPE, ShardIteratorType.AFTER_SEQUENCE_NUMBER.toString());
     streamConfigMap.put(KinesisConfig.ENDPOINT, LOCALSTACK_KINESIS_ENDPOINT);
     streamConfigMap.put(KinesisConfig.ACCESS_KEY, getLocalAWSCredentials().resolveCredentials().accessKeyId());
     streamConfigMap.put(KinesisConfig.SECRET_KEY, getLocalAWSCredentials().resolveCredentials().secretAccessKey());
     streamConfigMap.put(StreamConfigProperties.SEGMENT_FLUSH_THRESHOLD_ROWS, Integer.toString(200));
     streamConfigMap.put(StreamConfigProperties.constructStreamProperty(streamType,
         StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA), "smallest");
     return streamConfigMap;
   }

   public void startKinesis()
       throws Exception {
     final LocalstackDockerConfiguration dockerConfig = PROCESSOR.process(this.getClass());
     StopAllLocalstackDockerCommand stopAllLocalstackDockerCommand = new StopAllLocalstackDockerCommand();
     stopAllLocalstackDockerCommand.execute();
     _localstackDocker.startup(dockerConfig);

     _kinesisClient = KinesisClient.builder().httpClient(new ApacheSdkHttpService().createHttpClientBuilder()
             .buildWithDefaults(
                 AttributeMap.builder().put(SdkHttpConfigurationOption.TRUST_ALL_CERTIFICATES, Boolean.TRUE).build()))
         .credentialsProvider(getLocalAWSCredentials()).region(Region.of(REGION))
         .endpointOverride(new URI(LOCALSTACK_KINESIS_ENDPOINT)).build();

     _kinesisClient.createStream(CreateStreamRequest.builder().streamName(STREAM_NAME).shardCount(NUM_SHARDS).build());

     TestUtils.waitForCondition(new Function<Void, Boolean>() {
       @Nullable
       @Override
       public Boolean apply(@Nullable Void aVoid) {
         try {
           String kinesisStreamStatus =
               _kinesisClient.describeStream(DescribeStreamRequest.builder().streamName(STREAM_NAME).build())
                   .streamDescription().streamStatusAsString();

           return kinesisStreamStatus.contentEquals("ACTIVE");
         } catch (Exception e) {
           LOGGER.warn("Could not fetch kinesis stream status", e);
           return null;
         }
       }
     }, 1000L, 30000, "Kinesis stream " + STREAM_NAME + " is not created or is not in active state", true);
   }

   public void stopKinesis() {
     if (_localstackDocker.isRunning()) {
       _localstackDocker.stop();
     }
   }

   private void publishRecordsToKinesis() {
     try {
       StringBuilder params = new StringBuilder("?");
       for (int i = 0; i < _h2FieldNameAndTypes.size() - 1; i++) {
         params.append(",?");
       }
       PreparedStatement h2Statement =
           _h2Connection.prepareStatement("INSERT INTO " + getTableName() + " VALUES (" + params.toString() + ")");

       InputStream inputStream =
           RealtimeKinesisIntegrationTest.class.getClassLoader().getResourceAsStream(DATA_FILE_PATH);

       try (BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
         String line;
         while ((line = br.readLine()) != null) {
           JsonNode data = JsonUtils.stringToJsonNode(line);

           PutRecordRequest putRecordRequest =
               PutRecordRequest.builder().streamName(STREAM_NAME).data(SdkBytes.fromUtf8String(line))
                   .partitionKey(data.get("Origin").textValue()).build();
           PutRecordResponse putRecordResponse = _kinesisClient.putRecord(putRecordRequest);
           if (putRecordResponse.sdkHttpResponse().statusCode() == 200) {
             if (StringUtils.isNotBlank(putRecordResponse.sequenceNumber()) && StringUtils.isNotBlank(
                 putRecordResponse.shardId())) {
               _totalRecordsPushedInStream++;

               int fieldIndex = 1;
               for (String fieldNameAndDatatype : _h2FieldNameAndTypes) {
                 String[] fieldNameAndDatatypeList = fieldNameAndDatatype.split(" ");
                 String fieldName = fieldNameAndDatatypeList[0];
                 String h2DataType = fieldNameAndDatatypeList[1];
                 switch (h2DataType) {
                   case "int": {
                     h2Statement.setObject(fieldIndex++, data.get(fieldName).intValue());
                     break;
                   }
                   case "varchar(128)": {
                     h2Statement.setObject(fieldIndex++, data.get(fieldName).textValue());
                     break;
                   }
                   default:
                     break;
                 }
               }
               h2Statement.execute();
             }
           }
         }
       }

       inputStream.close();
     } catch (Exception e) {
       throw new RuntimeException("Could not publish records to Kinesis Stream", e);
     }
   }

   private static AwsCredentialsProvider getLocalAWSCredentials() {
     return StaticCredentialsProvider.create(AwsBasicCredentials.create("access", "secret"));
   }

   @Test(enabled = false)
   public void testRecords()
       throws Exception {
     Assert.assertNotEquals(_totalRecordsPushedInStream, 0);

     ResultSet pinotResultSet =
         getPinotConnection().execute("SELECT * FROM " + getTableName() + " ORDER BY Origin LIMIT 10000")
             .getResultSet(0);

     Assert.assertNotEquals(pinotResultSet.getRowCount(), 0);

     Statement h2statement =
         _h2Connection.createStatement(java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY);
     h2statement.execute("SELECT * FROM " + getTableName() + " ORDER BY Origin");
     java.sql.ResultSet h2ResultSet = h2statement.getResultSet();

     Assert.assertFalse(h2ResultSet.isLast());

     h2ResultSet.beforeFirst();
     int row = 0;
     Map<String, Integer> columnToIndex = new HashMap<>();
     for (int i = 0; i < _h2FieldNameAndTypes.size(); i++) {
       columnToIndex.put(pinotResultSet.getColumnName(i), i);
     }

     while (h2ResultSet.next()) {

       for (String fieldNameAndDatatype : _h2FieldNameAndTypes) {
         String[] fieldNameAndDatatypeList = fieldNameAndDatatype.split(" ");
         String fieldName = fieldNameAndDatatypeList[0];
         String h2DataType = fieldNameAndDatatypeList[1];
         switch (h2DataType) {
           case "int": {
             int expectedValue = h2ResultSet.getInt(fieldName);
             int actualValue = pinotResultSet.getInt(row, columnToIndex.get(fieldName));
             Assert.assertEquals(expectedValue, actualValue);
             break;
           }
           case "varchar(128)": {
             String expectedValue = h2ResultSet.getString(fieldName);
             String actualValue = pinotResultSet.getString(row, columnToIndex.get(fieldName));
             Assert.assertEquals(expectedValue, actualValue);
             break;
           }
           default:
             break;
         }
       }

       row++;

       if (row >= pinotResultSet.getRowCount()) {
         int cnt = 0;
         while (h2ResultSet.next()) {
           cnt++;
         }
         Assert.assertEquals(cnt, 0);
         break;
       }
     }
   }

   @Test(enabled = false)
   public void testCountRecords() {
     long count = getPinotConnection().execute("SELECT COUNT(*) FROM " + getTableName()).getResultSet(0).getLong(0);
     Assert.assertEquals(count, _totalRecordsPushedInStream);
   }

   public void createH2ConnectionAndTable()
       throws Exception {
     Assert.assertNull(_h2Connection);
     Class.forName("org.h2.Driver");
     _h2Connection = DriverManager.getConnection("jdbc:h2:mem:");
     _h2Connection.prepareCall("DROP TABLE IF EXISTS " + getTableName()).execute();
     _h2FieldNameAndTypes = new ArrayList<>();

     InputStream inputStream = RealtimeKinesisIntegrationTest.class.getClassLoader().getResourceAsStream(DATA_FILE_PATH);

     String line;
     try (BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
       while ((line = br.readLine()) != null) {
         break;
       }
     } finally {
       inputStream.close();
     }

     if (StringUtils.isNotBlank(line)) {
       JsonNode dataObject = JsonUtils.stringToJsonNode(line);

       Iterator<Map.Entry<String, JsonNode>> fieldIterator = dataObject.fields();
       while (fieldIterator.hasNext()) {
         Map.Entry<String, JsonNode> field = fieldIterator.next();
         String fieldName = field.getKey();
         JsonNodeType fieldDataType = field.getValue().getNodeType();

         String h2DataType;

         switch (fieldDataType) {
           case NUMBER: {
             h2DataType = "int";
             break;
           }
           case STRING: {
             h2DataType = "varchar(128)";
             break;
           }
           case BOOLEAN: {
             h2DataType = "boolean";
             break;
           }
           default: {
             throw new UnsupportedDataTypeException(
                 "Kinesis Integration test doesn't support datatype: " + fieldDataType.name());
           }
         }

         _h2FieldNameAndTypes.add(fieldName + " " + h2DataType);
       }
     }

     _h2Connection.prepareCall("CREATE TABLE " + getTableName() + "(" + StringUtil.join(",",
         _h2FieldNameAndTypes.toArray(new String[_h2FieldNameAndTypes.size()])) + ")").execute();
   }

   @AfterClass(enabled = false)
   public void tearDown()
       throws Exception {
     if (_skipTestNoDockerInstalled) {
       return;
     }

     dropRealtimeTable(getTableName());
     stopServer();
     stopBroker();
     stopController();
     stopZk();
     stopKinesis();
     FileUtils.deleteDirectory(_tempDir);
   }

   public static class StopAllLocalstackDockerCommand extends Command {

     public void execute() {
       String runningDockerContainers =
           dockerExe.execute(Arrays.asList("ps", "-a", "-q", "-f", "ancestor=localstack/localstack"));
       if (StringUtils.isNotBlank(runningDockerContainers) && !runningDockerContainers.toLowerCase().contains("error")) {
         String[] containerList = runningDockerContainers.split("\n");

         for (String containerId : containerList) {
           dockerExe.execute(Arrays.asList("stop", containerId));
         }
       }
     }
   }

   public static class DockerInfoCommand extends Command {

     public void execute() {
       String dockerInfo = dockerExe.execute(Collections.singletonList("info"));

       if (dockerInfo.toLowerCase().contains("error")) {
         throw new IllegalStateException("Docker daemon is not running!");
       }
     }
   }
 }
	/**
	* 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.pinot.integration.tests;

	import cloud.localstack.Localstack;
	import cloud.localstack.ServiceName;
	import cloud.localstack.docker.annotation.LocalstackDockerAnnotationProcessor;
	import cloud.localstack.docker.annotation.LocalstackDockerConfiguration;
	import cloud.localstack.docker.annotation.LocalstackDockerProperties;
	import cloud.localstack.docker.command.Command;
	import com.fasterxml.jackson.databind.JsonNode;
	import com.fasterxml.jackson.databind.node.JsonNodeType;
	import com.google.common.base.Function;
	import java.io.BufferedReader;
	import java.io.File;
	import java.io.InputStream;
	import java.io.InputStreamReader;
	import java.net.URI;
	import java.net.URL;
	import java.nio.charset.StandardCharsets;
	import java.sql.DriverManager;
	import java.sql.PreparedStatement;
	import java.sql.Statement;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import javax.activation.UnsupportedDataTypeException;
	import javax.annotation.Nullable;
	import org.apache.commons.io.FileUtils;
	import org.apache.commons.lang3.StringUtils;
	import org.apache.pinot.client.ResultSet;
	import org.apache.pinot.plugin.stream.kinesis.KinesisConfig;
	import org.apache.pinot.plugin.stream.kinesis.KinesisConsumerFactory;
	import org.apache.pinot.spi.config.table.TableConfig;
	import org.apache.pinot.spi.config.table.TableType;
	import org.apache.pinot.spi.data.Schema;
	import org.apache.pinot.spi.stream.StreamConfig;
	import org.apache.pinot.spi.stream.StreamConfigProperties;
	import org.apache.pinot.spi.utils.JsonUtils;
	import org.apache.pinot.spi.utils.StringUtil;
	import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
	import org.apache.pinot.util.TestUtils;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;
	import org.testng.Assert;
	import org.testng.SkipException;
	import org.testng.annotations.AfterClass;
	import org.testng.annotations.BeforeClass;
	import org.testng.annotations.Test;
	import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
	import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
	import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
	import software.amazon.awssdk.core.SdkBytes;
	import software.amazon.awssdk.http.SdkHttpConfigurationOption;
	import software.amazon.awssdk.http.apache.ApacheSdkHttpService;
	import software.amazon.awssdk.regions.Region;
	import software.amazon.awssdk.services.kinesis.KinesisClient;
	import software.amazon.awssdk.services.kinesis.model.CreateStreamRequest;
	import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest;
	import software.amazon.awssdk.services.kinesis.model.PutRecordRequest;
	import software.amazon.awssdk.services.kinesis.model.PutRecordResponse;
	import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
	import software.amazon.awssdk.utils.AttributeMap;


	@LocalstackDockerProperties(services = {ServiceName.KINESIS}, imageTag = "0.12.15")
	public class RealtimeKinesisIntegrationTest extends BaseClusterIntegrationTestSet {
	private static final Logger LOGGER = LoggerFactory.getLogger(RealtimeKinesisIntegrationTest.class);

	private static final LocalstackDockerAnnotationProcessor PROCESSOR = new LocalstackDockerAnnotationProcessor();
	private static final String STREAM_NAME = "kinesis-test";
	private static final String STREAM_TYPE = "kinesis";
	public static final int MAX_RECORDS_TO_FETCH = Integer.MAX_VALUE;

	public static final String REGION = "us-east-1";
	public static final String LOCALSTACK_KINESIS_ENDPOINT = "http://localhost:4566";
	public static final int NUM_SHARDS = 10;

	// Localstack Kinesis doesn't support large rows.
	// So, this airlineStats data file consists of only few fields and rows from the original data
	public static final String SCHEMA_FILE_PATH = "kinesis/airlineStats_data_reduced.schema";
	public static final String DATA_FILE_PATH = "kinesis/airlineStats_data_reduced.json";

	private final Localstack _localstackDocker = Localstack.INSTANCE;

	private static KinesisClient _kinesisClient = null;

	private long _totalRecordsPushedInStream = 0;

	List<String> _h2FieldNameAndTypes = new ArrayList<>();

	private boolean _skipTestNoDockerInstalled = false;

	@BeforeClass(enabled = false)
	public void setUp()
	throws Exception {
	try {
	DockerInfoCommand dockerInfoCommand = new DockerInfoCommand();
	dockerInfoCommand.execute();
	} catch (IllegalStateException e) {
	_skipTestNoDockerInstalled = true;
	LOGGER.warn("Skipping test! Docker is not found running", e);
	throw new SkipException(e.getMessage());
	}

	TestUtils.ensureDirectoriesExistAndEmpty(_tempDir);

	// Start the Pinot cluster
	startZk();
	startController();
	startBroker();
	startServer();

	// Start Kinesis
	startKinesis();

	// Create and upload the schema and table config
	addSchema(createKinesisSchema());
	addTableConfig(createKinesisTableConfig());

	createH2ConnectionAndTable();

	// Push data into Kinesis
	publishRecordsToKinesis();

	// Wait for all documents loaded
	waitForAllDocsLoadedKinesis(120_000L);
	}

	public Schema createKinesisSchema()
	throws Exception {
	URL resourceUrl = BaseClusterIntegrationTest.class.getClassLoader().getResource(SCHEMA_FILE_PATH);
	Assert.assertNotNull(resourceUrl);
	return Schema.fromFile(new File(resourceUrl.getFile()));
	}

	protected void waitForAllDocsLoadedKinesis(long timeoutMs)
	throws Exception {
	waitForAllDocsLoadedKinesis(timeoutMs, true);
	}

	protected void waitForAllDocsLoadedKinesis(long timeoutMs, boolean raiseError) {
	TestUtils.waitForCondition(new Function<Void, Boolean>() {
	@Nullable
	@Override
	public Boolean apply(@Nullable Void aVoid) {
	try {
	return getCurrentCountStarResult() >= _totalRecordsPushedInStream;
	} catch (Exception e) {
	LOGGER.warn("Could not fetch current number of rows in pinot table " + getTableName(), e);
	return null;
	}
	}
	}, 1000L, timeoutMs, "Failed to load " + _totalRecordsPushedInStream + " documents", raiseError);
	}

	public TableConfig createKinesisTableConfig() {
	return new TableConfigBuilder(TableType.REALTIME).setTableName(getTableName()).setSchemaName(getTableName())
	.setTimeColumnName("DaysSinceEpoch").setFieldConfigList(getFieldConfigs()).setNumReplicas(getNumReplicas())
	.setSegmentVersion(getSegmentVersion()).setLoadMode(getLoadMode()).setTaskConfig(getTaskConfig())
	.setBrokerTenant(getBrokerTenant()).setServerTenant(getServerTenant()).setIngestionConfig(getIngestionConfig())
	.setLLC(true).setStreamConfigs(createKinesisStreamConfig()).setNullHandlingEnabled(getNullHandlingEnabled())
	.build();
	}

	public Map<String, String> createKinesisStreamConfig() {
	Map<String, String> streamConfigMap = new HashMap<>();
	String streamType = "kinesis";
	streamConfigMap.put(StreamConfigProperties.STREAM_TYPE, streamType);

	streamConfigMap.put(
	StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_TOPIC_NAME),
	STREAM_NAME);

	streamConfigMap.put(
	StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_FETCH_TIMEOUT_MILLIS),
	"30000");
	streamConfigMap.put(
	StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_CONSUMER_TYPES),
	StreamConfig.ConsumerType.LOWLEVEL.toString());
	streamConfigMap.put(StreamConfigProperties.constructStreamProperty(STREAM_TYPE,
	StreamConfigProperties.STREAM_CONSUMER_FACTORY_CLASS), KinesisConsumerFactory.class.getName());
	streamConfigMap.put(
	StreamConfigProperties.constructStreamProperty(STREAM_TYPE, StreamConfigProperties.STREAM_DECODER_CLASS),
	"org.apache.pinot.plugin.inputformat.json.JSONMessageDecoder");
	streamConfigMap.put(KinesisConfig.REGION, REGION);
	streamConfigMap.put(KinesisConfig.MAX_RECORDS_TO_FETCH, String.valueOf(MAX_RECORDS_TO_FETCH));
	streamConfigMap.put(KinesisConfig.SHARD_ITERATOR_TYPE, ShardIteratorType.AFTER_SEQUENCE_NUMBER.toString());
	streamConfigMap.put(KinesisConfig.ENDPOINT, LOCALSTACK_KINESIS_ENDPOINT);
	streamConfigMap.put(KinesisConfig.ACCESS_KEY, getLocalAWSCredentials().resolveCredentials().accessKeyId());
	streamConfigMap.put(KinesisConfig.SECRET_KEY, getLocalAWSCredentials().resolveCredentials().secretAccessKey());
	streamConfigMap.put(StreamConfigProperties.SEGMENT_FLUSH_THRESHOLD_ROWS, Integer.toString(200));
	streamConfigMap.put(StreamConfigProperties.constructStreamProperty(streamType,
	StreamConfigProperties.STREAM_CONSUMER_OFFSET_CRITERIA), "smallest");
	return streamConfigMap;
	}

	public void startKinesis()
	throws Exception {
	final LocalstackDockerConfiguration dockerConfig = PROCESSOR.process(this.getClass());
	StopAllLocalstackDockerCommand stopAllLocalstackDockerCommand = new StopAllLocalstackDockerCommand();
	stopAllLocalstackDockerCommand.execute();
	_localstackDocker.startup(dockerConfig);

	_kinesisClient = KinesisClient.builder().httpClient(new ApacheSdkHttpService().createHttpClientBuilder()
	.buildWithDefaults(
	AttributeMap.builder().put(SdkHttpConfigurationOption.TRUST_ALL_CERTIFICATES, Boolean.TRUE).build()))
	.credentialsProvider(getLocalAWSCredentials()).region(Region.of(REGION))
	.endpointOverride(new URI(LOCALSTACK_KINESIS_ENDPOINT)).build();

	_kinesisClient.createStream(CreateStreamRequest.builder().streamName(STREAM_NAME).shardCount(NUM_SHARDS).build());

	TestUtils.waitForCondition(new Function<Void, Boolean>() {
	@Nullable
	@Override
	public Boolean apply(@Nullable Void aVoid) {
	try {
	String kinesisStreamStatus =
	_kinesisClient.describeStream(DescribeStreamRequest.builder().streamName(STREAM_NAME).build())
	.streamDescription().streamStatusAsString();

	return kinesisStreamStatus.contentEquals("ACTIVE");
	} catch (Exception e) {
	LOGGER.warn("Could not fetch kinesis stream status", e);
	return null;
	}
	}
	}, 1000L, 30000, "Kinesis stream " + STREAM_NAME + " is not created or is not in active state", true);
	}

	public void stopKinesis() {
	if (_localstackDocker.isRunning()) {
	_localstackDocker.stop();
	}
	}

	private void publishRecordsToKinesis() {
	try {
	StringBuilder params = new StringBuilder("?");
	for (int i = 0; i < _h2FieldNameAndTypes.size() - 1; i++) {
	params.append(",?");
	}
	PreparedStatement h2Statement =
	_h2Connection.prepareStatement("INSERT INTO " + getTableName() + " VALUES (" + params.toString() + ")");

	InputStream inputStream =
	RealtimeKinesisIntegrationTest.class.getClassLoader().getResourceAsStream(DATA_FILE_PATH);

	try (BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
	String line;
	while ((line = br.readLine()) != null) {
	JsonNode data = JsonUtils.stringToJsonNode(line);

	PutRecordRequest putRecordRequest =
	PutRecordRequest.builder().streamName(STREAM_NAME).data(SdkBytes.fromUtf8String(line))
	.partitionKey(data.get("Origin").textValue()).build();
	PutRecordResponse putRecordResponse = _kinesisClient.putRecord(putRecordRequest);
	if (putRecordResponse.sdkHttpResponse().statusCode() == 200) {
	if (StringUtils.isNotBlank(putRecordResponse.sequenceNumber()) && StringUtils.isNotBlank(
	putRecordResponse.shardId())) {
	_totalRecordsPushedInStream++;

	int fieldIndex = 1;
	for (String fieldNameAndDatatype : _h2FieldNameAndTypes) {
	String[] fieldNameAndDatatypeList = fieldNameAndDatatype.split(" ");
	String fieldName = fieldNameAndDatatypeList[0];
	String h2DataType = fieldNameAndDatatypeList[1];
	switch (h2DataType) {
	case "int": {
	h2Statement.setObject(fieldIndex++, data.get(fieldName).intValue());
	break;
	}
	case "varchar(128)": {
	h2Statement.setObject(fieldIndex++, data.get(fieldName).textValue());
	break;
	}
	default:
	break;
	}
	}
	h2Statement.execute();
	}
	}
	}
	}

	inputStream.close();
	} catch (Exception e) {
	throw new RuntimeException("Could not publish records to Kinesis Stream", e);
	}
	}

	private static AwsCredentialsProvider getLocalAWSCredentials() {
	return StaticCredentialsProvider.create(AwsBasicCredentials.create("access", "secret"));
	}

	@Test(enabled = false)
	public void testRecords()
	throws Exception {
	Assert.assertNotEquals(_totalRecordsPushedInStream, 0);

	ResultSet pinotResultSet =
	getPinotConnection().execute("SELECT * FROM " + getTableName() + " ORDER BY Origin LIMIT 10000")
	.getResultSet(0);

	Assert.assertNotEquals(pinotResultSet.getRowCount(), 0);

	Statement h2statement =
	_h2Connection.createStatement(java.sql.ResultSet.TYPE_FORWARD_ONLY, java.sql.ResultSet.CONCUR_READ_ONLY);
	h2statement.execute("SELECT * FROM " + getTableName() + " ORDER BY Origin");
	java.sql.ResultSet h2ResultSet = h2statement.getResultSet();

	Assert.assertFalse(h2ResultSet.isLast());

	h2ResultSet.beforeFirst();
	int row = 0;
	Map<String, Integer> columnToIndex = new HashMap<>();
	for (int i = 0; i < _h2FieldNameAndTypes.size(); i++) {
	columnToIndex.put(pinotResultSet.getColumnName(i), i);
	}

	while (h2ResultSet.next()) {

	for (String fieldNameAndDatatype : _h2FieldNameAndTypes) {
	String[] fieldNameAndDatatypeList = fieldNameAndDatatype.split(" ");
	String fieldName = fieldNameAndDatatypeList[0];
	String h2DataType = fieldNameAndDatatypeList[1];
	switch (h2DataType) {
	case "int": {
	int expectedValue = h2ResultSet.getInt(fieldName);
	int actualValue = pinotResultSet.getInt(row, columnToIndex.get(fieldName));
	Assert.assertEquals(expectedValue, actualValue);
	break;
	}
	case "varchar(128)": {
	String expectedValue = h2ResultSet.getString(fieldName);
	String actualValue = pinotResultSet.getString(row, columnToIndex.get(fieldName));
	Assert.assertEquals(expectedValue, actualValue);
	break;
	}
	default:
	break;
	}
	}

	row++;

	if (row >= pinotResultSet.getRowCount()) {
	int cnt = 0;
	while (h2ResultSet.next()) {
	cnt++;
	}
	Assert.assertEquals(cnt, 0);
	break;
	}
	}
	}

	@Test(enabled = false)
	public void testCountRecords() {
	long count = getPinotConnection().execute("SELECT COUNT(*) FROM " + getTableName()).getResultSet(0).getLong(0);
	Assert.assertEquals(count, _totalRecordsPushedInStream);
	}

	public void createH2ConnectionAndTable()
	throws Exception {
	Assert.assertNull(_h2Connection);
	Class.forName("org.h2.Driver");
	_h2Connection = DriverManager.getConnection("jdbc:h2:mem:");
	_h2Connection.prepareCall("DROP TABLE IF EXISTS " + getTableName()).execute();
	_h2FieldNameAndTypes = new ArrayList<>();

	InputStream inputStream = RealtimeKinesisIntegrationTest.class.getClassLoader().getResourceAsStream(DATA_FILE_PATH);

	String line;
	try (BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
	while ((line = br.readLine()) != null) {
	break;
	}
	} finally {
	inputStream.close();
	}

	if (StringUtils.isNotBlank(line)) {
	JsonNode dataObject = JsonUtils.stringToJsonNode(line);

	Iterator<Map.Entry<String, JsonNode>> fieldIterator = dataObject.fields();
	while (fieldIterator.hasNext()) {
	Map.Entry<String, JsonNode> field = fieldIterator.next();
	String fieldName = field.getKey();
	JsonNodeType fieldDataType = field.getValue().getNodeType();

	String h2DataType;

	switch (fieldDataType) {
	case NUMBER: {
	h2DataType = "int";
	break;
	}
	case STRING: {
	h2DataType = "varchar(128)";
	break;
	}
	case BOOLEAN: {
	h2DataType = "boolean";
	break;
	}
	default: {
	throw new UnsupportedDataTypeException(
	"Kinesis Integration test doesn't support datatype: " + fieldDataType.name());
	}
	}

	_h2FieldNameAndTypes.add(fieldName + " " + h2DataType);
	}
	}

	_h2Connection.prepareCall("CREATE TABLE " + getTableName() + "(" + StringUtil.join(",",
	_h2FieldNameAndTypes.toArray(new String[_h2FieldNameAndTypes.size()])) + ")").execute();
	}

	@AfterClass(enabled = false)
	public void tearDown()
	throws Exception {
	if (_skipTestNoDockerInstalled) {
	return;
	}

	dropRealtimeTable(getTableName());
	stopServer();
	stopBroker();
	stopController();
	stopZk();
	stopKinesis();
	FileUtils.deleteDirectory(_tempDir);
	}

	public static class StopAllLocalstackDockerCommand extends Command {

	public void execute() {
	String runningDockerContainers =
	dockerExe.execute(Arrays.asList("ps", "-a", "-q", "-f", "ancestor=localstack/localstack"));
	if (StringUtils.isNotBlank(runningDockerContainers) && !runningDockerContainers.toLowerCase().contains("error")) {
	String[] containerList = runningDockerContainers.split("\n");

	for (String containerId : containerList) {
	dockerExe.execute(Arrays.asList("stop", containerId));
	}
	}
	}
	}

	public static class DockerInfoCommand extends Command {

	public void execute() {
	String dockerInfo = dockerExe.execute(Collections.singletonList("info"));

	if (dockerInfo.toLowerCase().contains("error")) {
	throw new IllegalStateException("Docker daemon is not running!");
	}
	}
	}
	}