/*
 * 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.beam.sdk.extensions.sql.meta.provider.pubsub;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.beam.sdk.testing.JsonMatcher.jsonBytesLike;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasEntry;
import static org.hamcrest.Matchers.hasProperty;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.generic.GenericRecordBuilder;
import org.apache.beam.sdk.coders.AvroCoder;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.extensions.protobuf.PayloadMessages;
import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv;
import org.apache.beam.sdk.extensions.sql.impl.JdbcConnection;
import org.apache.beam.sdk.extensions.sql.impl.JdbcDriver;
import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils;
import org.apache.beam.sdk.extensions.sql.meta.provider.SchemaIOTableProviderWrapper;
import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
import org.apache.beam.sdk.extensions.sql.meta.store.InMemoryMetaStore;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
import org.apache.beam.sdk.io.gcp.pubsub.TestPubsub;
import org.apache.beam.sdk.io.gcp.pubsub.TestPubsubSignal;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.SchemaCoder;
import org.apache.beam.sdk.schemas.utils.AvroUtils;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.util.common.ReflectHelpers;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.calcite.v1_26_0.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.calcite.v1_26_0.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.calcite.v1_26_0.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.calcite.v1_26_0.org.apache.calcite.jdbc.CalciteConnection;
import org.hamcrest.Matcher;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
import org.junit.runners.Parameterized.Parameter;
import org.junit.runners.Parameterized.Parameters;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@RunWith(Parameterized.class)
@SuppressWarnings({
  "keyfor",
})
public class PubsubTableProviderIT implements Serializable {

  private static final Logger LOG = LoggerFactory.getLogger(PubsubTableProviderIT.class);

  private static final Schema PAYLOAD_SCHEMA =
      Schema.builder()
          .addNullableField("id", Schema.FieldType.INT32)
          .addNullableField("name", Schema.FieldType.STRING)
          .build();

  @Rule public transient TestPubsub eventsTopic = TestPubsub.create();
  @Rule public transient TestPubsub filteredEventsTopic = TestPubsub.create();
  @Rule public transient TestPubsub dlqTopic = TestPubsub.create();
  @Rule public transient TestPubsubSignal resultSignal = TestPubsubSignal.create();
  @Rule public transient TestPipeline pipeline = TestPipeline.create();
  @Rule public transient TestPipeline filterPipeline = TestPipeline.create();
  private final SchemaIOTableProviderWrapper tableProvider = new PubsubTableProvider();

  /** How long to wait on the result signal. */
  private final Duration timeout = Duration.standardMinutes(10);

  @Parameters
  public static Collection<Object[]> data() {
    return Arrays.asList(
        new Object[][] {
          {new PubsubJsonObjectProvider()},
          {new PubsubAvroObjectProvider()},
          {new PubsubProtoObjectProvider()}
        });
  }

  @Parameter public PubsubObjectProvider objectsProvider;

  /**
   * HACK: we need an objectmapper to turn pipelineoptions back into a map. We need to use
   * ReflectHelpers to get the extra PipelineOptions.
   */
  private static final ObjectMapper MAPPER =
      new ObjectMapper()
          .registerModules(ObjectMapper.findModules(ReflectHelpers.findClassLoader()));

  @Test
  public void testSQLSelectsPayloadContent() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "attributes MAP<VARCHAR, VARCHAR>, \n"
                + "payload ROW< \n"
                + "             id INTEGER, \n"
                + "             name VARCHAR \n"
                + "           > \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES '{ "
                + "%s"
                + "\"protoClass\" : \"%s\", "
                + "\"timestampAttributeKey\" : \"ts\" }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            PayloadMessages.SimpleMessage.class.getName());

    String queryString = "SELECT message.payload.id, message.payload.name from message";

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    // Apply the PTransform to query the pubsub topic
    PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);

    // Observe the query results and send success signal after seeing the expected messages
    queryOutput.apply(
        "waitForSuccess",
        resultSignal.signalSuccessWhen(
            SchemaCoder.of(PAYLOAD_SCHEMA),
            observedRows ->
                observedRows.equals(
                    ImmutableSet.of(
                        row(PAYLOAD_SCHEMA, 3, "foo"),
                        row(PAYLOAD_SCHEMA, 5, "bar"),
                        row(PAYLOAD_SCHEMA, 7, "baz")))));

    // Start the pipeline
    pipeline.run();

    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    // Start publishing the messages when main pipeline is started and signaling topic is ready
    eventsTopic.publish(
        ImmutableList.of(
            objectsProvider.messageIdName(ts(1), 3, "foo"),
            objectsProvider.messageIdName(ts(2), 5, "bar"),
            objectsProvider.messageIdName(ts(3), 7, "baz")));

    // Poll the signaling topic for success message
    resultSignal.waitForSuccess(timeout);
  }

  @Ignore("https://issues.apache.org/jira/browse/BEAM-12320")
  @Test
  public void testSQLSelectsArrayAttributes() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "attributes ARRAY<ROW<key VARCHAR, `value` VARCHAR>>, \n"
                + "payload ROW< \n"
                + "             id INTEGER, \n"
                + "             name VARCHAR \n"
                + "           > \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES '{ "
                + "%s"
                + "\"protoClass\" : \"%s\", "
                + "\"timestampAttributeKey\" : \"ts\" }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            PayloadMessages.SimpleMessage.class.getName());

    String queryString =
        "SELECT message.payload.id, attributes[1].key AS a1, attributes[2].key AS a2 FROM message";

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    // Apply the PTransform to query the pubsub topic
    PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);

    // Observe the query results and send success signal after seeing the expected messages
    queryOutput.apply(
        "waitForSuccess",
        resultSignal.signalSuccessWhen(
            SchemaCoder.of(PAYLOAD_SCHEMA),
            observedRows -> {
              Map<Integer, String> entries = new HashMap<>();
              for (Row row : observedRows) {
                if ("ts".equals(row.getString("a1"))) {
                  entries.put(row.getInt32("id"), row.getString("a2"));
                } else {
                  entries.put(row.getInt32("id"), row.getString("a1"));
                }
              }

              LOG.info("Entries: {}", entries);

              return entries.equals(ImmutableMap.of(3, "foo", 5, "bar", 7, "baz"));
            }));

    // Start the pipeline
    pipeline.run();

    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    // Start publishing the messages when main pipeline is started and signaling topic is ready
    eventsTopic.publish(
        ImmutableList.of(
            objectsProvider.messageIdName(ts(1), 3, "foo"),
            objectsProvider.messageIdName(ts(2), 5, "bar"),
            objectsProvider.messageIdName(ts(3), 7, "baz")));

    // Poll the signaling topic for success message
    resultSignal.waitForSuccess(timeout);
  }

  @Test
  public void testSQLWithBytePayload() throws Exception {

    // Prepare messages to send later
    List<PubsubMessage> messages =
        ImmutableList.of(
            objectsProvider.messageIdName(ts(1), 3, "foo"),
            objectsProvider.messageIdName(ts(2), 5, "bar"),
            objectsProvider.messageIdName(ts(3), 7, "baz"));

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "attributes MAP<VARCHAR, VARCHAR>, \n"
                + "payload VARBINARY \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES '{ "
                + "\"protoClass\" : \"%s\", "
                + "\"timestampAttributeKey\" : \"ts\" }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            PayloadMessages.SimpleMessage.class.getName());

    String queryString = "SELECT message.payload AS some_bytes FROM message";

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    // Apply the PTransform to query the pubsub topic
    PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);

    // Observe the query results and send success signal after seeing the expected messages
    Schema justBytesSchema =
        Schema.builder().addField("some_bytes", FieldType.BYTES.withNullable(true)).build();
    Row expectedRow0 = row(justBytesSchema, (Object) messages.get(0).getPayload());
    Row expectedRow1 = row(justBytesSchema, (Object) messages.get(1).getPayload());
    Row expectedRow2 = row(justBytesSchema, (Object) messages.get(2).getPayload());
    Set<Row> expected = ImmutableSet.of(expectedRow0, expectedRow1, expectedRow2);
    queryOutput.apply(
        "waitForSuccess",
        resultSignal.signalSuccessWhen(
            SchemaCoder.of(justBytesSchema), observedRows -> observedRows.equals(expected)));

    // Start the pipeline
    pipeline.run();

    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    // Start publishing the messages when main pipeline is started and signaling topic is ready
    eventsTopic.publish(messages);

    // Poll the signaling topic for success message
    resultSignal.waitForSuccess(timeout);
  }

  @Test
  @SuppressWarnings("unchecked")
  public void testUsesDlq() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "attributes MAP<VARCHAR, VARCHAR>, \n"
                + "payload ROW< \n"
                + "             id INTEGER, \n"
                + "             name VARCHAR \n"
                + "           > \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES "
                + "    '{ "
                + "       %s"
                + "       \"timestampAttributeKey\" : \"ts\", "
                + "       \"deadLetterQueue\" : \"%s\", "
                + "       \"protoClass\" : \"%s\" "
                + "     }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            dlqTopic.topicPath(),
            PayloadMessages.SimpleMessage.class.getName());

    String queryString = "SELECT message.payload.id, message.payload.name from message";

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    // Apply the PTransform to query the pubsub topic
    PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);

    // Observe the query results and send success signal after seeing the expected messages
    queryOutput.apply(
        "waitForSuccess",
        resultSignal.signalSuccessWhen(
            SchemaCoder.of(PAYLOAD_SCHEMA),
            observedRows ->
                observedRows.equals(
                    ImmutableSet.of(
                        row(PAYLOAD_SCHEMA, 3, "foo"),
                        row(PAYLOAD_SCHEMA, 5, "bar"),
                        row(PAYLOAD_SCHEMA, 7, "baz")))));

    // Start the pipeline
    pipeline.run();

    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    // Start publishing the messages when main pipeline is started and signaling topics are ready
    eventsTopic.publish(
        ImmutableList.of(
            objectsProvider.messageIdName(ts(1), 3, "foo"),
            objectsProvider.messageIdName(ts(2), 5, "bar"),
            objectsProvider.messageIdName(ts(3), 7, "baz"),
            messagePayload(ts(4), "{ - }", ImmutableMap.of()), // invalid message, will go to DLQ
            messagePayload(ts(5), "{ + }", ImmutableMap.of()))); // invalid message, will go to DLQ

    // Poll the signaling topic for success message
    resultSignal.waitForSuccess(timeout);
    dlqTopic
        .assertThatTopicEventuallyReceives(
            matcherPayload(ts(4), "{ - }"), matcherPayload(ts(5), "{ + }"))
        .waitForUpTo(Duration.standardSeconds(40));
  }

  @Test
  @SuppressWarnings({"unchecked", "rawtypes"})
  public void testSQLLimit() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "attributes MAP<VARCHAR, VARCHAR>, \n"
                + "payload ROW< \n"
                + "             id INTEGER, \n"
                + "             name VARCHAR \n"
                + "           > \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES "
                + "    '{ "
                + "       %s"
                + "       \"timestampAttributeKey\" : \"ts\", "
                + "       \"deadLetterQueue\" : \"%s\", "
                + "       \"protoClass\" : \"%s\" "
                + "     }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            dlqTopic.topicPath(),
            PayloadMessages.SimpleMessage.class.getName());

    List<PubsubMessage> messages =
        ImmutableList.of(
            objectsProvider.messageIdName(ts(1), 3, "foo"),
            objectsProvider.messageIdName(ts(2), 5, "bar"),
            objectsProvider.messageIdName(ts(3), 7, "baz"),
            objectsProvider.messageIdName(ts(4), 9, "ba2"),
            objectsProvider.messageIdName(ts(5), 10, "ba3"),
            objectsProvider.messageIdName(ts(6), 13, "ba4"),
            objectsProvider.messageIdName(ts(7), 15, "ba5"));

    // We need the default options on the schema to include the project passed in for the
    // integration test
    CalciteConnection connection = connect(pipeline.getOptions(), new PubsubTableProvider());

    Statement statement = connection.createStatement();
    statement.execute(createTableString);

    // Because Pubsub only allow new subscription receives message after the subscription is
    // created, eventsTopic.publish(messages) can only be called after statement.executeQuery.
    // However, because statement.executeQuery is a blocking call, it has to be put into a
    // seperate thread to execute.
    ExecutorService pool = Executors.newFixedThreadPool(1);
    Future<List<String>> queryResult =
        pool.submit(
            (Callable)
                () -> {
                  ResultSet resultSet =
                      statement.executeQuery("SELECT message.payload.id FROM message LIMIT 3");
                  ImmutableList.Builder<String> result = ImmutableList.builder();
                  while (resultSet.next()) {
                    result.add(resultSet.getString(1));
                  }
                  return result.build();
                });

    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    eventsTopic.publish(messages);

    assertThat(queryResult.get(2, TimeUnit.MINUTES).size(), equalTo(3));
    pool.shutdown();
  }

  @Test
  public void testSQLSelectsPayloadContentFlat() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "id INTEGER, \n"
                + "name VARCHAR \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES "
                + "    '{ "
                + "       %s"
                + "       \"protoClass\" : \"%s\", "
                + "       \"timestampAttributeKey\" : \"ts\" "
                + "     }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            PayloadMessages.SimpleMessage.class.getName());

    String queryString = "SELECT message.id, message.name from message";

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    // Apply the PTransform to query the pubsub topic
    PCollection<Row> queryOutput = query(sqlEnv, pipeline, queryString);

    // Observe the query results and send success signal after seeing the expected messages
    queryOutput.apply(
        "waitForSuccess",
        resultSignal.signalSuccessWhen(
            SchemaCoder.of(PAYLOAD_SCHEMA),
            observedRows ->
                observedRows.equals(
                    ImmutableSet.of(
                        row(PAYLOAD_SCHEMA, 3, "foo"),
                        row(PAYLOAD_SCHEMA, 5, "bar"),
                        row(PAYLOAD_SCHEMA, 7, "baz")))));

    // Start the pipeline
    pipeline.run();

    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    // Start publishing the messages when main pipeline is started and signaling topic is ready
    eventsTopic.publish(
        ImmutableList.of(
            objectsProvider.messageIdName(ts(1), 3, "foo"),
            objectsProvider.messageIdName(ts(2), 5, "bar"),
            objectsProvider.messageIdName(ts(3), 7, "baz")));

    // Poll the signaling topic for success message
    resultSignal.waitForSuccess(timeout);
  }

  @Test
  @SuppressWarnings("unchecked")
  public void testSQLInsertRowsToPubsubFlat() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "name VARCHAR, \n"
                + "height INTEGER, \n"
                + "knows_javascript BOOLEAN \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES "
                + "    '{ "
                + "       %s"
                + "       \"protoClass\" : \"%s\", "
                + "       \"deadLetterQueue\" : \"%s\""
                + "     }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            PayloadMessages.NameHeightKnowsJSMessage.class.getName(),
            dlqTopic.topicPath());

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    // TODO(BEAM-8741): Ideally we could write this query without specifying a column list, because
    //   it shouldn't be possible to write to event_timestamp when it's mapped to  publish time.
    String queryString =
        "INSERT INTO message (name, height, knows_javascript) \n"
            + "VALUES \n"
            + "('person1', 80, TRUE), \n"
            + "('person2', 70, FALSE)";

    // Apply the PTransform to insert the rows
    query(sqlEnv, pipeline, queryString);

    pipeline.run().waitUntilFinish(Duration.standardMinutes(5));

    eventsTopic
        .assertThatTopicEventuallyReceives(
            objectsProvider.matcherNameHeightKnowsJS("person1", 80, true),
            objectsProvider.matcherNameHeightKnowsJS("person2", 70, false))
        .waitForUpTo(Duration.standardSeconds(40));
  }

  @Test
  @SuppressWarnings("unchecked")
  public void testSQLInsertRowsToPubsubWithTimestampAttributeFlat() throws Exception {

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE message (\n"
                + "  event_timestamp TIMESTAMP, \n"
                + "  name VARCHAR, \n"
                + "  height INTEGER, \n"
                + "  knows_javascript BOOLEAN \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "TBLPROPERTIES "
                + "  '{ "
                + "     %s "
                + "     \"protoClass\" : \"%s\", "
                + "     \"deadLetterQueue\" : \"%s\","
                + "     \"timestampAttributeKey\" : \"ts\""
                + "   }'",
            tableProvider.getTableType(),
            eventsTopic.topicPath(),
            payloadFormatParam(),
            PayloadMessages.NameHeightKnowsJSMessage.class.getName(),
            dlqTopic.topicPath());

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);

    String queryString =
        "INSERT INTO message "
            + "VALUES "
            + "(TIMESTAMP '1970-01-01 00:00:00.002', 'person2', 70, FALSE)";
    query(sqlEnv, pipeline, queryString);

    pipeline.run().waitUntilFinish(Duration.standardMinutes(5));

    eventsTopic
        .assertThatTopicEventuallyReceives(matcherTsNameHeightKnowsJS(ts(2), "person2", 70, false))
        .waitForUpTo(Duration.standardSeconds(40));
  }

  @Test
  @SuppressWarnings("unchecked")
  public void testSQLReadAndWriteWithSameFlatTableDefinition() throws Exception {
    // This test verifies that the same pubsub table definition can be used for both reading and
    // writing
    // pipeline: Use SQL to insert data into `people`
    // filterPipeline: Use SQL to read from `people`, filter the rows, and write to
    // `javascript_people`

    String tblProperties =
        objectsProvider.getPayloadFormat() == null
            ? ""
            : String.format(
                "TBLPROPERTIES '{ \"protoClass\" : \"%s\", \"format\": \"%s\" }'",
                PayloadMessages.NameHeightKnowsJSMessage.class.getName(),
                objectsProvider.getPayloadFormat());

    String createTableString =
        String.format(
            "CREATE EXTERNAL TABLE people (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "name VARCHAR, \n"
                + "height INTEGER, \n"
                + "knows_javascript BOOLEAN \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "%s",
            tableProvider.getTableType(), eventsTopic.topicPath(), tblProperties);

    String filteredTblProperties =
        objectsProvider.getPayloadFormat() == null
            ? ""
            : String.format(
                "TBLPROPERTIES '{ \"protoClass\" : \"%s\", \"format\": \"%s\" }'",
                PayloadMessages.NameHeightMessage.class.getName(),
                objectsProvider.getPayloadFormat());

    String createFilteredTableString =
        String.format(
            "CREATE EXTERNAL TABLE javascript_people (\n"
                + "event_timestamp TIMESTAMP, \n"
                + "name VARCHAR, \n"
                + "height INTEGER \n"
                + ") \n"
                + "TYPE '%s' \n"
                + "LOCATION '%s' \n"
                + "%s",
            tableProvider.getTableType(), filteredEventsTopic.topicPath(), filteredTblProperties);

    // Initialize SQL environment and create the pubsub table
    BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new PubsubTableProvider());
    sqlEnv.executeDdl(createTableString);
    sqlEnv.executeDdl(createFilteredTableString);

    // TODO(BEAM-8741): Ideally we could write these queries without specifying a column list,
    // because
    //   it shouldn't be possible to write to event_timestamp when it's mapped to  publish time.
    String filterQueryString =
        "INSERT INTO javascript_people (name, height) (\n"
            + "  SELECT \n"
            + "    name, \n"
            + "    height \n"
            + "  FROM people \n"
            + "  WHERE knows_javascript \n"
            + ")";

    String injectQueryString =
        "INSERT INTO people (name, height, knows_javascript) VALUES \n"
            + "('person1', 80, TRUE),  \n"
            + "('person2', 70, FALSE), \n"
            + "('person3', 60, TRUE),  \n"
            + "('person4', 50, FALSE), \n"
            + "('person5', 40, TRUE)";

    // Apply the PTransform to do the filtering
    query(sqlEnv, filterPipeline, filterQueryString);

    // Apply the PTransform to inject the input data
    query(sqlEnv, pipeline, injectQueryString);

    // Start the filter pipeline and wait until it has started.
    filterPipeline.run();

    // Block until a subscription for this topic exists
    eventsTopic.assertSubscriptionEventuallyCreated(
        pipeline.getOptions().as(GcpOptions.class).getProject(), Duration.standardMinutes(5));

    // .. then run the injector pipeline
    pipeline.run().waitUntilFinish(Duration.standardMinutes(5));

    filteredEventsTopic
        .assertThatTopicEventuallyReceives(
            objectsProvider.matcherNameHeight("person1", 80),
            objectsProvider.matcherNameHeight("person3", 60),
            objectsProvider.matcherNameHeight("person5", 40))
        .waitForUpTo(Duration.standardMinutes(5));
  }

  @SuppressWarnings("unchecked")
  private CalciteConnection connect(PipelineOptions options, TableProvider... tableProviders) {
    // HACK: PipelineOptions should expose a prominent method to do this reliably
    // The actual options are in the "options" field of the converted map
    Map<String, String> argsMap =
        ((Map<String, Object>) MAPPER.convertValue(pipeline.getOptions(), Map.class).get("options"))
            .entrySet().stream()
                .filter(
                    (entry) -> {
                      if (entry.getValue() instanceof List) {
                        if (!((List) entry.getValue()).isEmpty()) {
                          throw new IllegalArgumentException("Cannot encode list arguments");
                        }
                        // We can encode empty lists, just omit them.
                        return false;
                      }
                      return true;
                    })
                .collect(Collectors.toMap(Map.Entry::getKey, entry -> toArg(entry.getValue())));

    InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
    for (TableProvider tableProvider : tableProviders) {
      inMemoryMetaStore.registerProvider(tableProvider);
    }

    JdbcConnection connection = JdbcDriver.connect(inMemoryMetaStore, options);
    connection.setPipelineOptionsMap(argsMap);
    return connection;
  }

  private static String toArg(Object o) {
    try {
      String jsonRepr = MAPPER.writeValueAsString(o);

      // String and enums are expected to be unquoted on the command line
      if (jsonRepr.startsWith("\"") && jsonRepr.endsWith("\"")) {
        return jsonRepr.substring(1, jsonRepr.length() - 1);
      } else {
        return jsonRepr;
      }
    } catch (JsonProcessingException e) {
      throw new RuntimeException(e);
    }
  }

  private String payloadFormatParam() {
    return objectsProvider.getPayloadFormat() == null
        ? ""
        : String.format("\"format\" : \"%s\", ", objectsProvider.getPayloadFormat());
  }

  private PCollection<Row> query(BeamSqlEnv sqlEnv, TestPipeline pipeline, String queryString) {

    return BeamSqlRelUtils.toPCollection(pipeline, sqlEnv.parseQuery(queryString));
  }

  private Row row(Schema schema, Object... values) {
    return Row.withSchema(schema).addValues(values).build();
  }

  private static PubsubMessage message(
      Instant timestamp, byte[] payload, Map<String, String> attributes) {
    return new PubsubMessage(
        payload,
        ImmutableMap.<String, String>builder()
            .putAll(attributes)
            .put("ts", String.valueOf(timestamp.getMillis()))
            .build());
  }

  private Matcher<PubsubMessage> matcherTsNameHeightKnowsJS(
      Instant ts, String name, int height, boolean knowsJS) throws Exception {
    return allOf(
        objectsProvider.matcherNameHeightKnowsJS(name, height, knowsJS),
        hasProperty("attributeMap", hasEntry("ts", String.valueOf(ts.getMillis()))));
  }

  private Matcher<PubsubMessage> matcherPayload(Instant timestamp, String payload) {
    return allOf(
        hasProperty("payload", equalTo(payload.getBytes(StandardCharsets.US_ASCII))),
        hasProperty("attributeMap", hasEntry("ts", String.valueOf(timestamp.getMillis()))));
  }

  private Instant ts(long millis) {
    return Instant.ofEpochMilli(millis);
  }

  private PubsubMessage messagePayload(
      Instant timestamp, String payload, Map<String, String> attributes) {
    return message(timestamp, payload.getBytes(StandardCharsets.US_ASCII), attributes);
  }

  private abstract static class PubsubObjectProvider implements Serializable {
    protected abstract String getPayloadFormat();

    protected abstract PubsubMessage messageIdName(Instant timestamp, int id, String name)
        throws Exception;

    protected abstract Matcher<PubsubMessage> matcherNames(String name) throws Exception;

    protected abstract Matcher<PubsubMessage> matcherNameHeightKnowsJS(
        String name, int height, boolean knowsJS) throws Exception;

    protected abstract Matcher<PubsubMessage> matcherNameHeight(String name, int height)
        throws Exception;
  }

  private static class PubsubProtoObjectProvider extends PubsubObjectProvider {

    @Override
    protected String getPayloadFormat() {
      return "proto";
    }

    @Override
    protected PubsubMessage messageIdName(Instant timestamp, int id, String name) {

      PayloadMessages.SimpleMessage.Builder simpleMessage =
          PayloadMessages.SimpleMessage.newBuilder().setId(id).setName(name);

      return PubsubTableProviderIT.message(
          timestamp,
          simpleMessage.build().toByteArray(),
          ImmutableMap.of(name, Integer.toString(id)));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNames(String name) throws IOException {

      PayloadMessages.NameMessage.Builder nameMessage =
          PayloadMessages.NameMessage.newBuilder().setName(name);

      return hasProperty("payload", equalTo(nameMessage.build().toByteArray()));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNameHeightKnowsJS(
        String name, int height, boolean knowsJS) throws IOException {

      PayloadMessages.NameHeightKnowsJSMessage.Builder nameHeightKnowsJSMessage =
          PayloadMessages.NameHeightKnowsJSMessage.newBuilder()
              .setHeight(height)
              .setName(name)
              .setKnowsJavascript(knowsJS);

      return hasProperty("payload", equalTo(nameHeightKnowsJSMessage.build().toByteArray()));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNameHeight(String name, int height) throws IOException {

      PayloadMessages.NameHeightMessage.Builder nameHeightMessage =
          PayloadMessages.NameHeightMessage.newBuilder().setName(name).setHeight(height);

      return hasProperty("payload", equalTo(nameHeightMessage.build().toByteArray()));
    }
  }

  private static class PubsubJsonObjectProvider extends PubsubObjectProvider {

    // Pubsub table provider should default to json
    @Override
    protected String getPayloadFormat() {
      return null;
    }

    @Override
    protected PubsubMessage messageIdName(Instant timestamp, int id, String name) {
      String jsonString = "{ \"id\" : " + id + ", \"name\" : \"" + name + "\" }";
      return message(timestamp, jsonString, ImmutableMap.of(name, Integer.toString(id)));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNames(String name) throws IOException {
      return hasProperty("payload", toJsonByteLike(String.format("{\"name\":\"%s\"}", name)));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNameHeightKnowsJS(
        String name, int height, boolean knowsJS) throws IOException {
      String jsonString =
          String.format(
              "{\"name\":\"%s\", \"height\": %s, \"knows_javascript\": %s}", name, height, knowsJS);

      return hasProperty("payload", toJsonByteLike(jsonString));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNameHeight(String name, int height) throws IOException {
      String jsonString = String.format("{\"name\":\"%s\", \"height\": %s}", name, height);
      return hasProperty("payload", toJsonByteLike(jsonString));
    }

    private PubsubMessage message(
        Instant timestamp, String jsonPayload, Map<String, String> attributes) {
      return PubsubTableProviderIT.message(timestamp, jsonPayload.getBytes(UTF_8), attributes);
    }

    private Matcher<byte[]> toJsonByteLike(String jsonString) throws IOException {
      return jsonBytesLike(jsonString);
    }
  }

  private static class PubsubAvroObjectProvider extends PubsubObjectProvider {
    private static final Schema NAME_HEIGHT_KNOWS_JS_SCHEMA =
        Schema.builder()
            .addNullableField("name", Schema.FieldType.STRING)
            .addNullableField("height", Schema.FieldType.INT32)
            .addNullableField("knows_javascript", Schema.FieldType.BOOLEAN)
            .build();

    private static final Schema NAME_HEIGHT_SCHEMA =
        Schema.builder()
            .addNullableField("name", Schema.FieldType.STRING)
            .addNullableField("height", Schema.FieldType.INT32)
            .build();

    @Override
    protected String getPayloadFormat() {
      return "avro";
    }

    @Override
    protected PubsubMessage messageIdName(Instant timestamp, int id, String name)
        throws IOException {
      byte[] encodedRecord =
          createEncodedGenericRecord(
              PAYLOAD_SCHEMA,
              org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.of(
                  id, name));
      return message(timestamp, encodedRecord, ImmutableMap.of(name, Integer.toString(id)));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNames(String name) throws IOException {
      Schema schema = Schema.builder().addStringField("name").build();
      byte[] encodedRecord =
          createEncodedGenericRecord(
              schema,
              org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.of(
                  name));
      return hasProperty("payload", equalTo(encodedRecord));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNameHeight(String name, int height) throws IOException {
      byte[] encodedRecord =
          createEncodedGenericRecord(
              NAME_HEIGHT_SCHEMA,
              org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.of(
                  name, height));
      return hasProperty("payload", equalTo(encodedRecord));
    }

    @Override
    protected Matcher<PubsubMessage> matcherNameHeightKnowsJS(
        String name, int height, boolean knowsJS) throws IOException {
      byte[] encodedRecord =
          createEncodedGenericRecord(
              NAME_HEIGHT_KNOWS_JS_SCHEMA,
              org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.of(
                  name, height, knowsJS));
      return hasProperty("payload", equalTo(encodedRecord));
    }

    private byte[] createEncodedGenericRecord(Schema beamSchema, List<Object> values)
        throws IOException {
      org.apache.avro.Schema avroSchema = AvroUtils.toAvroSchema(beamSchema);
      GenericRecordBuilder builder = new GenericRecordBuilder(avroSchema);
      List<org.apache.avro.Schema.Field> fields = avroSchema.getFields();
      for (int i = 0; i < fields.size(); ++i) {
        builder.set(fields.get(i), values.get(i));
      }
      AvroCoder<GenericRecord> coder = AvroCoder.of(avroSchema);
      ByteArrayOutputStream out = new ByteArrayOutputStream();

      coder.encode(builder.build(), out);
      return out.toByteArray();
    }
  }
}