public static <T extends Message> ImmutableList<T>
parseAsList(Collection<String> jsons, Class<T> protoClass) {
  return jsons
      .stream()
      .map(json -> ProtoJsonConverter.parseJson(json, protoClass))
      .filter(Objects::nonNull)
      .collect(toImmutableList());
} @SuppressWarnings("unchecked")
public static <T extends Message> T parseJson(String json, Class<T> protoClass) {
  try {
    Message.Builder builder =
        (Message.Builder)
            protoClass.getMethod("newBuilder")
                .invoke(null);

    JsonFormat.parser().merge(json, builder);
    return (T) builder.build();
  } catch (Exception protoException) {
    logger.atSevere().withCause(protoException).atMostEvery(1, TimeUnit.MINUTES)
        .log("error converting json:\n%s", json);
    return null;
  }
} public abstract class BigQueryTableCreator {

  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

  /**
   * Matches the given String as a Legacy or Standard Table name. If no match found, it returns a
   *
   * @param bigQueryTableName the table name in legacy or standard SQL format.
   * @return Table information with valid values or if no match found then tableName set as the
   * input String or empty table name.
   */
  public static BigQueryTableEntity usingBestEffort(String bigQueryTableName) {
    if (bigQueryTableName != null && bigQueryTableName.startsWith("$")) {
      return BigQueryTableEntity.create(null, null, bigQueryTableName);
    }

    for (String pattern : ImmutableList
        .of(BQ_LEGACY_STANDARD_TABLE_NAME_FORMAT, BQ_RESOURCE_FORMAT, BQ_LINKED_RESOURCE_FORMAT)) {
      try {
        return extractInformation(pattern, bigQueryTableName);
      } catch (IllegalArgumentException aex) {
        logger.atInfo().atMostEvery(1, TimeUnit.MINUTES).withCause(aex)
            .log("error parsing %s", bigQueryTableName);
      }
    }

    throw new IllegalArgumentException(
        "Couldn't convert into any known types: (" + bigQueryTableName + ")");
  }

  /**
   * Returns a parsed TableEntity from the legacy SQL form (<project-id>:<dataset-id>.<table-id>) of
   * a BigQuery table.
   */
  public static BigQueryTableEntity fromLegacyTableName(String legacyName) {
    return extractInformation(LEGACY_TABLE_FORMAT, legacyName);
  }

  public static BigQueryTableEntity fromSqlResource(String sqlResource) {
    return extractInformation(SQL_RESOURCE_FORMAT, sqlResource);
  }

  public static BigQueryTableEntity fromBigQueryResource(String resource) {
    return extractInformation(BQ_RESOURCE_FORMAT, resource);
  }

  public static BigQueryTableEntity fromLinkedResource(String linkedResource) {
    return extractInformation(BQ_LINKED_RESOURCE_FORMAT, linkedResource);
  }

  private static final String PROJECT_ID_TAG = "projectId";
  private static final String DATASET_ID_TAG = "dataset";
  private static final String TABLE_ID_TAG = "table";

  private static final String PROJECT_PATTERN = "[a-zA-Z0-9\\.\\-\\:]+";
  private static final String DATASET_PATTERN = "[a-zA-Z_][a-zA-Z0-9\\_]+";
  private static final String TABLE_PATTERN = "[a-zA-Z][a-zA-Z0-9\\_]+";


  private static final String LEGACY_TABLE_FORMAT =
      String.format(
          "^(?<%s>%s)\\:(?<%s>%s)\\.(?<%s>%s)$",
          PROJECT_ID_TAG, PROJECT_PATTERN, DATASET_ID_TAG, DATASET_PATTERN, TABLE_ID_TAG,
          TABLE_PATTERN);

  private static final String SQL_RESOURCE_FORMAT =
      String.format(
          "^bigquery\\.table\\.(?<%s>%s)\\.(?<%s>%s)\\.(?<%s>%s)$",
          PROJECT_ID_TAG, PROJECT_PATTERN, DATASET_ID_TAG, DATASET_PATTERN, TABLE_ID_TAG,
          TABLE_PATTERN);

  private static final String BQ_RESOURCE_FORMAT =
      String.format(
          "^projects/(?<%s>%s)/datasets/(?<%s>%s)/tables/(?<%s>%s)(\\$[0-9]+)?$",
          PROJECT_ID_TAG, PROJECT_PATTERN, DATASET_ID_TAG, DATASET_PATTERN, TABLE_ID_TAG,
          TABLE_PATTERN);

  private static final String BQ_LINKED_RESOURCE_FORMAT =
      String.format(
          "^//bigquery.googleapis.com/projects/(?<%s>%s)/datasets/(?<%s>%s)/tables/(?<%s>%s)$",
          PROJECT_ID_TAG, PROJECT_PATTERN, DATASET_ID_TAG, DATASET_PATTERN, TABLE_ID_TAG,
          TABLE_PATTERN);

  private static final String BQ_LEGACY_STANDARD_TABLE_NAME_FORMAT =
      String.format(
          "^(?<%s>%s)[:\\.](?<%s>%s)\\.(?<%s>%s)$",
          PROJECT_ID_TAG, PROJECT_PATTERN, DATASET_ID_TAG, DATASET_PATTERN, TABLE_ID_TAG,
          TABLE_PATTERN);

  private static BigQueryTableEntity extractInformation(String pattern, String resource) {
    Matcher matcher = Pattern.compile(pattern).matcher(resource);

    if (!matcher.find()) {
      throw new IllegalArgumentException(
          "input (" + resource + ") not in correct format (" + pattern + ")");
    }

    return BigQueryTableEntity.builder()
        .setProjectId(matcher.group(PROJECT_ID_TAG))
        .setDataset(matcher.group(DATASET_ID_TAG))
        .setTable(matcher.group(TABLE_ID_TAG))
        .build();
  }

  private BigQueryTableCreator() {
  }
} @AutoValue
public abstract class BigQueryTableEntity implements DataEntityConvertible, Serializable {

  public abstract String getProjectId();

  public abstract String getDataset();

  public abstract String getTable();

  /**
   * Returns {@code true} if the table is a temporary table.
   * <p> It uses rule dataset name starts with '_' or table name starts with '_' or 'anon'.
   */
  public final boolean isTempTable() {
    return getDataset().startsWith("_")
        || getTable().startsWith("_")
        || getTable().startsWith("anon");
  }

  public static Builder builder() {
    return new AutoValue_BigQueryTableEntity.Builder();
  }

  public static BigQueryTableEntity create(String projectId, String dataset, String table) {
    return builder()
        .setProjectId(projectId)
        .setDataset(dataset)
        .setTable(table)
        .build();
  }

  @Override
  public DataEntity dataEntity() {
    return DataEntity.newBuilder()
        .setKind(DataEntityTypes.BIGQUERY_TABLE)
        .setLinkedResource(
            String.format("//bigquery.googleapis.com/projects/%s/datasets/%s/tables/%s",
                getProjectId(),
                getDataset(),
                getTable()))
        .setSqlResource(
            String.format("bigquery.table.%s.%s.%s", getProjectId(), getDataset(), getTable()))
        .build();
  }

  public String getLegacySqlName() {
    return String.format("%s:%s.%s", getProjectId(), getDataset(), getTable());
  }

  public String getStandSqlName() {
    return String.format("%s.%s.%s", getProjectId(), getDataset(), getTable());
  }

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder setProjectId(String projectId);

    public abstract Builder setDataset(String dataset);

    public abstract Builder setTable(String table);

    public abstract BigQueryTableEntity build();
  }
} @AutoValue
@DefaultSchema(AutoValueSchema.class)
public abstract class TagsForCatalog {

  public abstract String getEntryId();

  public abstract ImmutableList<String> getTagsJson();

  public static TagsForCatalog empty() {
    return create("", ImmutableList.of());
  }

  public final boolean isEmpty() {
    return this.equals(empty());
  }

  public final ImmutableList<Tag> parsedTags() {
    return ProtoJsonConverter.parseAsList(getTagsJson(), Tag.class);
  }

  @SchemaCreate
  public static TagsForCatalog create(String entryId, Collection<String> tagsJson) {
    return builder()
        .setEntryId(entryId)
        .setTagsJson(ImmutableList.copyOf(tagsJson))
        .build();
  }

  public static Builder builder() {
    return new AutoValue_TagsForCatalog.Builder();
  }

  public static Builder forTags(Collection<Tag> tags) {
    return builder()
        .setTagsJson(
            tags.stream()
                .map(ProtoJsonConverter::asJsonString)
                .collect(toImmutableList()));
  }

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder setEntryId(String newEntryId);

    public Builder setEntry(Entry entry) {
      return setEntryId(entry.getName());
    }

    public abstract Builder setTagsJson(ImmutableList<String> newTagsJson);

    public abstract TagsForCatalog build();
  }
} public final class BigQueryTableLoadService implements Serializable {

  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

  private final BigQueryServiceFactory bqServiceFactory;
  private static final Cache<BigQueryTableEntity, Table> LOCAL_CACHE = buildCache();

  public BigQueryTableLoadService(
      BigQueryServiceFactory bqServiceFactory) {
    this.bqServiceFactory = bqServiceFactory;
  }

  public static BigQueryTableLoadService usingServiceFactory(
      BigQueryServiceFactory bqServiceFactory) {
    return new BigQueryTableLoadService(bqServiceFactory);
  }

  private static Cache<BigQueryTableEntity, Table> buildCache() {
    return CacheBuilder.newBuilder()
        .expireAfterWrite(Duration.ofMinutes(5))
        .maximumSize(10000L)
        .build();
  }

  /**
   * Loads a single table information.
   *
   * @param tableName the fully qualified table name.
   */
  public Table loadTable(String tableName) {
    return loadTable(BigQueryTableCreator.usingBestEffort(tableName));
  }

  /**
   * Loads a single table information.
   *
   * @param table the fully qualified table name.
   */
  public Table loadTable(BigQueryTableEntity table) {
    try {
      //check in cache
      return LOCAL_CACHE.get(table, () -> loadSingleTableFromServer(table));
    } catch (ExecutionException executionException) {
      logger.atWarning()
          .withCause(executionException)
          .atMostEvery(1, TimeUnit.MINUTES)
          .log(String.format("Unable to load table %s", table));

      throw new BigQueryOperationException(table, executionException);
    }
  }

  /**
   * Loads multiple tables' information.
   *
   * @param tableNames multiple fully qualified table names.
   */
  public ImmutableSet<Table> loadTables(String... tableNames) {
    return loadTables(
        Arrays.stream(tableNames)
            .distinct()
            .map(BigQueryTableCreator::usingBestEffort)
            .collect(toImmutableSet()));
  }

  /**
   * Loads multiple tables from BigQuery, with Cache-through.
   *
   * <b>It does NOT use Batch API.</b>
   *
   * @param tables the list of table names to load
   * @return all the loaded tables. It can throw runtime exception if table is not found.
   */
  public ImmutableSet<Table> loadTables(ImmutableSet<BigQueryTableEntity> tables) {
    ImmutableMap<BigQueryTableEntity, Table> allCachedData =
        LOCAL_CACHE.getAllPresent(tables);
    ImmutableSet<BigQueryTableEntity> tablesToFetch
        = Sets.difference(tables, allCachedData.keySet()).immutableCopy();

    return Sets.union(ImmutableSet.copyOf(allCachedData.values()),
        bulkFetchFromServer(tablesToFetch)).immutableCopy();
  }

  private Table loadSingleTableFromServer(BigQueryTableEntity tableSpec) throws IOException {
    return bqServiceFactory.buildService()
        .tables()
        .get(tableSpec.getProjectId(), tableSpec.getDataset(), tableSpec.getTable())
        .execute();
  }

  private ImmutableSet<Table> bulkFetchFromServer(ImmutableSet<BigQueryTableEntity> tablesToFetch) {
    return tablesToFetch.stream()
        .map(this::loadTable)
        .collect(toImmutableSet());
  }

  @VisibleForTesting
  static void clearLocalCache() {
    LOCAL_CACHE.invalidateAll();
    LOCAL_CACHE.cleanUp();
  }
} public final class BigQueryZetaSqlSchemaLoader implements ZetaSqlSchemaLoader {

  private final BigQueryTableLoadService bqTableLoader;

  public BigQueryZetaSqlSchemaLoader(BigQueryTableLoadService bqTableLoader) {
    this.bqTableLoader = bqTableLoader;
  }

  @Override
  public ImmutableSet<SimpleTable> loadSchemas(String... tableNames) {
    return bqTableLoader.loadTables(tableNames).stream()
        .map(BigQuerySchemaConverter::convert)
        .collect(toImmutableSet());
  }

  @Override
  public ImmutableSet<SimpleTable> loadSchemas(ImmutableSet<String> tableNames) {
    return loadSchemas(tableNames.asList().toArray(new String[0]));
  }
} public interface ZetaSqlSchemaLoaderFactory extends Serializable {

  ZetaSqlSchemaLoader newLoader();

  /**
   * Returns a SchemaLoader which always returns  {@code null} schema.
   */
  static ZetaSqlSchemaLoaderFactory emptyLoaderFactory() {
    return () -> null;
  }
} public final class FakeBigQueryServiceFactory implements BigQueryServiceFactory {

  private final String[] tableSchemas;

  public FakeBigQueryServiceFactory(String[] tableSchemas) {
    this.tableSchemas = tableSchemas;
  }

  public static FakeBigQueryServiceFactory forTableSchemas(String... tableSchemas) {
    return new FakeBigQueryServiceFactory(tableSchemas);
  }

  public static BigQueryServiceFactory forStub(FakeBigquery fakeService) {
    return () -> fakeService;
  }

  @Override
  public Bigquery buildService() {
    return FakeBigquery.forTableSchemas(tableSchemas);
  }
} public final class FakeDataCatalogStub extends DataCatalogStub implements Serializable {

  private boolean shutdown;
  private boolean terminated;

  private final ImmutableMap<String, Entry> predefinedEntries;
  private final ImmutableMultimap<String, Tag> predefinedTags;

  public FakeDataCatalogStub(
      ImmutableCollection<Entry> predefinedEntries,
      ImmutableMultimap<String, Tag> predefinedTags) {
    this.predefinedEntries = predefinedEntries.stream()
        .collect(toImmutableMap(Entry::getLinkedResource, identity()));

    this.predefinedTags = predefinedTags;
  }

  public static FakeDataCatalogStub buildWithTestData(
      List<String> entryResourcesNames,
      List<String> tagsResourceNames) {

    ImmutableList<Entry> entries =
        entryResourcesNames.stream()
            .map(TestResourceLoader::load)
            .map(json -> parseJson(json, Entry.class))
            .collect(toImmutableList());

    ImmutableListMultimap<String, Tag> entityTags =
        tagsResourceNames.stream()
            .map(TestResourceLoader::load)
            .map(tagsJson -> parseAsList(tagsJson, Tag.class))
            .flatMap(List::stream)
            .collect(toImmutableListMultimap(TagUtility::extractParent, identity()));

    return new FakeDataCatalogStub(entries, entityTags);
  }


  @Override
  public UnaryCallable<LookupEntryRequest, Entry> lookupEntryCallable() {
    checkState(!shutdown, "Stub shutdown");

    return new UnaryCallable<LookupEntryRequest, Entry>() {

      @Override
      public ApiFuture<Entry> futureCall(LookupEntryRequest lookupEntryRequest,
          ApiCallContext apiCallContext) {
        return new FakeDataCatalogLookupEntryResponse(
            predefinedEntries.get(lookupEntryRequest.getLinkedResource()),
            lookupEntryRequest);
      }
    };
  }

  @Override
  public UnaryCallable<ListTagsRequest, ListTagsPagedResponse> listTagsPagedCallable() {
    checkState(!shutdown, "Stub shutdown");

    return new UnaryCallable<ListTagsRequest, ListTagsPagedResponse>() {
      @Override
      public ApiFuture<ListTagsPagedResponse> futureCall(ListTagsRequest listTagsRequest,
          ApiCallContext apiCallContext) {
        return
            new FakeDataCatalogPagesListTagsResponse(
                listTagsRequest, apiCallContext, predefinedTags.get(listTagsRequest.getParent()));
      }
    };
  }

  @Override
  public void close() {
    // Do nothing because this is a Fake and doesn't implement an actual gRPC operation.
  }

  @Override
  public void shutdown() {
    shutdown = true;
  }

  @Override
  public boolean isShutdown() {
    return shutdown;
  }

  @Override
  public boolean isTerminated() {
    return terminated;
  }

  @Override
  public void shutdownNow() {
    shutdown();
  }

  @Override
  public boolean awaitTermination(long l, TimeUnit timeUnit) {
    terminated = true;
    return true;
  }


} public final class JsonTransforms {

  public static <T> String stringify(T obj) {
    try {
      return new ObjectMapper().writeValueAsString(obj);
    } catch (JsonProcessingException jsonProcessingException) {
      return "";
    }
  }

  private JsonTransforms() {
  }
} public final class TestResourceLoader {

  private static final String TEST_RESOURCE_FOLDER = "test";

  public static String load(String resourceFileName) {
    try {
      byte[] bytes = Files
          .readAllBytes(Paths.get("src", TEST_RESOURCE_FOLDER, "resources", resourceFileName));
      return new String(bytes, StandardCharsets.UTF_8);
    } catch (IOException ioException) {
      return "";
    }
  }

  private TestResourceLoader() {
  }
} @AutoValue
public abstract class CompositeLineageToTagTransformation extends
    PTransform<PCollection<CompositeLineage>, PCollection<TagsForCatalog>> {

  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

  public abstract String lineageTagTemplateId();

  @Nullable
  public abstract DataCatalogStub catalogStub();

  @Override
  public PCollection<TagsForCatalog> expand(PCollection<CompositeLineage> input) {
    return input.apply("convert to Data Catalog Tags", ParDo.of(
        new DoFn<CompositeLineage, TagsForCatalog>() {
          @ProcessElement
          public void convertToTag(
              @Element CompositeLineage lineage,
              OutputReceiver<TagsForCatalog> out) {

            try {
              DataEntity targetTable = lineage.getTableLineage().getTarget();

              DataCatalogService.usingStub(catalogStub())
                  .lookupEntry(targetTable)
                  .ifPresent(tableEntry ->
                      out.output(
                          TagsForCatalog
                              .forTags(LineageToTagConverterFactory
                                  .forTemplateId(lineageTagTemplateId())
                                  .converterFor(lineage)
                                  .buildTags())
                              .setEntryId(tableEntry.getName())
                              .build()));

            } catch (Exception exception) {
              logger.atWarning().withCause(exception)
                  .atMostEvery(1, TimeUnit.MINUTES)
                  .log("error converting CompositeLineage\n%s", lineage);
            }
          }
        }
    ));
  }

  public static CompositeLineageToTagTransformation withTagTemplateId(String tagTemplateId) {
    return builder().lineageTagTemplateId(tagTemplateId).build();
  }

  public static Builder builder() {
    return new AutoValue_CompositeLineageToTagTransformation.Builder();
  }

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder lineageTagTemplateId(String lineageTagTemplateId);

    public abstract Builder catalogStub(@Nullable DataCatalogStub catalogStub);

    public abstract CompositeLineageToTagTransformation build();
  }
} @AutoValue
public abstract class LineageExtractionTransform extends
    PTransform<PCollection<String>, PCollection<CompositeLineage>> {

  public abstract BigQueryTableEntity getOutputLineageTable();

  @Nullable
  public abstract Clock getClock();

  @Nullable
  public abstract ZetaSqlSchemaLoaderFactory getZetaSqlSchemaLoaderFactory();

  @Override
  public PCollection<CompositeLineage> expand(PCollection<String> auditLogMessages) {
    return auditLogMessages
        .apply("Validate Log Events", Filter.by(AuditLogValidator.create(getOutputLineageTable())))
        .apply("Extract Lineage", ParDo.of(new IdentifyAndExtract()))
        .apply("Validate Lineage", Filter.by(checkNonEmptyLineage()));
  }

  /**
   * Returns true if the {@link CompositeLineage} contains at least jobInformation and TableLineage
   * and the target table is not the Lineage Output BigQuery table.
   */
  private SerializableFunction<CompositeLineage, Boolean> checkNonEmptyLineage() {
    return input ->
        isNonEmptyMessage(input.getJobInformation()) &&
            isNonEmptyMessage(input.getTableLineage());
  }

  /**
   * Checks if the given Protobuf Message is non-null and not-default.
   *
   * @param message the Protobuf message to check
   * @return {@code true} if message is non-null and not-equal to defaultInstance()
   */
  private static boolean isNonEmptyMessage(Message message) {
    return message != null && !message.equals(message.getDefaultInstanceForType());
  }

  private class IdentifyAndExtract extends DoFn<String, CompositeLineage> {

    @ProcessElement
    public void extract(@Element String messageJson, OutputReceiver<CompositeLineage> out) {
      out.output(
          new InsertJobTableLineageExtractor(getClock(), JsonMessageParser.of(messageJson),
              getZetaSqlSchemaLoaderFactory())
              .extract());
    }
  }

  public static LineageExtractionTransform create(
      BigQueryTableEntity outputLineageTable,
      @Nullable Clock clock,
      @Nullable ZetaSqlSchemaLoaderFactory zetaSqlSchemaLoaderFactory) {
    return builder()
        .setOutputLineageTable(outputLineageTable)
        .setClock(clock)
        .setZetaSqlSchemaLoaderFactory(zetaSqlSchemaLoaderFactory)
        .build();
  }

  public static Builder builder() {
    return new AutoValue_LineageExtractionTransform.Builder()
        .setZetaSqlSchemaLoaderFactory(emptyLoaderFactory());
  }

  abstract Builder toBuilder();

  @AutoValue.Builder
  public abstract static class Builder {

    public abstract Builder setOutputLineageTable(BigQueryTableEntity newLineageTable);

    public abstract Builder setClock(@Nullable Clock newClock);

    public abstract Builder setZetaSqlSchemaLoaderFactory(
        @Nullable ZetaSqlSchemaLoaderFactory zetaSqlSchemaLoaderFactory);

    abstract LineageExtractionTransform autoBuild();

    public LineageExtractionTransform build() {
      LineageExtractionTransform transformObj = autoBuild();

      if (transformObj.getClock() == null) {
        transformObj = transformObj.toBuilder().setClock(Clock.systemUTC()).build();
      }

      if (transformObj.getZetaSqlSchemaLoaderFactory() == null) {
        transformObj = transformObj.toBuilder()
            .setZetaSqlSchemaLoaderFactory(emptyLoaderFactory()).build();
      }

      return transformObj;
    }
  }
}
import static com.google.cloud.solutions.datalineage.converter.ProtoJsonConverter.parseAsList;
import static com.google.cloud.solutions.datalineage.converter.ProtoJsonConverter.parseJson;
import static com.google.common.collect.ImmutableList.toImmutableList;
import com.google.cloud.datacatalog.v1beta1.Entry;
import com.google.cloud.datacatalog.v1beta1.Tag;
import com.google.cloud.datacatalog.v1beta1.TagField;
import com.google.cloud.solutions.datalineage.extractor.BigQueryTableCreator;
import com.google.cloud.solutions.datalineage.model.BigQueryTableEntity;
import com.google.cloud.solutions.datalineage.model.LineageMessages.CompositeLineage;
import com.google.cloud.solutions.datalineage.model.TagsForCatalog;
import com.google.cloud.solutions.datalineage.service.BigQueryTableLoadService;
import com.google.cloud.solutions.datalineage.service.BigQueryZetaSqlSchemaLoader;
import com.google.cloud.solutions.datalineage.service.ZetaSqlSchemaLoaderFactory;
import com.google.cloud.solutions.datalineage.testing.FakeBigQueryServiceFactory;
import com.google.cloud.solutions.datalineage.testing.FakeDataCatalogStub;
import com.google.cloud.solutions.datalineage.testing.JsonTransforms;
import com.google.cloud.solutions.datalineage.testing.TestResourceLoader;
import com.google.cloud.solutions.datalineage.transform.CompositeLineageToTagTransformation;
import com.google.cloud.solutions.datalineage.transform.LineageExtractionTransform;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.jayway.jsonpath.JsonPath;
import java.time.Clock;
import java.time.Instant;
import java.time.ZoneId;
import net.minidev.json.JSONArray;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.PCollection;
import org.apache.commons.lang3.StringUtils;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
/*
 * Copyright 2020 Google LLC
 *
 * Licensed 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 com.google.cloud.solutions.datalineage;



@RunWith(JUnit4.class)
public final class PipelineTest {

  @Rule
  public final transient TestPipeline p = TestPipeline.create();

  static ImmutableList<String> sampleMessages() {
    return JsonPath
        .parse(
            TestResourceLoader.load("message_sequence_for_bq_query_with_destination_table.json"))
        .<JSONArray>read("$[*]")
        .stream()
        .map(JsonTransforms::stringify)
        .filter(StringUtils::isNotBlank)
        .collect(toImmutableList());
  }

  @Test
  public void compositeLineageSerializable() {
    p.apply(Create.of(CompositeLineage.getDefaultInstance()));
    p.run();
  }

  @Test
  public void tagsForCatalog_empty_Serializable() {
    p.apply(Create.of(TagsForCatalog.empty()));
    p.run();
  }

  @Test
  public void tagsForCatalog_filled_Serializable() {
    TagsForCatalog sample =
        TagsForCatalog
            .forTags(ImmutableSet.of(
                Tag.newBuilder()
                    .setTemplate("templateId")
                    .putFields("test", TagField.newBuilder().setStringValue("string1").build())
                    .build()))
            .setEntry(Entry.newBuilder().setName("asddsdadasda").build())
            .build();

    p.apply(Create.of(sample));
    p.run();
  }

  @Test
  public void auditLogLineageParser_changesToLineageTable_discardChanges() {
    BigQueryTableEntity lineageTable =
        BigQueryTableCreator
            .fromLegacyTableName("myproject:demo.sample_table");

    PCollection<CompositeLineage> output = p
        .apply(Create.of(sampleMessages()))
        .setCoder(StringUtf8Coder.of())
        .apply(LineageExtractionTransform.builder().setOutputLineageTable(lineageTable).build());

    PAssert.that(output).empty();

    p.run();

  }

  @Test
  public void auditLogLineageParser_noSchemaLoader_onlyTableLineage() {
    Clock fixedClock = Clock.fixed(Instant.parse("2020-04-13T21:29:34.923Z"), ZoneId.of("UTC"));
    BigQueryTableEntity lineageTable =
        BigQueryTableCreator
            .fromLegacyTableName("projectId:dataset.table");

    PCollection<CompositeLineage> output = p
        .apply(Create.of(sampleMessages()))
        .setCoder(StringUtf8Coder.of())
        .apply(
            LineageExtractionTransform.builder()
                .setClock(fixedClock)
                .setOutputLineageTable(lineageTable)
                .setZetaSqlSchemaLoaderFactory(ZetaSqlSchemaLoaderFactory.emptyLoaderFactory())
                .build());

    PAssert.that(output)
        .containsInAnyOrder(
            parseJson(TestResourceLoader
                    .load("composite-lineages/only_table_level_lineage.json"),
                CompositeLineage.class));

    p.run();
  }

  @Test
  public void auditLogLineageParser_sqlParserProvided_containsColumnLineages() {
    Clock fixedClock = Clock.fixed(Instant.parse("2020-04-13T21:29:34.923Z"), ZoneId.of("UTC"));
    BigQueryTableEntity lineageTable =
        BigQueryTableCreator
            .fromLegacyTableName("projecId:dataset.table");
    ZetaSqlSchemaLoaderFactory fakeSchemaLoaderFactory =
        (ZetaSqlSchemaLoaderFactory) () ->
            new BigQueryZetaSqlSchemaLoader(