BaseDataPublisherTest.java

/*
 * 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 gobblin.publisher;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Properties;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.testng.Assert;
import org.testng.annotations.Test;

import com.google.common.collect.ImmutableList;
import com.google.common.io.Files;

import gobblin.configuration.ConfigurationKeys;
import gobblin.configuration.State;
import gobblin.configuration.WorkUnitState;
import gobblin.metadata.MetadataMerger;
import gobblin.metadata.types.GlobalMetadata;
import gobblin.util.ForkOperatorUtils;


/**
 * Tests for BaseDataPublisher
 */
public class BaseDataPublisherTest {
  /**
   * Test DATA_PUBLISHER_METADATA_STR: a user should be able to put an arbitrary metadata string in job configuration
   * and have that written out.
   */
  @Test
  public void testMetadataStrOneBranch()
      throws IOException {
    State s = buildDefaultState(1);

    WorkUnitState wuState = new WorkUnitState();
    wuState.setProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_STR, "foobar");
    addStateToWorkunit(s, wuState);

    BaseDataPublisher publisher = new BaseDataPublisher(s);
    publisher.publishMetadata(wuState);

    try (InputStream mdStream = new FileInputStream(openMetadataFile(s, 1, 0))) {
      String mdBytes = IOUtils.toString(mdStream, StandardCharsets.UTF_8);
      Assert.assertEquals(mdBytes, "foobar", "Expected to read back metadata from string");
    }
  }

  /**
   * Test that DATA_PUBLISHER_METADATA_STR functionality works across multiple branches.
   */
  @Test
  public void testMetadataStrMultipleWorkUnitsAndBranches()
      throws IOException {
    final int numBranches = 3;
    State s = buildDefaultState(numBranches);

    List<WorkUnitState> workUnits = new ArrayList<>();
    for (int i = 0; i < numBranches; i++) {
      WorkUnitState wuState = new WorkUnitState();
      wuState.setProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_STR, "foobar");
      addStateToWorkunit(s, wuState);
      workUnits.add(wuState);
    }

    BaseDataPublisher publisher = new BaseDataPublisher(s);
    publisher.publishMetadata(workUnits);

    for (int branch = 0; branch < numBranches; branch++) {
      try (InputStream mdStream = new FileInputStream(openMetadataFile(s, numBranches, branch))) {
        String mdBytes = IOUtils.toString(mdStream, StandardCharsets.UTF_8);
        Assert.assertEquals(mdBytes, "foobar", "Expected to read back metadata from string");
      }
    }
  }

  /**
   * Test that an exception is properly thrown if we configure a merger that doesn't actually implement
   * MetadataMerger
   */
  @Test(expectedExceptions = IllegalArgumentException.class)
  public void testBogusMetadataMerger()
      throws IOException {
    State s = buildDefaultState(1);
    s.setProp(ConfigurationKeys.DATA_PUBLISH_WRITER_METADATA_KEY, "true");
    s.setProp(ConfigurationKeys.DATA_PUBLISH_WRITER_METADATA_MERGER_NAME_KEY, "java.lang.String");
    s.setProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_STR, "foobar");
    WorkUnitState wuState = new WorkUnitState();
    addStateToWorkunit(s, wuState);

    BaseDataPublisher publisher = new BaseDataPublisher(s);
    publisher.publishMetadata(Collections.singletonList(wuState));
  }

  /**
   * This test is testing several things at once:
   *  1. That a merger is called properly for all workunits in a brach
   *  2. That different mergers can be instantiated per branch
   */
  @Test
  public void testMergedMetadata()
      throws IOException {
    final int numBranches = 2;
    final int numWorkUnits = 10;

    State s = buildDefaultState(numBranches);

    for (int i = 0; i < numBranches; i++) {
      String mdKeyName = ForkOperatorUtils
          .getPropertyNameForBranch(ConfigurationKeys.DATA_PUBLISH_WRITER_METADATA_KEY, numBranches, i);
      String mdMergerKeyName = ForkOperatorUtils
          .getPropertyNameForBranch(ConfigurationKeys.DATA_PUBLISH_WRITER_METADATA_MERGER_NAME_KEY, numBranches, i);

      s.setProp(mdKeyName, "true");
      s.setProp(mdMergerKeyName,
          (i % 2) == 0 ? TestAdditionMerger.class.getName() : TestMultiplicationMerger.class.getName());
    }

    // For each branch, metadata is (branchId+1*workUnitNumber+1) - adding 1 so we don't ever multiply by 0
    List<WorkUnitState> workUnits = new ArrayList<>();
    for (int workUnitId = 0; workUnitId < numWorkUnits; workUnitId++) {
      WorkUnitState wuState = new WorkUnitState();
      addStateToWorkunit(s, wuState);

      for (int branchId = 0; branchId < numBranches; branchId++) {
        String mdForBranchName =
            ForkOperatorUtils.getPropertyNameForBranch(ConfigurationKeys.WRITER_METADATA_KEY, numBranches, branchId);
        wuState.setProp(mdForBranchName, String.valueOf((branchId + 1) * (workUnitId + 1)));
      }

      workUnits.add(wuState);
    }

    BaseDataPublisher publisher = new BaseDataPublisher(s);
    publisher.publishMetadata(workUnits);

    for (int branch = 0; branch < numBranches; branch++) {
      int expectedSum = (branch % 2 == 0) ? 0 : 1;
      for (int i = 0; i < numWorkUnits; i++) {
        if (branch % 2 == 0) {
          expectedSum += (branch + 1) * (i + 1);
        } else {
          expectedSum *= (branch + 1) * (i + 1);
        }
      }

      try (InputStream mdStream = new FileInputStream(openMetadataFile(s, numBranches, branch))) {
        String mdBytes = IOUtils.toString(mdStream, StandardCharsets.UTF_8);
        Assert.assertEquals(mdBytes, String.valueOf(expectedSum), "Expected to read back correctly merged metadata from string");
      }
    }
  }

  @Test
  public void testNoOutputWhenDisabled()
      throws IOException {
    State s = buildDefaultState(1);

    WorkUnitState wuState = new WorkUnitState();
    addStateToWorkunit(s, wuState);

    wuState.setProp(ConfigurationKeys.WRITER_METADATA_KEY, "abcdefg");

    BaseDataPublisher publisher = new BaseDataPublisher(s);
    publisher.publishMetadata(Collections.singletonList(wuState));

    File mdFile = openMetadataFile(s, 1, 0);
    Assert.assertFalse(mdFile.exists(), "Internal metadata from writer should not be written out if no merger is set in config");
  }

  @Test
  public void testWithPartitionKey() throws IOException {
    File publishPath = Files.createTempDir();
    try {
      File part1 = new File(publishPath, "1-2-3-4");
      part1.mkdir();

      File part2 = new File(publishPath, "5-6-7-8");
      part2.mkdir();

      State s = buildDefaultState(1);
      String md = new GlobalMetadata().toJson();

      s.removeProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_OUTPUT_DIR);
      s.setProp(ConfigurationKeys.DATA_PUBLISH_WRITER_METADATA_KEY, "true");
      s.setProp(ConfigurationKeys.WRITER_METADATA_KEY, md);
      s.setProp(ConfigurationKeys.DATA_PUBLISHER_FINAL_DIR, publishPath.getAbsolutePath());
      s.setProp(ConfigurationKeys.DATA_PUBLISHER_APPEND_EXTRACT_TO_FINAL_DIR, "false");
      s.setProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_OUTPUT_FILE, "metadata.json");

      WorkUnitState wuState1 = new WorkUnitState();
      wuState1.setProp(ConfigurationKeys.WRITER_PARTITION_PATH_KEY, "1-2-3-4");
      wuState1.setProp(ConfigurationKeys.WRITER_METADATA_KEY, md);
      addStateToWorkunit(s, wuState1);

      WorkUnitState wuState2 = new WorkUnitState();
      wuState2.setProp(ConfigurationKeys.WRITER_PARTITION_PATH_KEY, "5-6-7-8");
      wuState2.setProp(ConfigurationKeys.WRITER_METADATA_KEY, md);
      addStateToWorkunit(s, wuState2);

      BaseDataPublisher publisher = new BaseDataPublisher(s);
      publisher.publishMetadata(ImmutableList.of(wuState1, wuState2));

      Assert.assertTrue(new File(part1, "metadata.json").exists());
      Assert.assertTrue(new File(part2, "metadata.json").exists());
    } finally {
      FileUtils.deleteDirectory(publishPath);
    }
  }

  public static class TestAdditionMerger implements MetadataMerger<String> {
    private int sum = 0;

    @Override
    public void update(String metadata) {
      sum += Integer.valueOf(metadata);
    }

    @Override
    public String getMergedMetadata() {
      return String.valueOf(sum);
    }
  }

  public static class TestMultiplicationMerger implements MetadataMerger<String> {
    private int product = 1;

    public TestMultiplicationMerger(Properties config) {
      // testing ctor call
    }

    @Override
    public void update(String metadata) {
      product *= Integer.valueOf(metadata);
    }

    @Override
    public String getMergedMetadata() {
      return String.valueOf(product);
    }
  }

  private void addStateToWorkunit(State s, WorkUnitState wuState) {
    for (Map.Entry<Object, Object> prop : s.getProperties().entrySet()) {
      wuState.setProp((String) prop.getKey(), prop.getValue());
    }
  }

  private File openMetadataFile(State state, int numBranches, int branchId) {
    String dir = state.getProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_OUTPUT_DIR);
    String fileName = state.getProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_OUTPUT_FILE);
    if (numBranches > 1) {
      fileName += "." + String.valueOf(branchId);
    }
    return new File(dir, fileName);
  }

  private State buildDefaultState(int numBranches)
      throws IOException {
    State state = new State();

    state.setProp(ConfigurationKeys.FORK_BRANCHES_KEY, numBranches);
    File tmpLocation = File.createTempFile("metadata", "");
    tmpLocation.delete();
    state.setProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_OUTPUT_DIR, tmpLocation.getParent());
    state.setProp(ConfigurationKeys.DATA_PUBLISHER_METADATA_OUTPUT_FILE, tmpLocation.getName());

    return state;
  }
}