Commit 9a605b66 by Ashutosh Mestry

ATLAS-3953: Export: ZipSink: Specify character endcoding when writing to ZIP file.

Change-Id: I2e3a2bad79f5cf4c91328ed9f64a454ded2acbf3
parent 545afb53
......@@ -59,7 +59,7 @@ public class HdfsPathEntityCreator {
}
public AtlasEntity.AtlasEntityWithExtInfo getCreateEntity(AtlasObjectId item) throws AtlasBaseException {
if(!item.getUniqueAttributes().containsKey(HDFS_PATH_ATTRIBUTE_NAME_PATH)) {
if(item.getUniqueAttributes() == null || !item.getUniqueAttributes().containsKey(HDFS_PATH_ATTRIBUTE_NAME_PATH)) {
return null;
}
......
......@@ -22,12 +22,13 @@ import org.apache.atlas.model.impexp.AtlasExportResult;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.type.AtlasType;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
......@@ -109,24 +110,7 @@ public class ZipSink {
}
private void writeBytes(String payload) throws IOException {
splitAndWriteBytes(payload, 10 * 1024 * 1024, zipOutputStream);
}
static void splitAndWriteBytes(String msg, int bufferSize, OutputStream os) throws IOException {
int numberOfSplits = (int) Math.ceil(((float) msg.length()) / bufferSize);
if (numberOfSplits == 0) {
numberOfSplits = 1;
} else {
if (LOG.isDebugEnabled()) {
LOG.info("ZipSink: number of splits: {}", numberOfSplits);
}
}
for (int i = 0, start = 0; i < numberOfSplits; i++, start += bufferSize) {
int end = bufferSize + start;
String s = StringUtils.substring(msg, start, end);
os.write(s.getBytes());
}
IOUtils.copy(IOUtils.toInputStream(payload, StandardCharsets.UTF_8), zipOutputStream);
}
public boolean hasEntity(String guid) {
......
......@@ -30,15 +30,16 @@ import org.testng.annotations.Test;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import static org.mockito.Mockito.mock;
import static org.testng.Assert.*;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertTrue;
public class ZipSinkTest {
private ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
......@@ -47,23 +48,6 @@ public class ZipSinkTest {
private AtlasExportResult defaultExportResult;
private String knownEntityGuidFormat = "111-222-333-%s";
private class MockOutputStream extends OutputStream {
List<byte[]> collected = new ArrayList<>();
@Override
public void write(int b) throws IOException {
}
@Override
public void write(byte[] bytes) {
collected.add(bytes);
}
public List<byte[]> getCollected() {
return collected;
}
};
private void initZipSinkWithExportOrder() throws AtlasBaseException {
zipSink = new ZipSink(byteArrayOutputStream);
zipSink.setExportOrder(defaultExportOrder);
......@@ -146,7 +130,6 @@ public class ZipSinkTest {
@Test
public void zipWithExactlyTwoEntries_ContentsVerified() throws AtlasBaseException, IOException {
ByteArrayOutputStream byteOutputStream = new ByteArrayOutputStream();
useZipSinkToCreateEntries(byteOutputStream);
......@@ -225,26 +208,4 @@ public class ZipSinkTest {
String json = AtlasType.toJson(defaultExportResult);
return json.equals(s);
}
@Test
public void splitTest() throws IOException {
assertSplit("ABCDEFGHIJKLMNOPQRSTUVWXYZ01", 7, new String[] {"ABCDEFG", "HIJKLMN", "OPQRSTU", "VWXYZ01"});
assertSplit("ABCDEFGHIJKLMNOPQRSTUVWXYZ", 7, new String[] {"ABCDEFG", "HIJKLMN", "OPQRSTU", "VWXYZ"});
}
private void assertSplit(String msg, int bufferSize, String[] splits) throws IOException {
MockOutputStream os = getOutputStream();
ZipSink.splitAndWriteBytes(msg, bufferSize, os);
assertEquals(os.getCollected().size(), splits.length);
for (int i = 0; i < os.collected.size(); i++) {
byte[] bytes = os.getCollected().get(i);
String s = new String(bytes);
assertEquals(s, splits[i]);
}
}
private MockOutputStream getOutputStream() {
return new MockOutputStream();
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment