mirror of
https://github.com/google/nomulus
synced 2026-09-19 22:44:26 +00:00
Load CommitLog into BEAM pipeline (#618)
* Load CommitLog into BEAM pipeline Created tools that can generate CommitLogs on a test datastore. Defined BEAM transforms that load from CommitLog files and added a few simple tests. This is work-in-progress. Next step is to build a consistent backup by merging Datastore exports and CommitLogs.
This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
// Copyright 2020 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.backup;
|
||||
|
||||
import com.google.apphosting.api.ApiProxy;
|
||||
import com.google.apphosting.api.ApiProxy.Environment;
|
||||
import com.google.common.collect.ImmutableMap;
|
||||
import java.io.Closeable;
|
||||
import java.lang.reflect.Method;
|
||||
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.
|
||||
*/
|
||||
public class AppEngineEnvironment implements Closeable {
|
||||
|
||||
private static final Environment PLACEHOLDER_ENV = createAppEngineEnvironment();
|
||||
|
||||
private boolean isPlaceHolderNeeded;
|
||||
|
||||
AppEngineEnvironment() {
|
||||
isPlaceHolderNeeded = ApiProxy.getCurrentEnvironment() == null;
|
||||
// isPlaceHolderNeeded may be true when we are invoked in a test with AppEngineRule.
|
||||
if (isPlaceHolderNeeded) {
|
||||
ApiProxy.setEnvironmentForCurrentThread(PLACEHOLDER_ENV);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
if (isPlaceHolderNeeded) {
|
||||
ApiProxy.setEnvironmentForCurrentThread(null);
|
||||
}
|
||||
}
|
||||
|
||||
/** Returns a placeholder {@link Environment} that can return hardcoded AppId and Attributes. */
|
||||
private static Environment createAppEngineEnvironment() {
|
||||
return (Environment)
|
||||
Proxy.newProxyInstance(
|
||||
Environment.class.getClassLoader(),
|
||||
new Class[] {Environment.class},
|
||||
(Object proxy, Method method, Object[] args) -> {
|
||||
switch (method.getName()) {
|
||||
case "getAppId":
|
||||
return "PlaceholderAppId";
|
||||
case "getAttributes":
|
||||
return ImmutableMap.<String, Object>of();
|
||||
default:
|
||||
throw new UnsupportedOperationException(method.getName());
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
// Copyright 2020 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.backup;
|
||||
|
||||
import static com.google.common.base.Preconditions.checkState;
|
||||
import static google.registry.backup.BackupUtils.createDeserializingIterator;
|
||||
|
||||
import com.google.common.collect.ImmutableList;
|
||||
import com.google.common.collect.Streams;
|
||||
import google.registry.model.ImmutableObject;
|
||||
import google.registry.model.ofy.CommitLogCheckpoint;
|
||||
import google.registry.model.ofy.CommitLogManifest;
|
||||
import google.registry.model.ofy.CommitLogMutation;
|
||||
import java.io.BufferedInputStream;
|
||||
import java.io.File;
|
||||
import java.io.FileInputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.nio.channels.Channels;
|
||||
import java.nio.channels.ReadableByteChannel;
|
||||
import java.util.Iterator;
|
||||
import java.util.stream.Stream;
|
||||
|
||||
/**
|
||||
* Helpers for reading CommitLog records from a file.
|
||||
*
|
||||
* <p>This class is adapted from {@link RestoreCommitLogsAction}, and will be used in the initial
|
||||
* population of the Cloud SQL database.
|
||||
*/
|
||||
public final class CommitLogImports {
|
||||
|
||||
private CommitLogImports() {}
|
||||
|
||||
/**
|
||||
* Returns entities in an {@code inputStream} (from a single CommitLog file) as an {@link
|
||||
* ImmutableList} of {@link VersionedEntity} records. Upon completion the {@code inputStream} is
|
||||
* closed.
|
||||
*
|
||||
* <p>The returned list may be empty, since CommitLogs are written at fixed intervals regardless
|
||||
* if actual changes exist.
|
||||
*
|
||||
* <p>A CommitLog file starts with a {@link CommitLogCheckpoint}, followed by (repeated)
|
||||
* subsequences of [{@link CommitLogManifest}, [{@link CommitLogMutation}] ...]. Each subsequence
|
||||
* represents the changes in one transaction. The {@code CommitLogManifest} contains deleted
|
||||
* entity keys, whereas each {@code CommitLogMutation} contains one whole entity.
|
||||
*/
|
||||
public static ImmutableList<VersionedEntity> loadEntities(InputStream inputStream) {
|
||||
try (AppEngineEnvironment appEngineEnvironment = new AppEngineEnvironment();
|
||||
InputStream input = new BufferedInputStream(inputStream)) {
|
||||
Iterator<ImmutableObject> commitLogs = createDeserializingIterator(input);
|
||||
checkState(commitLogs.hasNext());
|
||||
checkState(commitLogs.next() instanceof CommitLogCheckpoint);
|
||||
|
||||
return Streams.stream(commitLogs)
|
||||
.map(
|
||||
e ->
|
||||
e instanceof CommitLogManifest
|
||||
? VersionedEntity.fromManifest((CommitLogManifest) e)
|
||||
: Stream.of(VersionedEntity.fromMutation((CommitLogMutation) e)))
|
||||
.flatMap(s -> s)
|
||||
.collect(ImmutableList.toImmutableList());
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
/** Covenience method that adapts {@link #loadEntities(InputStream)} to a {@link File}. */
|
||||
public static ImmutableList<VersionedEntity> loadEntities(File commitLogFile) {
|
||||
try {
|
||||
return loadEntities(new FileInputStream(commitLogFile));
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Covenience method that adapts {@link #loadEntities(InputStream)} to a {@link
|
||||
* ReadableByteChannel}.
|
||||
*/
|
||||
public static ImmutableList<VersionedEntity> loadEntities(ReadableByteChannel channel) {
|
||||
return loadEntities(Channels.newInputStream(channel));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
// Copyright 2020 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.backup;
|
||||
|
||||
import com.google.appengine.api.datastore.Entity;
|
||||
import com.google.appengine.api.datastore.EntityTranslator;
|
||||
import com.google.appengine.api.datastore.Key;
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.auto.value.extension.memoized.Memoized;
|
||||
import google.registry.model.ofy.CommitLogManifest;
|
||||
import google.registry.model.ofy.CommitLogMutation;
|
||||
import java.io.Serializable;
|
||||
import java.util.Arrays;
|
||||
import java.util.Optional;
|
||||
import java.util.stream.Stream;
|
||||
import javax.annotation.Nullable;
|
||||
|
||||
/**
|
||||
* A Datastore {@link Entity Entity's} timestamped state.
|
||||
*
|
||||
* <p>For a new or updated Entity, its ProtocolBuffer bytes are stored along with its {@link Key}.
|
||||
* For a deleted entity, only its {@link Key} is stored, and the {@link #entityProtoBytes} is left
|
||||
* as null.
|
||||
*
|
||||
* <p>Note that {@link Optional java.util.Optional} is not serializable, therefore cannot be used as
|
||||
* property type in this class.
|
||||
*/
|
||||
@AutoValue
|
||||
public abstract class VersionedEntity implements Serializable {
|
||||
|
||||
private static final long serialVersionUID = 1L;
|
||||
|
||||
public abstract long commitTimeMills();
|
||||
|
||||
/** The {@link Key} of the {@link Entity}. */
|
||||
public abstract Key key();
|
||||
|
||||
/** Serialized form of the {@link Entity}. This property is {@code null} for a deleted Entity. */
|
||||
@Nullable
|
||||
abstract ImmutableBytes entityProtoBytes();
|
||||
|
||||
@Memoized
|
||||
public Optional<Entity> getEntity() {
|
||||
return Optional.ofNullable(entityProtoBytes())
|
||||
.map(ImmutableBytes::getBytes)
|
||||
.map(EntityTranslator::createFromPbBytes);
|
||||
}
|
||||
|
||||
public boolean isDelete() {
|
||||
return entityProtoBytes() == null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Converts deleted entity keys in {@code manifest} into a {@link Stream} of {@link
|
||||
* VersionedEntity VersionedEntities}. See {@link CommitLogImports#loadEntities} for more
|
||||
* information.
|
||||
*/
|
||||
public static Stream<VersionedEntity> fromManifest(CommitLogManifest manifest) {
|
||||
long commitTimeMillis = manifest.getCommitTime().getMillis();
|
||||
return manifest.getDeletions().stream()
|
||||
.map(com.googlecode.objectify.Key::getRaw)
|
||||
.map(key -> builder().commitTimeMills(commitTimeMillis).key(key).build());
|
||||
}
|
||||
|
||||
/* Converts a {@link CommitLogMutation} to a {@link VersionedEntity}. */
|
||||
public static VersionedEntity fromMutation(CommitLogMutation mutation) {
|
||||
return from(
|
||||
com.googlecode.objectify.Key.create(mutation).getParent().getId(),
|
||||
mutation.getEntityProtoBytes());
|
||||
}
|
||||
|
||||
public static VersionedEntity from(long commitTimeMillis, byte[] entityProtoBytes) {
|
||||
return builder()
|
||||
.entityProtoBytes(entityProtoBytes)
|
||||
.key(EntityTranslator.createFromPbBytes(entityProtoBytes).getKey())
|
||||
.commitTimeMills(commitTimeMillis)
|
||||
.build();
|
||||
}
|
||||
|
||||
static Builder builder() {
|
||||
return new AutoValue_VersionedEntity.Builder();
|
||||
}
|
||||
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
|
||||
public abstract Builder commitTimeMills(long commitTimeMillis);
|
||||
|
||||
abstract Builder entityProtoBytes(ImmutableBytes bytes);
|
||||
|
||||
public abstract Builder key(Key key);
|
||||
|
||||
public abstract VersionedEntity build();
|
||||
|
||||
public Builder entityProtoBytes(byte[] bytes) {
|
||||
return entityProtoBytes(new ImmutableBytes(bytes));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps a byte array and prevents it from being modified by its original owner.
|
||||
*
|
||||
* <p>While this class seems an overkill, it exists for two reasons:
|
||||
*
|
||||
* <ul>
|
||||
* <li>It is easier to override the {@link #equals} method here (for value-equivalence check)
|
||||
* than to override the AutoValue-generated {@code equals} method.
|
||||
* <li>To appease the style checker, which forbids arrays as AutoValue property.
|
||||
* </ul>
|
||||
*/
|
||||
static final class ImmutableBytes implements Serializable {
|
||||
|
||||
private static final long serialVersionUID = 1L;
|
||||
|
||||
private final byte[] bytes;
|
||||
|
||||
ImmutableBytes(byte[] bytes) {
|
||||
this.bytes = Arrays.copyOf(bytes, bytes.length);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the saved byte array. Invocation is restricted to trusted callers, who must not
|
||||
* modify the array.
|
||||
*/
|
||||
byte[] getBytes() {
|
||||
return bytes;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean equals(Object o) {
|
||||
if (this == o) {
|
||||
return true;
|
||||
}
|
||||
if (!(o instanceof ImmutableBytes)) {
|
||||
return false;
|
||||
}
|
||||
ImmutableBytes that = (ImmutableBytes) o;
|
||||
// Do not use Objects.equals, which checks reference identity instead of data in array.
|
||||
return Arrays.equals(bytes, that.bytes);
|
||||
}
|
||||
|
||||
@Override
|
||||
public int hashCode() {
|
||||
// Do not use Objects.hashCode, which hashes the reference, not the data in array.
|
||||
return Arrays.hashCode(bytes);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -17,8 +17,7 @@ package google.registry.beam.initsql;
|
||||
import static com.google.common.base.Preconditions.checkArgument;
|
||||
import static com.google.common.base.Strings.isNullOrEmpty;
|
||||
|
||||
import java.util.regex.Matcher;
|
||||
import java.util.regex.Pattern;
|
||||
import org.joda.time.DateTime;
|
||||
|
||||
/**
|
||||
* Helpers for determining the fully qualified paths to Nomulus backup files. A backup consists of a
|
||||
@@ -30,13 +29,9 @@ public final class BackupPaths {
|
||||
|
||||
private static final String WILDCARD_CHAR = "*";
|
||||
private static final String EXPORT_PATTERN_TEMPLATE = "%s/all_namespaces/kind_%s/input-%s";
|
||||
/**
|
||||
* Regex pattern that captures the kind string in a file name. Datastore places no restrictions on
|
||||
* what characters may be used in a kind string.
|
||||
*/
|
||||
private static final String FILENAME_TO_KIND_REGEX = ".+/all_namespaces/kind_(.+)/input-.+";
|
||||
|
||||
private static final Pattern FILENAME_TO_KIND_PATTERN = Pattern.compile(FILENAME_TO_KIND_REGEX);
|
||||
public static final String COMMIT_LOG_NAME_PREFIX = "commit_diff_until_";
|
||||
private static final String COMMIT_LOG_PATTERN_TEMPLATE = "%s/" + COMMIT_LOG_NAME_PREFIX + "*";
|
||||
|
||||
/**
|
||||
* Returns a regex pattern that matches all Datastore export files of a given {@code kind}.
|
||||
@@ -65,22 +60,15 @@ public final class BackupPaths {
|
||||
return String.format(EXPORT_PATTERN_TEMPLATE, exportDir, kind, Integer.toString(shard));
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the 'kind' of entity stored in a file based on the file name.
|
||||
*
|
||||
* <p>This method poses low risk and greatly simplifies the implementation of some transforms in
|
||||
* {@link ExportLoadingTransforms}.
|
||||
*
|
||||
* @see ExportLoadingTransforms
|
||||
*/
|
||||
public static String getKindFromFileName(String fileName) {
|
||||
public static String getCommitLogFileNamePattern(String commitLogDir) {
|
||||
return String.format(COMMIT_LOG_PATTERN_TEMPLATE, commitLogDir);
|
||||
}
|
||||
|
||||
/** Gets the Commit timestamp from a CommitLog file name. */
|
||||
public static DateTime getCommitLogTimestamp(String fileName) {
|
||||
checkArgument(!isNullOrEmpty(fileName), "Null or empty fileName.");
|
||||
Matcher matcher = FILENAME_TO_KIND_PATTERN.matcher(fileName);
|
||||
checkArgument(
|
||||
matcher.matches(),
|
||||
"Illegal file name %s, should match %s.",
|
||||
fileName,
|
||||
FILENAME_TO_KIND_REGEX);
|
||||
return matcher.group(1);
|
||||
int start = fileName.lastIndexOf(COMMIT_LOG_NAME_PREFIX);
|
||||
checkArgument(start >= 0, "Illegal file name %s.", fileName);
|
||||
return DateTime.parse(fileName.substring(start + COMMIT_LOG_NAME_PREFIX.length()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
// Copyright 2020 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.beam.initsql;
|
||||
|
||||
import static com.google.common.base.Preconditions.checkArgument;
|
||||
import static com.google.common.base.Preconditions.checkNotNull;
|
||||
import static google.registry.beam.initsql.BackupPaths.getCommitLogFileNamePattern;
|
||||
import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
|
||||
import static google.registry.beam.initsql.Transforms.processFiles;
|
||||
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
|
||||
|
||||
import google.registry.backup.CommitLogImports;
|
||||
import google.registry.backup.VersionedEntity;
|
||||
import java.io.IOException;
|
||||
import org.apache.beam.sdk.coders.StringUtf8Coder;
|
||||
import org.apache.beam.sdk.io.FileIO.ReadableFile;
|
||||
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.PTransform;
|
||||
import org.apache.beam.sdk.transforms.ParDo;
|
||||
import org.apache.beam.sdk.values.PBegin;
|
||||
import org.apache.beam.sdk.values.PCollection;
|
||||
import org.joda.time.DateTime;
|
||||
|
||||
/**
|
||||
* {@link org.apache.beam.sdk.transforms.PTransform Pipeline transforms} for loading from Nomulus
|
||||
* CommitLog files. They are all part of a transformation that loads raw records from a sequence of
|
||||
* Datastore CommitLog files, and are broken apart for testing.
|
||||
*/
|
||||
public class CommitLogTransforms {
|
||||
|
||||
/**
|
||||
* Returns a {@link PTransform transform} that can generate a collection of patterns that match
|
||||
* all Datastore CommitLog files.
|
||||
*/
|
||||
public static PTransform<PBegin, PCollection<String>> getCommitLogFilePatterns(
|
||||
String commitLogDir) {
|
||||
return Create.of(getCommitLogFileNamePattern(commitLogDir)).withCoder(StringUtf8Coder.of());
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns files with timestamps between {@code fromTime} (inclusive) and {@code endTime}
|
||||
* (exclusive).
|
||||
*/
|
||||
public static PTransform<PCollection<? extends String>, PCollection<String>>
|
||||
filterCommitLogsByTime(DateTime fromTime, DateTime toTime) {
|
||||
checkNotNull(fromTime, "fromTime");
|
||||
checkNotNull(toTime, "toTime");
|
||||
checkArgument(
|
||||
fromTime.isBefore(toTime),
|
||||
"Invalid time range: fromTime (%s) is before endTime (%s)",
|
||||
fromTime,
|
||||
toTime);
|
||||
return ParDo.of(new FilterCommitLogFileByTime(fromTime, toTime));
|
||||
}
|
||||
|
||||
/** Returns a {@link PTransform} from file {@link Metadata} to {@link VersionedEntity}. */
|
||||
public static PTransform<PCollection<Metadata>, PCollection<VersionedEntity>>
|
||||
loadCommitLogsFromFiles() {
|
||||
return processFiles(new LoadOneCommitLogsFile());
|
||||
}
|
||||
|
||||
static class FilterCommitLogFileByTime extends DoFn<String, String> {
|
||||
private final DateTime fromTime;
|
||||
private final DateTime toTime;
|
||||
|
||||
public FilterCommitLogFileByTime(DateTime fromTime, DateTime toTime) {
|
||||
this.fromTime = fromTime;
|
||||
this.toTime = toTime;
|
||||
}
|
||||
|
||||
@ProcessElement
|
||||
public void processElement(@Element String fileName, OutputReceiver<String> out) {
|
||||
DateTime timestamp = getCommitLogTimestamp(fileName);
|
||||
if (isBeforeOrAt(fromTime, timestamp) && timestamp.isBefore(toTime)) {
|
||||
out.output(fileName);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads a CommitLog file and converts its content into {@link VersionedEntity VersionedEntities}.
|
||||
*/
|
||||
static class LoadOneCommitLogsFile extends DoFn<ReadableFile, VersionedEntity> {
|
||||
|
||||
@ProcessElement
|
||||
public void processElement(@Element ReadableFile file, OutputReceiver<VersionedEntity> out) {
|
||||
try {
|
||||
CommitLogImports.loadEntities(file.open()).forEach(out::output);
|
||||
} catch (IOException e) {
|
||||
// Let the pipeline retry the whole file.
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -15,38 +15,26 @@
|
||||
package google.registry.beam.initsql;
|
||||
|
||||
import static google.registry.beam.initsql.BackupPaths.getExportFileNamePattern;
|
||||
import static google.registry.beam.initsql.BackupPaths.getKindFromFileName;
|
||||
import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
|
||||
import static org.apache.beam.sdk.values.TypeDescriptors.strings;
|
||||
import static google.registry.beam.initsql.Transforms.processFiles;
|
||||
|
||||
import com.google.common.collect.ImmutableList;
|
||||
import google.registry.backup.VersionedEntity;
|
||||
import google.registry.tools.LevelDbLogReader;
|
||||
import java.io.IOException;
|
||||
import java.util.Collection;
|
||||
import org.apache.beam.sdk.coders.StringUtf8Coder;
|
||||
import org.apache.beam.sdk.io.Compression;
|
||||
import org.apache.beam.sdk.io.FileIO;
|
||||
import org.apache.beam.sdk.io.FileIO.ReadableFile;
|
||||
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.MapElements;
|
||||
import org.apache.beam.sdk.transforms.PTransform;
|
||||
import org.apache.beam.sdk.transforms.ParDo;
|
||||
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.TypeDescriptor;
|
||||
|
||||
/**
|
||||
* {@link PTransform Pipeline transforms} for loading from Datastore export files. They are all part
|
||||
* of a transformation that loads raw records from a Datastore export, and are broken apart for
|
||||
* testing.
|
||||
*
|
||||
* <p>We drop the 'kind' information in {@link #getDatastoreExportFilePatterns} and recover it later
|
||||
* using the file paths. Although we could have kept it by passing around {@link KV key-value
|
||||
* pairs}, the code would be more complicated, especially in {@link #loadDataFromFiles()}.
|
||||
*/
|
||||
public class ExportLoadingTransforms {
|
||||
|
||||
@@ -63,51 +51,28 @@ public class ExportLoadingTransforms {
|
||||
.withCoder(StringUtf8Coder.of());
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a {@link PTransform} from file name patterns to file {@link Metadata Metadata records}.
|
||||
*/
|
||||
public static PTransform<PCollection<String>, PCollection<Metadata>> getFilesByPatterns() {
|
||||
return new PTransform<PCollection<String>, PCollection<Metadata>>() {
|
||||
@Override
|
||||
public PCollection<Metadata> expand(PCollection<String> input) {
|
||||
return input.apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.DISALLOW));
|
||||
}
|
||||
};
|
||||
/** Returns a {@link PTransform} from file {@link Metadata} to {@link VersionedEntity}. */
|
||||
public static PTransform<PCollection<Metadata>, PCollection<VersionedEntity>>
|
||||
loadExportDataFromFiles() {
|
||||
return processFiles(new LoadOneExportShard());
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a {@link PTransform} from file {@link Metadata} to {@link KV Key-Value pairs} with
|
||||
* entity 'kind' as key and raw record as value.
|
||||
*/
|
||||
public static PTransform<PCollection<Metadata>, PCollection<KV<String, byte[]>>>
|
||||
loadDataFromFiles() {
|
||||
return new PTransform<PCollection<Metadata>, PCollection<KV<String, byte[]>>>() {
|
||||
@Override
|
||||
public PCollection<KV<String, byte[]>> expand(PCollection<Metadata> input) {
|
||||
return input
|
||||
.apply(FileIO.readMatches().withCompression(Compression.UNCOMPRESSED))
|
||||
.apply(
|
||||
MapElements.into(kvs(strings(), TypeDescriptor.of(ReadableFile.class)))
|
||||
.via(file -> KV.of(getKindFromFileName(file.getMetadata().toString()), file)))
|
||||
.apply("Load One LevelDb File", ParDo.of(new LoadOneFile()));
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads a LevelDb file and converts each raw record into a {@link KV pair} of kind and bytes.
|
||||
* Reads a LevelDb file and converts each raw record into a {@link VersionedEntity}. All such
|
||||
* entities use {@link Long#MIN_VALUE} as timestamp, so that they go before data from CommitLogs.
|
||||
*
|
||||
* <p>LevelDb files are not seekable because a large object may span multiple blocks. If a
|
||||
* sequential read fails, the file needs to be retried from the beginning.
|
||||
*/
|
||||
private static class LoadOneFile extends DoFn<KV<String, ReadableFile>, KV<String, byte[]>> {
|
||||
private static class LoadOneExportShard extends DoFn<ReadableFile, VersionedEntity> {
|
||||
|
||||
private static final long TIMESTAMP = Long.MIN_VALUE;
|
||||
|
||||
@ProcessElement
|
||||
public void processElement(
|
||||
@Element KV<String, ReadableFile> kv, OutputReceiver<KV<String, byte[]>> output) {
|
||||
public void processElement(@Element ReadableFile file, OutputReceiver<VersionedEntity> output) {
|
||||
try {
|
||||
LevelDbLogReader.from(kv.getValue().open())
|
||||
.forEachRemaining(record -> output.output(KV.of(kv.getKey(), record)));
|
||||
LevelDbLogReader.from(file.open())
|
||||
.forEachRemaining(record -> output.output(VersionedEntity.from(TIMESTAMP, record)));
|
||||
} catch (IOException e) {
|
||||
// Let the pipeline retry the whole file.
|
||||
throw new RuntimeException(e);
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
// Copyright 2020 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.beam.initsql;
|
||||
|
||||
import google.registry.backup.VersionedEntity;
|
||||
import org.apache.beam.sdk.io.Compression;
|
||||
import org.apache.beam.sdk.io.FileIO;
|
||||
import org.apache.beam.sdk.io.FileIO.ReadableFile;
|
||||
import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
|
||||
import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
|
||||
import org.apache.beam.sdk.transforms.DoFn;
|
||||
import org.apache.beam.sdk.transforms.PTransform;
|
||||
import org.apache.beam.sdk.transforms.ParDo;
|
||||
import org.apache.beam.sdk.values.PCollection;
|
||||
|
||||
/**
|
||||
* Common {@link PTransform pipeline transforms} used in pipelines that load from both Datastore
|
||||
* export files and Nomulus CommitLog files.
|
||||
*/
|
||||
public class Transforms {
|
||||
|
||||
/**
|
||||
* Returns a {@link PTransform} from file name patterns to file {@link Metadata Metadata records}.
|
||||
*/
|
||||
public static PTransform<PCollection<String>, PCollection<Metadata>> getFilesByPatterns() {
|
||||
return new PTransform<PCollection<String>, PCollection<Metadata>>() {
|
||||
@Override
|
||||
public PCollection<Metadata> expand(PCollection<String> input) {
|
||||
return input.apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.DISALLOW));
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a {@link PTransform} from file {@link Metadata} to {@link VersionedEntity} using
|
||||
* caller-provided {@code transformer}.
|
||||
*/
|
||||
public static PTransform<PCollection<Metadata>, PCollection<VersionedEntity>> processFiles(
|
||||
DoFn<ReadableFile, VersionedEntity> transformer) {
|
||||
return new PTransform<PCollection<Metadata>, PCollection<VersionedEntity>>() {
|
||||
@Override
|
||||
public PCollection<VersionedEntity> expand(PCollection<Metadata> input) {
|
||||
return input
|
||||
.apply(FileIO.readMatches().withCompression(Compression.UNCOMPRESSED))
|
||||
.apply(transformer.getClass().getSimpleName(), ParDo.of(transformer));
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user