mirrors/nomulus
1
0
Fork 0
mirror of https://github.com/google/nomulus synced 2026-09-03 06:36:56 +00:00
Code Issues Releases Wiki Activity

Load Datastore snapshot from backup files (#660)

Browse Source 
* Load Datastore snapshot from backup files

Defined a composite transform that loads from a Datastore export and
concurrent CommitLog files, identify entities that still exist at the
end of the time window, and resolve their latest states in the window.
This commit is contained in:
Weimin Yu
2020-07-01 09:58:42 -04:00
committed by GitHub
parent d065ff63fc
commit 1961a5759d
7 changed files with 385 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/src/main/java/google/registry/backup/AppEngineEnvironment.java
+2 -1
View File
@@ -23,7 +23,8 @@ import java.lang.reflect.Proxy;
/**
* Sets up a placeholder {@link Environment} on a non-AppEngine platform so that Datastore Entities
* can be deserialized. See {@code DatastoreEntityExtension} in test source for more information.
* can be converted from/to Objectify entities. See {@code DatastoreEntityExtension} in test source
* for more information.
*/
public class AppEngineEnvironment implements Closeable {
core/src/main/java/google/registry/beam/initsql/Transforms.java
+133 -8
View File
@@ -16,18 +16,27 @@ package google.registry.beam.initsql;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
import static google.registry.util.DateTimeUtils.START_OF_TIME;
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
import static java.util.Comparator.comparing;
import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
import static org.apache.beam.sdk.values.TypeDescriptors.strings;
import avro.shaded.com.google.common.collect.Iterators;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Streams;
import google.registry.backup.CommitLogImports;
import google.registry.backup.VersionedEntity;
import google.registry.tools.LevelDbLogReader;
import java.util.Collection;
import java.util.Iterator;
import java.util.Optional;
import java.util.Set;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.Compression;
import org.apache.beam.sdk.io.FileIO;
@@ -36,11 +45,20 @@ import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.ProcessFunction;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.joda.time.DateTime;
/**
@@ -57,6 +75,110 @@ public final class Transforms {
*/
@VisibleForTesting static final long EXPORT_ENTITY_TIME_STAMP = START_OF_TIME.getMillis();
/**
* Returns a {@link TupleTag} that can be used to retrieve entities of the given {@code kind} from
* the Datastore snapshot returned by {@link #loadDatastoreSnapshot}.
*/
public static TupleTag<VersionedEntity> createTagForKind(String kind) {
// When used with PCollectionTuple the result must retain generic type information.
// Both the Generic param and the empty bracket below are important.
return new TupleTag<VersionedEntity>(Transforms.class.getSimpleName() + ":" + kind) {};
}
/**
* Composite {@link PTransform transform} that loads the Datastore snapshot at {@code
* commitLogToTime} for caller specified {@code kinds}.
*
* <p>Caller must provide the location of a Datastore export that started AFTER {@code
* commitLogFromTime} and completed BEFORE {@code commitLogToTime}, as well as the root directory
* of all CommitLog files.
*
* <p>Selection of {@code commitLogFromTime} and {@code commitLogToTime} should follow the
* guidelines below to ensure that all incremental changes concurrent with the export are covered:
*
* <ul>
* <li>Two or more CommitLogs should exist between {@code commitLogFromTime} and the starting
* time of the Datastore export. This ensures that the earlier CommitLog file was complete
* before the export started.
* <li>Two or more CommitLogs should exit between the export completion time and {@code
* commitLogToTime}.
* </ul>
*
* <p>The output from the returned transform is a {@link PCollectionTuple} consisting of {@link
* VersionedEntity VersionedEntities} grouped into {@link PCollection PCollections} by {@code
* kind}.
*/
public static PTransform<PBegin, PCollectionTuple> loadDatastoreSnapshot(
String exportDir,
String commitLogDir,
DateTime commitLogFromTime,
DateTime commitLogToTime,
Set<String> kinds) {
checkArgument(kinds != null && !kinds.isEmpty(), "At least one kind is expected.");
// Create tags to collect entities by kind in final step.
final ImmutableMap<String, TupleTag<VersionedEntity>> outputTags =
kinds.stream()
.collect(ImmutableMap.toImmutableMap(kind -> kind, Transforms::createTagForKind));
// Arbitrarily select one tag as mainOutTag and put the remaining ones in a TupleTagList.
// This separation is required by ParDo's config API.
Iterator<TupleTag<VersionedEntity>> tagsIt = outputTags.values().iterator();
final TupleTag<VersionedEntity> mainOutputTag = tagsIt.next();
final TupleTagList additionalTags = TupleTagList.of(ImmutableList.copyOf(tagsIt));
return new PTransform<PBegin, PCollectionTuple>() {
@Override
public PCollectionTuple expand(PBegin input) {
PCollection<VersionedEntity> exportedEntities =
input
.apply("Get export file patterns", getDatastoreExportFilePatterns(exportDir, kinds))
.apply("Find export files", getFilesByPatterns())
.apply("Load export data", loadExportDataFromFiles());
PCollection<VersionedEntity> commitLogEntities =
input
.apply("Get commitlog file patterns", getCommitLogFilePatterns(commitLogDir))
.apply("Find commitlog files", getFilesByPatterns())
.apply(
"Filter commitLog by time",
filterCommitLogsByTime(commitLogFromTime, commitLogToTime))
.apply("Load commitlog data", loadCommitLogsFromFiles(kinds));
return PCollectionList.of(exportedEntities)
.and(commitLogEntities)
.apply("Merge exports and CommitLogs", Flatten.pCollections())
.apply(
"Key entities by Datastore Keys",
// Converting to KV<String, VE> instead of KV<Key, VE> b/c default coder for Key
// (SerializableCoder) is not deterministic and cannot be used with GroupBy.
MapElements.into(kvs(strings(), TypeDescriptor.of(VersionedEntity.class)))
.via((VersionedEntity e) -> KV.of(e.key().toString(), e)))
.apply("Gather entities by key", GroupByKey.create())
.apply(
"Output latest version per entity",
ParDo.of(
new DoFn<KV<String, Iterable<VersionedEntity>>, VersionedEntity>() {
@ProcessElement
public void processElement(
@Element KV<String, Iterable<VersionedEntity>> kv,
MultiOutputReceiver out) {
Optional<VersionedEntity> latest =
Streams.stream(kv.getValue())
.sorted(comparing(VersionedEntity::commitTimeMills).reversed())
.findFirst();
// Throw to abort (after default retries). Investigate, fix, and rerun.
checkState(
latest.isPresent(), "Unexpected key with no data", kv.getKey());
if (latest.get().isDelete()) {
return;
}
String kind = latest.get().getEntity().get().getKind();
out.get(outputTags.get(kind)).output(latest.get());
}
})
.withOutputTags(mainOutputTag, additionalTags));
}
};
}
/**
* Returns a {@link PTransform transform} that can generate a collection of patterns that match
* all Datastore CommitLog files.
@@ -96,7 +218,7 @@ public final class Transforms {
* Returns CommitLog files with timestamps between {@code fromTime} (inclusive) and {@code
* endTime} (exclusive).
*/
public static PTransform<PCollection<? extends String>, PCollection<String>>
public static PTransform<PCollection<? extends Metadata>, PCollection<Metadata>>
filterCommitLogsByTime(DateTime fromTime, DateTime toTime) {
return ParDo.of(new FilterCommitLogFileByTime(fromTime, toTime));
}
@@ -114,9 +236,13 @@ public final class Transforms {
/** Returns a {@link PTransform} from file {@link Metadata} to {@link VersionedEntity}. */
public static PTransform<PCollection<Metadata>, PCollection<VersionedEntity>>
loadCommitLogsFromFiles() {
loadCommitLogsFromFiles(Set<String> kinds) {
return processFiles(
new BackupFileReader(file -> CommitLogImports.loadEntities(file.open()).iterator()));
new BackupFileReader(
file ->
CommitLogImports.loadEntities(file.open()).stream()
.filter(e -> kinds.contains(e.key().getKind()))
.iterator()));
}
/**
@@ -139,12 +265,11 @@ public final class Transforms {
return input
.apply(FileIO.readMatches().withCompression(Compression.UNCOMPRESSED))
.apply(transformer.getClass().getSimpleName(), ParDo.of(transformer));
// TODO(weiminyu): reshuffle to enable dynamic work rebalance per beam dev guide
}
};
}
private static class FilterCommitLogFileByTime extends DoFn<String, String> {
private static class FilterCommitLogFileByTime extends DoFn<Metadata, Metadata> {
private final DateTime fromTime;
private final DateTime toTime;
@@ -161,10 +286,10 @@ public final class Transforms {
}
@ProcessElement
public void processElement(@Element String fileName, OutputReceiver<String> out) {
DateTime timestamp = getCommitLogTimestamp(fileName);
public void processElement(@Element Metadata fileMeta, OutputReceiver<Metadata> out) {
DateTime timestamp = getCommitLogTimestamp(fileMeta.resourceId().toString());
if (isBeforeOrAt(fromTime, timestamp) && timestamp.isBefore(toTime)) {
out.output(fileName);
out.output(fileMeta);
}
}
}