Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
RecordAccumulator builds a set of datastore Entity records from a set of leveldb logfiles in a directory tree (which is how we receive them for database backup/restore testing). This CL also refactors some of the logfile test code out of LevelDbLogReaderTest so that we can reuse it for building test logs. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=167313553
- Loading branch information
Showing
5 changed files
with
291 additions
and
43 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
// Copyright 2017 The Nomulus Authors. All Rights Reserved. | ||
// | ||
// 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 google.registry.tools; | ||
|
||
import com.google.appengine.api.datastore.EntityTranslator; | ||
import com.google.common.collect.ImmutableSet; | ||
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; | ||
import java.io.File; | ||
import java.io.FileInputStream; | ||
import java.io.IOException; | ||
|
||
/** Utility class that accumulates Entity records from level db files. */ | ||
class RecordAccumulator { | ||
private final LevelDbLogReader reader = new LevelDbLogReader(); | ||
|
||
/** Recursively reads all records in the directory. */ | ||
public final RecordAccumulator readDirectory(File dir) { | ||
for (File child : dir.listFiles()) { | ||
if (child.isDirectory()) { | ||
readDirectory(child); | ||
} else if (child.isFile()) { | ||
try { | ||
reader.readFrom(new FileInputStream(child)); | ||
} catch (IOException e) { | ||
throw new RuntimeException("IOException reading from file: " + child, e); | ||
} | ||
} | ||
} | ||
|
||
return this; | ||
} | ||
|
||
/** Creates an entity set from the current set of raw records. */ | ||
ImmutableSet<ComparableEntity> getComparableEntitySet() { | ||
ImmutableSet.Builder<ComparableEntity> builder = new ImmutableSet.Builder<>(); | ||
for (byte[] rawRecord : reader.getRecords()) { | ||
// Parse the entity proto and create an Entity object from it. | ||
EntityProto proto = new EntityProto(); | ||
proto.parseFrom(rawRecord); | ||
ComparableEntity entity = new ComparableEntity(EntityTranslator.createFromPb(proto)); | ||
|
||
builder.add(entity); | ||
} | ||
|
||
return builder.build(); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,87 @@ | ||
// Copyright 2017 The Nomulus Authors. All Rights Reserved. | ||
// | ||
// 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 google.registry.tools; | ||
|
||
import static com.google.common.base.Preconditions.checkArgument; | ||
import google.registry.tools.LevelDbLogReader.ChunkType; | ||
|
||
class LevelDbUtil { | ||
|
||
public static final int MAX_RECORD = LevelDbLogReader.BLOCK_SIZE - LevelDbLogReader.HEADER_SIZE; | ||
|
||
/** Adds a new record header to "bytes" at "pos", returns the new position. */ | ||
private static int addRecordHeader(byte[] bytes, int pos, ChunkType type, int size) { | ||
// Write a bogus checksum. | ||
for (int i = 0; i < 4; ++i) { | ||
bytes[pos++] = -1; | ||
} | ||
|
||
// Write size and type. | ||
bytes[pos++] = (byte) size; | ||
bytes[pos++] = (byte) (size >> 8); | ||
bytes[pos++] = (byte) type.getCode(); | ||
|
||
return pos; | ||
} | ||
|
||
/** | ||
* Adds a record of repeating bytes of 'val' of the given size to bytes at pos. | ||
* | ||
* <p>Writes the two least-significant bytes of 'val', alternating the order of them. So if the | ||
* value of 'val' is 0x1234, writes 0x12 0x34 0x34 0x12 0x12 ... If the value is greater than | ||
* 0xffff, it will be truncated to 16 bits. | ||
* | ||
* <p>This currently doesn't write a real checksum since we're not doing anything with that in the | ||
* leveldb reader. | ||
* | ||
* <p>Returns the new offset for the next block. | ||
*/ | ||
static int addRecord(byte[] bytes, int pos, ChunkType type, int size, int val) { | ||
pos = addRecordHeader(bytes, pos, type, size); | ||
|
||
// Write "size" bytes of data. | ||
for (int i = 0; i < size; ++i) { | ||
bytes[pos + i] = (byte) val; | ||
|
||
// Swap the least significant bytes in val so we can have more than 256 different same-sized | ||
// records. | ||
val = ((val >> 8) & 0xff) | ((val & 0xff) << 8); | ||
} | ||
|
||
return pos + size; | ||
} | ||
|
||
/** | ||
* Adds a record containing "data" to "bytes". | ||
* | ||
* <p>This currently doesn't write a real checksum since we're not doing anything with that in the | ||
* leveldb reader. | ||
* | ||
* <p>Returns the new offset for the next block. | ||
*/ | ||
static int addRecord(byte[] bytes, int pos, ChunkType type, byte[] data) { | ||
checkArgument( | ||
data.length < MAX_RECORD, "Record length (%d) > max record size (%d)", data.length, | ||
MAX_RECORD); | ||
pos = addRecordHeader(bytes, pos, type, data.length); | ||
|
||
// Write the contents of "data". | ||
for (int i = 0; i < data.length; ++i) { | ||
bytes[pos + i] = data[i]; | ||
} | ||
|
||
return pos + data.length; | ||
} | ||
} |
140 changes: 140 additions & 0 deletions
140
javatests/google/registry/tools/RecordAccumulatorTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,140 @@ | ||
// Copyright 2017 The Nomulus Authors. All Rights Reserved. | ||
// | ||
// 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 google.registry.tools; | ||
|
||
import static com.google.common.truth.Truth.assertThat; | ||
import static google.registry.tools.LevelDbLogReader.BLOCK_SIZE; | ||
|
||
import com.google.appengine.api.datastore.Entity; | ||
import com.google.appengine.api.datastore.EntityTranslator; | ||
import com.google.auto.value.AutoValue; | ||
import com.google.common.collect.ImmutableSet; | ||
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; | ||
import google.registry.testing.AppEngineRule; | ||
import google.registry.tools.LevelDbLogReader.ChunkType; | ||
import java.io.File; | ||
import java.io.FileNotFoundException; | ||
import java.io.FileOutputStream; | ||
import java.io.IOException; | ||
import org.junit.Rule; | ||
import org.junit.Test; | ||
import org.junit.rules.TemporaryFolder; | ||
import org.junit.runner.RunWith; | ||
import org.junit.runners.JUnit4; | ||
|
||
@RunWith(JUnit4.class) | ||
public class RecordAccumulatorTest { | ||
|
||
private static final int BASE_ID = 1001; | ||
private static final String TEST_ENTITY_KIND = "TestEntity"; | ||
|
||
@Rule public final TemporaryFolder tempFs = new TemporaryFolder(); | ||
@Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build(); | ||
|
||
@Test | ||
public void testReadDirectory() throws IOException { | ||
File subdir = tempFs.newFolder("folder"); | ||
LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(subdir, "data1")); | ||
|
||
// Note that we need to specify property values as "Long" for property comparisons to work | ||
// correctly because that's how they are deserialized from protos. | ||
ComparableEntity e1 = | ||
builder.addEntityProto( | ||
BASE_ID, | ||
Property.create("eeny", 100L), | ||
Property.create("meeny", 200L), | ||
Property.create("miney", 300L)); | ||
ComparableEntity e2 = | ||
builder.addEntityProto( | ||
BASE_ID + 1, | ||
Property.create("eeny", 100L), | ||
Property.create("meeny", 200L), | ||
Property.create("miney", 300L)); | ||
builder.build(); | ||
|
||
builder = new LevelDbFileBuilder(new File(subdir, "data2")); | ||
|
||
// Duplicate of the record in the other file. | ||
builder.addEntityProto( | ||
BASE_ID, | ||
Property.create("eeny", 100L), | ||
Property.create("meeny", 200L), | ||
Property.create("miney", 300L)); | ||
|
||
ComparableEntity e3 = | ||
builder.addEntityProto( | ||
BASE_ID + 2, | ||
Property.create("moxy", 100L), | ||
Property.create("fruvis", 200L), | ||
Property.create("cortex", 300L)); | ||
builder.build(); | ||
|
||
ImmutableSet<ComparableEntity> entities = | ||
new RecordAccumulator().readDirectory(subdir).getComparableEntitySet(); | ||
assertThat(entities).containsExactly(e1, e2, e3); | ||
} | ||
|
||
/** Utility class for building a leveldb logfile. */ | ||
private static final class LevelDbFileBuilder { | ||
private final FileOutputStream out; | ||
private byte[] currentBlock = new byte[BLOCK_SIZE]; | ||
|
||
// Write position in the current block. | ||
private int currentPos = 0; | ||
|
||
LevelDbFileBuilder(File file) throws FileNotFoundException { | ||
out = new FileOutputStream(file); | ||
} | ||
|
||
/** | ||
* Adds a record containing a new entity protobuf to the file. | ||
* | ||
* <p>Returns the ComparableEntity object rather than "this" so that we can check for the | ||
* presence of the entity in the result set. | ||
*/ | ||
private ComparableEntity addEntityProto(int id, Property... properties) throws IOException { | ||
Entity entity = new Entity(TEST_ENTITY_KIND, id); | ||
for (Property prop : properties) { | ||
entity.setProperty(prop.name(), prop.value()); | ||
} | ||
EntityProto proto = EntityTranslator.convertToPb(entity); | ||
byte[] protoBytes = proto.toByteArray(); | ||
if (protoBytes.length > BLOCK_SIZE - currentPos) { | ||
out.write(currentBlock); | ||
currentBlock = new byte[BLOCK_SIZE]; | ||
} | ||
|
||
currentPos = LevelDbUtil.addRecord(currentBlock, currentPos, ChunkType.FULL, protoBytes); | ||
return new ComparableEntity(entity); | ||
} | ||
|
||
/** Writes all remaining data and closes the block. */ | ||
void build() throws IOException { | ||
out.write(currentBlock); | ||
out.close(); | ||
} | ||
} | ||
|
||
@AutoValue | ||
abstract static class Property { | ||
static Property create(String name, Object value) { | ||
return new AutoValue_RecordAccumulatorTest_Property(name, value); | ||
} | ||
|
||
abstract String name(); | ||
|
||
abstract Object value(); | ||
} | ||
} |