zydronium/google-nomulus
1
0
Fork 0
mirror of https://github.com/google/nomulus.git synced 2025-04-29 19:47:51 +02:00
Code Issues Projects Releases Packages Wiki Activity

Delete DatastoreTM and most other references to Datastore (#1681)

Browse source 
This includes:
- deletion of helper DB methods in tests
- deletion of various old Datastore-only classes and removal of any
  endpoints
- removal of the dual-database test concept
- removal of 'ofy' from the AppEngineExtension
This commit is contained in:
gbrodman 2022-07-01 13:33:38 -04:00 • committed by GitHub
parent cb1957f01a
commit f62732547f
336 changed files with 3452 additions and 12054 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
core/build.gradle
View file

@ -705,14 +705,6 @@ createToolTask(
createToolTask(
'jpaDemoPipeline', 'google.registry.beam.common.JpaDemoPipeline')
// Caller must provide projectId, GCP region, runner, and the kinds to delete
// (comma-separated kind names or '*' for all). E.g.:
// nom_build :core:bulkDeleteDatastore --args="--project=domain-registry-crash \
// --region=us-central1 --runner=DataflowRunner --kindsToDelete=*"
createToolTask(
'bulkDeleteDatastore',
'google.registry.beam.datastore.BulkDeleteDatastorePipeline')
project.tasks.create('generateSqlSchema', JavaExec) {
classpath = sourceSets.nonprod.runtimeClasspath
main = 'google.registry.tools.DevTool'
@ -757,11 +749,6 @@ createUberJar(
// User should install gcloud and login to GCP before invoking this tasks.
if (environment == 'alpha') {
def pipelines = [
bulkDeleteDatastore:
[
mainClass: 'google.registry.beam.datastore.BulkDeleteDatastorePipeline',
metaData : 'google/registry/beam/bulk_delete_datastore_pipeline_metadata.json'
],
spec11 :
[
mainClass: 'google.registry.beam.spec11.Spec11Pipeline',

15
core/src/main/java/google/registry/batch/DeleteExpiredDomainsAction.java
View file

@ -18,7 +18,6 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
import static google.registry.flows.FlowUtils.marshalWithLenientRetry;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.DateTimeUtils.END_OF_TIME;
import static google.registry.util.ResourceUtils.readResourceUtf8;
import static java.nio.charset.StandardCharsets.UTF_8;
@ -129,15 +128,13 @@ public class DeleteExpiredDomainsAction implements Runnable {
logger.atInfo().log(
"Deleting non-renewing domains with autorenew end times up through %s.", runTime);
// Note: in Datastore, this query is (and must be) non-transactional, and thus, is only
// eventually consistent.
ImmutableList<DomainBase> domainsToDelete =
transactIfJpaTm(
() ->
tm().createQueryComposer(DomainBase.class)
.where("autorenewEndTime", Comparator.LTE, runTime)
.where("deletionTime", Comparator.EQ, END_OF_TIME)
.list());
tm().transact(
() ->
tm().createQueryComposer(DomainBase.class)
.where("autorenewEndTime", Comparator.LTE, runTime)
.where("deletionTime", Comparator.EQ, END_OF_TIME)
.list());
if (domainsToDelete.isEmpty()) {
logger.atInfo().log("Found 0 domains to delete.");
response.setPayload("Found 0 domains to delete.");

2
core/src/main/java/google/registry/batch/DeleteProberDataAction.java
View file

@ -42,7 +42,6 @@ import google.registry.model.domain.DomainHistory;
import google.registry.model.tld.Registry.TldType;
import google.registry.request.Action;
import google.registry.request.Parameter;
import google.registry.request.Response;
import google.registry.request.auth.Auth;
import java.util.concurrent.atomic.AtomicInteger;
import javax.inject.Inject;
@ -111,7 +110,6 @@ public class DeleteProberDataAction implements Runnable {
@Config("registryAdminClientId")
String registryAdminRegistrarId;
@Inject Response response;
@Inject DeleteProberDataAction() {}
@Override

11
core/src/main/java/google/registry/batch/ExpandRecurringBillingEventsAction.java
View file

@ -25,7 +25,6 @@ import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_AUTORENEW
import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.CollectionUtils.union;
import static google.registry.util.DateTimeUtils.START_OF_TIME;
import static google.registry.util.DateTimeUtils.earliestOf;
@ -96,11 +95,11 @@ public class ExpandRecurringBillingEventsAction implements Runnable {
public void run() {
DateTime executeTime = clock.nowUtc();
DateTime persistedCursorTime =
transactIfJpaTm(
() ->
tm().loadByKeyIfPresent(Cursor.createGlobalVKey(RECURRING_BILLING))
.orElse(Cursor.createGlobal(RECURRING_BILLING, START_OF_TIME))
.getCursorTime());
tm().transact(
() ->
tm().loadByKeyIfPresent(Cursor.createGlobalVKey(RECURRING_BILLING))
.orElse(Cursor.createGlobal(RECURRING_BILLING, START_OF_TIME))
.getCursorTime());
DateTime cursorTime = cursorTimeParam.orElse(persistedCursorTime);
checkArgument(
cursorTime.isBefore(executeTime), "Cursor time must be earlier than execution time.");

5
core/src/main/java/google/registry/batch/RelockDomainAction.java
View file

@ -120,10 +120,7 @@ public class RelockDomainAction implements Runnable {
* for more details on retry behavior. */
response.setStatus(SC_NO_CONTENT);
response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
// nb: DomainLockUtils relies on the JPA transaction being the outermost transaction
// if we have Datastore as the primary DB (if SQL is the primary DB, it's irrelevant)
jpaTm().transact(() -> tm().transact(this::relockDomain));
tm().transact(this::relockDomain);
}
private void relockDomain() {

122
core/src/main/java/google/registry/batch/WipeoutDatastoreAction.java
View file

@ -1,122 +0,0 @@
// Copyright 2021 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.batch;
import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
import static google.registry.beam.BeamUtils.createJobName;
import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
import static javax.servlet.http.HttpServletResponse.SC_OK;
import com.google.api.services.dataflow.Dataflow;
import com.google.api.services.dataflow.model.LaunchFlexTemplateParameter;
import com.google.api.services.dataflow.model.LaunchFlexTemplateRequest;
import com.google.api.services.dataflow.model.LaunchFlexTemplateResponse;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.flogger.FluentLogger;
import google.registry.config.RegistryConfig.Config;
import google.registry.config.RegistryEnvironment;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.request.Action;
import google.registry.request.Response;
import google.registry.request.auth.Auth;
import google.registry.util.Clock;
import javax.inject.Inject;
/**
* Wipes out all Cloud Datastore data in a Nomulus GCP environment.
*
* <p>This class is created for the QA environment, where migration testing with production data
* will happen. A regularly scheduled wipeout is a prerequisite to using production data there.
*/
@Action(
service = Action.Service.BACKEND,
path = "/_dr/task/wipeOutDatastore",
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
@DeleteAfterMigration
public class WipeoutDatastoreAction implements Runnable {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
private static final String PIPELINE_NAME = "bulk_delete_datastore_pipeline";
private static final ImmutableSet<RegistryEnvironment> FORBIDDEN_ENVIRONMENTS =
ImmutableSet.of(RegistryEnvironment.PRODUCTION, RegistryEnvironment.SANDBOX);
private final String projectId;
private final String jobRegion;
private final Response response;
private final Dataflow dataflow;
private final String stagingBucketUrl;
private final Clock clock;
@Inject
WipeoutDatastoreAction(
@Config("projectId") String projectId,
@Config("defaultJobRegion") String jobRegion,
@Config("beamStagingBucketUrl") String stagingBucketUrl,
Clock clock,
Response response,
Dataflow dataflow) {
this.projectId = projectId;
this.jobRegion = jobRegion;
this.stagingBucketUrl = stagingBucketUrl;
this.clock = clock;
this.response = response;
this.dataflow = dataflow;
}
@Override
public void run() {
response.setContentType(PLAIN_TEXT_UTF_8);
if (FORBIDDEN_ENVIRONMENTS.contains(RegistryEnvironment.get())) {
response.setStatus(SC_FORBIDDEN);
response.setPayload("Wipeout is not allowed in " + RegistryEnvironment.get());
return;
}
try {
LaunchFlexTemplateParameter parameters =
new LaunchFlexTemplateParameter()
.setJobName(createJobName("bulk-delete-datastore-", clock))
.setContainerSpecGcsPath(
String.format("%s/%s_metadata.json", stagingBucketUrl, PIPELINE_NAME))
.setParameters(
ImmutableMap.of(
"kindsToDelete",
"*",
"registryEnvironment",
RegistryEnvironment.get().name()));
LaunchFlexTemplateResponse launchResponse =
dataflow
.projects()
.locations()
.flexTemplates()
.launch(
projectId,
jobRegion,
new LaunchFlexTemplateRequest().setLaunchParameter(parameters))
.execute();
response.setStatus(SC_OK);
response.setPayload("Launched " + launchResponse.getJob().getName());
} catch (Exception e) {
String msg = String.format("Failed to launch %s.", PIPELINE_NAME);
logger.atSevere().withCause(e).log(msg);
response.setStatus(SC_INTERNAL_SERVER_ERROR);
response.setPayload(msg);
}
}
}

336
core/src/main/java/google/registry/beam/datastore/BulkDeleteDatastorePipeline.java
View file

@ -1,336 +0,0 @@
// 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.datastore;
import static com.google.common.base.Preconditions.checkState;
import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
import static org.apache.beam.sdk.values.TypeDescriptors.strings;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedSet;
import com.google.common.flogger.FluentLogger;
import com.google.datastore.v1.Entity;
import google.registry.config.RegistryEnvironment;
import google.registry.model.annotations.DeleteAfterMigration;
import java.util.Iterator;
import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.io.gcp.datastore.DatastoreIO;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.Validation;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
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.Reshuffle;
import org.apache.beam.sdk.transforms.View;
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.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
/**
* A BEAM pipeline that deletes Datastore entities in bulk.
*
* <p>This pipeline provides an alternative to the <a
* href="https://cloud.google.com/datastore/docs/bulk-delete">GCP builtin template</a> that performs
* the same task. It solves the following performance and usability problems in the builtin
* template:
*
* <ul>
* <li>When deleting all data (by using the {@code select __key__} or {@code select *} queries),
* the builtin template cannot parallelize the query, therefore has to query with a single
* worker.
* <li>When deleting all data, the builtin template also attempts to delete Datastore internal
* tables which would cause permission-denied errors, which in turn MAY cause the pipeline to
* abort before all data has been deleted.
* <li>With the builtin template, it is possible to delete multiple entity types in one pipeline
* ONLY if the user can come up with a single literal query that covers all of them. This is
* not the case with most Nomulus entity types.
* </ul>
*
* <p>A user of this pipeline must specify the types of entities to delete using the {@code
* --kindsToDelete} command line argument. To delete specific entity types, give a comma-separated
* string of their kind names; to delete all data, give {@code "*"}.
*
* <p>When deleting all data, it is recommended for the user to specify the number of user entity
* types in the Datastore using the {@code --numOfKindsHint} argument. If the default value for this
* parameter is too low, performance will suffer.
*/
@DeleteAfterMigration
public class BulkDeleteDatastorePipeline {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
// This tool is not for use in our critical projects.
private static final ImmutableSet<String> FORBIDDEN_PROJECTS =
ImmutableSet.of("domain-registry", "domain-registry-sandbox");
private final BulkDeletePipelineOptions options;
BulkDeleteDatastorePipeline(BulkDeletePipelineOptions options) {
this.options = options;
}
public void run() {
Pipeline pipeline = Pipeline.create(options);
setupPipeline(pipeline);
pipeline.run();
}
@SuppressWarnings("deprecation") // org.apache.beam.sdk.transforms.Reshuffle
private void setupPipeline(Pipeline pipeline) {
checkState(
!FORBIDDEN_PROJECTS.contains(options.getProject()),
"Bulk delete is forbidden in %s",
options.getProject());
// Pre-allocated tags to label entities by kind. In the case of delete-all, we must use a guess.
TupleTagList deletionTags;
PCollection<String> kindsToDelete;
if (options.getKindsToDelete().equals("*")) {
deletionTags = getDeletionTags(options.getNumOfKindsHint());
kindsToDelete =
pipeline.apply("DiscoverEntityKinds", discoverEntityKinds(options.getProject()));
} else {
ImmutableList<String> kindsToDeleteParam = parseKindsToDelete(options);
checkState(
!kindsToDeleteParam.contains("*"),
"The --kindsToDelete argument should not contain both '*' and other kinds.");
deletionTags = getDeletionTags(kindsToDeleteParam.size());
kindsToDelete = pipeline.apply("UseProvidedKinds", Create.of(kindsToDeleteParam));
}
// Map each kind to a tag. The "SplitByKind" stage below will group entities by kind using
// this mapping. In practice, this has been effective at avoiding entity group contentions.
PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping =
mapKindsToDeletionTags(kindsToDelete, deletionTags).apply("GetKindsToTagMap", View.asMap());
PCollectionTuple entities =
kindsToDelete
.apply("GenerateQueries", ParDo.of(new GenerateQueries()))
.apply("ReadEntities", DatastoreV1.read().withProjectId(options.getProject()))
.apply(
"SplitByKind",
ParDo.of(new SplitEntities(kindToTagMapping))
.withSideInputs(kindToTagMapping)
.withOutputTags(getOneDeletionTag("placeholder"), deletionTags));
for (TupleTag<?> tag : deletionTags.getAll()) {
entities
.get((TupleTag<Entity>) tag)
// Reshuffle calls GroupByKey which is one way to trigger load rebalance in the pipeline.
// Using the deprecated "Reshuffle" for convenience given the short life of this tool.
.apply("RebalanceLoad", Reshuffle.viaRandomKey())
.apply(
"DeleteEntities_" + tag.getId(),
DatastoreIO.v1().deleteEntity().withProjectId(options.getProject()));
}
}
private static String toKeyOnlyQueryForKind(String kind) {
return "select __key__ from `" + kind + "`";
}
/**
* Returns a {@link TupleTag} that retains the generic type parameter and may be used in a
* multi-output {@link ParDo} (e.g. {@link SplitEntities}).
*
* <p>This method is NOT needed in tests when creating tags for assertions. Simply create them
* with {@code new TupleTag<Entity>(String)}.
*/
@VisibleForTesting
static TupleTag<Entity> getOneDeletionTag(String id) {
// The trailing {} is needed to retain generic param type.
return new TupleTag<Entity>(id) {};
}
@VisibleForTesting
static ImmutableList<String> parseKindsToDelete(BulkDeletePipelineOptions options) {
return ImmutableList.copyOf(
Splitter.on(",").omitEmptyStrings().trimResults().split(options.getKindsToDelete().trim()));
}
/**
* Returns a list of {@code n} {@link TupleTag TupleTags} numbered from {@code 0} to {@code n-1}.
*/
@VisibleForTesting
static TupleTagList getDeletionTags(int n) {
ImmutableList.Builder<TupleTag<?>> builder = new ImmutableList.Builder<>();
for (int i = 0; i < n; i++) {
builder.add(getOneDeletionTag(String.valueOf(i)));
}
return TupleTagList.of(builder.build());
}
/** Returns a {@link PTransform} that finds all entity kinds in Datastore. */
@VisibleForTesting
static PTransform<PBegin, PCollection<String>> discoverEntityKinds(String project) {
return new PTransform<PBegin, PCollection<String>>() {
@Override
public PCollection<String> expand(PBegin input) {
// Use the __kind__ table to discover entity kinds. Data in the more informational
// __Stat_Kind__ table may be up to 48-hour stale.
return input
.apply(
"LoadEntityMetaData",
DatastoreIO.v1()
.read()
.withProjectId(project)
.withLiteralGqlQuery("select * from __kind__"))
.apply(
"GetKindNames",
ParDo.of(
new DoFn<Entity, String>() {
@ProcessElement
public void processElement(
@Element Entity entity, OutputReceiver<String> out) {
String kind = entity.getKey().getPath(0).getName();
if (kind.startsWith("_")) {
return;
}
out.output(kind);
}
}));
}
};
}
@VisibleForTesting
static PCollection<KV<String, TupleTag<Entity>>> mapKindsToDeletionTags(
PCollection<String> kinds, TupleTagList tags) {
// The first two stages send all strings in the 'kinds' PCollection to one worker which
// performs the mapping in the last stage.
return kinds
.apply(
"AssignSingletonKeyToKinds",
MapElements.into(kvs(strings(), strings())).via(kind -> KV.of("", kind)))
.apply("GatherKindsIntoCollection", GroupByKey.create())
.apply("MapKindsToTag", ParDo.of(new MapKindsToTags(tags)));
}
/** Transforms each {@code kind} string into a Datastore query for that kind. */
@VisibleForTesting
static class GenerateQueries extends DoFn<String, String> {
@ProcessElement
public void processElement(@Element String kind, OutputReceiver<String> out) {
out.output(toKeyOnlyQueryForKind(kind));
}
}
private static class MapKindsToTags
extends DoFn<KV<String, Iterable<String>>, KV<String, TupleTag<Entity>>> {
private final TupleTagList tupleTags;
MapKindsToTags(TupleTagList tupleTags) {
this.tupleTags = tupleTags;
}
@ProcessElement
public void processElement(
@Element KV<String, Iterable<String>> kv,
OutputReceiver<KV<String, TupleTag<Entity>>> out) {
// Sort kinds so that mapping is deterministic.
ImmutableSortedSet<String> sortedKinds = ImmutableSortedSet.copyOf(kv.getValue());
Iterator<String> kinds = sortedKinds.iterator();
Iterator<TupleTag<?>> tags = tupleTags.getAll().iterator();
while (kinds.hasNext() && tags.hasNext()) {
out.output(KV.of(kinds.next(), (TupleTag<Entity>) tags.next()));
}
if (kinds.hasNext()) {
logger.atWarning().log(
"There are more kinds to delete (%s) than our estimate (%s). "
+ "Performance may suffer.",
sortedKinds.size(), tupleTags.size());
}
// Round robin assignment so that mapping is deterministic
while (kinds.hasNext()) {
tags = tupleTags.getAll().iterator();
while (kinds.hasNext() && tags.hasNext()) {
out.output(KV.of(kinds.next(), (TupleTag<Entity>) tags.next()));
}
}
}
}
/**
* {@link DoFn} that splits one {@link PCollection} of mixed kinds into multiple single-kind
* {@code PCollections}.
*/
@VisibleForTesting
static class SplitEntities extends DoFn<Entity, Entity> {
private final PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping;
SplitEntities(PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping) {
super();
this.kindToTagMapping = kindToTagMapping;
}
@ProcessElement
public void processElement(ProcessContext context) {
Entity entity = context.element();
com.google.datastore.v1.Key entityKey = entity.getKey();
String kind = entityKey.getPath(entityKey.getPathCount() - 1).getKind();
TupleTag<Entity> tag = context.sideInput(kindToTagMapping).get(kind);
context.output(tag, entity);
}
}
public static void main(String[] args) {
BulkDeletePipelineOptions options =
PipelineOptionsFactory.fromArgs(args).withValidation().as(BulkDeletePipelineOptions.class);
BulkDeleteDatastorePipeline pipeline = new BulkDeleteDatastorePipeline(options);
pipeline.run();
System.exit(0);
}
public interface BulkDeletePipelineOptions extends GcpOptions {
@Description("The Registry environment.")
RegistryEnvironment getRegistryEnvironment();
void setRegistryEnvironment(RegistryEnvironment environment);
@Description(
"The Datastore KINDs to be deleted. The format may be:\n"
+ "\t- The list of kinds to be deleted as a comma-separated string, or\n"
+ "\t- '*', which causes all kinds to be deleted.")
@Validation.Required
String getKindsToDelete();
void setKindsToDelete(String kinds);
@Description(
"An estimate of the number of KINDs to be deleted. "
+ "This is recommended if --kindsToDelete is '*' and the default value is too low.")
@Default.Integer(30)
int getNumOfKindsHint();
void setNumOfKindsHint(int numOfKindsHint);
}
}

768
core/src/main/java/google/registry/beam/datastore/DatastoreV1.java
View file

@ -1,768 +0,0 @@
// 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.
// This class is adapted from the Apache BEAM SDK. The original license may
// be found at <a href="https://github.com/apache/beam/blob/master/LICENSE">
// this link</a>.
package google.registry.beam.datastore;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Verify.verify;
import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL;
import static com.google.datastore.v1.PropertyOrder.Direction.DESCENDING;
import static com.google.datastore.v1.QueryResultBatch.MoreResultsType.NOT_FINISHED;
import static com.google.datastore.v1.client.DatastoreHelper.makeAndFilter;
import static com.google.datastore.v1.client.DatastoreHelper.makeFilter;
import static com.google.datastore.v1.client.DatastoreHelper.makeOrder;
import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
import com.google.api.client.http.HttpRequestInitializer;
import com.google.auth.Credentials;
import com.google.auth.http.HttpCredentialsAdapter;
import com.google.auto.value.AutoValue;
import com.google.cloud.hadoop.util.ChainingHttpRequestInitializer;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.flogger.FluentLogger;
import com.google.datastore.v1.Entity;
import com.google.datastore.v1.EntityResult;
import com.google.datastore.v1.GqlQuery;
import com.google.datastore.v1.PartitionId;
import com.google.datastore.v1.Query;
import com.google.datastore.v1.QueryResultBatch;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
import com.google.datastore.v1.client.Datastore;
import com.google.datastore.v1.client.DatastoreException;
import com.google.datastore.v1.client.DatastoreFactory;
import com.google.datastore.v1.client.DatastoreHelper;
import com.google.datastore.v1.client.DatastoreOptions;
import com.google.datastore.v1.client.QuerySplitter;
import com.google.protobuf.Int32Value;
import com.google.rpc.Code;
import google.registry.model.annotations.DeleteAfterMigration;
import java.io.Serializable;
import java.util.List;
import java.util.NoSuchElementException;
import javax.annotation.Nullable;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
import org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.PipelineOptions;
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.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.HasDisplayData;
import org.apache.beam.sdk.util.BackOff;
import org.apache.beam.sdk.util.BackOffUtils;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.joda.time.Duration;
/**
* Contains an adaptation of {@link org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read}. See
* {@link MultiRead} for details.
*/
@DeleteAfterMigration
public class DatastoreV1 {
// A package-private constructor to prevent direct instantiation from outside of this package
DatastoreV1() {}
/**
* Non-retryable errors. See https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes .
*/
private static final ImmutableSet<Code> NON_RETRYABLE_ERRORS =
ImmutableSet.of(
Code.FAILED_PRECONDITION,
Code.INVALID_ARGUMENT,
Code.PERMISSION_DENIED,
Code.UNAUTHENTICATED);
/**
* Returns an empty {@link MultiRead} builder. Configure the source {@code projectId}, {@code
* query}, and optionally {@code namespace} and {@code numQuerySplits} using {@link
* MultiRead#withProjectId}, {@link MultiRead#withNamespace}, {@link
* MultiRead#withNumQuerySplits}.
*/
public static MultiRead read() {
return new AutoValue_DatastoreV1_MultiRead.Builder().setNumQuerySplits(0).build();
}
/**
* A {@link PTransform} that executes every Cloud SQL queries in a {@link PCollection } and reads
* their result rows as {@code Entity} objects.
*
* <p>This class is adapted from {@link org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read}. It
* uses literal GQL queries in the input {@link PCollection} instead of a constant query provided
* to the builder. Only the {@link #expand} method is modified from the original. Everything else
* including comments have been copied verbatim.
*/
@AutoValue
public abstract static class MultiRead
extends PTransform<PCollection<String>, PCollection<Entity>> {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
/** An upper bound on the number of splits for a query. */
public static final int NUM_QUERY_SPLITS_MAX = 50000;
/** A lower bound on the number of splits for a query. */
static final int NUM_QUERY_SPLITS_MIN = 12;
/** Default bundle size of 64MB. */
static final long DEFAULT_BUNDLE_SIZE_BYTES = 64L * 1024L * 1024L;
/**
* Maximum number of results to request per query.
*
* <p>Must be set, or it may result in an I/O error when querying Cloud Datastore.
*/
static final int QUERY_BATCH_LIMIT = 500;
public abstract @Nullable String getProjectId();
public abstract @Nullable String getNamespace();
public abstract int getNumQuerySplits();
public abstract @Nullable String getLocalhost();
@Override
public abstract String toString();
abstract Builder toBuilder();
@AutoValue.Builder
abstract static class Builder {
abstract Builder setProjectId(String projectId);
abstract Builder setNamespace(String namespace);
abstract Builder setNumQuerySplits(int numQuerySplits);
abstract Builder setLocalhost(String localhost);
abstract MultiRead build();
}
/**
* Computes the number of splits to be performed on the given query by querying the estimated
* size from Cloud Datastore.
*/
static int getEstimatedNumSplits(Datastore datastore, Query query, @Nullable String namespace) {
int numSplits;
try {
long estimatedSizeBytes = getEstimatedSizeBytes(datastore, query, namespace);
logger.atInfo().log("Estimated size for the query is %d bytes.", estimatedSizeBytes);
numSplits =
(int)
Math.min(
NUM_QUERY_SPLITS_MAX,
Math.round(((double) estimatedSizeBytes) / DEFAULT_BUNDLE_SIZE_BYTES));
} catch (Exception e) {
logger.atWarning().withCause(e).log(
"Failed the fetch estimatedSizeBytes for query: %s", query);
// Fallback in case estimated size is unavailable.
numSplits = NUM_QUERY_SPLITS_MIN;
}
return Math.max(numSplits, NUM_QUERY_SPLITS_MIN);
}
/**
* Cloud Datastore system tables with statistics are periodically updated. This method fetches
* the latest timestamp (in microseconds) of statistics update using the {@code __Stat_Total__}
* table.
*/
private static long queryLatestStatisticsTimestamp(
Datastore datastore, @Nullable String namespace) throws DatastoreException {
Query.Builder query = Query.newBuilder();
// Note: namespace either being null or empty represents the default namespace, in which
// case we treat it as not provided by the user.
if (Strings.isNullOrEmpty(namespace)) {
query.addKindBuilder().setName("__Stat_Total__");
} else {
query.addKindBuilder().setName("__Stat_Ns_Total__");
}
query.addOrder(makeOrder("timestamp", DESCENDING));
query.setLimit(Int32Value.newBuilder().setValue(1));
RunQueryRequest request = makeRequest(query.build(), namespace);
RunQueryResponse response = datastore.runQuery(request);
QueryResultBatch batch = response.getBatch();
if (batch.getEntityResultsCount() == 0) {
throw new NoSuchElementException("Datastore total statistics unavailable");
}
Entity entity = batch.getEntityResults(0).getEntity();
return entity.getProperties().get("timestamp").getTimestampValue().getSeconds() * 1000000;
}
/** Retrieve latest table statistics for a given kind, namespace, and datastore. */
private static Entity getLatestTableStats(
String ourKind, @Nullable String namespace, Datastore datastore) throws DatastoreException {
long latestTimestamp = queryLatestStatisticsTimestamp(datastore, namespace);
logger.atInfo().log("Latest stats timestamp for kind %s is %s.", ourKind, latestTimestamp);
Query.Builder queryBuilder = Query.newBuilder();
if (Strings.isNullOrEmpty(namespace)) {
queryBuilder.addKindBuilder().setName("__Stat_Kind__");
} else {
queryBuilder.addKindBuilder().setName("__Stat_Ns_Kind__");
}
queryBuilder.setFilter(
makeAndFilter(
makeFilter("kind_name", EQUAL, makeValue(ourKind).build()).build(),
makeFilter("timestamp", EQUAL, makeValue(latestTimestamp).build()).build()));
RunQueryRequest request = makeRequest(queryBuilder.build(), namespace);
long now = System.currentTimeMillis();
RunQueryResponse response = datastore.runQuery(request);
logger.atFine().log(
"Query for per-kind statistics took %d ms.", System.currentTimeMillis() - now);
QueryResultBatch batch = response.getBatch();
if (batch.getEntityResultsCount() == 0) {
throw new NoSuchElementException(
"Datastore statistics for kind " + ourKind + " unavailable");
}
return batch.getEntityResults(0).getEntity();
}
/**
* Get the estimated size of the data returned by the given query.
*
* <p>Cloud Datastore provides no way to get a good estimate of how large the result of a query
* entity kind being queried, using the __Stat_Kind__ system table, assuming exactly 1 kind is
* specified in the query.
*
* <p>See https://cloud.google.com/datastore/docs/concepts/stats.
*/
static long getEstimatedSizeBytes(Datastore datastore, Query query, @Nullable String namespace)
throws DatastoreException {
String ourKind = query.getKind(0).getName();
Entity entity = getLatestTableStats(ourKind, namespace, datastore);
return entity.getProperties().get("entity_bytes").getIntegerValue();
}
private static PartitionId.Builder forNamespace(@Nullable String namespace) {
PartitionId.Builder partitionBuilder = PartitionId.newBuilder();
// Namespace either being null or empty represents the default namespace.
// Datastore Client libraries expect users to not set the namespace proto field in
// either of these cases.
if (!Strings.isNullOrEmpty(namespace)) {
partitionBuilder.setNamespaceId(namespace);
}
return partitionBuilder;
}
/** Builds a {@link RunQueryRequest} from the {@code query} and {@code namespace}. */
static RunQueryRequest makeRequest(Query query, @Nullable String namespace) {
return RunQueryRequest.newBuilder()
.setQuery(query)
.setPartitionId(forNamespace(namespace))
.build();
}
/** Builds a {@link RunQueryRequest} from the {@code GqlQuery} and {@code namespace}. */
private static RunQueryRequest makeRequest(GqlQuery gqlQuery, @Nullable String namespace) {
return RunQueryRequest.newBuilder()
.setGqlQuery(gqlQuery)
.setPartitionId(forNamespace(namespace))
.build();
}
/**
* A helper function to get the split queries, taking into account the optional {@code
* namespace}.
*/
private static List<Query> splitQuery(
Query query,
@Nullable String namespace,
Datastore datastore,
QuerySplitter querySplitter,
int numSplits)
throws DatastoreException {
// If namespace is set, include it in the split request so splits are calculated accordingly.
return querySplitter.getSplits(query, forNamespace(namespace).build(), numSplits, datastore);
}
/**
* Translates a Cloud Datastore gql query string to {@link Query}.
*
* <p>Currently, the only way to translate a gql query string to a Query is to run the query
* against Cloud Datastore and extract the {@code Query} from the response. To prevent reading
* any data, we set the {@code LIMIT} to 0 but if the gql query already has a limit set, we
* catch the exception with {@code INVALID_ARGUMENT} error code and retry the translation
* without the zero limit.
*
* <p>Note: This may result in reading actual data from Cloud Datastore but the service has a
* cap on the number of entities returned for a single rpc request, so this should not be a
* problem in practice.
*/
private static Query translateGqlQueryWithLimitCheck(
String gql, Datastore datastore, String namespace) throws DatastoreException {
String gqlQueryWithZeroLimit = gql + " LIMIT 0";
try {
Query translatedQuery = translateGqlQuery(gqlQueryWithZeroLimit, datastore, namespace);
// Clear the limit that we set.
return translatedQuery.toBuilder().clearLimit().build();
} catch (DatastoreException e) {
// Note: There is no specific error code or message to detect if the query already has a
// limit, so we just check for INVALID_ARGUMENT and assume that that the query might have
// a limit already set.
if (e.getCode() == Code.INVALID_ARGUMENT) {
logger.atWarning().log(
"Failed to translate Gql query '%s': %s", gqlQueryWithZeroLimit, e.getMessage());
logger.atWarning().log(
"User query might have a limit already set, so trying without zero limit.");
// Retry without the zero limit.
return translateGqlQuery(gql, datastore, namespace);
} else {
throw e;
}
}
}
/** Translates a gql query string to {@link Query}. */
private static Query translateGqlQuery(String gql, Datastore datastore, String namespace)
throws DatastoreException {
logger.atInfo().log("Translating gql %s", gql);
GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build();
RunQueryRequest req = makeRequest(gqlQuery, namespace);
return datastore.runQuery(req).getQuery();
}
/**
* Returns a new {@link MultiRead} that reads from the Cloud Datastore for the specified
* project.
*/
public MultiRead withProjectId(String projectId) {
checkArgument(projectId != null, "projectId can not be null");
return toBuilder().setProjectId(projectId).build();
}
/** Returns a new {@link MultiRead} that reads from the given namespace. */
public MultiRead withNamespace(String namespace) {
return toBuilder().setNamespace(namespace).build();
}
/**
* Returns a new {@link MultiRead} that reads by splitting the given {@code query} into {@code
* numQuerySplits}.
*
* <p>The semantics for the query splitting is defined below:
*
* <ul>
* <li>Any value less than or equal to 0 will be ignored, and the number of splits will be
* chosen dynamically at runtime based on the query data size.
* <li>Any value greater than {@link MultiRead#NUM_QUERY_SPLITS_MAX} will be capped at {@code
* NUM_QUERY_SPLITS_MAX}.
* <li>If the {@code query} has a user limit set, then {@code numQuerySplits} will be ignored
* and no split will be performed.
* <li>Under certain cases Cloud Datastore is unable to split query to the requested number of
* splits. In such cases we just use whatever the Cloud Datastore returns.
* </ul>
*/
public MultiRead withNumQuerySplits(int numQuerySplits) {
return toBuilder()
.setNumQuerySplits(Math.min(Math.max(numQuerySplits, 0), NUM_QUERY_SPLITS_MAX))
.build();
}
/**
* Returns a new {@link MultiRead} that reads from a Datastore Emulator running at the given
* localhost address.
*/
public MultiRead withLocalhost(String localhost) {
return toBuilder().setLocalhost(localhost).build();
}
/** Returns Number of entities available for reading. */
public long getNumEntities(
PipelineOptions options, String ourKind, @Nullable String namespace) {
try {
V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
V1DatastoreFactory datastoreFactory = new V1DatastoreFactory();
Datastore datastore =
datastoreFactory.getDatastore(
options, v1Options.getProjectId(), v1Options.getLocalhost());
Entity entity = getLatestTableStats(ourKind, namespace, datastore);
return entity.getProperties().get("count").getIntegerValue();
} catch (Exception e) {
return -1;
}
}
@Override
public PCollection<Entity> expand(PCollection<String> gqlQueries) {
checkArgument(getProjectId() != null, "projectId cannot be null");
V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
/*
* This composite transform involves the following steps:
* 1. Apply a {@link ParDo} that translates each query in {@code gqlQueries} into a {@code
* query}.
*
* 2. A {@link ParDo} splits the resulting query into {@code numQuerySplits} and
* assign each split query a unique {@code Integer} as the key. The resulting output is
* of the type {@code PCollection<KV<Integer, Query>>}.
*
* If the value of {@code numQuerySplits} is less than or equal to 0, then the number of
* splits will be computed dynamically based on the size of the data for the {@code query}.
*
* 3. The resulting {@code PCollection} is sharded using a {@link GroupByKey} operation. The
* queries are extracted from they {@code KV<Integer, Iterable<Query>>} and flattened to
* output a {@code PCollection<Query>}.
*
* 4. In the third step, a {@code ParDo} reads entities for each query and outputs
* a {@code PCollection<Entity>}.
*/
PCollection<Query> inputQuery =
gqlQueries.apply(ParDo.of(new GqlQueryTranslateFn(v1Options)));
return inputQuery
.apply("Split", ParDo.of(new SplitQueryFn(v1Options, getNumQuerySplits())))
.apply("Reshuffle", Reshuffle.viaRandomKey())
.apply("Read", ParDo.of(new ReadFn(v1Options)));
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
builder
.addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("ProjectId"))
.addIfNotNull(DisplayData.item("namespace", getNamespace()).withLabel("Namespace"));
}
private static class V1Options implements HasDisplayData, Serializable {
private final String project;
private final @Nullable String namespace;
private final @Nullable String localhost;
private V1Options(String project, @Nullable String namespace, @Nullable String localhost) {
this.project = project;
this.namespace = namespace;
this.localhost = localhost;
}
public static V1Options from(
String projectId, @Nullable String namespace, @Nullable String localhost) {
return new V1Options(projectId, namespace, localhost);
}
public String getProjectId() {
return project;
}
public @Nullable String getNamespace() {
return namespace;
}
public @Nullable String getLocalhost() {
return localhost;
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
builder
.addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("ProjectId"))
.addIfNotNull(DisplayData.item("namespace", getNamespace()).withLabel("Namespace"));
}
}
/** A DoFn that translates a Cloud Datastore gql query string to {@code Query}. */
static class GqlQueryTranslateFn extends DoFn<String, Query> {
private final V1Options v1Options;
private transient Datastore datastore;
private final V1DatastoreFactory datastoreFactory;
GqlQueryTranslateFn(V1Options options) {
this(options, new V1DatastoreFactory());
}
GqlQueryTranslateFn(V1Options options, V1DatastoreFactory datastoreFactory) {
this.v1Options = options;
this.datastoreFactory = datastoreFactory;
}
@StartBundle
public void startBundle(StartBundleContext c) {
datastore =
datastoreFactory.getDatastore(
c.getPipelineOptions(), v1Options.getProjectId(), v1Options.getLocalhost());
}
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
String gqlQuery = c.element();
logger.atInfo().log("User query: '%s'.", gqlQuery);
Query query =
translateGqlQueryWithLimitCheck(gqlQuery, datastore, v1Options.getNamespace());
logger.atInfo().log("User gql query translated to Query(%s).", query);
c.output(query);
}
}
/**
* A {@link DoFn} that splits a given query into multiple sub-queries, assigns them unique keys
* and outputs them as {@link KV}.
*/
private static class SplitQueryFn extends DoFn<Query, Query> {
private final V1Options options;
// number of splits to make for a given query
private final int numSplits;
private final V1DatastoreFactory datastoreFactory;
// Datastore client
private transient Datastore datastore;
// Query splitter
private transient QuerySplitter querySplitter;
public SplitQueryFn(V1Options options, int numSplits) {
this(options, numSplits, new V1DatastoreFactory());
}
private SplitQueryFn(V1Options options, int numSplits, V1DatastoreFactory datastoreFactory) {
this.options = options;
this.numSplits = numSplits;
this.datastoreFactory = datastoreFactory;
}
@StartBundle
public void startBundle(StartBundleContext c) {
datastore =
datastoreFactory.getDatastore(
c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
querySplitter = datastoreFactory.getQuerySplitter();
}
@ProcessElement
public void processElement(ProcessContext c) {
Query query = c.element();
// If query has a user set limit, then do not split.
if (query.hasLimit()) {
c.output(query);
return;
}
int estimatedNumSplits;
// Compute the estimated numSplits if numSplits is not specified by the user.
if (numSplits <= 0) {
estimatedNumSplits = getEstimatedNumSplits(datastore, query, options.getNamespace());
} else {
estimatedNumSplits = numSplits;
}
logger.atInfo().log("Splitting the query into %d splits.", estimatedNumSplits);
List<Query> querySplits;
try {
querySplits =
splitQuery(
query, options.getNamespace(), datastore, querySplitter, estimatedNumSplits);
} catch (Exception e) {
logger.atWarning().log("Unable to parallelize the given query: %s", query, e);
querySplits = ImmutableList.of(query);
}
// assign unique keys to query splits.
for (Query subquery : querySplits) {
c.output(subquery);
}
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
builder.include("options", options);
if (numSplits > 0) {
builder.add(
DisplayData.item("numQuerySplits", numSplits)
.withLabel("Requested number of Query splits"));
}
}
}
/** A {@link DoFn} that reads entities from Cloud Datastore for each query. */
private static class ReadFn extends DoFn<Query, Entity> {
private final V1Options options;
private final V1DatastoreFactory datastoreFactory;
// Datastore client
private transient Datastore datastore;
private final Counter rpcErrors = Metrics.counter(ReadFn.class, "datastoreRpcErrors");
private final Counter rpcSuccesses = Metrics.counter(ReadFn.class, "datastoreRpcSuccesses");
private static final int MAX_RETRIES = 5;
private static final FluentBackoff RUNQUERY_BACKOFF =
FluentBackoff.DEFAULT
.withMaxRetries(MAX_RETRIES)
.withInitialBackoff(Duration.standardSeconds(5));
public ReadFn(V1Options options) {
this(options, new V1DatastoreFactory());
}
private ReadFn(V1Options options, V1DatastoreFactory datastoreFactory) {
this.options = options;
this.datastoreFactory = datastoreFactory;
}
@StartBundle
public void startBundle(StartBundleContext c) {
datastore =
datastoreFactory.getDatastore(
c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
}
private RunQueryResponse runQueryWithRetries(RunQueryRequest request) throws Exception {
Sleeper sleeper = Sleeper.DEFAULT;
BackOff backoff = RUNQUERY_BACKOFF.backoff();
while (true) {
try {
RunQueryResponse response = datastore.runQuery(request);
rpcSuccesses.inc();
return response;
} catch (DatastoreException exception) {
rpcErrors.inc();
if (NON_RETRYABLE_ERRORS.contains(exception.getCode())) {
throw exception;
}
if (!BackOffUtils.next(sleeper, backoff)) {
logger.atSevere().log("Aborting after %d retries.", MAX_RETRIES);
throw exception;
}
}
}
}
/** Read and output entities for the given query. */
@ProcessElement
public void processElement(ProcessContext context) throws Exception {
Query query = context.element();
String namespace = options.getNamespace();
int userLimit = query.hasLimit() ? query.getLimit().getValue() : Integer.MAX_VALUE;
boolean moreResults = true;
QueryResultBatch currentBatch = null;
while (moreResults) {
Query.Builder queryBuilder = query.toBuilder();
queryBuilder.setLimit(
Int32Value.newBuilder().setValue(Math.min(userLimit, QUERY_BATCH_LIMIT)));
if (currentBatch != null && !currentBatch.getEndCursor().isEmpty()) {
queryBuilder.setStartCursor(currentBatch.getEndCursor());
}
RunQueryRequest request = makeRequest(queryBuilder.build(), namespace);
RunQueryResponse response = runQueryWithRetries(request);
currentBatch = response.getBatch();
// MORE_RESULTS_AFTER_LIMIT is not implemented yet:
// https://groups.google.com/forum/#!topic/gcd-discuss/iNs6M1jA2Vw, so
// use result count to determine if more results might exist.
int numFetch = currentBatch.getEntityResultsCount();
if (query.hasLimit()) {
verify(
userLimit >= numFetch,
"Expected userLimit %s >= numFetch %s, because query limit %s must be <= userLimit",
userLimit,
numFetch,
query.getLimit());
userLimit -= numFetch;
}
// output all the entities from the current batch.
for (EntityResult entityResult : currentBatch.getEntityResultsList()) {
context.output(entityResult.getEntity());
}
// Check if we have more entities to be read.
moreResults =
// User-limit does not exist (so userLimit == MAX_VALUE) and/or has not been satisfied
(userLimit > 0)
// All indications from the API are that there are/may be more results.
&& ((numFetch == QUERY_BATCH_LIMIT)
|| (currentBatch.getMoreResults() == NOT_FINISHED));
}
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
builder.include("options", options);
}
}
}
/**
* A wrapper factory class for Cloud Datastore singleton classes {@link DatastoreFactory} and
* {@link QuerySplitter}
*
* <p>{@link DatastoreFactory} and {@link QuerySplitter} are not java serializable, hence wrapping
* them under this class, which implements {@link Serializable}.
*/
private static class V1DatastoreFactory implements Serializable {
/** Builds a Cloud Datastore client for the given pipeline options and project. */
public Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) {
return getDatastore(pipelineOptions, projectId, null);
}
/**
* Builds a Cloud Datastore client for the given pipeline options, project and an optional
* locahost.
*/
public Datastore getDatastore(
PipelineOptions pipelineOptions, String projectId, @Nullable String localhost) {
Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential();
HttpRequestInitializer initializer;
if (credential != null) {
initializer =
new ChainingHttpRequestInitializer(
new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer());
} else {
initializer = new RetryHttpRequestInitializer();
}
DatastoreOptions.Builder builder =
new DatastoreOptions.Builder().projectId(projectId).initializer(initializer);
if (localhost != null) {
builder.localHost(localhost);
} else {
builder.host("batch-datastore.googleapis.com");
}
return DatastoreFactory.get().create(builder.build());
}
/** Builds a Cloud Datastore {@link QuerySplitter}. */
public QuerySplitter getQuerySplitter() {
return DatastoreHelper.getQuerySplitter();
}
}
}

10
core/src/main/java/google/registry/beam/rde/RdeIO.java
View file

@ -18,7 +18,6 @@ import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Verify.verify;
import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.rde.RdeModule.BRDA_QUEUE;
import static google.registry.rde.RdeModule.RDE_UPLOAD_QUEUE;
import static java.nio.charset.StandardCharsets.UTF_8;
@ -270,16 +269,15 @@ public class RdeIO {
@ProcessElement
public void processElement(
@Element KV<PendingDeposit, Integer> input, PipelineOptions options) {
tm().transact(
() -> {
PendingDeposit key = input.getKey();
Registry registry = Registry.get(key.tld());
Optional<Cursor> cursor =
transactIfJpaTm(
() ->
tm().loadByKeyIfPresent(
Cursor.createScopedVKey(key.cursor(), registry)));
tm().transact(
() ->
tm().loadByKeyIfPresent(
Cursor.createScopedVKey(key.cursor(), registry)));
DateTime position = getCursorTimeOrStartOfTime(cursor);
checkState(key.interval() != null, "Interval must be present");
DateTime newPosition = key.watermark().plus(key.interval());

14
core/src/main/java/google/registry/config/RegistryConfig.java
View file

@ -1331,19 +1331,7 @@ public final class RegistryConfig {
}
/**
* Returns the Google Cloud Storage bucket for storing Datastore backups.
*
* @see google.registry.export.BackupDatastoreAction
*/
public static String getDatastoreBackupsBucket() {
return "gs://" + getProjectId() + "-datastore-backups";
}
/**
* Returns the length of time before commit logs should be deleted from Datastore.
*
* <p>The only reason you'll want to retain this commit logs in Datastore is for performing
* point-in-time restoration queries for subsystems like RDE.
* Returns the length of time before commit logs should be deleted from the database.
*
* @see google.registry.tools.server.GenerateZoneFilesAction
*/

36
core/src/main/java/google/registry/env/common/backend/WEB-INF/web.xml vendored
View file

@ -169,36 +169,6 @@
<url-pattern>/_dr/dnsRefresh</url-pattern>
</servlet-mapping>
<!-- Exports a Datastore backup snapshot to GCS. -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
<url-pattern>/_dr/task/backupDatastore</url-pattern>
</servlet-mapping>
<!-- Checks the completion of a Datastore backup snapshot. -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
<url-pattern>/_dr/task/checkDatastoreBackup</url-pattern>
</servlet-mapping>
<!-- Loads a Datastore backup snapshot into BigQuery. -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
<url-pattern>/_dr/task/uploadDatastoreBackup</url-pattern>
</servlet-mapping>
<!-- Updates a view to point at a certain snapshot in BigQuery. -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
<url-pattern>/_dr/task/updateSnapshotView</url-pattern>
</servlet-mapping>
<!-- Polls state of jobs in Bigquery -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
<url-pattern>/_dr/task/pollBigqueryJob</url-pattern>
</servlet-mapping>
<!-- Fans out a cron task over an adjustable range of TLDs. -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
@ -323,12 +293,6 @@ have been in the database for a certain period of time. -->
<url-pattern>/_dr/task/wipeOutCloudSql</url-pattern>
</servlet-mapping>
<!-- Action to wipeout Cloud Datastore data -->
<servlet-mapping>
<servlet-name>backend-servlet</servlet-name>
<url-pattern>/_dr/task/wipeOutDatastore</url-pattern>
</servlet-mapping>
<!-- Security config -->
<security-constraint>
<web-resource-collection>

64
core/src/main/java/google/registry/export/AnnotatedEntities.java
View file

@ -1,64 +0,0 @@
// 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.export;
import static com.google.common.collect.ImmutableSortedSet.toImmutableSortedSet;
import static google.registry.util.TypeUtils.hasAnnotation;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Ordering;
import com.googlecode.objectify.Key;
import google.registry.model.EntityClasses;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.model.annotations.InCrossTld;
import google.registry.model.annotations.NotBackedUp;
import google.registry.model.annotations.ReportedOn;
import google.registry.model.annotations.VirtualEntity;
/** Constants related to export code. */
@DeleteAfterMigration
public final class AnnotatedEntities {
/** Returns the names of kinds to include in Datastore backups. */
public static ImmutableSet<String> getBackupKinds() {
// Back up all entity classes that aren't annotated with @VirtualEntity (never even persisted
// to Datastore, so they can't be backed up) or @NotBackedUp (intentionally omitted).
return EntityClasses.ALL_CLASSES
.stream()
.filter(hasAnnotation(VirtualEntity.class).negate())
.filter(hasAnnotation(NotBackedUp.class).negate())
.map(Key::getKind)
.collect(toImmutableSortedSet(Ordering.natural()));
}
/** Returns the names of kinds to import into reporting tools (e.g. BigQuery). */
public static ImmutableSet<String> getReportingKinds() {
return EntityClasses.ALL_CLASSES
.stream()
.filter(hasAnnotation(ReportedOn.class))
.filter(hasAnnotation(VirtualEntity.class).negate())
.map(Key::getKind)
.collect(toImmutableSortedSet(Ordering.natural()));
}
/** Returns the names of kinds that are in the cross-TLD entity group. */
public static ImmutableSet<String> getCrossTldKinds() {
return EntityClasses.ALL_CLASSES.stream()
.filter(hasAnnotation(InCrossTld.class))
.filter(hasAnnotation(VirtualEntity.class).negate())
.map(Key::getKind)
.collect(toImmutableSortedSet(Ordering.natural()));
}
}

105
core/src/main/java/google/registry/export/BackupDatastoreAction.java
View file

@ -1,105 +0,0 @@
// Copyright 2018 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.export;
import static google.registry.request.Action.Method.POST;
import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.flogger.FluentLogger;
import google.registry.config.RegistryConfig;
import google.registry.export.datastore.DatastoreAdmin;
import google.registry.export.datastore.Operation;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.request.Action;
import google.registry.request.Action.Service;
import google.registry.request.HttpException.InternalServerErrorException;
import google.registry.request.Response;
import google.registry.request.auth.Auth;
import google.registry.util.Clock;
import google.registry.util.CloudTasksUtils;
import javax.inject.Inject;
/**
* Action to trigger a Datastore backup job that writes a snapshot to Google Cloud Storage.
*
* <p>This is the first step of a four step workflow for exporting snapshots, with each step calling
* the next upon successful completion:
*
* <ol>
* <li>The snapshot is exported to Google Cloud Storage (this action).
* <li>The {@link CheckBackupAction} polls until the export is completed.
* <li>The {@link UploadDatastoreBackupAction} uploads the data from GCS to BigQuery.
* <li>The {@link UpdateSnapshotViewAction} updates the view in latest_datastore_export.
* </ol>
*/
@Action(
service = Action.Service.BACKEND,
path = BackupDatastoreAction.PATH,
method = POST,
automaticallyPrintOk = true,
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
@DeleteAfterMigration
public class BackupDatastoreAction implements Runnable {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
/** Queue to use for enqueuing the task that will actually launch the backup. */
static final String QUEUE = "export-snapshot"; // See queue.xml.
static final String PATH = "/_dr/task/backupDatastore"; // See web.xml.
@Inject DatastoreAdmin datastoreAdmin;
@Inject Response response;
@Inject Clock clock;
@Inject CloudTasksUtils cloudTasksUtils;
@Inject
BackupDatastoreAction() {}
@Override
public void run() {
try {
Operation backup =
datastoreAdmin
.export(
RegistryConfig.getDatastoreBackupsBucket(), AnnotatedEntities.getBackupKinds())
.execute();
String backupName = backup.getName();
// Enqueue a poll task to monitor the backup for completion and load reporting-related kinds
// into bigquery.
cloudTasksUtils.enqueue(
CheckBackupAction.QUEUE,
cloudTasksUtils.createPostTaskWithDelay(
CheckBackupAction.PATH,
Service.BACKEND.toString(),
ImmutableMultimap.of(
CheckBackupAction.CHECK_BACKUP_NAME_PARAM,
backupName,
CheckBackupAction.CHECK_BACKUP_KINDS_TO_LOAD_PARAM,
Joiner.on(',').join(AnnotatedEntities.getReportingKinds())),
CheckBackupAction.POLL_COUNTDOWN));
String message =
String.format(
"Datastore backup started with name: %s\nSaving to %s",
backupName, backup.getExportFolderUrl());
logger.atInfo().log(message);
response.setPayload(message);
} catch (Throwable e) {
throw new InternalServerErrorException("Exception occurred while backing up Datastore", e);
}
}
}

129
core/src/main/java/google/registry/export/BigqueryPollJobAction.java
View file

@ -1,129 +0,0 @@
// 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.export;
import static google.registry.bigquery.BigqueryUtils.toJobReferenceString;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobReference;
import com.google.cloud.tasks.v2.Task;
import com.google.common.flogger.FluentLogger;
import com.google.protobuf.ByteString;
import dagger.Lazy;
import google.registry.request.Action;
import google.registry.request.Header;
import google.registry.request.HttpException.BadRequestException;
import google.registry.request.HttpException.NotModifiedException;
import google.registry.request.Payload;
import google.registry.request.auth.Auth;
import google.registry.util.CloudTasksUtils;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import javax.inject.Inject;
import org.joda.time.Duration;
/**
* An action which polls the state of a bigquery job. If it is completed then it will log its
* completion state; otherwise it will return a failure code so that the task will be retried.
*
* <p>Note that this is AUTH_INTERNAL_ONLY: we don't allow "admin" for this to mitigate a
* vulnerability, see b/177308043.
*/
@Action(
service = Action.Service.BACKEND,
path = BigqueryPollJobAction.PATH,
method = {Action.Method.GET, Action.Method.POST},
automaticallyPrintOk = true,
auth = Auth.AUTH_INTERNAL_ONLY)
public class BigqueryPollJobAction implements Runnable {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
static final String QUEUE = "export-bigquery-poll"; // See queue.xml
static final String PATH = "/_dr/task/pollBigqueryJob"; // See web.xml
static final String CHAINED_TASK_QUEUE_HEADER = "X-DomainRegistry-ChainedTaskQueue";
static final String PROJECT_ID_HEADER = "X-DomainRegistry-ProjectId";
static final String JOB_ID_HEADER = "X-DomainRegistry-JobId";
static final Duration POLL_COUNTDOWN = Duration.standardSeconds(20);
@Inject Bigquery bigquery;
@Inject CloudTasksUtils cloudTasksUtils;
@Inject @Header(CHAINED_TASK_QUEUE_HEADER) Lazy<String> chainedQueueName;
@Inject @Header(PROJECT_ID_HEADER) String projectId;
@Inject @Header(JOB_ID_HEADER) String jobId;
@Inject @Payload ByteString payload;
@Inject BigqueryPollJobAction() {}
@Override
public void run() {
boolean jobOutcome =
checkJobOutcome(); // Throws a NotModifiedException if the job hasn't completed.
// If the job failed, do not enqueue the next step.
if (!jobOutcome || payload == null || payload.size() == 0) {
return;
}
// If there is a payload, it's a chained task, so enqueue it.
Task task;
try {
task =
(Task)
new ObjectInputStream(new ByteArrayInputStream(payload.toByteArray())).readObject();
} catch (ClassNotFoundException | IOException e) {
throw new BadRequestException("Cannot deserialize task from payload", e);
}
Task enqueuedTask = cloudTasksUtils.enqueue(chainedQueueName.get(), task);
logger.atInfo().log(
"Added chained task %s for %s to queue %s: %s",
enqueuedTask.getName(),
enqueuedTask.getAppEngineHttpRequest().getRelativeUri(),
chainedQueueName.get(),
enqueuedTask);
}
/**
* Returns true if the provided job succeeded, false if it failed, and throws an exception if it
* is still pending.
*/
private boolean checkJobOutcome() {
Job job = null;
String jobRefString =
toJobReferenceString(new JobReference().setProjectId(projectId).setJobId(jobId));
try {
job = bigquery.jobs().get(projectId, jobId).execute();
} catch (IOException e) {
// We will throw a new exception because done==false, but first log this exception.
logger.atWarning().withCause(e).log("Error checking outcome of BigQuery job %s.", jobId);
}
// If job is not yet done, then throw an exception so that we'll return a failing HTTP status
// code and the task will be retried.
if (job == null || !job.getStatus().getState().equals("DONE")) {
throw new NotModifiedException(jobRefString);
}
// Check if the job ended with an error.
if (job.getStatus().getErrorResult() != null) {
logger.atSevere().log("Bigquery job failed - %s - %s.", jobRefString, job);
return false;
}
logger.atInfo().log("Bigquery job succeeded - %s.", jobRefString);
return true;
}
}

195
core/src/main/java/google/registry/export/CheckBackupAction.java
View file

@ -1,195 +0,0 @@
// Copyright 2018 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.export;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.collect.Sets.intersection;
import static google.registry.request.Action.Method.GET;
import static google.registry.request.Action.Method.POST;
import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.common.base.Joiner;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import com.google.common.flogger.FluentLogger;
import google.registry.export.datastore.DatastoreAdmin;
import google.registry.export.datastore.Operation;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.request.Action;
import google.registry.request.Action.Service;
import google.registry.request.HttpException;
import google.registry.request.HttpException.BadRequestException;
import google.registry.request.HttpException.InternalServerErrorException;
import google.registry.request.HttpException.NoContentException;
import google.registry.request.HttpException.NotModifiedException;
import google.registry.request.Parameter;
import google.registry.request.RequestMethod;
import google.registry.request.Response;
import google.registry.request.auth.Auth;
import google.registry.util.Clock;
import google.registry.util.CloudTasksUtils;
import java.io.IOException;
import java.util.Set;
import javax.inject.Inject;
import org.joda.time.Duration;
import org.joda.time.PeriodType;
import org.joda.time.format.PeriodFormat;
/**
* Action that checks the status of a snapshot, and if complete, trigger loading it into BigQuery.
*/
@Action(
service = Action.Service.BACKEND,
path = CheckBackupAction.PATH,
method = {POST, GET},
automaticallyPrintOk = true,
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
@DeleteAfterMigration
public class CheckBackupAction implements Runnable {
/** Parameter names for passing parameters into this action. */
static final String CHECK_BACKUP_NAME_PARAM = "name";
static final String CHECK_BACKUP_KINDS_TO_LOAD_PARAM = "kindsToLoad";
/** Action-specific details needed for enqueuing tasks against itself. */
static final String QUEUE = "export-snapshot-poll"; // See queue.xml.
static final String PATH = "/_dr/task/checkDatastoreBackup"; // See web.xml.
static final Duration POLL_COUNTDOWN = Duration.standardMinutes(2);
/** The maximum amount of time we allow a backup to run before abandoning it. */
static final Duration MAXIMUM_BACKUP_RUNNING_TIME = Duration.standardHours(20);
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@Inject DatastoreAdmin datastoreAdmin;
@Inject Clock clock;
@Inject Response response;
@Inject CloudTasksUtils cloudTasksUtils;
@Inject @RequestMethod Action.Method requestMethod;
@Inject
@Parameter(CHECK_BACKUP_NAME_PARAM)
String backupName;
@Inject
@Parameter(CHECK_BACKUP_KINDS_TO_LOAD_PARAM)
String kindsToLoadParam;
@Inject
CheckBackupAction() {}
@Override
public void run() {
try {
if (requestMethod == POST) {
checkAndLoadBackupIfComplete();
} else {
// This is a GET request.
// TODO(weiminyu): consider moving this functionality to Registry tool.
response.setPayload(getExportStatus().toPrettyString());
}
} catch (HttpException e) {
// Rethrow and let caller propagate status code and error message to the response.
// See google.registry.request.RequestHandler#handleRequest.
throw e;
} catch (Throwable e) {
throw new InternalServerErrorException(
"Exception occurred while checking datastore exports.", e);
}
}
private Operation getExportStatus() throws IOException {
try {
return datastoreAdmin.get(backupName).execute();
} catch (GoogleJsonResponseException e) {
if (e.getStatusCode() == SC_NOT_FOUND) {
String message = String.format("Bad backup name %s: %s", backupName, e.getMessage());
// TODO(b/19081569): Ideally this would return a 2XX error so the task would not be
// retried but we might abandon backups that start late and haven't yet written to
// Datastore. We could fix that by replacing this with a two-phase polling strategy.
throw new BadRequestException(message, e);
}
throw e;
}
}
private void checkAndLoadBackupIfComplete() throws IOException {
Set<String> kindsToLoad = ImmutableSet.copyOf(Splitter.on(',').split(kindsToLoadParam));
Operation backup = getExportStatus();
checkArgument(backup.isExport(), "Expecting an export operation: [%s].", backupName);
if (backup.isProcessing()
&& backup.getRunningTime(clock).isShorterThan(MAXIMUM_BACKUP_RUNNING_TIME)) {
// Backup might still be running, so send a 304 to have the task retry.
throw new NotModifiedException(
String.format(
"Datastore backup %s still in progress: %s", backupName, backup.getProgress()));
}
if (!backup.isSuccessful()) {
// Declare the backup a lost cause, and send 204 No Content so the task will
// not be retried.
String message =
String.format(
"Datastore backup %s abandoned - not complete after %s. Progress: %s",
backupName,
PeriodFormat.getDefault()
.print(
backup
.getRunningTime(clock)
.toPeriod()
.normalizedStandard(PeriodType.dayTime().withMillisRemoved())),
backup.getProgress());
throw new NoContentException(message);
}
String backupId = backup.getExportId();
// Log a warning if kindsToLoad is not a subset of the exported kinds.
if (!backup.getKinds().containsAll(kindsToLoad)) {
logger.atWarning().log(
"Kinds to load included non-exported kinds: %s",
Sets.difference(kindsToLoad, backup.getKinds()));
}
// Load kinds from the backup, limited to those also in kindsToLoad (if it's present).
ImmutableSet<String> exportedKindsToLoad =
ImmutableSet.copyOf(intersection(backup.getKinds(), kindsToLoad));
String message = String.format("Datastore backup %s complete - ", backupName);
if (exportedKindsToLoad.isEmpty()) {
message += "no kinds to load into BigQuery.";
} else {
/** Enqueue a task for starting a backup load. */
cloudTasksUtils.enqueue(
UploadDatastoreBackupAction.QUEUE,
cloudTasksUtils.createPostTask(
UploadDatastoreBackupAction.PATH,
Service.BACKEND.toString(),
ImmutableMultimap.of(
UploadDatastoreBackupAction.UPLOAD_BACKUP_ID_PARAM,
backupId,
UploadDatastoreBackupAction.UPLOAD_BACKUP_FOLDER_PARAM,
backup.getExportFolderUrl(),
UploadDatastoreBackupAction.UPLOAD_BACKUP_KINDS_PARAM,
Joiner.on(',').join(exportedKindsToLoad))));
message += "BigQuery load task enqueued.";
}
logger.atInfo().log(message);
response.setPayload(message);
}
}

113
core/src/main/java/google/registry/export/ExportRequestModule.java
View file

@ -1,113 +0,0 @@
// 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.export;
import static google.registry.export.BigqueryPollJobAction.CHAINED_TASK_QUEUE_HEADER;
import static google.registry.export.BigqueryPollJobAction.JOB_ID_HEADER;
import static google.registry.export.BigqueryPollJobAction.PROJECT_ID_HEADER;
import static google.registry.export.CheckBackupAction.CHECK_BACKUP_KINDS_TO_LOAD_PARAM;
import static google.registry.export.CheckBackupAction.CHECK_BACKUP_NAME_PARAM;
import static google.registry.export.UpdateSnapshotViewAction.UPDATE_SNAPSHOT_DATASET_ID_PARAM;
import static google.registry.export.UpdateSnapshotViewAction.UPDATE_SNAPSHOT_KIND_PARAM;
import static google.registry.export.UpdateSnapshotViewAction.UPDATE_SNAPSHOT_TABLE_ID_PARAM;
import static google.registry.export.UpdateSnapshotViewAction.UPDATE_SNAPSHOT_VIEWNAME_PARAM;
import static google.registry.export.UploadDatastoreBackupAction.UPLOAD_BACKUP_FOLDER_PARAM;
import static google.registry.export.UploadDatastoreBackupAction.UPLOAD_BACKUP_ID_PARAM;
import static google.registry.export.UploadDatastoreBackupAction.UPLOAD_BACKUP_KINDS_PARAM;
import static google.registry.request.RequestParameters.extractRequiredHeader;
import static google.registry.request.RequestParameters.extractRequiredParameter;
import dagger.Module;
import dagger.Provides;
import google.registry.request.Header;
import google.registry.request.Parameter;
import javax.servlet.http.HttpServletRequest;
/** Dagger module for data export tasks. */
@Module
public final class ExportRequestModule {
@Provides
@Parameter(UPDATE_SNAPSHOT_DATASET_ID_PARAM)
static String provideUpdateSnapshotDatasetId(HttpServletRequest req) {
return extractRequiredParameter(req, UPDATE_SNAPSHOT_DATASET_ID_PARAM);
}
@Provides
@Parameter(UPDATE_SNAPSHOT_TABLE_ID_PARAM)
static String provideUpdateSnapshotTableId(HttpServletRequest req) {
return extractRequiredParameter(req, UPDATE_SNAPSHOT_TABLE_ID_PARAM);
}
@Provides
@Parameter(UPDATE_SNAPSHOT_KIND_PARAM)
static String provideUpdateSnapshotKind(HttpServletRequest req) {
return extractRequiredParameter(req, UPDATE_SNAPSHOT_KIND_PARAM);
}
@Provides
@Parameter(UPDATE_SNAPSHOT_VIEWNAME_PARAM)
static String provideUpdateSnapshotViewName(HttpServletRequest req) {
return extractRequiredParameter(req, UPDATE_SNAPSHOT_VIEWNAME_PARAM);
}
@Provides
@Parameter(UPLOAD_BACKUP_FOLDER_PARAM)
static String provideSnapshotUrlPrefix(HttpServletRequest req) {
return extractRequiredParameter(req, UPLOAD_BACKUP_FOLDER_PARAM);
}
@Provides
@Parameter(UPLOAD_BACKUP_ID_PARAM)
static String provideLoadSnapshotId(HttpServletRequest req) {
return extractRequiredParameter(req, UPLOAD_BACKUP_ID_PARAM);
}
@Provides
@Parameter(UPLOAD_BACKUP_KINDS_PARAM)
static String provideLoadSnapshotKinds(HttpServletRequest req) {
return extractRequiredParameter(req, UPLOAD_BACKUP_KINDS_PARAM);
}
@Provides
@Parameter(CHECK_BACKUP_NAME_PARAM)
static String provideCheckSnapshotName(HttpServletRequest req) {
return extractRequiredParameter(req, CHECK_BACKUP_NAME_PARAM);
}
@Provides
@Parameter(CHECK_BACKUP_KINDS_TO_LOAD_PARAM)
static String provideCheckSnapshotKindsToLoad(HttpServletRequest req) {
return extractRequiredParameter(req, CHECK_BACKUP_KINDS_TO_LOAD_PARAM);
}
@Provides
@Header(CHAINED_TASK_QUEUE_HEADER)
static String provideChainedTaskQueue(HttpServletRequest req) {
return extractRequiredHeader(req, CHAINED_TASK_QUEUE_HEADER);
}
@Provides
@Header(JOB_ID_HEADER)
static String provideJobId(HttpServletRequest req) {
return extractRequiredHeader(req, JOB_ID_HEADER);
}
@Provides
@Header(PROJECT_ID_HEADER)
static String provideProjectId(HttpServletRequest req) {
return extractRequiredHeader(req, PROJECT_ID_HEADER);
}
}

146
core/src/main/java/google/registry/export/UpdateSnapshotViewAction.java
View file

@ -1,146 +0,0 @@
// 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.export;
import static google.registry.request.Action.Method.POST;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.Table;
import com.google.api.services.bigquery.model.TableReference;
import com.google.api.services.bigquery.model.ViewDefinition;
import com.google.common.flogger.FluentLogger;
import google.registry.bigquery.CheckedBigquery;
import google.registry.config.RegistryConfig.Config;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.request.Action;
import google.registry.request.HttpException.InternalServerErrorException;
import google.registry.request.Parameter;
import google.registry.request.auth.Auth;
import google.registry.util.SqlTemplate;
import java.io.IOException;
import javax.inject.Inject;
/** Update a well-known view to point at a certain Datastore snapshot table in BigQuery. */
@Action(
service = Action.Service.BACKEND,
path = UpdateSnapshotViewAction.PATH,
method = POST,
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
@DeleteAfterMigration
public class UpdateSnapshotViewAction implements Runnable {
/** Headers for passing parameters into the servlet. */
static final String UPDATE_SNAPSHOT_DATASET_ID_PARAM = "dataset";
static final String UPDATE_SNAPSHOT_TABLE_ID_PARAM = "table";
static final String UPDATE_SNAPSHOT_KIND_PARAM = "kind";
static final String UPDATE_SNAPSHOT_VIEWNAME_PARAM = "viewname";
/** Servlet-specific details needed for enqueuing tasks against itself. */
// For now this queue is shared by the backup workflows started by BackupDatastoreAction.
// TODO(weiminyu): update queue name (snapshot->backup) after ExportSnapshot flow is removed.
static final String QUEUE = "export-snapshot-update-view"; // See queue.xml.
static final String PATH = "/_dr/task/updateSnapshotView"; // See web.xml.
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@Inject
@Parameter(UPDATE_SNAPSHOT_DATASET_ID_PARAM)
String datasetId;
@Inject
@Parameter(UPDATE_SNAPSHOT_TABLE_ID_PARAM)
String tableId;
@Inject
@Parameter(UPDATE_SNAPSHOT_KIND_PARAM)
String kindName;
@Inject
@Parameter(UPDATE_SNAPSHOT_VIEWNAME_PARAM)
String viewName;
@Inject
@Config("projectId")
String projectId;
@Inject CheckedBigquery checkedBigquery;
@Inject
UpdateSnapshotViewAction() {}
@Override
public void run() {
try {
SqlTemplate sqlTemplate =
SqlTemplate.create(
"#standardSQL\nSELECT * FROM `%PROJECT%.%SOURCE_DATASET%.%SOURCE_TABLE%`");
updateSnapshotView(datasetId, tableId, kindName, viewName, sqlTemplate);
} catch (Throwable e) {
throw new InternalServerErrorException(
String.format("Could not update snapshot view %s for table %s", viewName, tableId), e);
}
}
private void updateSnapshotView(
String sourceDatasetId,
String sourceTableId,
String kindName,
String viewDataset,
SqlTemplate viewQueryTemplate)
throws IOException {
Bigquery bigquery = checkedBigquery.ensureDataSetExists(projectId, viewDataset);
updateTable(
bigquery,
new Table()
.setTableReference(
new TableReference()
.setProjectId(projectId)
.setDatasetId(viewDataset)
.setTableId(kindName))
.setView(
new ViewDefinition()
.setUseLegacySql(false)
.setQuery(
viewQueryTemplate
.put("PROJECT", projectId)
.put("SOURCE_DATASET", sourceDatasetId)
.put("SOURCE_TABLE", sourceTableId)
.build())));
logger.atInfo().log(
"Updated view [%s:%s.%s] to point at snapshot table [%s:%s.%s].",
projectId, viewDataset, kindName, projectId, sourceDatasetId, sourceTableId);
}
private static void updateTable(Bigquery bigquery, Table table) throws IOException {
TableReference ref = table.getTableReference();
try {
bigquery
.tables()
.update(ref.getProjectId(), ref.getDatasetId(), ref.getTableId(), table)
.execute();
} catch (GoogleJsonResponseException e) {
if (e.getDetails() != null && e.getDetails().getCode() == 404) {
bigquery.tables().insert(ref.getProjectId(), ref.getDatasetId(), table).execute();
} else {
logger.atWarning().withCause(e).log("UpdateSnapshotViewAction errored out.");
}
}
}
}

215
core/src/main/java/google/registry/export/UploadDatastoreBackupAction.java
View file

@ -1,215 +0,0 @@
// Copyright 2018 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.export;
import static com.google.common.base.MoreObjects.firstNonNull;
import static google.registry.request.Action.Method.POST;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobConfiguration;
import com.google.api.services.bigquery.model.JobConfigurationLoad;
import com.google.api.services.bigquery.model.JobReference;
import com.google.api.services.bigquery.model.TableReference;
import com.google.cloud.tasks.v2.Task;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Joiner;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.flogger.FluentLogger;
import com.google.common.net.HttpHeaders;
import com.google.common.net.MediaType;
import com.google.protobuf.ByteString;
import com.google.protobuf.util.Timestamps;
import google.registry.bigquery.BigqueryUtils.SourceFormat;
import google.registry.bigquery.BigqueryUtils.WriteDisposition;
import google.registry.bigquery.CheckedBigquery;
import google.registry.config.RegistryConfig.Config;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.request.Action;
import google.registry.request.Action.Service;
import google.registry.request.HttpException.BadRequestException;
import google.registry.request.HttpException.InternalServerErrorException;
import google.registry.request.Parameter;
import google.registry.request.auth.Auth;
import google.registry.util.Clock;
import google.registry.util.CloudTasksUtils;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import javax.inject.Inject;
/** Action to load a Datastore backup from Google Cloud Storage into BigQuery. */
@Action(
service = Action.Service.BACKEND,
path = UploadDatastoreBackupAction.PATH,
method = POST,
auth = Auth.AUTH_INTERNAL_OR_ADMIN)
@DeleteAfterMigration
public class UploadDatastoreBackupAction implements Runnable {
/** Parameter names for passing parameters into the servlet. */
static final String UPLOAD_BACKUP_ID_PARAM = "id";
static final String UPLOAD_BACKUP_FOLDER_PARAM = "folder";
static final String UPLOAD_BACKUP_KINDS_PARAM = "kinds";
static final String BACKUP_DATASET = "datastore_backups";
/** Servlet-specific details needed for enqueuing tasks against itself. */
static final String QUEUE = "export-snapshot"; // See queue.xml.
static final String LATEST_BACKUP_VIEW_NAME = "latest_datastore_export";
static final String PATH = "/_dr/task/uploadDatastoreBackup"; // See web.xml.
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@Inject CheckedBigquery checkedBigquery;
@Inject CloudTasksUtils cloudTasksUtils;
@Inject Clock clock;
@Inject @Config("projectId") String projectId;
@Inject
@Parameter(UPLOAD_BACKUP_FOLDER_PARAM)
String backupFolderUrl;
@Inject
@Parameter(UPLOAD_BACKUP_ID_PARAM)
String backupId;
@Inject
@Parameter(UPLOAD_BACKUP_KINDS_PARAM)
String backupKinds;
@Inject
UploadDatastoreBackupAction() {}
@Override
public void run() {
try {
String message = uploadBackup(backupId, backupFolderUrl, Splitter.on(',').split(backupKinds));
logger.atInfo().log("Loaded backup successfully: %s", message);
} catch (Throwable e) {
logger.atSevere().withCause(e).log("Error loading backup.");
if (e instanceof IllegalArgumentException) {
throw new BadRequestException("Error calling load backup: " + e.getMessage(), e);
} else {
throw new InternalServerErrorException(
"Error loading backup: " + firstNonNull(e.getMessage(), e.toString()));
}
}
}
private String uploadBackup(String backupId, String backupFolderUrl, Iterable<String> kinds)
throws IOException {
Bigquery bigquery = checkedBigquery.ensureDataSetExists(projectId, BACKUP_DATASET);
String loadMessage =
String.format("Loading Datastore backup %s from %s...", backupId, backupFolderUrl);
logger.atInfo().log(loadMessage);
String sanitizedBackupId = sanitizeForBigquery(backupId);
StringBuilder builder = new StringBuilder(loadMessage + "\n");
builder.append("Load jobs:\n");
for (String kindName : kinds) {
String jobId = String.format("load-backup-%s-%s", sanitizedBackupId, kindName);
JobReference jobRef = new JobReference().setProjectId(projectId).setJobId(jobId);
String sourceUri = getBackupInfoFileForKind(backupFolderUrl, kindName);
String tableId = String.format("%s_%s", sanitizedBackupId, kindName);
// Launch the load job.
Job job = makeLoadJob(jobRef, sourceUri, tableId);
bigquery.jobs().insert(projectId, job).execute();
// Serialize the chainedTask into a byte array to put in the task payload.
ByteArrayOutputStream taskBytes = new ByteArrayOutputStream();
new ObjectOutputStream(taskBytes)
.writeObject(
cloudTasksUtils.createPostTask(
UpdateSnapshotViewAction.PATH,
Service.BACKEND.toString(),
ImmutableMultimap.of(
UpdateSnapshotViewAction.UPDATE_SNAPSHOT_DATASET_ID_PARAM,
BACKUP_DATASET,
UpdateSnapshotViewAction.UPDATE_SNAPSHOT_TABLE_ID_PARAM,
tableId,
UpdateSnapshotViewAction.UPDATE_SNAPSHOT_KIND_PARAM,
kindName,
UpdateSnapshotViewAction.UPDATE_SNAPSHOT_VIEWNAME_PARAM,
LATEST_BACKUP_VIEW_NAME)));
// Enqueues a task to poll for the success or failure of the referenced BigQuery job and to
// launch the provided task in the specified queue if the job succeeds.
cloudTasksUtils.enqueue(
BigqueryPollJobAction.QUEUE,
Task.newBuilder()
.setAppEngineHttpRequest(
cloudTasksUtils
.createPostTask(BigqueryPollJobAction.PATH, Service.BACKEND.toString(), null)
.getAppEngineHttpRequest()
.toBuilder()
.putHeaders(BigqueryPollJobAction.PROJECT_ID_HEADER, jobRef.getProjectId())
.putHeaders(BigqueryPollJobAction.JOB_ID_HEADER, jobRef.getJobId())
.putHeaders(
BigqueryPollJobAction.CHAINED_TASK_QUEUE_HEADER,
UpdateSnapshotViewAction.QUEUE)
// need to include CONTENT_TYPE in header when body is not empty
.putHeaders(HttpHeaders.CONTENT_TYPE, MediaType.FORM_DATA.toString())
.setBody(ByteString.copyFrom(taskBytes.toByteArray()))
.build())
.setScheduleTime(
Timestamps.fromMillis(
clock.nowUtc().plus(BigqueryPollJobAction.POLL_COUNTDOWN).getMillis()))
.build());
builder.append(String.format(" - %s:%s\n", projectId, jobId));
logger.atInfo().log("Submitted load job %s:%s.", projectId, jobId);
}
return builder.toString();
}
private static String sanitizeForBigquery(String backupId) {
return backupId.replaceAll("[^a-zA-Z0-9_]", "_");
}
@VisibleForTesting
static String getBackupInfoFileForKind(String backupFolderUrl, String kindName) {
return Joiner.on('/')
.join(
backupFolderUrl,
"all_namespaces",
String.format("kind_%s", kindName),
String.format("all_namespaces_kind_%s.%s", kindName, "export_metadata"));
}
private Job makeLoadJob(JobReference jobRef, String sourceUri, String tableId) {
TableReference tableReference =
new TableReference()
.setProjectId(jobRef.getProjectId())
.setDatasetId(BACKUP_DATASET)
.setTableId(tableId);
return new Job()
.setJobReference(jobRef)
.setConfiguration(new JobConfiguration()
.setLoad(new JobConfigurationLoad()
.setWriteDisposition(WriteDisposition.WRITE_EMPTY.toString())
.setSourceFormat(SourceFormat.DATASTORE_BACKUP.toString())
.setSourceUris(ImmutableList.of(sourceUri))
.setDestinationTable(tableReference)));
}
}

279
core/src/main/java/google/registry/export/datastore/DatastoreAdmin.java
View file

@ -1,279 +0,0 @@
// Copyright 2018 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.export.datastore;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import com.google.api.client.googleapis.services.json.AbstractGoogleJsonClient;
import com.google.api.client.googleapis.services.json.AbstractGoogleJsonClientRequest;
import com.google.api.client.http.HttpRequestInitializer;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.util.Key;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import google.registry.model.annotations.DeleteAfterMigration;
import java.util.Collection;
import java.util.Optional;
/**
* Java client to <a href="https://cloud.google.com/datastore/docs/reference/admin/rest/">Cloud
* Datastore Admin REST API</a>.
*/
@DeleteAfterMigration
public class DatastoreAdmin extends AbstractGoogleJsonClient {
private static final String ROOT_URL = "https://datastore.googleapis.com/v1/";
private static final String SERVICE_PATH = "";
// GCP project that this instance is associated with.
private final String projectId;
protected DatastoreAdmin(Builder builder) {
super(builder);
this.projectId = checkNotNull(builder.projectId, "GCP projectId missing.");
}
/**
* Returns an {@link Export} request that starts exporting all Cloud Datastore databases owned by
* the GCP project identified by {@link #projectId}.
*
* <p>Typical usage is:
*
* <pre>
* {@code Export export = datastoreAdmin.export(parameters ...);}
* {@code Operation operation = export.execute();}
* {@code while (!operation.isSuccessful()) { ...}}
* </pre>
*
* <p>Please see the <a
* href="https://cloud.google.com/datastore/docs/reference/admin/rest/v1/projects/export">API
* specification of the export method for details</a>.
*
* <p>The following undocumented behaviors with regard to {@code outputUrlPrefix} have been
* observed:
*
* <ul>
* <li>If outputUrlPrefix refers to a GCS bucket, exported data will be nested deeper in the
* bucket with a timestamped path. This is useful when periodical backups are desired
* <li>If outputUrlPrefix is a already a nested path in a GCS bucket, exported data will be put
* under this path. This means that a nested path is not reusable, since the export process
* by default would not overwrite existing files.
* </ul>
*
* @param outputUrlPrefix the full resource URL of the external storage location
* @param kinds the datastore 'kinds' to be exported. If empty, all kinds will be exported
*/
public Export export(String outputUrlPrefix, Collection<String> kinds) {
return new Export(new ExportRequest(outputUrlPrefix, kinds));
}
/**
* Imports the entire backup specified by {@code backupUrl} back to Cloud Datastore.
*
* <p>A successful backup restores deleted entities and reverts updates to existing entities since
* the backup time. However, it does not affect newly added entities.
*/
public Import importBackup(String backupUrl) {
return new Import(new ImportRequest(backupUrl, ImmutableList.of()));
}
/**
* Imports the backup specified by {@code backupUrl} back to Cloud Datastore. Only entities whose
* types are included in {@code kinds} are imported.
*
* @see #importBackup(String)
*/
public Import importBackup(String backupUrl, Collection<String> kinds) {
return new Import(new ImportRequest(backupUrl, kinds));
}
/**
* Returns a {@link Get} request that retrieves the details of an export or import {@link
* Operation}.
*
* @param operationName name of the {@code Operation} as returned by an export or import request
*/
public Get get(String operationName) {
return new Get(operationName);
}
/**
* Returns a {@link ListOperations} request that retrieves all export or import {@link Operation
* operations} matching {@code filter}.
*
* <p>Sample usage: find all operations started after 2018-10-31 00:00:00 UTC and has stopped:
*
* <pre>
* {@code String filter = "metadata.common.startTime>\"2018-10-31T0:0:0Z\" AND done=true";}
* {@code List<Operation> operations = datastoreAdmin.list(filter);}
* </pre>
*
* <p>Please refer to {@link Operation} for how to reference operation properties.
*/
public ListOperations list(String filter) {
checkArgument(!Strings.isNullOrEmpty(filter), "Filter must not be null or empty.");
return new ListOperations(Optional.of(filter));
}
/**
* Returns a {@link ListOperations} request that retrieves all export or import {@link Operation *
* operations}.
*/
public ListOperations listAll() {
return new ListOperations(Optional.empty());
}
/** Builder for {@link DatastoreAdmin}. */
public static class Builder extends AbstractGoogleJsonClient.Builder {
private String projectId;
public Builder(
HttpTransport httpTransport,
JsonFactory jsonFactory,
HttpRequestInitializer httpRequestInitializer) {
super(httpTransport, jsonFactory, ROOT_URL, SERVICE_PATH, httpRequestInitializer, false);
}
@Override
public Builder setApplicationName(String applicationName) {
return (Builder) super.setApplicationName(applicationName);
}
/** Sets the GCP project ID of the Cloud Datastore databases being managed. */
public Builder setProjectId(String projectId) {
this.projectId = projectId;
return this;
}
@Override
public DatastoreAdmin build() {
return new DatastoreAdmin(this);
}
}
/** A request to export Cloud Datastore databases. */
public class Export extends DatastoreAdminRequest<Operation> {
Export(ExportRequest exportRequest) {
super(
DatastoreAdmin.this,
"POST",
"projects/{projectId}:export",
exportRequest,
Operation.class);
set("projectId", projectId);
}
}
/** A request to restore an backup to a Cloud Datastore database. */
public class Import extends DatastoreAdminRequest<Operation> {
Import(ImportRequest importRequest) {
super(
DatastoreAdmin.this,
"POST",
"projects/{projectId}:import",
importRequest,
Operation.class);
set("projectId", projectId);
}
}
/** A request to retrieve details of an export or import operation. */
public class Get extends DatastoreAdminRequest<Operation> {
Get(String operationName) {
super(DatastoreAdmin.this, "GET", operationName, null, Operation.class);
}
}
/** A request to retrieve all export or import operations matching a given filter. */
public class ListOperations extends DatastoreAdminRequest<Operation.OperationList> {
ListOperations(Optional<String> filter) {
super(
DatastoreAdmin.this,
"GET",
"projects/{projectId}/operations",
null,
Operation.OperationList.class);
set("projectId", projectId);
filter.ifPresent(f -> set("filter", f));
}
}
/** Base class of all DatastoreAdmin requests. */
abstract static class DatastoreAdminRequest<T> extends AbstractGoogleJsonClientRequest<T> {
/**
* @param client Google JSON client
* @param requestMethod HTTP Method
* @param uriTemplate URI template for the path relative to the base URL. If it starts with a
* "/" the base path from the base URL will be stripped out. The URI template can also be a
* full URL. URI template expansion is done using {@link
* com.google.api.client.http.UriTemplate#expand(String, String, Object, boolean)}
* @param jsonContent POJO that can be serialized into JSON content or {@code null} for none
* @param responseClass response class to parse into
*/
protected DatastoreAdminRequest(
DatastoreAdmin client,
String requestMethod,
String uriTemplate,
Object jsonContent,
Class<T> responseClass) {
super(client, requestMethod, uriTemplate, jsonContent, responseClass);
}
}
/**
* Model object that describes the JSON content in an export request.
*
* <p>Please note that some properties defined in the API are excluded, e.g., {@code databaseId}
* (not supported by Cloud Datastore) and labels (not used by Domain Registry).
*/
@SuppressWarnings("unused")
static class ExportRequest extends GenericJson {
@Key private final String outputUrlPrefix;
@Key private final EntityFilter entityFilter;
ExportRequest(String outputUrlPrefix, Collection<String> kinds) {
checkNotNull(outputUrlPrefix, "outputUrlPrefix");
checkArgument(!kinds.isEmpty(), "kinds must not be empty");
this.outputUrlPrefix = outputUrlPrefix;
this.entityFilter = new EntityFilter(kinds);
}
}
/**
* Model object that describes the JSON content in an export request.
*
* <p>Please note that some properties defined in the API are excluded, e.g., {@code databaseId}
* (not supported by Cloud Datastore) and labels (not used by Domain Registry).
*/
@SuppressWarnings("unused")
static class ImportRequest extends GenericJson {
@Key private final String inputUrl;
@Key private final EntityFilter entityFilter;
ImportRequest(String inputUrl, Collection<String> kinds) {
checkNotNull(inputUrl, "outputUrlPrefix");
this.inputUrl = inputUrl;
this.entityFilter = new EntityFilter(kinds);
}
}
}

43
core/src/main/java/google/registry/export/datastore/DatastoreAdminModule.java
View file

@ -1,43 +0,0 @@
// Copyright 2018 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.export.datastore;
import dagger.Module;
import dagger.Provides;
import google.registry.config.CredentialModule;
import google.registry.config.RegistryConfig;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.util.GoogleCredentialsBundle;
import javax.inject.Singleton;
/** Dagger module that configures provision of {@link DatastoreAdmin}. */
@Module
@DeleteAfterMigration
public abstract class DatastoreAdminModule {
@Singleton
@Provides
static DatastoreAdmin provideDatastoreAdmin(
@CredentialModule.DefaultCredential GoogleCredentialsBundle credentialsBundle,
@RegistryConfig.Config("projectId") String projectId) {
return new DatastoreAdmin.Builder(
credentialsBundle.getHttpTransport(),
credentialsBundle.getJsonFactory(),
credentialsBundle.getHttpRequestInitializer())
.setApplicationName(projectId)
.setProjectId(projectId)
.build();
}
}

49
core/src/main/java/google/registry/export/datastore/EntityFilter.java
View file

@ -1,49 +0,0 @@
// Copyright 2018 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.export.datastore;
import static com.google.common.base.Preconditions.checkNotNull;
import com.google.api.client.json.GenericJson;
import com.google.api.client.util.Key;
import com.google.common.collect.ImmutableList;
import google.registry.model.annotations.DeleteAfterMigration;
import java.util.Collection;
import java.util.List;
/**
* Model object that describes the Cloud Datastore 'kinds' to be exported or imported. The JSON form
* of this type is found in export/import requests and responses.
*
* <p>Please note that properties not used by Domain Registry are not included, e.g., {@code
* namespaceIds}.
*/
@DeleteAfterMigration
public class EntityFilter extends GenericJson {
@Key private List<String> kinds = ImmutableList.of();
/** For JSON deserialization. */
public EntityFilter() {}
EntityFilter(Collection<String> kinds) {
checkNotNull(kinds, "kinds");
this.kinds = ImmutableList.copyOf(kinds);
}
List<String> getKinds() {
return ImmutableList.copyOf(kinds);
}
}

250
core/src/main/java/google/registry/export/datastore/Operation.java
View file

@ -1,250 +0,0 @@
// Copyright 2018 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.export.datastore;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Strings.isNullOrEmpty;
import com.google.api.client.json.GenericJson;
import com.google.api.client.util.Key;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import google.registry.export.datastore.DatastoreAdmin.Get;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.util.Clock;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nullable;
import org.joda.time.DateTime;
import org.joda.time.Duration;
/**
* Model object that describes the details of an export or import operation in Cloud Datastore.
*
* <p>{@link Operation} instances are parsed from the JSON payload in Datastore response messages.
*/
@DeleteAfterMigration
public class Operation extends GenericJson {
private static final String STATE_SUCCESS = "SUCCESSFUL";
private static final String STATE_PROCESSING = "PROCESSING";
@Key private String name;
@Key private Metadata metadata;
@Key private boolean done;
/** For JSON deserialization. */
public Operation() {}
/** Returns the name of this operation, which may be used in a {@link Get} request. */
public String getName() {
checkState(name != null, "Name must not be null.");
return name;
}
public boolean isExport() {
return !isNullOrEmpty(getExportFolderUrl());
}
public boolean isImport() {
return !isNullOrEmpty(getMetadata().getInputUrl());
}
public boolean isDone() {
return done;
}
private String getState() {
return getMetadata().getCommonMetadata().getState();
}
public boolean isSuccessful() {
return getState().equals(STATE_SUCCESS);
}
public boolean isProcessing() {
return getState().equals(STATE_PROCESSING);
}
/**
* Returns the elapsed time since starting if this operation is still running, or the total
* running time if this operation has completed.
*/
public Duration getRunningTime(Clock clock) {
return new Duration(
getStartTime(), getMetadata().getCommonMetadata().getEndTime().orElse(clock.nowUtc()));
}
public DateTime getStartTime() {
return getMetadata().getCommonMetadata().getStartTime();
}
public ImmutableSet<String> getKinds() {
return ImmutableSet.copyOf(getMetadata().getEntityFilter().getKinds());
}
/**
* Returns the URL to the GCS folder that holds the exported data. This folder is created by
* Datastore and is under the {@code outputUrlPrefix} set to {@linkplain
* DatastoreAdmin#export(String, java.util.Collection) the export request}.
*
* @throws IllegalStateException if this is not an export operation
*/
public String getExportFolderUrl() {
return getMetadata().getOutputUrlPrefix();
}
/**
* Returns the last segment of the {@linkplain #getExportFolderUrl() export folder URL} which can
* be used as unique identifier of this export operation. This is a better ID than the {@linkplain
* #getName() operation name}, which is opaque.
*
* @throws IllegalStateException if this is not an export operation
*/
public String getExportId() {
String exportFolderUrl = getExportFolderUrl();
return exportFolderUrl.substring(exportFolderUrl.lastIndexOf('/') + 1);
}
public String getProgress() {
StringBuilder result = new StringBuilder();
Progress progress = getMetadata().getProgressBytes();
if (progress != null) {
result.append(
String.format(" [%s/%s bytes]", progress.workCompleted, progress.workEstimated));
}
progress = getMetadata().getProgressEntities();
if (progress != null) {
result.append(
String.format(" [%s/%s entities]", progress.workCompleted, progress.workEstimated));
}
if (result.length() == 0) {
return "Progress: N/A";
}
return "Progress:" + result;
}
private Metadata getMetadata() {
checkState(metadata != null, "Response metadata missing.");
return metadata;
}
/** Models the common metadata properties of all operations. */
public static class CommonMetadata extends GenericJson {
@Key private String startTime;
@Key @Nullable private String endTime;
@Key private String operationType;
@Key private String state;
public CommonMetadata() {}
String getOperationType() {
checkState(!isNullOrEmpty(operationType), "operationType may not be null or empty");
return operationType;
}
String getState() {
checkState(!isNullOrEmpty(state), "state may not be null or empty");
return state;
}
DateTime getStartTime() {
checkState(startTime != null, "StartTime missing.");
return DateTime.parse(startTime);
}
Optional<DateTime> getEndTime() {
return Optional.ofNullable(endTime).map(DateTime::parse);
}
}
/** Models the metadata of a Cloud Datatore export or import operation. */
public static class Metadata extends GenericJson {
@Key("common")
private CommonMetadata commonMetadata;
@Key private Progress progressEntities;
@Key private Progress progressBytes;
@Key private EntityFilter entityFilter;
@Key private String inputUrl;
@Key private String outputUrlPrefix;
public Metadata() {}
CommonMetadata getCommonMetadata() {
checkState(commonMetadata != null, "CommonMetadata field is null.");
return commonMetadata;
}
public Progress getProgressEntities() {
return progressEntities;
}
public Progress getProgressBytes() {
return progressBytes;
}
public EntityFilter getEntityFilter() {
return entityFilter;
}
public String getInputUrl() {
return checkUrls().inputUrl;
}
public String getOutputUrlPrefix() {
return checkUrls().outputUrlPrefix;
}
Metadata checkUrls() {
checkState(
isNullOrEmpty(inputUrl) || isNullOrEmpty(outputUrlPrefix),
"inputUrl and outputUrlPrefix must not be both present");
checkState(
!isNullOrEmpty(inputUrl) || !isNullOrEmpty(outputUrlPrefix),
"inputUrl and outputUrlPrefix must not be both missing");
return this;
}
}
/** Progress of an export or import operation. */
public static class Progress extends GenericJson {
@Key private long workCompleted;
@Key private long workEstimated;
public Progress() {}
long getWorkCompleted() {
return workCompleted;
}
public long getWorkEstimated() {
return workEstimated;
}
}
/** List of {@link Operation Operations}. */
public static class OperationList extends GenericJson {
@Key private List<Operation> operations;
/** For JSON deserialization. */
public OperationList() {}
ImmutableList<Operation> toList() {
return ImmutableList.copyOf(operations);
}
}
}

4
core/src/main/java/google/registry/export/sheet/SyncRegistrarsSheet.java
View file

@ -25,7 +25,6 @@ import static google.registry.model.registrar.RegistrarPoc.Type.MARKETING;
import static google.registry.model.registrar.RegistrarPoc.Type.TECH;
import static google.registry.model.registrar.RegistrarPoc.Type.WHOIS;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.DateTimeUtils.START_OF_TIME;
import com.google.common.base.Joiner;
@ -63,8 +62,7 @@ class SyncRegistrarsSheet {
*/
boolean wereRegistrarsModified() {
Optional<Cursor> cursor =
transactIfJpaTm(
() -> tm().loadByKeyIfPresent(Cursor.createGlobalVKey(SYNC_REGISTRAR_SHEET)));
tm().transact(() -> tm().loadByKeyIfPresent(Cursor.createGlobalVKey(SYNC_REGISTRAR_SHEET)));
DateTime lastUpdateTime = !cursor.isPresent() ? START_OF_TIME : cursor.get().getCursorTime();
for (Registrar registrar : Registrar.loadAllCached()) {
if (DateTimeUtils.isAtOrAfter(registrar.getLastUpdateTime(), lastUpdateTime)) {

2
core/src/main/java/google/registry/flows/FlowUtils.java
View file

@ -53,7 +53,7 @@ public final class FlowUtils {
}
}
/** Persists the saves and deletes in an {@link EntityChanges} to Datastore. */
/** Persists the saves and deletes in an {@link EntityChanges} to the DB. */
public static void persistEntityChanges(EntityChanges entityChanges) {
tm().putAll(entityChanges.getSaves());
tm().delete(entityChanges.getDeletes());

3
core/src/main/java/google/registry/flows/ResourceFlowUtils.java
View file

@ -19,7 +19,6 @@ import static google.registry.model.EppResourceUtils.isLinked;
import static google.registry.model.EppResourceUtils.loadByForeignKey;
import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
@ -176,7 +175,7 @@ public final class ResourceFlowUtils {
throw new BadAuthInfoForResourceException();
}
// Check the authInfo against the contact.
verifyAuthInfo(authInfo, transactIfJpaTm(() -> tm().loadByKey(foundContact.get())));
verifyAuthInfo(authInfo, tm().transact(() -> tm().loadByKey(foundContact.get())));
}
/** Check that the given {@link AuthInfo} is valid for the given contact. */

4
core/src/main/java/google/registry/flows/TransactionalFlow.java
View file

@ -17,7 +17,7 @@ package google.registry.flows;
/**
* Interface for a {@link Flow} that needs to be run transactionally.
*
* <p>Any flow that mutates Datastore should implement this so that {@link FlowRunner} will know how
* to run it.
* <p>Any flow that mutates the DB should implement this so that {@link FlowRunner} will know how to
* run it.
*/
public interface TransactionalFlow extends Flow {}

5
core/src/main/java/google/registry/flows/domain/DomainInfoFlow.java
View file

@ -22,7 +22,6 @@ import static google.registry.flows.domain.DomainFlowUtils.handleFeeRequest;
import static google.registry.flows.domain.DomainFlowUtils.loadForeignKeyedDesignatedContacts;
import static google.registry.model.EppResourceUtils.loadByForeignKey;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
@ -110,7 +109,7 @@ public final class DomainInfoFlow implements Flow {
.setRepoId(domain.getRepoId())
.setCurrentSponsorClientId(domain.getCurrentSponsorRegistrarId())
.setRegistrant(
transactIfJpaTm(() -> tm().loadByKey(domain.getRegistrant())).getContactId());
tm().transact(() -> tm().loadByKey(domain.getRegistrant())).getContactId());
// If authInfo is non-null, then the caller is authorized to see the full information since we
// will have already verified the authInfo is valid.
if (registrarId.equals(domain.getCurrentSponsorRegistrarId()) || authInfo.isPresent()) {
@ -118,7 +117,7 @@ public final class DomainInfoFlow implements Flow {
infoBuilder
.setStatusValues(domain.getStatusValues())
.setContacts(
transactIfJpaTm(() -> loadForeignKeyedDesignatedContacts(domain.getContacts())))
tm().transact(() -> loadForeignKeyedDesignatedContacts(domain.getContacts())))
.setNameservers(hostsRequest.requestDelegated() ? domain.loadNameserverHostNames() : null)
.setSubordinateHosts(
hostsRequest.requestSubordinate() ? domain.getSubordinateHosts() : null)

2
core/src/main/java/google/registry/flows/domain/DomainTransferRejectFlow.java
View file

@ -63,7 +63,7 @@ import org.joda.time.DateTime;
* transfer is automatically approved. Within that window, this flow allows the losing client to
* reject the transfer request.
*
* <p>When the transfer was requested, poll messages and billing events were saved to Datastore with
* <p>When the transfer was requested, poll messages and billing events were saved to SQL with
* timestamps such that they only would become active when the transfer period passed. In this flow,
* those speculative objects are deleted.
*

3
core/src/main/java/google/registry/flows/domain/token/AllocationTokenFlowUtils.java
View file

@ -16,7 +16,6 @@ package google.registry.flows.domain.token;
import static com.google.common.base.Preconditions.checkArgument;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
@ -161,7 +160,7 @@ public class AllocationTokenFlowUtils {
throw new InvalidAllocationTokenException();
}
Optional<AllocationToken> maybeTokenEntity =
transactIfJpaTm(() -> tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token)));
tm().transact(() -> tm().loadByKeyIfPresent(VKey.create(AllocationToken.class, token)));
if (!maybeTokenEntity.isPresent()) {
throw new InvalidAllocationTokenException();
}

5
core/src/main/java/google/registry/flows/host/HostInfoFlow.java
View file

@ -19,7 +19,6 @@ import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
import static google.registry.flows.host.HostFlowUtils.validateHostName;
import static google.registry.model.EppResourceUtils.isLinked;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.collect.ImmutableSet;
import google.registry.flows.EppException;
@ -78,8 +77,8 @@ public final class HostInfoFlow implements Flow {
// there is no superordinate domain, the host's own values for these fields will be correct.
if (host.isSubordinate()) {
DomainBase superordinateDomain =
transactIfJpaTm(
() -> tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now));
tm().transact(
() -> tm().loadByKey(host.getSuperordinateDomain()).cloneProjectedAtTime(now));
hostInfoDataBuilder
.setCurrentSponsorClientId(superordinateDomain.getCurrentSponsorRegistrarId())
.setLastTransferTime(host.computeLastTransferTime(superordinateDomain));

7
core/src/main/java/google/registry/flows/poll/PollFlowUtils.java
View file

@ -18,7 +18,6 @@ import static com.google.common.base.Preconditions.checkArgument;
import static google.registry.persistence.transaction.QueryComposer.Comparator.EQ;
import static google.registry.persistence.transaction.QueryComposer.Comparator.LTE;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
import google.registry.model.poll.PollMessage;
@ -31,13 +30,13 @@ public final class PollFlowUtils {
/** Returns the number of poll messages for the given registrar that are not in the future. */
public static int getPollMessageCount(String registrarId, DateTime now) {
return transactIfJpaTm(() -> createPollMessageQuery(registrarId, now).count()).intValue();
return tm().transact(() -> createPollMessageQuery(registrarId, now).count()).intValue();
}
/** Returns the first (by event time) poll message not in the future for this registrar. */
public static Optional<PollMessage> getFirstPollMessage(String registrarId, DateTime now) {
return transactIfJpaTm(
() -> createPollMessageQuery(registrarId, now).orderBy("eventTime").first());
return tm().transact(
() -> createPollMessageQuery(registrarId, now).orderBy("eventTime").first());
}
/**

3
core/src/main/java/google/registry/model/EppResourceUtils.java
View file

@ -19,7 +19,6 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static google.registry.model.ofy.ObjectifyService.auditedOfy;
import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.DateTimeUtils.START_OF_TIME;
import static google.registry.util.DateTimeUtils.isAtOrAfter;
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
@ -163,7 +162,7 @@ public final class EppResourceUtils {
T resource =
useCache
? EppResource.loadCached(fki.getResourceKey())
: transactIfJpaTm(() -> tm().loadByKeyIfPresent(fki.getResourceKey()).orElse(null));
: tm().transact(() -> tm().loadByKeyIfPresent(fki.getResourceKey()).orElse(null));
if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
return Optional.empty();
}

1
core/src/main/java/google/registry/model/UnsafeSerializable.java
View file

@ -21,7 +21,6 @@ import java.io.Serializable;
* entities are made {@link Serializable} so that they can be passed between JVMs. The intended use
* case is BEAM pipeline-based cross-database data validation between Datastore and Cloud SQL during
* the migration. Note that only objects loaded from the SQL database need serialization support.
* Objects exported from Datastore can already be serialized as protocol buffers.
*
* <p>All entities implementing this interface take advantage of the fact that all Java collection
* classes we use, either directly or indirectly, including those in Java libraries, Guava,

3
core/src/main/java/google/registry/model/UpdateAutoTimestamp.java
View file

@ -19,7 +19,6 @@ import static google.registry.util.DateTimeUtils.START_OF_TIME;
import com.googlecode.objectify.annotation.Ignore;
import com.googlecode.objectify.annotation.OnLoad;
import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
import google.registry.util.DateTimeUtils;
import java.time.ZonedDateTime;
import java.util.Optional;
@ -34,8 +33,6 @@ import org.joda.time.DateTime;
/**
* A timestamp that auto-updates on each save to Datastore/Cloud SQL.
*
* @see UpdateAutoTimestampTranslatorFactory
*/
@Embeddable
public class UpdateAutoTimestamp extends ImmutableObject implements UnsafeSerializable {

3
core/src/main/java/google/registry/model/billing/BillingEvent.java
View file

@ -19,7 +19,6 @@ 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.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.CollectionUtils.forceEmptyToNull;
import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
import static google.registry.util.DateTimeUtils.END_OF_TIME;
@ -909,7 +908,7 @@ public abstract class BillingEvent extends ImmutableObject
checkNotNull(instance.reason);
checkNotNull(instance.eventRef);
BillingEvent.OneTime billingEvent =
transactIfJpaTm(() -> tm().loadByKey(VKey.from(instance.eventRef)));
tm().transact(() -> tm().loadByKey(VKey.from(instance.eventRef)));
checkArgument(
Objects.equals(instance.cost.getCurrencyUnit(), billingEvent.cost.getCurrencyUnit()),
"Referenced billing event is in a different currency");

11
core/src/main/java/google/registry/model/domain/DomainContent.java
View file

@ -24,7 +24,6 @@ import static com.google.common.collect.Sets.intersection;
import static google.registry.model.EppResourceUtils.projectResourceOntoBuilderAtTime;
import static google.registry.model.EppResourceUtils.setAutomaticTransferSuccessProperties;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.CollectionUtils.forceEmptyToNull;
import static google.registry.util.CollectionUtils.nullToEmpty;
import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
@ -604,11 +603,11 @@ public class DomainContent extends EppResource
/** Loads and returns the fully qualified host names of all linked nameservers. */
public ImmutableSortedSet<String> loadNameserverHostNames() {
return transactIfJpaTm(
() ->
tm().loadByKeys(getNameservers()).values().stream()
.map(HostResource::getHostName)
.collect(toImmutableSortedSet(Ordering.natural())));
return tm().transact(
() ->
tm().loadByKeys(getNameservers()).values().stream()
.map(HostResource::getHostName)
.collect(toImmutableSortedSet(Ordering.natural())));
}
/** A key to the registrant who registered this domain. */

509
core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java
View file

@ -1,509 +0,0 @@
// Copyright 2019 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.model.ofy;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
import static google.registry.model.ofy.ObjectifyService.auditedOfy;
import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
import com.google.common.base.Functions;
import com.google.common.collect.ImmutableCollection;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Streams;
import com.googlecode.objectify.Key;
import com.googlecode.objectify.Result;
import com.googlecode.objectify.cmd.Query;
import google.registry.model.ImmutableObject;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.model.annotations.InCrossTld;
import google.registry.model.contact.ContactHistory;
import google.registry.model.domain.DomainHistory;
import google.registry.model.host.HostHistory;
import google.registry.model.reporting.HistoryEntry;
import google.registry.persistence.VKey;
import google.registry.persistence.transaction.QueryComposer;
import google.registry.persistence.transaction.TransactionManager;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.function.Supplier;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;
import javax.annotation.Nullable;
import javax.persistence.NoResultException;
import javax.persistence.NonUniqueResultException;
import org.joda.time.DateTime;
/** Datastore implementation of {@link TransactionManager}. */
@DeleteAfterMigration
public class DatastoreTransactionManager implements TransactionManager {
private Ofy injectedOfy;
/** Constructs an instance. */
public DatastoreTransactionManager(Ofy injectedOfy) {
this.injectedOfy = injectedOfy;
}
private Ofy getOfy() {
return injectedOfy == null ? auditedOfy() : injectedOfy;
}
@Override
public boolean inTransaction() {
return getOfy().inTransaction();
}
@Override
public void assertInTransaction() {
getOfy().assertInTransaction();
}
@Override
public <T> T transact(Supplier<T> work) {
return getOfy().transact(work);
}
@Override
public void transact(Runnable work) {
getOfy().transact(work);
}
@Override
public <T> T transactNew(Supplier<T> work) {
return getOfy().transactNew(work);
}
@Override
public void transactNew(Runnable work) {
getOfy().transactNew(work);
}
@Override
public <R> R transactNewReadOnly(Supplier<R> work) {
return getOfy().transactNewReadOnly(work);
}
@Override
public void transactNewReadOnly(Runnable work) {
getOfy().transactNewReadOnly(work);
}
@Override
public <R> R doTransactionless(Supplier<R> work) {
return getOfy().doTransactionless(work);
}
@Override
public DateTime getTransactionTime() {
return getOfy().getTransactionTime();
}
@Override
public void insert(Object entity) {
put(entity);
}
@Override
public void insertAll(ImmutableCollection<?> entities) {
putAll(entities);
}
@Override
public void insertAll(ImmutableObject... entities) {
putAll(entities);
}
@Override
public void insertWithoutBackup(ImmutableObject entity) {
putWithoutBackup(entity);
}
@Override
public void insertAllWithoutBackup(ImmutableCollection<?> entities) {
putAllWithoutBackup(entities);
}
@Override
public void put(Object entity) {
saveEntity(entity);
}
@Override
public void putAll(ImmutableObject... entities) {
syncIfTransactionless(
getOfy().save().entities(toDatastoreEntities(ImmutableList.copyOf(entities))));
}
@Override
public void putAll(ImmutableCollection<?> entities) {
syncIfTransactionless(getOfy().save().entities(toDatastoreEntities(entities)));
}
@Override
public void putWithoutBackup(ImmutableObject entity) {
syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntity(entity)));
}
@Override
public void putAllWithoutBackup(ImmutableCollection<?> entities) {
syncIfTransactionless(getOfy().saveWithoutBackup().entities(toDatastoreEntities(entities)));
}
@Override
public void update(Object entity) {
put(entity);
}
@Override
public void updateAll(ImmutableCollection<?> entities) {
putAll(entities);
}
@Override
public void updateAll(ImmutableObject... entities) {
updateAll(ImmutableList.copyOf(entities));
}
@Override
public void updateWithoutBackup(ImmutableObject entity) {
putWithoutBackup(entity);
}
@Override
public void updateAllWithoutBackup(ImmutableCollection<?> entities) {
putAllWithoutBackup(entities);
}
@Override
public boolean exists(Object entity) {
return getOfy().load().key(Key.create(toDatastoreEntity(entity))).now() != null;
}
@Override
public <T> boolean exists(VKey<T> key) {
return loadNullable(key) != null;
}
// TODO: add tests for these methods. They currently have some degree of test coverage because
// they are used when retrieving the nameservers which require these, as they are now loaded by
// VKey instead of by ofy Key. But ideally, there should be one set of TransactionManager
// interface tests that are applied to both the datastore and SQL implementations.
@Override
public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
return Optional.ofNullable(loadNullable(key));
}
@Override
public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
Iterable<? extends VKey<? extends T>> keys) {
// Keep track of the Key -> VKey mapping so we can translate them back.
ImmutableMap<Key<T>, VKey<? extends T>> keyMap =
StreamSupport.stream(keys.spliterator(), false)
.distinct()
.collect(toImmutableMap(key -> (Key<T>) key.getOfyKey(), Functions.identity()));
return getOfy().load().keys(keyMap.keySet()).entrySet().stream()
.collect(
toImmutableMap(
entry -> keyMap.get(entry.getKey()), entry -> toSqlEntity(entry.getValue())));
}
@SuppressWarnings("unchecked")
@Override
public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
return getOfy()
.load()
.entities(toDatastoreEntities(ImmutableList.copyOf(entities)))
.values()
.stream()
.map(DatastoreTransactionManager::toSqlEntity)
.map(entity -> (T) entity)
.collect(toImmutableList());
}
@Override
public <T> T loadByKey(VKey<T> key) {
T result = loadNullable(key);
if (result == null) {
throw new NoSuchElementException(key.toString());
}
return result;
}
@Override
public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
Iterable<? extends VKey<? extends T>> keys) {
ImmutableMap<VKey<? extends T>, T> result = loadByKeysIfPresent(keys);
ImmutableSet<? extends VKey<? extends T>> missingKeys =
Streams.stream(keys).filter(k -> !result.containsKey(k)).collect(toImmutableSet());
if (!missingKeys.isEmpty()) {
// Ofy ignores nonexistent keys but the method contract specifies to throw if nonexistent
throw new NoSuchElementException(
String.format("Failed to load nonexistent entities for keys: %s", missingKeys));
}
return result;
}
@SuppressWarnings("unchecked")
@Override
public <T> T loadByEntity(T entity) {
return (T) toSqlEntity(auditedOfy().load().entity(toDatastoreEntity(entity)).now());
}
@Override
public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
ImmutableList<T> result = loadByEntitiesIfPresent(entities);
if (result.size() != Iterables.size(entities)) {
throw new NoSuchElementException(
String.format("Attempted to load entities, some of which are missing: %s", entities));
}
return result;
}
@Override
public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
return ImmutableList.copyOf(getPossibleAncestorQuery(clazz));
}
@Override
public <T> Stream<T> loadAllOfStream(Class<T> clazz) {
return Streams.stream(getPossibleAncestorQuery(clazz));
}
@Override
public <T> Optional<T> loadSingleton(Class<T> clazz) {
List<T> elements = getPossibleAncestorQuery(clazz).limit(2).list();
checkArgument(
elements.size() <= 1,
"Expected at most one entity of type %s, found at least two",
clazz.getSimpleName());
return elements.stream().findFirst();
}
@Override
public void delete(VKey<?> key) {
syncIfTransactionless(getOfy().delete().key(key.getOfyKey()));
}
@Override
public void delete(Iterable<? extends VKey<?>> vKeys) {
// We have to create a list to work around the wildcard capture issue here.
// See https://docs.oracle.com/javase/tutorial/java/generics/capture.html
ImmutableList<Key<?>> list =
StreamSupport.stream(vKeys.spliterator(), false)
.map(VKey::getOfyKey)
.collect(toImmutableList());
syncIfTransactionless(getOfy().delete().keys(list));
}
@Override
public <T> T delete(T entity) {
syncIfTransactionless(getOfy().delete().entity(toDatastoreEntity(entity)));
return entity;
}
@Override
public void deleteWithoutBackup(VKey<?> key) {
syncIfTransactionless(getOfy().deleteWithoutBackup().key(key.getOfyKey()));
}
@Override
public void deleteWithoutBackup(Iterable<? extends VKey<?>> keys) {
syncIfTransactionless(
getOfy()
.deleteWithoutBackup()
.keys(Streams.stream(keys).map(VKey::getOfyKey).collect(toImmutableList())));
}
@Override
public void deleteWithoutBackup(Object entity) {
syncIfTransactionless(getOfy().deleteWithoutBackup().entity(toDatastoreEntity(entity)));
}
@Override
public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
return new DatastoreQueryComposerImpl<>(entity);
}
@Override
public void clearSessionCache() {
getOfy().clearSessionCache();
}
@Override
public boolean isOfy() {
return true;
}
/**
* Executes the given {@link Result} instance synchronously if not in a transaction.
*
* <p>The {@link Result} instance contains a task that will be executed by Objectify
* asynchronously. If it is in a transaction, we don't need to execute the task immediately
* because it is guaranteed to be done by the end of the transaction. However, if it is not in a
* transaction, we need to execute it in case the following code expects that happens before
* themselves.
*/
private void syncIfTransactionless(Result<?> result) {
if (!inTransaction()) {
result.now();
}
}
/**
* The following three methods exist due to the migration to Cloud SQL.
*
* <p>In Cloud SQL, {@link HistoryEntry} objects are represented instead as {@link DomainHistory},
* {@link ContactHistory}, and {@link HostHistory} objects. During the migration, we do not wish
* to change the Datastore schema so all of these objects are stored in Datastore as HistoryEntry
* objects. They are converted to/from the appropriate classes upon retrieval, and converted to
* HistoryEntry on save. See go/r3.0-history-objects for more details.
*/
private void saveEntity(Object entity) {
checkArgumentNotNull(entity, "entity must be specified");
syncIfTransactionless(getOfy().save().entity(toDatastoreEntity(entity)));
}
@Nullable
private <T> T loadNullable(VKey<T> key) {
return toSqlEntity(getOfy().load().key(key.getOfyKey()).now());
}
/** Converts a possible {@link HistoryEntry} child to a {@link HistoryEntry}. */
private static Object toDatastoreEntity(@Nullable Object obj) {
if (obj instanceof HistoryEntry) {
return ((HistoryEntry) obj).asHistoryEntry();
}
return obj;
}
private static ImmutableList<Object> toDatastoreEntities(ImmutableCollection<?> collection) {
return collection.stream()
.map(DatastoreTransactionManager::toDatastoreEntity)
.collect(toImmutableList());
}
/**
* Converts an object to the corresponding child {@link HistoryEntry} if necessary and possible.
*
* <p>This should be used when returning objects from Datastore to make sure they reflect the most
* recent type of the object in question.
*/
@SuppressWarnings("unchecked")
public static <T> T toSqlEntity(@Nullable T obj) {
if (obj instanceof HistoryEntry) {
return (T) ((HistoryEntry) obj).toChildHistoryEntity();
}
return obj;
}
/** A query for returning any/all results of an object, with an ancestor if possible. */
private <T> Query<T> getPossibleAncestorQuery(Class<T> clazz) {
Query<T> query = getOfy().load().type(clazz);
// If the entity is in the cross-TLD entity group, then we can take advantage of an ancestor
// query to give us strong transactional consistency.
if (clazz.isAnnotationPresent(InCrossTld.class)) {
query = query.ancestor(getCrossTldKey());
}
return query;
}
private static class DatastoreQueryComposerImpl<T> extends QueryComposer<T> {
DatastoreQueryComposerImpl(Class<T> entityClass) {
super(entityClass);
}
Query<T> buildQuery() {
checkOnlyOneInequalityField();
Query<T> result = auditedOfy().load().type(entityClass);
for (WhereClause pred : predicates) {
String comparatorString = pred.comparator.getDatastoreString();
if (comparatorString == null) {
throw new UnsupportedOperationException(
String.format("The %s operation is not supported on Datastore.", pred.comparator));
}
result = result.filter(pred.fieldName + comparatorString, pred.value);
}
if (orderBy != null) {
result = result.order(orderBy);
}
return result;
}
@Override
public Optional<T> first() {
return Optional.ofNullable(buildQuery().limit(1).first().now());
}
@Override
public T getSingleResult() {
List<T> results = buildQuery().limit(2).list();
if (results.size() == 0) {
// The exception text here is the same as what we get for JPA queries.
throw new NoResultException("No entity found for query");
} else if (results.size() > 1) {
throw new NonUniqueResultException("More than one result found for getSingleResult query");
}
return results.get(0);
}
@Override
public Stream<T> stream() {
return Streams.stream(buildQuery());
}
@Override
public long count() {
// Objectify provides a count() function, but unfortunately that doesn't work if there are
// more than 1000 (the default response page size?) entries in the result set. We also use
// chunkAll() here as it provides a nice performance boost.
//
// There is some information on this issue on SO, see:
// https://stackoverflow.com/questions/751124/how-does-one-get-a-count-of-rows-in-a-datastore-model-in-google-app-engine
return Iterables.size(buildQuery().chunkAll().keys());
}
@Override
public ImmutableList<T> list() {
return ImmutableList.copyOf(buildQuery().list());
}
private void checkOnlyOneInequalityField() {
// Datastore inequality queries are limited to one property, see
// https://cloud.google.com/appengine/docs/standard/go111/datastore/query-restrictions#inequality_filters_are_limited_to_at_most_one_property
long numInequalityFields =
predicates.stream()
.filter(pred -> !pred.comparator.equals(Comparator.EQ))
.map(pred -> pred.fieldName)
.distinct()
.count();
checkArgument(
numInequalityFields <= 1,
"Datastore cannot handle inequality queries on multiple fields, we found %s fields.",
numInequalityFields);
}
}
}

4
core/src/main/java/google/registry/model/ofy/ObjectifyService.java
View file

@ -44,7 +44,6 @@ import google.registry.model.translators.DurationTranslatorFactory;
import google.registry.model.translators.EppHistoryVKeyTranslatorFactory;
import google.registry.model.translators.InetAddressTranslatorFactory;
import google.registry.model.translators.ReadableInstantUtcTranslatorFactory;
import google.registry.model.translators.UpdateAutoTimestampTranslatorFactory;
import google.registry.model.translators.VKeyTranslatorFactory;
/**
@ -130,8 +129,7 @@ public class ObjectifyService {
new InetAddressTranslatorFactory(),
new MoneyStringTranslatorFactory(),
new ReadableInstantUtcTranslatorFactory(),
new VKeyTranslatorFactory(),
new UpdateAutoTimestampTranslatorFactory())) {
new VKeyTranslatorFactory())) {
factory().getTranslators().add(translatorFactory);
}
}

5
core/src/main/java/google/registry/model/rde/RdeRevision.java
View file

@ -17,7 +17,6 @@ package google.registry.model.rde;
import static com.google.common.base.Preconditions.checkArgument;
import static google.registry.model.rde.RdeNamingUtils.makePartialName;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.base.VerifyException;
import com.googlecode.objectify.Key;
@ -97,8 +96,8 @@ public final class RdeRevision extends BackupGroupRoot {
RdeRevisionId sqlKey = RdeRevisionId.create(tld, date.toLocalDate(), mode);
Key<RdeRevision> ofyKey = Key.create(RdeRevision.class, id);
Optional<RdeRevision> revisionOptional =
transactIfJpaTm(
() -> tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey)));
tm().transact(
() -> tm().loadByKeyIfPresent(VKey.create(RdeRevision.class, sqlKey, ofyKey)));
return revisionOptional.map(rdeRevision -> rdeRevision.revision + 1).orElse(0);
}

7
core/src/main/java/google/registry/model/registrar/Registrar.java
View file

@ -34,7 +34,6 @@ import static google.registry.model.ofy.ObjectifyService.auditedOfy;
import static google.registry.model.tld.Registries.assertTldsExist;
import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
import static google.registry.util.CollectionUtils.nullToEmptyImmutableSortedCopy;
import static google.registry.util.DateTimeUtils.START_OF_TIME;
@ -820,7 +819,7 @@ public class Registrar extends ImmutableObject
.collect(toImmutableSet());
Set<VKey<Registry>> missingTldKeys =
Sets.difference(
newTldKeys, transactIfJpaTm(() -> tm().loadByKeysIfPresent(newTldKeys)).keySet());
newTldKeys, tm().transact(() -> tm().loadByKeysIfPresent(newTldKeys)).keySet());
checkArgument(missingTldKeys.isEmpty(), "Trying to set nonexisting TLDs: %s", missingTldKeys);
getInstance().allowedTlds = ImmutableSortedSet.copyOf(allowedTlds);
return this;
@ -1023,7 +1022,7 @@ public class Registrar extends ImmutableObject
/** Loads all registrar entities directly from Datastore. */
public static Iterable<Registrar> loadAll() {
return transactIfJpaTm(() -> tm().loadAllOf(Registrar.class));
return tm().transact(() -> tm().loadAllOf(Registrar.class));
}
/** Loads all registrar entities using an in-memory cache. */
@ -1041,7 +1040,7 @@ public class Registrar extends ImmutableObject
/** Loads and returns a registrar entity by its id directly from Datastore. */
public static Optional<Registrar> loadByRegistrarId(String registrarId) {
checkArgument(!Strings.isNullOrEmpty(registrarId), "registrarId must be specified");
return transactIfJpaTm(() -> tm().loadByKeyIfPresent(createVKey(registrarId)));
return tm().transact(() -> tm().loadByKeyIfPresent(createVKey(registrarId)));
}
/**

55
core/src/main/java/google/registry/model/translators/UpdateAutoTimestampTranslatorFactory.java
View file

@ -1,55 +0,0 @@
// 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.model.translators;
import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
import static org.joda.time.DateTimeZone.UTC;
import google.registry.model.UpdateAutoTimestamp;
import java.util.Date;
import org.joda.time.DateTime;
/** Saves {@link UpdateAutoTimestamp} as the current time. */
public class UpdateAutoTimestampTranslatorFactory
extends AbstractSimpleTranslatorFactory<UpdateAutoTimestamp, Date> {
public UpdateAutoTimestampTranslatorFactory() {
super(UpdateAutoTimestamp.class);
}
@Override
SimpleTranslator<UpdateAutoTimestamp, Date> createTranslator() {
return new SimpleTranslator<UpdateAutoTimestamp, Date>() {
/**
* Load an existing timestamp. It can be assumed to be non-null since if the field is null in
* Datastore then Objectify will skip this translator and directly load a null.
*/
@Override
public UpdateAutoTimestamp loadValue(Date datastoreValue) {
// Load an existing timestamp, or treat it as START_OF_TIME if none exists.
return UpdateAutoTimestamp.create(new DateTime(datastoreValue, UTC));
}
/** Save a timestamp, setting it to the current time. */
@Override
public Date saveValue(UpdateAutoTimestamp pojoValue) {
return UpdateAutoTimestamp.autoUpdateEnabled()
? ofyTm().getTransactionTime().toDate()
: pojoValue.getTimestamp().toDate();
}
};
}
}

4
core/src/main/java/google/registry/module/backend/BackendComponent.java
View file

@ -24,7 +24,6 @@ import google.registry.config.CredentialModule;
import google.registry.config.RegistryConfig.ConfigModule;
import google.registry.dns.writer.VoidDnsWriterModule;
import google.registry.export.DriveModule;
import google.registry.export.datastore.DatastoreAdminModule;
import google.registry.export.sheet.SheetsServiceModule;
import google.registry.flows.ServerTridProviderModule;
import google.registry.flows.custom.CustomLogicFactoryModule;
@ -40,7 +39,6 @@ import google.registry.monitoring.whitebox.StackdriverModule;
import google.registry.persistence.PersistenceModule;
import google.registry.privileges.secretmanager.SecretManagerModule;
import google.registry.rde.JSchModule;
import google.registry.request.Modules.DatastoreServiceModule;
import google.registry.request.Modules.Jackson2Module;
import google.registry.request.Modules.NetHttpTransportModule;
import google.registry.request.Modules.UrlConnectionServiceModule;
@ -63,8 +61,6 @@ import javax.inject.Singleton;
CloudTasksUtilsModule.class,
CredentialModule.class,
CustomLogicFactoryModule.class,
DatastoreAdminModule.class,
DatastoreServiceModule.class,
DirectoryModule.class,
DummyKeyringModule.class,
DriveModule.class,

24
core/src/main/java/google/registry/module/backend/BackendRequestComponent.java
View file

@ -27,7 +27,6 @@ import google.registry.batch.ResaveEntityAction;
import google.registry.batch.SendExpiringCertificateNotificationEmailAction;
import google.registry.batch.WipeOutCloudSqlAction;
import google.registry.batch.WipeOutContactHistoryPiiAction;
import google.registry.batch.WipeoutDatastoreAction;
import google.registry.cron.CronModule;
import google.registry.cron.TldFanoutAction;
import google.registry.dns.DnsModule;
@ -38,16 +37,10 @@ import google.registry.dns.writer.VoidDnsWriterModule;
import google.registry.dns.writer.clouddns.CloudDnsWriterModule;
import google.registry.dns.writer.dnsupdate.DnsUpdateConfigModule;
import google.registry.dns.writer.dnsupdate.DnsUpdateWriterModule;
import google.registry.export.BackupDatastoreAction;
import google.registry.export.BigqueryPollJobAction;
import google.registry.export.CheckBackupAction;
import google.registry.export.ExportDomainListsAction;
import google.registry.export.ExportPremiumTermsAction;
import google.registry.export.ExportRequestModule;
import google.registry.export.ExportReservedTermsAction;
import google.registry.export.SyncGroupMembersAction;
import google.registry.export.UpdateSnapshotViewAction;
import google.registry.export.UploadDatastoreBackupAction;
import google.registry.export.sheet.SheetModule;
import google.registry.export.sheet.SyncRegistrarsSheetAction;
import google.registry.flows.FlowComponent;
@ -95,7 +88,6 @@ import google.registry.tmch.TmchSmdrlAction;
DnsModule.class,
DnsUpdateConfigModule.class,
DnsUpdateWriterModule.class,
ExportRequestModule.class,
IcannReportingModule.class,
RdeModule.class,
ReportingModule.class,
@ -108,14 +100,8 @@ import google.registry.tmch.TmchSmdrlAction;
})
interface BackendRequestComponent {
BackupDatastoreAction backupDatastoreAction();
BigqueryPollJobAction bigqueryPollJobAction();
BrdaCopyAction brdaCopyAction();
CheckBackupAction checkBackupAction();
CopyDetailReportsAction copyDetailReportAction();
DeleteExpiredDomainsAction deleteExpiredDomainsAction();
@ -148,6 +134,8 @@ interface BackendRequestComponent {
PublishDnsUpdatesAction publishDnsUpdatesAction();
PublishInvoicesAction uploadInvoicesAction();
PublishSpec11ReportAction publishSpec11ReportAction();
ReadDnsQueueAction readDnsQueueAction();
@ -182,18 +170,10 @@ interface BackendRequestComponent {
TmchSmdrlAction tmchSmdrlAction();
UploadDatastoreBackupAction uploadDatastoreBackupAction();
UpdateRegistrarRdapBaseUrlsAction updateRegistrarRdapBaseUrlsAction();
UpdateSnapshotViewAction updateSnapshotViewAction();
PublishInvoicesAction uploadInvoicesAction();
WipeOutCloudSqlAction wipeOutCloudSqlAction();
WipeoutDatastoreAction wipeoutDatastoreAction();
WipeOutContactHistoryPiiAction wipeOutContactHistoryPiiAction();
@Subcomponent.Builder

2
core/src/main/java/google/registry/module/tools/ToolsComponent.java
View file

@ -33,7 +33,6 @@ import google.registry.keyring.kms.KmsModule;
import google.registry.module.tools.ToolsRequestComponent.ToolsRequestComponentModule;
import google.registry.monitoring.whitebox.StackdriverModule;
import google.registry.privileges.secretmanager.SecretManagerModule;
import google.registry.request.Modules.DatastoreServiceModule;
import google.registry.request.Modules.Jackson2Module;
import google.registry.request.Modules.NetHttpTransportModule;
import google.registry.request.Modules.UserServiceModule;
@ -50,7 +49,6 @@ import javax.inject.Singleton;
CredentialModule.class,
CustomLogicFactoryModule.class,
CloudTasksUtilsModule.class,
DatastoreServiceModule.class,
DirectoryModule.class,
DummyKeyringModule.class,
DriveModule.class,

17
core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java
View file

@ -22,7 +22,6 @@ import com.google.appengine.api.utils.SystemProperty.Environment.Value;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Suppliers;
import google.registry.config.RegistryEnvironment;
import google.registry.model.ofy.DatastoreTransactionManager;
import google.registry.persistence.DaggerPersistenceComponent;
import google.registry.tools.RegistryToolEnvironment;
import google.registry.util.NonFinalForTesting;
@ -33,8 +32,6 @@ import java.util.function.Supplier;
// TODO: Rename this to PersistenceFactory and move to persistence package.
public final class TransactionManagerFactory {
private static final DatastoreTransactionManager ofyTm = createTransactionManager();
/** Optional override to manually set the transaction manager per-test. */
private static Optional<TransactionManager> tmForTest = Optional.empty();
@ -67,10 +64,6 @@ public final class TransactionManagerFactory {
}
}
private static DatastoreTransactionManager createTransactionManager() {
return new DatastoreTransactionManager(null);
}
/**
* This function uses App Engine API to determine if the current runtime environment is App
* Engine.
@ -87,8 +80,8 @@ public final class TransactionManagerFactory {
/**
* Returns the {@link TransactionManager} instance.
*
* <p>Returns the {@link JpaTransactionManager} or {@link DatastoreTransactionManager} based on
* the migration schedule or the manually specified per-test transaction manager.
* <p>Returns the {@link JpaTransactionManager} or replica based on the possible manually
* specified per-test transaction manager.
*/
public static TransactionManager tm() {
return tmForTest.orElseGet(TransactionManagerFactory::jpaTm);
@ -119,12 +112,6 @@ public final class TransactionManagerFactory {
return tm().isOfy() ? tm() : replicaJpaTm();
}
/** Returns {@link DatastoreTransactionManager} instance. */
@VisibleForTesting
public static DatastoreTransactionManager ofyTm() {
return ofyTm;
}
/** Sets the return of {@link #jpaTm()} to the given instance of {@link JpaTransactionManager}. */
public static void setJpaTm(Supplier<JpaTransactionManager> jpaTmSupplier) {
checkArgumentNotNull(jpaTmSupplier, "jpaTmSupplier");

76
core/src/main/java/google/registry/persistence/transaction/TransactionManagerUtil.java
View file

@ -1,76 +0,0 @@
// 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.persistence.transaction;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import google.registry.model.ofy.DatastoreTransactionManager;
import java.util.function.Supplier;
/** Utility class that provides supplementary methods for {@link TransactionManager}. */
public class TransactionManagerUtil {
/**
* Returns the result of the given {@link Supplier}.
*
* <p>If {@link TransactionManagerFactory#tm()} returns a {@link JpaTransactionManager} instance,
* the {@link Supplier} is executed in a transaction.
*/
public static <T> T transactIfJpaTm(Supplier<T> supplier) {
if (tm() instanceof JpaTransactionManager) {
return tm().transact(supplier);
} else {
return supplier.get();
}
}
/**
* Executes the given {@link Runnable}.
*
* <p>If {@link TransactionManagerFactory#tm()} returns a {@link JpaTransactionManager} instance,
* the {@link Runnable} is executed in a transaction.
*/
public static void transactIfJpaTm(Runnable runnable) {
transactIfJpaTm(
() -> {
runnable.run();
return null;
});
}
/**
* Executes the given {@link Runnable} if {@link TransactionManagerFactory#tm()} returns a {@link
* DatastoreTransactionManager} instance, otherwise does nothing.
*/
public static void ofyTmOrDoNothing(Runnable ofyRunnable) {
if (tm() instanceof DatastoreTransactionManager) {
ofyRunnable.run();
}
}
/**
* Returns the result from the given {@link Supplier} if {@link TransactionManagerFactory#tm()}
* returns a {@link DatastoreTransactionManager} instance, otherwise returns null.
*/
public static <T> T ofyTmOrDoNothing(Supplier<T> ofySupplier) {
if (tm() instanceof DatastoreTransactionManager) {
return ofySupplier.get();
} else {
return null;
}
}
private TransactionManagerUtil() {}
}

3
core/src/main/java/google/registry/rdap/RdapEntityAction.java
View file

@ -15,7 +15,6 @@
package google.registry.rdap;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.rdap.RdapUtils.getRegistrarByIanaIdentifier;
import static google.registry.rdap.RdapUtils.getRegistrarByName;
import static google.registry.request.Action.Method.GET;
@ -72,7 +71,7 @@ public class RdapEntityAction extends RdapActionBase {
if (ROID_PATTERN.matcher(pathSearchString).matches()) {
VKey<ContactResource> contactVKey = VKey.create(ContactResource.class, pathSearchString);
Optional<ContactResource> contactResource =
transactIfJpaTm(() -> tm().loadByKeyIfPresent(contactVKey));
tm().transact(() -> tm().loadByKeyIfPresent(contactVKey));
// As per Andy Newton on the regext mailing list, contacts by themselves have no role, since
// they are global, and might have different roles for different domains.
if (contactResource.isPresent() && isAuthorized(contactResource.get())) {

11
core/src/main/java/google/registry/rdap/RdapEntitySearchAction.java
View file

@ -17,7 +17,6 @@ package google.registry.rdap;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static google.registry.persistence.transaction.TransactionManagerFactory.replicaJpaTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.rdap.RdapUtils.getRegistrarByIanaIdentifier;
import static google.registry.request.Action.Method.GET;
import static google.registry.request.Action.Method.HEAD;
@ -325,11 +324,11 @@ public class RdapEntitySearchAction extends RdapSearchActionBase {
contactResourceList = ImmutableList.of();
} else {
Optional<ContactResource> contactResource =
transactIfJpaTm(
() ->
tm().loadByKeyIfPresent(
VKey.create(
ContactResource.class, partialStringQuery.getInitialString())));
tm().transact(
() ->
tm().loadByKeyIfPresent(
VKey.create(
ContactResource.class, partialStringQuery.getInitialString())));
contactResourceList =
(contactResource.isPresent() && shouldBeVisible(contactResource.get()))
? ImmutableList.of(contactResource.get())

19
core/src/main/java/google/registry/rdap/RdapJsonFormatter.java
View file

@ -22,7 +22,6 @@ import static com.google.common.collect.ImmutableSetMultimap.toImmutableSetMulti
import static google.registry.model.EppResourceUtils.isLinked;
import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.rdap.RdapIcannStandardInformation.CONTACT_REDACTED_VALUE;
import static google.registry.util.CollectionUtils.union;
@ -358,11 +357,11 @@ public class RdapJsonFormatter {
// Kick off the database loads of the nameservers that we will need, so it can load
// asynchronously while we load and process the contacts.
ImmutableSet<HostResource> loadedHosts =
transactIfJpaTm(
() -> ImmutableSet.copyOf(tm().loadByKeys(domainBase.getNameservers()).values()));
tm().transact(
() -> ImmutableSet.copyOf(tm().loadByKeys(domainBase.getNameservers()).values()));
// Load the registrant and other contacts and add them to the data.
ImmutableMap<VKey<? extends ContactResource>, ContactResource> loadedContacts =
transactIfJpaTm(() -> tm().loadByKeysIfPresent(domainBase.getReferencedContacts()));
tm().transact(() -> tm().loadByKeysIfPresent(domainBase.getReferencedContacts()));
// RDAP Response Profile 2.7.3, A domain MUST have the REGISTRANT, ADMIN, TECH roles and MAY
// have others. We also add the BILLING.
//
@ -441,12 +440,12 @@ public class RdapJsonFormatter {
statuses.add(StatusValue.LINKED);
}
if (hostResource.isSubordinate()
&& transactIfJpaTm(
() ->
tm().loadByKey(hostResource.getSuperordinateDomain())
.cloneProjectedAtTime(getRequestTime())
.getStatusValues()
.contains(StatusValue.PENDING_TRANSFER))) {
&& tm().transact(
() ->
tm().loadByKey(hostResource.getSuperordinateDomain())
.cloneProjectedAtTime(getRequestTime())
.getStatusValues()
.contains(StatusValue.PENDING_TRANSFER))) {
statuses.add(StatusValue.PENDING_TRANSFER);
}
builder

2
core/src/main/java/google/registry/rde/DepositFragment.java
View file

@ -17,7 +17,7 @@ package google.registry.rde;
import com.google.auto.value.AutoValue;
import java.io.Serializable;
/** Container of Datastore resource marshalled by {@link RdeMarshaller}. */
/** Container of RDE resource marshalled by {@link RdeMarshaller}. */
@AutoValue
public abstract class DepositFragment implements Serializable {

5
core/src/main/java/google/registry/rde/DomainBaseToXjcConverter.java
View file

@ -16,7 +16,6 @@ package google.registry.rde;
import static com.google.common.base.Preconditions.checkState;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.base.Ascii;
import com.google.common.base.Strings;
@ -173,7 +172,7 @@ final class DomainBaseToXjcConverter {
if (registrant == null) {
logger.atWarning().log("Domain %s has no registrant contact.", domainName);
} else {
ContactResource registrantContact = transactIfJpaTm(() -> tm().loadByKey(registrant));
ContactResource registrantContact = tm().transact(() -> tm().loadByKey(registrant));
checkState(
registrantContact != null,
"Registrant contact %s on domain %s does not exist",
@ -305,7 +304,7 @@ final class DomainBaseToXjcConverter {
"Contact key for type %s is null on domain %s",
model.getType(),
domainName);
ContactResource contact = transactIfJpaTm(() -> tm().loadByKey(model.getContactKey()));
ContactResource contact = tm().transact(() -> tm().loadByKey(model.getContactKey()));
checkState(
contact != null,
"Contact %s on domain %s does not exist",

3
core/src/main/java/google/registry/rde/EscrowTaskRunner.java
View file

@ -15,7 +15,6 @@
package google.registry.rde;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.flogger.FluentLogger;
import google.registry.model.common.Cursor;
@ -91,7 +90,7 @@ class EscrowTaskRunner {
logger.atInfo().log("Performing escrow for TLD '%s'.", registry.getTld());
DateTime startOfToday = clock.nowUtc().withTimeAtStartOfDay();
DateTime nextRequiredRun =
transactIfJpaTm(
tm().transact(
() -> tm().loadByKeyIfPresent(Cursor.createScopedVKey(cursorType, registry)))
.map(Cursor::getCursorTime)
.orElse(startOfToday);

5
core/src/main/java/google/registry/rde/PendingDepositChecker.java
View file

@ -16,7 +16,6 @@ package google.registry.rde;
import static com.google.common.base.Preconditions.checkArgument;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
import com.google.common.collect.ImmutableSetMultimap;
@ -91,8 +90,8 @@ public final class PendingDepositChecker {
}
// Avoid creating a transaction unless absolutely necessary.
Optional<Cursor> maybeCursor =
transactIfJpaTm(
() -> tm().loadByKeyIfPresent(Cursor.createScopedVKey(cursorType, registry)));
tm().transact(
() -> tm().loadByKeyIfPresent(Cursor.createScopedVKey(cursorType, registry)));
DateTime cursorValue = maybeCursor.map(Cursor::getCursorTime).orElse(startingPoint);
if (isBeforeOrAt(cursorValue, now)) {
DateTime watermark =

31
core/src/main/java/google/registry/reporting/spec11/Spec11EmailUtils.java
View file

@ -19,7 +19,6 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.io.Resources.getResource;
import static google.registry.persistence.transaction.QueryComposer.Comparator;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
@ -131,21 +130,21 @@ public class Spec11EmailUtils {
private RegistrarThreatMatches filterOutNonPublishedMatches(
RegistrarThreatMatches registrarThreatMatches) {
ImmutableList<ThreatMatch> filteredMatches =
transactIfJpaTm(
() -> {
return registrarThreatMatches.threatMatches().stream()
.filter(
threatMatch ->
tm()
.createQueryComposer(DomainBase.class)
.where(
"fullyQualifiedDomainName",
Comparator.EQ,
threatMatch.fullyQualifiedDomainName())
.stream()
.anyMatch(DomainBase::shouldPublishToDns))
.collect(toImmutableList());
});
tm().transact(
() -> {
return registrarThreatMatches.threatMatches().stream()
.filter(
threatMatch ->
tm()
.createQueryComposer(DomainBase.class)
.where(
"fullyQualifiedDomainName",
Comparator.EQ,
threatMatch.fullyQualifiedDomainName())
.stream()
.anyMatch(DomainBase::shouldPublishToDns))
.collect(toImmutableList());
});
return RegistrarThreatMatches.create(registrarThreatMatches.clientId(), filteredMatches);
}

14
core/src/main/java/google/registry/request/Modules.java
View file

@ -14,15 +14,12 @@
package google.registry.request;
import static com.google.appengine.api.datastore.DatastoreServiceFactory.getDatastoreService;
import com.google.api.client.extensions.appengine.http.UrlFetchTransport;
import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.http.javanet.NetHttpTransport;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.jackson2.JacksonFactory;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.urlfetch.URLFetchService;
import com.google.appengine.api.urlfetch.URLFetchServiceFactory;
import com.google.appengine.api.users.UserService;
@ -35,17 +32,6 @@ import javax.inject.Singleton;
/** Dagger modules for App Engine services and other vendor classes. */
public final class Modules {
/** Dagger module for {@link DatastoreService}. */
@Module
public static final class DatastoreServiceModule {
private static final DatastoreService datastoreService = getDatastoreService();
@Provides
static DatastoreService provideDatastoreService() {
return datastoreService;
}
}
/** Dagger module for {@link UrlConnectionService}. */
@Module
public static final class UrlConnectionServiceModule {

25
core/src/main/java/google/registry/request/auth/AuthenticatedRegistrarAccessor.java
View file

@ -18,7 +18,6 @@ import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkNotNull;
import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.google.appengine.api.users.User;
import com.google.common.annotations.VisibleForTesting;
@ -331,18 +330,18 @@ public class AuthenticatedRegistrarAccessor {
// Admins have ADMIN access to all registrars, and also OWNER access to the registry registrar
// and all non-REAL or non-live registrars.
if (isAdmin) {
transactIfJpaTm(
() ->
tm().loadAllOf(Registrar.class)
.forEach(
registrar -> {
if (registrar.getType() != Registrar.Type.REAL
|| !registrar.isLive()
|| registrar.getRegistrarId().equals(registryAdminClientId)) {
builder.put(registrar.getRegistrarId(), Role.OWNER);
}
builder.put(registrar.getRegistrarId(), Role.ADMIN);
}));
tm().transact(
() ->
tm().loadAllOf(Registrar.class)
.forEach(
registrar -> {
if (registrar.getType() != Registrar.Type.REAL
|| !registrar.isLive()
|| registrar.getRegistrarId().equals(registryAdminClientId)) {
builder.put(registrar.getRegistrarId(), Role.OWNER);
}
builder.put(registrar.getRegistrarId(), Role.ADMIN);
}));
}
return builder.build();

2
core/src/main/java/google/registry/tmch/ClaimsListParser.java
View file

@ -36,7 +36,7 @@ public class ClaimsListParser {
/**
* Converts the lines from the DNL CSV file into a {@link ClaimsList} object.
*
* <p>Please note that this does <b>not</b> insert the object into Datastore.
* <p>Please note that this does <b>not</b> insert the object into the DB.
*/
public static ClaimsList parse(List<String> lines) {
ImmutableMap.Builder<String, String> builder = new ImmutableMap.Builder<>();

2
core/src/main/java/google/registry/tmch/TmchData.java
View file

@ -24,7 +24,7 @@ import org.bouncycastle.openpgp.PGPPublicKey;
import org.bouncycastle.openpgp.PGPUtil;
import org.bouncycastle.openpgp.bc.BcPGPPublicKeyRing;
/** Helper class for common data loaded from the jar and Datastore at runtime. */
/** Helper class for common data loaded from the jar and SQL at runtime. */
public final class TmchData {
private static final String BEGIN_ENCODED_SMD = "-----BEGIN ENCODED SMD-----";

13
core/src/main/java/google/registry/tools/CountDomainsCommand.java
View file

@ -17,7 +17,6 @@ package google.registry.tools;
import static google.registry.model.tld.Registries.assertTldsExist;
import static google.registry.persistence.transaction.QueryComposer.Comparator;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
@ -47,11 +46,11 @@ final class CountDomainsCommand implements CommandWithRemoteApi {
}
private long getCountForTld(String tld, DateTime now) {
return transactIfJpaTm(
() ->
tm().createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.where("deletionTime", Comparator.GT, now)
.count());
return tm().transact(
() ->
tm().createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.where("deletionTime", Comparator.GT, now)
.count());
}
}

13
core/src/main/java/google/registry/tools/DeleteTldCommand.java
View file

@ -16,7 +16,6 @@ package google.registry.tools;
import static com.google.common.base.Preconditions.checkState;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
@ -78,11 +77,11 @@ final class DeleteTldCommand extends ConfirmingCommand implements CommandWithRem
}
private boolean tldContainsDomains(String tld) {
return transactIfJpaTm(
() ->
tm().createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.first()
.isPresent());
return tm().transact(
() ->
tm().createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.first()
.isPresent());
}
}

13
core/src/main/java/google/registry/tools/GenerateAllocationTokensCommand.java
View file

@ -21,7 +21,6 @@ import static google.registry.model.billing.BillingEvent.RenewalPriceBehavior.DE
import static google.registry.model.domain.token.AllocationToken.TokenType.SINGLE_USE;
import static google.registry.model.domain.token.AllocationToken.TokenType.UNLIMITED_USE;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.CollectionUtils.nullToEmpty;
import static google.registry.util.StringGenerator.DEFAULT_PASSWORD_LENGTH;
import static java.nio.charset.StandardCharsets.UTF_8;
@ -278,7 +277,7 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
if (dryRun) {
savedTokens = tokens;
} else {
transactIfJpaTm(() -> tm().transact(() -> tm().putAll(tokens)));
tm().transact(() -> tm().transact(() -> tm().putAll(tokens)));
savedTokens = tm().transact(() -> tm().loadByEntities(tokens));
}
savedTokens.forEach(
@ -307,10 +306,10 @@ class GenerateAllocationTokensCommand implements CommandWithRemoteApi {
candidates.stream()
.map(input -> VKey.create(AllocationToken.class, input))
.collect(toImmutableSet());
return transactIfJpaTm(
() ->
tm().loadByKeysIfPresent(existingTokenKeys).values().stream()
.map(AllocationToken::getToken)
.collect(toImmutableSet()));
return tm().transact(
() ->
tm().loadByKeysIfPresent(existingTokenKeys).values().stream()
.map(AllocationToken::getToken)
.collect(toImmutableSet()));
}
}

14
core/src/main/java/google/registry/tools/GenerateDnsReportCommand.java
View file

@ -18,7 +18,6 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.io.BaseEncoding.base16;
import static google.registry.model.tld.Registries.assertTldExists;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.DateTimeUtils.isBeforeOrAt;
import static java.nio.charset.StandardCharsets.US_ASCII;
@ -75,11 +74,11 @@ final class GenerateDnsReportCommand implements CommandWithRemoteApi {
result.append("[\n");
List<DomainBase> domains =
transactIfJpaTm(
() ->
tm().createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.list());
tm().transact(
() ->
tm().createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.list());
for (DomainBase domain : domains) {
// Skip deleted domains and domains that don't get published to DNS.
if (isBeforeOrAt(domain.getDeletionTime(), now) || !domain.shouldPublishToDns()) {
@ -88,8 +87,7 @@ final class GenerateDnsReportCommand implements CommandWithRemoteApi {
write(domain);
}
Iterable<HostResource> nameservers =
transactIfJpaTm(() -> tm().loadAllOf(HostResource.class));
Iterable<HostResource> nameservers = tm().transact(() -> tm().loadAllOf(HostResource.class));
for (HostResource nameserver : nameservers) {
// Skip deleted hosts and external hosts.
if (isBeforeOrAt(nameserver.getDeletionTime(), now)

17
core/src/main/java/google/registry/tools/GenerateLordnCommand.java
View file

@ -16,7 +16,6 @@ package google.registry.tools;
import static google.registry.persistence.transaction.QueryComposer.Comparator;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static java.nio.charset.StandardCharsets.UTF_8;
import com.beust.jcommander.Parameter;
@ -63,14 +62,14 @@ final class GenerateLordnCommand implements CommandWithRemoteApi {
DateTime now = clock.nowUtc();
ImmutableList.Builder<String> claimsCsv = new ImmutableList.Builder<>();
ImmutableList.Builder<String> sunriseCsv = new ImmutableList.Builder<>();
transactIfJpaTm(
() ->
tm()
.createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.orderBy("repoId")
.stream()
.forEach(domain -> processDomain(claimsCsv, sunriseCsv, domain)));
tm().transact(
() ->
tm()
.createQueryComposer(DomainBase.class)
.where("tld", Comparator.EQ, tld)
.orderBy("repoId")
.stream()
.forEach(domain -> processDomain(claimsCsv, sunriseCsv, domain)));
ImmutableList<String> claimsRows = claimsCsv.build();
ImmutableList<String> claimsAll =
new ImmutableList.Builder<String>()

47
core/src/main/java/google/registry/tools/GetOperationStatusCommand.java
View file

@ -1,47 +0,0 @@
// Copyright 2019 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 com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.jackson2.JacksonFactory;
import com.google.common.base.Strings;
import google.registry.export.datastore.DatastoreAdmin;
import java.util.List;
import javax.inject.Inject;
/** Command to get the status of a Datastore operation, e.g., an import or export. */
@Parameters(separators = " =", commandDescription = "Get status of a Datastore operation.")
public class GetOperationStatusCommand implements Command {
private static final JsonFactory JSON_FACTORY = JacksonFactory.getDefaultInstance();
@Parameter(description = "Name of the Datastore import or export operation.")
private List<String> mainParameters;
@Inject DatastoreAdmin datastoreAdmin;
@Override
public void run() throws Exception {
checkArgument(
mainParameters.size() == 1, "Requires exactly one argument: the name of the operation.");
String operationName = mainParameters.get(0);
checkArgument(!Strings.isNullOrEmpty(operationName), "Missing operation name.");
System.out.println(JSON_FACTORY.toPrettyString(datastoreAdmin.get(operationName).execute()));
}
}

129
core/src/main/java/google/registry/tools/ImportDatastoreCommand.java
View file

@ -1,129 +0,0 @@
// Copyright 2019 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 com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.common.base.Ascii;
import com.google.common.collect.ImmutableList;
import google.registry.export.datastore.DatastoreAdmin;
import google.registry.export.datastore.Operation;
import google.registry.model.annotations.DeleteAfterMigration;
import java.util.List;
import java.util.concurrent.TimeUnit;
import javax.annotation.Nullable;
import javax.inject.Inject;
import org.joda.time.Duration;
/**
* Command that imports an earlier backup into Datastore.
*
* <p>This command is part of the Datastore restore process. Please refer to <a
* href="http://playbooks/domain_registry/procedures/backup-restore-testing.md">the playbook</a> for
* the entire process.
*/
@DeleteAfterMigration
@Parameters(separators = " =", commandDescription = "Imports a backup of the Datastore.")
public class ImportDatastoreCommand extends ConfirmingCommand {
@Parameter(names = "--backup_url", description = "URL to the backup on GCS to be imported.")
private String backupUrl;
@Nullable
@Parameter(
names = "--kinds",
description = "List of entity kinds to be imported. Default is to import all.")
private List<String> kinds = ImmutableList.of();
@Parameter(
names = "--async",
description = "If true, command will launch import operation and quit.")
private boolean async;
@Parameter(
names = "--poll_interval",
description =
"Polling interval while waiting for completion synchronously. "
+ "Value is in ISO-8601 format, e.g., PT10S for 10 seconds.")
private Duration pollingInterval = Duration.standardSeconds(30);
@Parameter(
names = "--confirm_production_import",
description = "Set this option to 'PRODUCTION' to confirm import in production environment.")
private String confirmProductionImport = "";
@Inject DatastoreAdmin datastoreAdmin;
@Override
protected String execute() throws Exception {
RegistryToolEnvironment currentEnvironment = RegistryToolEnvironment.get();
// Extra confirmation for running in production
checkArgument(
!currentEnvironment.equals(RegistryToolEnvironment.PRODUCTION)
|| confirmProductionImport.equals("PRODUCTION"),
"The confirm_production_import option must be set when restoring production environment.");
Operation importOperation = datastoreAdmin.importBackup(backupUrl, kinds).execute();
String statusCommand =
String.format(
"nomulus -e %s get_operation_status %s",
Ascii.toLowerCase(currentEnvironment.name()), importOperation.getName());
if (async) {
return String.format(
"Datastore import started. Run this command to check its progress:\n%s",
statusCommand);
}
System.out.println(
"Waiting for import to complete.\n"
+ "You may press Ctrl-C at any time, and use this command to check progress:\n"
+ statusCommand);
while (importOperation.isProcessing()) {
waitInteractively(pollingInterval);
importOperation = datastoreAdmin.get(importOperation.getName()).execute();
System.out.printf("\n%s\n", importOperation.getProgress());
}
return String.format(
"\nDatastore import %s %s.",
importOperation.getName(), importOperation.isSuccessful() ? "succeeded" : "failed");
}
@Override
protected String prompt() {
return "\nThis command is an intermediate step in the Datastore restore process.\n\n"
+ "Please read and understand the playbook entry at\n"
+ " http://playbooks/domain_registry/procedures/backup-restore-testing.md\n"
+ "before proceeding.\n";
}
/** Prints dots to console at regular interval while waiting. */
private static void waitInteractively(Duration pollingInterval) throws InterruptedException {
int sleepSeconds = 2;
long iterations = (pollingInterval.getStandardSeconds() + sleepSeconds - 1) / sleepSeconds;
for (int i = 0; i < iterations; i++) {
TimeUnit.SECONDS.sleep(sleepSeconds);
System.out.print('.');
System.out.flush();
}
}
}

221
core/src/main/java/google/registry/tools/LevelDbLogReader.java
View file

@ -1,221 +0,0 @@
// 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.checkState;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.ReadableByteChannel;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Optional;
/**
* Iterator that incrementally parses binary data in LevelDb format into records.
*
* <p>The input source is automatically closed when all data have been read.
*
* <p>There are several other implementations of this, none of which appeared suitable for our use
* case: <a href="https://github.com/google/leveldb">The original C++ implementation</a>. <a
* href="https://cloud.google.com/appengine/docs/standard/java/javadoc/com/google/appengine/api/files/RecordReadChannel">
* com.google.appengine.api.files.RecordReadChannel</a> - Exactly what we need but deprecated. The
* referenced replacement: <a
* href="https://github.com/GoogleCloudPlatform/appengine-gcs-client.git">The App Engine GCS
* Client</a> - Does not appear to have any support for working with LevelDB. *
*
* <p>See <a
* href="https://github.com/google/leveldb/blob/master/doc/log_format.md">log_format.md</a>
*/
public final class LevelDbLogReader implements Iterator<byte[]> {
@VisibleForTesting static final int BLOCK_SIZE = 32 * 1024;
@VisibleForTesting static final int HEADER_SIZE = 7;
private final ByteArrayOutputStream recordContents = new ByteArrayOutputStream();
private final LinkedList<byte[]> recordList = Lists.newLinkedList();
private final ByteBuffer byteBuffer = ByteBuffer.allocate(BLOCK_SIZE);
private final ReadableByteChannel channel;
LevelDbLogReader(ReadableByteChannel channel) {
this.channel = channel;
}
@Override
public boolean hasNext() {
while (recordList.isEmpty()) {
try {
Optional<byte[]> block = readFromChannel();
if (!block.isPresent()) {
return false;
}
if (block.get().length != BLOCK_SIZE) {
throw new IllegalStateException("Data size is not multiple of " + BLOCK_SIZE);
}
processBlock(block.get());
} catch (IOException e) {
throw new RuntimeException(e);
}
}
return true;
}
@Override
public byte[] next() {
checkState(hasNext(), "The next() method called on empty iterator.");
return recordList.removeFirst();
}
/**
* Returns the next {@link #BLOCK_SIZE} bytes from the input channel, or {@link Optional#empty()}
* if there is no more data.
*/
// TODO(weiminyu): use ByteBuffer directly.
private Optional<byte[]> readFromChannel() throws IOException {
while (channel.isOpen()) {
int bytesRead = channel.read(byteBuffer);
if (!byteBuffer.hasRemaining() || bytesRead < 0) {
byteBuffer.flip();
if (!byteBuffer.hasRemaining()) {
channel.close();
return Optional.empty();
}
byte[] result = new byte[byteBuffer.remaining()];
byteBuffer.get(result);
byteBuffer.clear();
return Optional.of(result);
}
}
return Optional.empty();
}
/** Read a complete block, which must be exactly 32 KB. */
private void processBlock(byte[] block) {
// Read records from the block until there is no longer enough space for a record (i.e. until
// we're at HEADER_SIZE - 1 bytes from the end of the block).
int i = 0;
while (i < BLOCK_SIZE - (HEADER_SIZE - 1)) {
RecordHeader recordHeader = readRecordHeader(block, i);
if (recordHeader.type == ChunkType.END) {
// A type of zero indicates that we've reached the padding zeroes at the end of the block.
break;
}
// Copy the contents of the record into recordContents.
recordContents.write(block, i + HEADER_SIZE, recordHeader.size);
// If this is the last (or only) chunk in the record, store the full contents into the List.
if (recordHeader.type == ChunkType.FULL || recordHeader.type == ChunkType.LAST) {
recordList.add(recordContents.toByteArray());
recordContents.reset();
}
i += recordHeader.size + HEADER_SIZE;
}
}
/**
* Gets a byte from "block" as an unsigned value.
*
* <p>Java bytes are signed, which doesn't work very well for our bit-shifting operations.
*/
private int getUnsignedByte(byte[] block, int pos) {
return block[pos] & 0xFF;
}
/** Reads the 7 byte record header. */
private RecordHeader readRecordHeader(byte[] block, int pos) {
// Read checksum (4 bytes, LE).
int checksum =
getUnsignedByte(block, pos)
| (getUnsignedByte(block, pos + 1) << 8)
| (getUnsignedByte(block, pos + 2) << 16)
| (getUnsignedByte(block, pos + 3) << 24);
// Read size (2 bytes, LE).
int size = getUnsignedByte(block, pos + 4) | (getUnsignedByte(block, pos + 5) << 8);
// Read type (1 byte).
int type = getUnsignedByte(block, pos + 6);
return new RecordHeader(checksum, size, ChunkType.fromCode(type));
}
/** Returns a {@link LevelDbLogReader} over a {@link ReadableByteChannel}. */
public static LevelDbLogReader from(ReadableByteChannel channel) {
return new LevelDbLogReader(channel);
}
/** Returns a {@link LevelDbLogReader} over an {@link InputStream}. */
public static LevelDbLogReader from(InputStream source) {
return new LevelDbLogReader(Channels.newChannel(source));
}
/** Returns a {@link LevelDbLogReader} over a file specified by {@link Path}. */
public static LevelDbLogReader from(Path path) throws IOException {
return from(Files.newInputStream(path));
}
/** Returns a {@link LevelDbLogReader} over a file specified by {@code filename}. */
public static LevelDbLogReader from(String filename) throws IOException {
return from(FileSystems.getDefault().getPath(filename));
}
/** Aggregates the fields in a record header. */
private static final class RecordHeader {
final int checksum;
final int size;
final ChunkType type;
public RecordHeader(int checksum, int size, ChunkType type) {
this.checksum = checksum;
this.size = size;
this.type = type;
}
}
@VisibleForTesting
enum ChunkType {
// Warning: these values must map to their array indices. If this relationship is broken,
// you'll need to change fromCode() to not simply index into values().
END(0),
FULL(1),
FIRST(2),
MIDDLE(3),
LAST(4);
private final int code;
ChunkType(int code) {
this.code = code;
}
int getCode() {
return code;
}
/** Construct a record type from the numeric record type code. */
static ChunkType fromCode(int code) {
return values()[code];
}
}
}

3
core/src/main/java/google/registry/tools/ListCursorsCommand.java
View file

@ -16,7 +16,6 @@ package google.registry.tools;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
@ -59,7 +58,7 @@ final class ListCursorsCommand implements CommandWithRemoteApi {
.filter(r -> !filterEscrowEnabled || r.getEscrowEnabled())
.collect(toImmutableMap(r -> r, r -> Cursor.createScopedVKey(cursorType, r)));
ImmutableMap<VKey<? extends Cursor>, Cursor> cursors =
transactIfJpaTm(() -> tm().loadByKeysIfPresent(registries.values()));
tm().transact(() -> tm().loadByKeysIfPresent(registries.values()));
if (!registries.isEmpty()) {
String header = String.format(OUTPUT_FMT, "TLD", "Cursor Time", "Last Update Time");
System.out.printf("%s\n%s\n", header, Strings.repeat("-", header.length()));

65
core/src/main/java/google/registry/tools/ListDatastoreOperationsCommand.java
View file

@ -1,65 +0,0 @@
// Copyright 2019 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.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.jackson2.JacksonFactory;
import google.registry.export.datastore.DatastoreAdmin;
import google.registry.export.datastore.DatastoreAdmin.ListOperations;
import google.registry.model.annotations.DeleteAfterMigration;
import google.registry.util.Clock;
import java.util.Optional;
import javax.annotation.Nullable;
import javax.inject.Inject;
import org.joda.time.DateTime;
import org.joda.time.Duration;
/** Command that lists Datastore operations. */
@DeleteAfterMigration
@Parameters(separators = " =", commandDescription = "List Datastore operations.")
public class ListDatastoreOperationsCommand implements Command {
private static final JsonFactory JSON_FACTORY = JacksonFactory.getDefaultInstance();
@Nullable
@Parameter(
names = "--start_time_filter",
description =
"Duration relative to current time, used to filter operations by start time. "
+ "Value is in ISO-8601 format, e.g., PT10S for 10 seconds.")
private Duration startTimeFilter;
@Inject DatastoreAdmin datastoreAdmin;
@Inject Clock clock;
@Override
public void run() throws Exception {
ListOperations listOperations =
getQueryFilter().map(datastoreAdmin::list).orElseGet(() -> datastoreAdmin.listAll());
System.out.println(JSON_FACTORY.toPrettyString(listOperations.execute()));
}
private Optional<String> getQueryFilter() {
if (startTimeFilter == null) {
return Optional.empty();
}
DateTime earliestStartingTime = clock.nowUtc().minus(startTimeFilter);
return Optional.of(
String.format("metadata.common.startTime>\"%s\"", earliestStartingTime));
}
}

131
core/src/main/java/google/registry/tools/LoadSnapshotCommand.java
View file

@ -1,131 +0,0 @@
// 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.util.concurrent.Futures.addCallback;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import google.registry.bigquery.BigqueryUtils.SourceFormat;
import google.registry.export.AnnotatedEntities;
import google.registry.model.annotations.DeleteAfterMigration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
/** Command to load Datastore snapshots into Bigquery. */
@DeleteAfterMigration
@Parameters(separators = " =", commandDescription = "Load Datastore snapshot into Bigquery")
final class LoadSnapshotCommand extends BigqueryCommand {
@Parameter(
names = "--snapshot",
description = "Common filename prefix of the specific snapshot series to import.")
private String snapshotPrefix = null;
@Parameter(
names = "--gcs_bucket",
description = "Name of the GCS bucket from which to import Datastore snapshots.")
private String snapshotGcsBucket = "domain-registry/snapshots/testing";
@Parameter(
names = "--kinds",
description = "List of Datastore kinds for which to import snapshot data.")
private List<String> kindNames = new ArrayList<>(AnnotatedEntities.getReportingKinds());
/** Runs the main snapshot import logic. */
@Override
public void runWithBigquery() throws Exception {
kindNames.removeAll(ImmutableList.of("")); // Filter out any empty kind names.
if (snapshotPrefix == null || kindNames.isEmpty()) {
System.err.println("Nothing to import; specify --snapshot and at least one kind.");
return;
}
Map<String, ListenableFuture<?>> loadJobs = loadSnapshotKinds(kindNames);
waitForLoadJobs(loadJobs);
}
/**
* Starts load jobs for the given snapshot kinds, and returns a map of kind name to
* ListenableFuture representing the result of the load job for that kind.
*/
private Map<String, ListenableFuture<?>> loadSnapshotKinds(List<String> kindNames) {
ImmutableMap.Builder<String, ListenableFuture<?>> builder = new ImmutableMap.Builder<>();
for (String kind : kindNames) {
String filename = String.format(
"gs://%s/%s.%s.backup_info", snapshotGcsBucket, snapshotPrefix, kind);
builder.put(kind, loadSnapshotFile(filename, kind));
System.err.println("Started load job for kind: " + kind);
}
return builder.build();
}
/** Starts a load job for the specified kind name, sourcing data from the given GCS file. */
private ListenableFuture<?> loadSnapshotFile(String filename, String kindName) {
return bigquery()
.startLoad(
bigquery()
.buildDestinationTable(kindName)
.description("Datastore snapshot import for " + kindName + ".")
.build(),
SourceFormat.DATASTORE_BACKUP,
ImmutableList.of(filename));
}
/**
* Block on the completion of the load jobs in the provided map, printing out information on
* each job's success or failure.
*/
private void waitForLoadJobs(Map<String, ListenableFuture<?>> loadJobs) throws Exception {
final long startTime = System.currentTimeMillis();
System.err.println("Waiting for load jobs...");
// Add callbacks to each load job that print information on successful completion or failure.
for (final String jobId : loadJobs.keySet()) {
final String jobName = "load-" + jobId;
addCallback(
loadJobs.get(jobId),
new FutureCallback<Object>() {
private double elapsedSeconds() {
return (System.currentTimeMillis() - startTime) / 1000.0;
}
@Override
public void onSuccess(Object unused) {
System.err.printf("Job %s succeeded (%.3fs)\n", jobName, elapsedSeconds());
}
@Override
public void onFailure(Throwable error) {
System.err.printf(
"Job %s failed (%.3fs): %s\n", jobName, elapsedSeconds(), error.getMessage());
}
},
directExecutor());
}
// Block on the completion of all the load jobs.
List<?> results = Futures.successfulAsList(loadJobs.values()).get();
int numSucceeded = (int) results.stream().filter(Objects::nonNull).count();
System.err.printf(
"All load jobs have terminated: %d/%d successful.\n",
numSucceeded, loadJobs.size());
}
}

3
core/src/main/java/google/registry/tools/RegistryCli.java
View file

@ -15,7 +15,6 @@
package google.registry.tools;
import static com.google.common.base.Preconditions.checkState;
import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
import static google.registry.tools.Injector.injectReflectively;
import static java.nio.charset.StandardCharsets.UTF_8;
@ -264,7 +263,7 @@ final class RegistryCli implements AutoCloseable, CommandRunner {
ObjectifyService.initOfy();
// Make sure we start the command with a clean cache, so that any previous command won't
// interfere with this one.
ofyTm().clearSessionCache();
ObjectifyService.ofy().clearSessionCache();
// Enable Cloud SQL for command that needs remote API as they will very likely use
// Cloud SQL after the database migration. Note that the DB password is stored in Datastore

6
core/src/main/java/google/registry/tools/RegistryTool.java
View file

@ -17,7 +17,6 @@ package google.registry.tools;
import com.google.common.collect.ImmutableMap;
import google.registry.tools.javascrap.BackfillRegistrarBillingAccountsCommand;
import google.registry.tools.javascrap.CompareEscrowDepositsCommand;
import google.registry.tools.javascrap.HardDeleteHostCommand;
/** Container class to create and run remote commands against a Datastore instance. */
public final class RegistryTool {
@ -70,7 +69,6 @@ public final class RegistryTool {
.put("get_history_entries", GetHistoryEntriesCommand.class)
.put("get_host", GetHostCommand.class)
.put("get_keyring_secret", GetKeyringSecretCommand.class)
.put("get_operation_status", GetOperationStatusCommand.class)
.put("get_premium_list", GetPremiumListCommand.class)
.put("get_registrar", GetRegistrarCommand.class)
.put("get_reserved_list", GetReservedListCommand.class)
@ -80,18 +78,14 @@ public final class RegistryTool {
.put("get_sql_credential", GetSqlCredentialCommand.class)
.put("get_tld", GetTldCommand.class)
.put("ghostryde", GhostrydeCommand.class)
.put("hard_delete_host", HardDeleteHostCommand.class)
.put("hash_certificate", HashCertificateCommand.class)
.put("import_datastore", ImportDatastoreCommand.class)
.put("list_cursors", ListCursorsCommand.class)
.put("list_datastore_operations", ListDatastoreOperationsCommand.class)
.put("list_domains", ListDomainsCommand.class)
.put("list_hosts", ListHostsCommand.class)
.put("list_premium_lists", ListPremiumListsCommand.class)
.put("list_registrars", ListRegistrarsCommand.class)
.put("list_reserved_lists", ListReservedListsCommand.class)
.put("list_tlds", ListTldsCommand.class)
.put("load_snapshot", LoadSnapshotCommand.class)
.put("load_test", LoadTestCommand.class)
.put("lock_domain", LockDomainCommand.class)
.put("login", LoginCommand.class)

15
core/src/main/java/google/registry/tools/RegistryToolComponent.java
View file

@ -26,7 +26,6 @@ import google.registry.config.RegistryConfig.ConfigModule;
import google.registry.dns.writer.VoidDnsWriterModule;
import google.registry.dns.writer.clouddns.CloudDnsWriterModule;
import google.registry.dns.writer.dnsupdate.DnsUpdateWriterModule;
import google.registry.export.datastore.DatastoreAdminModule;
import google.registry.keyring.KeyringModule;
import google.registry.keyring.api.DummyKeyringModule;
import google.registry.keyring.api.KeyModule;
@ -37,7 +36,6 @@ import google.registry.persistence.PersistenceModule.ReadOnlyReplicaJpaTm;
import google.registry.persistence.transaction.JpaTransactionManager;
import google.registry.privileges.secretmanager.SecretManagerModule;
import google.registry.rde.RdeModule;
import google.registry.request.Modules.DatastoreServiceModule;
import google.registry.request.Modules.Jackson2Module;
import google.registry.request.Modules.UrlConnectionServiceModule;
import google.registry.request.Modules.UrlFetchServiceModule;
@ -45,7 +43,6 @@ import google.registry.request.Modules.UserServiceModule;
import google.registry.tools.AuthModule.LocalCredentialModule;
import google.registry.tools.javascrap.BackfillRegistrarBillingAccountsCommand;
import google.registry.tools.javascrap.CompareEscrowDepositsCommand;
import google.registry.tools.javascrap.HardDeleteHostCommand;
import google.registry.util.UtilsModule;
import google.registry.whois.NonCachingWhoisModule;
import javax.annotation.Nullable;
@ -67,8 +64,6 @@ import javax.inject.Singleton;
ConfigModule.class,
CloudDnsWriterModule.class,
CloudTasksUtilsModule.class,
DatastoreAdminModule.class,
DatastoreServiceModule.class,
DummyKeyringModule.class,
DnsUpdateWriterModule.class,
Jackson2Module.class,
@ -125,22 +120,12 @@ interface RegistryToolComponent {
void inject(GetKeyringSecretCommand command);
void inject(GetOperationStatusCommand command);
void inject(GetSqlCredentialCommand command);
void inject(GhostrydeCommand command);
void inject(HardDeleteHostCommand command);
void inject(ImportDatastoreCommand command);
void inject(ListCursorsCommand command);
void inject(ListDatastoreOperationsCommand command);
void inject(LoadSnapshotCommand command);
void inject(LockDomainCommand command);
void inject(LoginCommand command);

25
core/src/main/java/google/registry/tools/UpdateAllocationTokensCommand.java
View file

@ -19,7 +19,6 @@ import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.collect.Iterables.partition;
import static com.google.common.collect.Streams.stream;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
@ -118,17 +117,19 @@ final class UpdateAllocationTokensCommand extends UpdateOrDeleteAllocationTokens
}
tokensToSave =
transactIfJpaTm(
() ->
tm().loadByKeys(getTokenKeys()).values().stream()
.collect(toImmutableMap(Function.identity(), this::updateToken))
.entrySet()
.stream()
.filter(
entry ->
!entry.getKey().equals(entry.getValue())) // only update changed tokens
.map(Map.Entry::getValue)
.collect(toImmutableSet()));
tm().transact(
() ->
tm().loadByKeys(getTokenKeys()).values().stream()
.collect(toImmutableMap(Function.identity(), this::updateToken))
.entrySet()
.stream()
.filter(
entry ->
!entry
.getKey()
.equals(entry.getValue())) // only update changed tokens
.map(Map.Entry::getValue)
.collect(toImmutableSet()));
}
@Override

13
core/src/main/java/google/registry/tools/UpdateDomainCommand.java
View file

@ -21,7 +21,6 @@ import static google.registry.model.domain.rgp.GracePeriodStatus.AUTO_RENEW;
import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
import static google.registry.model.eppcommon.StatusValue.SERVER_UPDATE_PROHIBITED;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.util.PreconditionsUtils.checkArgumentPresent;
import static java.util.function.Predicate.isEqual;
@ -344,11 +343,11 @@ final class UpdateDomainCommand extends CreateOrUpdateDomainCommand {
ImmutableSet<String> getContactsOfType(
DomainBase domainBase, final DesignatedContact.Type contactType) {
return transactIfJpaTm(
() ->
domainBase.getContacts().stream()
.filter(contact -> contact.getType().equals(contactType))
.map(contact -> tm().loadByKey(contact.getContactKey()).getContactId())
.collect(toImmutableSet()));
return tm().transact(
() ->
domainBase.getContacts().stream()
.filter(contact -> contact.getType().equals(contactType))
.map(contact -> tm().loadByKey(contact.getContactKey()).getContactId())
.collect(toImmutableSet()));
}
}

17
core/src/main/java/google/registry/tools/UpdateOrDeleteAllocationTokensCommand.java
View file

@ -18,7 +18,6 @@ import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import com.beust.jcommander.Parameter;
import com.google.common.collect.ImmutableSet;
@ -59,18 +58,18 @@ abstract class UpdateOrDeleteAllocationTokensCommand extends ConfirmingCommand
.map(token -> VKey.create(AllocationToken.class, token))
.collect(toImmutableSet());
ImmutableSet<VKey<AllocationToken>> nonexistentKeys =
transactIfJpaTm(
() -> keys.stream().filter(key -> !tm().exists(key)).collect(toImmutableSet()));
tm().transact(
() -> keys.stream().filter(key -> !tm().exists(key)).collect(toImmutableSet()));
checkState(nonexistentKeys.isEmpty(), "Tokens with keys %s did not exist.", nonexistentKeys);
return keys;
} else {
checkArgument(!prefix.isEmpty(), "Provided prefix should not be blank");
return transactIfJpaTm(
() ->
tm().loadAllOf(AllocationToken.class).stream()
.filter(token -> token.getToken().startsWith(prefix))
.map(AllocationToken::createVKey)
.collect(toImmutableSet()));
return tm().transact(
() ->
tm().loadAllOf(AllocationToken.class).stream()
.filter(token -> token.getToken().startsWith(prefix))
.map(AllocationToken::createVKey)
.collect(toImmutableSet()));
}
}
}

99
core/src/main/java/google/registry/tools/javascrap/HardDeleteHostCommand.java
View file

@ -1,99 +0,0 @@
// Copyright 2021 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.javascrap;
import static com.google.common.base.Verify.verify;
import static google.registry.model.ofy.ObjectifyService.auditedOfy;
import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.common.collect.ImmutableList;
import com.googlecode.objectify.Key;
import google.registry.model.host.HostResource;
import google.registry.model.index.EppResourceIndex;
import google.registry.model.index.ForeignKeyIndex;
import google.registry.tools.CommandWithRemoteApi;
import google.registry.tools.ConfirmingCommand;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
/**
* Deletes a {@link HostResource} by its ROID.
*
* <p>This deletes the host itself, everything in the same entity group including all {@link
* google.registry.model.reporting.HistoryEntry}s and {@link
* google.registry.model.poll.PollMessage}s, the {@link EppResourceIndex}, and the {@link
* ForeignKeyIndex} (if it exists).
*
* <p>DO NOT use this to hard-delete a host that is still in use on a domain. Bad things will
* happen.
*/
@Parameters(separators = " =", commandDescription = "Delete a host by its ROID.")
public class HardDeleteHostCommand extends ConfirmingCommand implements CommandWithRemoteApi {
@Parameter(names = "--roid", description = "The ROID of the host to be deleted.")
String roid;
@Parameter(names = "--hostname", description = "The hostname, for verification.")
String hostname;
private ImmutableList<Key<Object>> toDelete;
@Override
protected void init() {
ofyTm()
.transact(
() -> {
Key<HostResource> targetKey = Key.create(HostResource.class, roid);
HostResource host = auditedOfy().load().key(targetKey).now();
verify(Objects.equals(host.getHostName(), hostname), "Hostname does not match");
List<Key<Object>> objectsInEntityGroup =
auditedOfy().load().ancestor(host).keys().list();
Optional<ForeignKeyIndex<HostResource>> fki =
Optional.ofNullable(
auditedOfy().load().key(ForeignKeyIndex.createKey(host)).now());
if (!fki.isPresent()) {
System.out.println(
"No ForeignKeyIndex exists, likely because resource is soft-deleted."
+ " Continuing.");
}
EppResourceIndex eppResourceIndex =
auditedOfy().load().entity(EppResourceIndex.create(targetKey)).now();
verify(eppResourceIndex.getKey().equals(targetKey), "Wrong EppResource Index loaded");
ImmutableList.Builder<Key<Object>> toDeleteBuilder =
new ImmutableList.Builder<Key<Object>>()
.addAll(objectsInEntityGroup)
.add(Key.create(eppResourceIndex));
fki.ifPresent(f -> toDeleteBuilder.add(Key.create(f)));
toDelete = toDeleteBuilder.build();
System.out.printf("\n\nAbout to delete %d entities with keys:\n", toDelete.size());
toDelete.forEach(System.out::println);
});
}
@Override
protected String execute() {
tm().transact(() -> auditedOfy().delete().keys(toDelete).now());
return "Done.";
}
}

3
core/src/main/java/google/registry/tools/server/ListHostsAction.java
View file

@ -16,7 +16,6 @@ package google.registry.tools.server;
import static com.google.common.collect.ImmutableSortedSet.toImmutableSortedSet;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.request.Action.Method.GET;
import static google.registry.request.Action.Method.POST;
import static java.util.Comparator.comparing;
@ -51,7 +50,7 @@ public final class ListHostsAction extends ListObjectsAction<HostResource> {
@Override
public ImmutableSet<HostResource> loadObjects() {
final DateTime now = clock.nowUtc();
return transactIfJpaTm(() -> tm().loadAllOf(HostResource.class)).stream()
return tm().transact(() -> tm().loadAllOf(HostResource.class)).stream()
.filter(host -> EppResourceUtils.isActive(host, now))
.collect(toImmutableSortedSet(comparing(HostResource::getHostName)));
}

29
core/src/main/resources/google/registry/beam/bulk_delete_datastore_pipeline_metadata.json
View file

@ -1,29 +0,0 @@
{
"name": "Bulk Delete Cloud Datastore",
"description": "An Apache Beam batch pipeline that deletes Cloud Datastore in bulk. This is easier to use than the GCP-provided template.",
"parameters": [
{
"name": "registryEnvironment",
"label": "The Registry environment.",
"helpText": "The Registry environment, required only because the worker initializer demands it.",
"is_optional": false,
"regexes": [
"^PRODUCTION|SANDBOX|CRASH|QA|ALPHA$"
]
},
{
"name": "kindsToDelete",
"label": "The data KINDs to delete.",
"helpText": "The Datastore KINDs to be deleted. The format may be: the list of kinds to be deleted as a comma-separated string; or '*', which causes all kinds to be deleted."
},
{
"name": "getNumOfKindsHint",
"label": "An estimate of the number of KINDs to be deleted.",
"helpText": "An estimate of the number of KINDs to be deleted. This is recommended if --kindsToDelete is '*' and the default value is too low.",
"is_optional": true,
"regexes": [
"^[1-9][0-9]*$"
]
}
]
}

2
core/src/test/java/google/registry/batch/AsyncTaskEnqueuerTest.java
View file

@ -55,7 +55,7 @@ public class AsyncTaskEnqueuerTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();

31
core/src/test/java/google/registry/batch/DeleteExpiredDomainsActionTest.java
View file

@ -19,7 +19,6 @@ import static com.google.common.truth.Truth.assertThat;
import static google.registry.model.eppcommon.StatusValue.PENDING_DELETE;
import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.testing.DatabaseHelper.createTld;
import static google.registry.testing.DatabaseHelper.loadByEntity;
import static google.registry.testing.DatabaseHelper.newDomainBase;
@ -41,30 +40,24 @@ import google.registry.model.poll.PollMessage;
import google.registry.model.reporting.HistoryEntry;
import google.registry.persistence.transaction.QueryComposer.Comparator;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeLockHandler;
import google.registry.testing.FakeResponse;
import google.registry.testing.InjectExtension;
import google.registry.testing.TestOfyAndSql;
import java.util.Optional;
import org.joda.time.DateTime;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link DeleteExpiredDomainsAction}. */
@DualDatabaseTest
class DeleteExpiredDomainsActionTest {
private final FakeClock clock = new FakeClock(DateTime.parse("2016-06-13T20:21:22Z"));
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder()
.withDatastoreAndCloudSql()
.withClock(clock)
.withTaskQueue()
.build();
AppEngineExtension.builder().withCloudSql().withClock(clock).withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();
@ -86,7 +79,7 @@ class DeleteExpiredDomainsActionTest {
eppController, "NewRegistrar", clock, new FakeLockHandler(true), response);
}
@TestOfyAndSql
@Test
void test_deletesOnlyExpiredDomain() {
// A normal, active autorenewing domain that shouldn't be touched.
DomainBase activeDomain = persistActiveDomain("foo.tld");
@ -131,7 +124,7 @@ class DeleteExpiredDomainsActionTest {
assertThat(reloadedExpiredDomain.getDeletionTime()).isEqualTo(clock.nowUtc().plusDays(35));
}
@TestOfyAndSql
@Test
void test_deletesThreeDomainsInOneRun() throws Exception {
DomainBase domain1 = persistNonAutorenewingDomain("ecck1.tld");
DomainBase domain2 = persistNonAutorenewingDomain("veee2.tld");
@ -143,14 +136,14 @@ class DeleteExpiredDomainsActionTest {
int maxRetries = 5;
while (true) {
ImmutableSet<String> matchingDomains =
transactIfJpaTm(
() ->
tm()
.createQueryComposer(DomainBase.class)
.where("autorenewEndTime", Comparator.LTE, clock.nowUtc())
.stream()
.map(DomainBase::getDomainName)
.collect(toImmutableSet()));
tm().transact(
() ->
tm()
.createQueryComposer(DomainBase.class)
.where("autorenewEndTime", Comparator.LTE, clock.nowUtc())
.stream()
.map(DomainBase::getDomainName)
.collect(toImmutableSet()));
if (matchingDomains.containsAll(ImmutableSet.of("ecck1.tld", "veee2.tld", "tarm3.tld"))) {
break;
}

38
core/src/test/java/google/registry/batch/DeleteProberDataActionTest.java
View file

@ -51,32 +51,25 @@ import google.registry.model.reporting.HistoryEntry;
import google.registry.model.tld.Registry;
import google.registry.model.tld.Registry.TldType;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeResponse;
import google.registry.testing.SystemPropertyExtension;
import google.registry.testing.TestOfyAndSql;
import java.util.Optional;
import java.util.Set;
import org.joda.money.Money;
import org.joda.time.DateTime;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link DeleteProberDataAction}. */
@DualDatabaseTest
class DeleteProberDataActionTest {
private static final DateTime DELETION_TIME = DateTime.parse("2010-01-01T00:00:00.000Z");
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder()
.withDatastoreAndCloudSql()
.withLocalModules()
.withTaskQueue()
.build();
AppEngineExtension.builder().withCloudSql().withLocalModules().withTaskQueue().build();
@RegisterExtension
final SystemPropertyExtension systemPropertyExtension = new SystemPropertyExtension();
@ -108,7 +101,6 @@ class DeleteProberDataActionTest {
private void resetAction() {
action = new DeleteProberDataAction();
action.dnsQueue = DnsQueue.createForTesting(new FakeClock());
action.response = new FakeResponse();
action.isDryRun = false;
action.tlds = ImmutableSet.of();
action.registryAdminRegistrarId = "TheRegistrar";
@ -120,7 +112,7 @@ class DeleteProberDataActionTest {
RegistryEnvironment.UNITTEST.setup(systemPropertyExtension);
}
@TestOfyAndSql
@Test
void test_deletesAllAndOnlyProberData() throws Exception {
Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
@ -135,7 +127,7 @@ class DeleteProberDataActionTest {
assertAllAbsent(oaEntities);
}
@TestOfyAndSql
@Test
void testSuccess_deletesAllAndOnlyGivenTlds() throws Exception {
Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
Set<ImmutableObject> exampleEntities = persistLotsOfDomains("example");
@ -151,7 +143,7 @@ class DeleteProberDataActionTest {
assertAllAbsent(ibEntities);
}
@TestOfyAndSql
@Test
void testFail_givenNonTestTld() {
action.tlds = ImmutableSet.of("not-test.test");
IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, action::run);
@ -160,7 +152,7 @@ class DeleteProberDataActionTest {
.contains("If tlds are given, they must all exist and be TEST tlds");
}
@TestOfyAndSql
@Test
void testFail_givenNonExistentTld() {
action.tlds = ImmutableSet.of("non-existent.test");
IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, action::run);
@ -169,7 +161,7 @@ class DeleteProberDataActionTest {
.contains("If tlds are given, they must all exist and be TEST tlds");
}
@TestOfyAndSql
@Test
void testFail_givenNonDotTestTldOnProd() {
action.tlds = ImmutableSet.of("example");
RegistryEnvironment.PRODUCTION.setup(systemPropertyExtension);
@ -179,7 +171,7 @@ class DeleteProberDataActionTest {
.contains("On production, can only work on TLDs that end with .test");
}
@TestOfyAndSql
@Test
void testSuccess_doesntDeleteNicDomainForProbers() throws Exception {
DomainBase nic = persistActiveDomain("nic.ib-any.test");
Set<ImmutableObject> ibEntities = persistLotsOfDomains("ib-any.test");
@ -188,7 +180,7 @@ class DeleteProberDataActionTest {
assertAllExist(ImmutableSet.of(nic));
}
@TestOfyAndSql
@Test
void testDryRun_doesntDeleteData() throws Exception {
Set<ImmutableObject> tldEntities = persistLotsOfDomains("tld");
Set<ImmutableObject> oaEntities = persistLotsOfDomains("oa-canary.test");
@ -198,7 +190,7 @@ class DeleteProberDataActionTest {
assertAllExist(oaEntities);
}
@TestOfyAndSql
@Test
void testSuccess_activeDomain_isSoftDeleted() throws Exception {
DomainBase domain =
persistResource(
@ -213,7 +205,7 @@ class DeleteProberDataActionTest {
assertDnsTasksEnqueued("blah.ib-any.test");
}
@TestOfyAndSql
@Test
void testSuccess_activeDomain_doubleMapSoftDeletes() throws Exception {
DomainBase domain = persistResource(
newDomainBase("blah.ib-any.test")
@ -230,7 +222,7 @@ class DeleteProberDataActionTest {
assertDnsTasksEnqueued("blah.ib-any.test");
}
@TestOfyAndSql
@Test
void test_recentlyCreatedDomain_isntDeletedYet() throws Exception {
persistResource(
newDomainBase("blah.ib-any.test")
@ -244,7 +236,7 @@ class DeleteProberDataActionTest {
assertThat(domain.get().getDeletionTime()).isEqualTo(END_OF_TIME);
}
@TestOfyAndSql
@Test
void testDryRun_doesntSoftDeleteData() throws Exception {
DomainBase domain =
persistResource(
@ -257,7 +249,7 @@ class DeleteProberDataActionTest {
assertThat(loadByEntity(domain).getDeletionTime()).isEqualTo(END_OF_TIME);
}
@TestOfyAndSql
@Test
void test_domainWithSubordinateHosts_isSkipped() throws Exception {
persistActiveHost("ns1.blah.ib-any.test");
DomainBase nakedDomain =
@ -275,7 +267,7 @@ class DeleteProberDataActionTest {
assertAllAbsent(ImmutableSet.of(nakedDomain));
}
@TestOfyAndSql
@Test
void testFailure_registryAdminClientId_isRequiredForSoftDeletion() {
persistResource(
newDomainBase("blah.ib-any.test")

86
core/src/test/java/google/registry/batch/ExpandRecurringBillingEventsActionTest.java
View file

@ -22,7 +22,6 @@ import static google.registry.model.domain.Period.Unit.YEARS;
import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_AUTORENEW;
import static google.registry.model.reporting.HistoryEntry.Type.DOMAIN_CREATE;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import static google.registry.persistence.transaction.TransactionManagerUtil.transactIfJpaTm;
import static google.registry.testing.DatabaseHelper.assertBillingEventsForResource;
import static google.registry.testing.DatabaseHelper.createTld;
import static google.registry.testing.DatabaseHelper.getHistoryEntriesOfType;
@ -55,29 +54,23 @@ import google.registry.model.reporting.DomainTransactionRecord.TransactionReport
import google.registry.model.reporting.HistoryEntry;
import google.registry.model.tld.Registry;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeResponse;
import google.registry.testing.TestOfyAndSql;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import org.joda.money.Money;
import org.joda.time.DateTime;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link ExpandRecurringBillingEventsAction}. */
@DualDatabaseTest
public class ExpandRecurringBillingEventsActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder()
.withDatastoreAndCloudSql()
.withLocalModules()
.withTaskQueue()
.build();
AppEngineExtension.builder().withCloudSql().withLocalModules().withTaskQueue().build();
private DateTime currentTestTime = DateTime.parse("1999-01-05T00:00:00Z");
private final FakeClock clock = new FakeClock(currentTestTime);
@ -139,8 +132,7 @@ public class ExpandRecurringBillingEventsActionTest {
}
private void assertCursorAt(DateTime expectedCursorTime) {
Cursor cursor =
transactIfJpaTm(() -> tm().loadByKey(Cursor.createGlobalVKey(RECURRING_BILLING)));
Cursor cursor = tm().transact(() -> tm().loadByKey(Cursor.createGlobalVKey(RECURRING_BILLING)));
assertThat(cursor).isNotNull();
assertThat(cursor.getCursorTime()).isEqualTo(expectedCursorTime);
}
@ -183,7 +175,7 @@ public class ExpandRecurringBillingEventsActionTest {
.setTargetId(domain.getDomainName());
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent() throws Exception {
persistResource(recurring);
action.cursorTimeParam = Optional.of(START_OF_TIME);
@ -197,7 +189,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_deletedDomain() throws Exception {
DateTime deletionTime = DateTime.parse("2000-08-01T00:00:00Z");
DomainBase deletedDomain = persistDeletedDomain("deleted.tld", deletionTime);
@ -240,7 +232,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_idempotentForDuplicateRuns() throws Exception {
persistResource(recurring);
action.cursorTimeParam = Optional.of(START_OF_TIME);
@ -258,7 +250,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertBillingEventsForResource(domain, expected, recurring);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_idempotentForExistingOneTime() throws Exception {
persistResource(recurring);
BillingEvent.OneTime persisted =
@ -272,7 +264,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertBillingEventsForResource(domain, persisted, recurring);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_notIdempotentForDifferentBillingTime() throws Exception {
persistResource(recurring);
action.cursorTimeParam = Optional.of(START_OF_TIME);
@ -295,7 +287,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertBillingEventsForResource(domain, persisted, expected, recurring);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_notIdempotentForDifferentRecurring() throws Exception {
persistResource(recurring);
BillingEvent.Recurring recurring2 = persistResource(recurring.asBuilder().setId(3L).build());
@ -321,7 +313,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertBillingEventsForResource(domain, persisted, expected, recurring, recurring2);
}
@TestOfyAndSql
@Test
void testSuccess_ignoreRecurringBeforeWindow() throws Exception {
recurring =
persistResource(
@ -338,7 +330,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_ignoreRecurringAfterWindow() throws Exception {
recurring =
persistResource(recurring.asBuilder().setEventTime(clock.nowUtc().plusYears(2)).build());
@ -349,7 +341,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertBillingEventsForResource(domain, recurring);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_billingTimeAtCursorTime() throws Exception {
persistResource(recurring);
action.cursorTimeParam = Optional.of(DateTime.parse("2000-02-19T00:00:00Z"));
@ -363,7 +355,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_cursorTimeBetweenEventAndBillingTime() throws Exception {
persistResource(recurring);
action.cursorTimeParam = Optional.of(DateTime.parse("2000-01-12T00:00:00Z"));
@ -377,7 +369,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_billingTimeAtExecutionTime() throws Exception {
clock.setTo(currentTestTime);
persistResource(recurring);
@ -393,7 +385,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_multipleYearCreate() throws Exception {
action.cursorTimeParam = Optional.of(recurring.getEventTime());
recurring =
@ -415,7 +407,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_withCursor() throws Exception {
persistResource(recurring);
saveCursor(START_OF_TIME);
@ -429,7 +421,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_withCursorPastExpected() throws Exception {
persistResource(recurring);
// Simulate a quick second run of the action (this should be a no-op).
@ -441,7 +433,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_recurrenceEndBeforeEvent() throws Exception {
// This can occur when a domain is transferred or deleted before a domain comes up for renewal.
recurring =
@ -458,7 +450,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_dryRun() throws Exception {
persistResource(recurring);
action.isDryRun = true;
@ -470,7 +462,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(START_OF_TIME); // Cursor doesn't move on a dry run.
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_multipleYears() throws Exception {
clock.setTo(clock.nowUtc().plusYears(5));
List<BillingEvent> expectedEvents = new ArrayList<>();
@ -497,7 +489,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_multipleYears_cursorInBetweenYears() throws Exception {
clock.setTo(clock.nowUtc().plusYears(5));
List<BillingEvent> expectedEvents = new ArrayList<>();
@ -524,7 +516,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_singleEvent_beforeRenewal() throws Exception {
// Need to restore to the time before the clock was advanced so that the commit log's timestamp
// is not inverted when the clock is later reverted.
@ -539,7 +531,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_singleEvent_afterRecurrenceEnd_inAutorenewGracePeriod() throws Exception {
// The domain creation date is 1999-01-05, and the first renewal date is thus 2000-01-05.
clock.setTo(DateTime.parse("2001-02-06T00:00:00Z"));
@ -568,7 +560,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_singleEvent_afterRecurrenceEnd_outsideAutorenewGracePeriod() throws Exception {
// The domain creation date is 1999-01-05, and the first renewal date is thus 2000-01-05.
clock.setTo(DateTime.parse("2001-02-06T00:00:00Z"));
@ -597,7 +589,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_billingTimeOnLeapYear() throws Exception {
recurring =
persistResource(
@ -618,7 +610,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandSingleEvent_billingTimeNotOnLeapYear() throws Exception {
recurring =
persistResource(
@ -640,7 +632,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandMultipleEvents() throws Exception {
persistResource(recurring);
DomainBase domain2 =
@ -735,7 +727,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandMultipleEvents_anchorTenant() throws Exception {
persistResource(
Registry.get("tld")
@ -783,7 +775,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_expandMultipleEvents_premiumDomain_internalRegistration() throws Exception {
persistResource(
Registry.get("tld")
@ -841,7 +833,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_premiumDomain() throws Exception {
persistResource(
Registry.get("tld")
@ -861,7 +853,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_premiumDomain_forAnchorTenant() throws Exception {
persistResource(
Registry.get("tld")
@ -881,7 +873,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_standardDomain_forAnchorTenant() throws Exception {
recurring = persistResource(recurring.asBuilder().setRenewalPriceBehavior(NONPREMIUM).build());
action.cursorTimeParam = Optional.of(START_OF_TIME);
@ -895,7 +887,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_premiumDomain_forInternalRegistration() throws Exception {
persistResource(
Registry.get("tld")
@ -921,7 +913,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_standardDomain_forInternalRegistration() throws Exception {
recurring =
persistResource(
@ -942,7 +934,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_varyingRenewPrices() throws Exception {
clock.setTo(currentTestTime);
persistResource(
@ -986,7 +978,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_varyingRenewPrices_anchorTenant() throws Exception {
clock.setTo(currentTestTime);
persistResource(
@ -1031,7 +1023,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testSuccess_varyingRenewPrices_internalRegistration() throws Exception {
clock.setTo(currentTestTime);
persistResource(
@ -1082,7 +1074,7 @@ public class ExpandRecurringBillingEventsActionTest {
assertCursorAt(currentTestTime);
}
@TestOfyAndSql
@Test
void testFailure_cursorAfterExecutionTime() {
action.cursorTimeParam = Optional.of(clock.nowUtc().plusYears(1));
IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, this::runAction);
@ -1091,7 +1083,7 @@ public class ExpandRecurringBillingEventsActionTest {
.contains("Cursor time must be earlier than execution time.");
}
@TestOfyAndSql
@Test
void testFailure_cursorAtExecutionTime() {
// The clock advances one milli on run.
action.cursorTimeParam = Optional.of(clock.nowUtc().plusMillis(1));

30
core/src/test/java/google/registry/batch/RelockDomainActionTest.java
View file

@ -43,10 +43,8 @@ import google.registry.testing.AppEngineExtension;
import google.registry.testing.CloudTasksHelper;
import google.registry.testing.CloudTasksHelper.TaskMatcher;
import google.registry.testing.DeterministicStringGenerator;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeResponse;
import google.registry.testing.TestOfyAndSql;
import google.registry.testing.UserInfo;
import google.registry.tools.DomainLockUtils;
import google.registry.util.EmailMessage;
@ -58,6 +56,7 @@ import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.Mock;
@ -65,7 +64,6 @@ import org.mockito.junit.jupiter.MockitoExtension;
/** Unit tests for {@link RelockDomainAction}. */
@ExtendWith(MockitoExtension.class)
@DualDatabaseTest
public class RelockDomainActionTest {
private static final String DOMAIN_NAME = "example.tld";
@ -84,7 +82,7 @@ public class RelockDomainActionTest {
@RegisterExtension
public final AppEngineExtension appEngineExtension =
AppEngineExtension.builder()
.withDatastoreAndCloudSql()
.withCloudSql()
.withTaskQueue()
.withUserService(UserInfo.create(POC_ID, "12345"))
.build();
@ -116,7 +114,7 @@ public class RelockDomainActionTest {
verifyNoMoreInteractions(sendEmailService);
}
@TestOfyAndSql
@Test
void testLock() {
action.run();
assertThat(loadByEntity(domain).getStatusValues())
@ -128,7 +126,7 @@ public class RelockDomainActionTest {
.isEqualTo(newLock);
}
@TestOfyAndSql
@Test
void testFailure_unknownCode() throws Exception {
action = createAction(12128675309L);
action.run();
@ -137,7 +135,7 @@ public class RelockDomainActionTest {
assertTaskEnqueued(1, 12128675309L, Duration.standardMinutes(10)); // should retry, transient
}
@TestOfyAndSql
@Test
void testFailure_pendingDelete() throws Exception {
persistResource(domain.asBuilder().setStatusValues(ImmutableSet.of(PENDING_DELETE)).build());
action.run();
@ -149,7 +147,7 @@ public class RelockDomainActionTest {
cloudTasksHelper.assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
}
@TestOfyAndSql
@Test
void testFailure_pendingTransfer() throws Exception {
persistResource(domain.asBuilder().setStatusValues(ImmutableSet.of(PENDING_TRANSFER)).build());
action.run();
@ -161,7 +159,7 @@ public class RelockDomainActionTest {
cloudTasksHelper.assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
}
@TestOfyAndSql
@Test
void testFailure_domainAlreadyLocked() {
domainLockUtils.administrativelyApplyLock(DOMAIN_NAME, CLIENT_ID, null, true);
action.run();
@ -171,7 +169,7 @@ public class RelockDomainActionTest {
cloudTasksHelper.assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
}
@TestOfyAndSql
@Test
void testFailure_domainDeleted() throws Exception {
persistDomainAsDeleted(domain, clock.nowUtc());
action.run();
@ -183,7 +181,7 @@ public class RelockDomainActionTest {
cloudTasksHelper.assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
}
@TestOfyAndSql
@Test
void testFailure_domainTransferred() throws Exception {
persistResource(
domain.asBuilder().setPersistedCurrentSponsorRegistrarId("NewRegistrar").build());
@ -198,7 +196,7 @@ public class RelockDomainActionTest {
cloudTasksHelper.assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
}
@TestOfyAndSql
@Test
public void testFailure_transientFailure_enqueuesTask() {
// Hard-delete the domain to simulate a DB failure
deleteTestDomain(domain, clock.nowUtc());
@ -209,7 +207,7 @@ public class RelockDomainActionTest {
assertTaskEnqueued(1);
}
@TestOfyAndSql
@Test
void testFailure_sufficientTransientFailures_sendsEmail() throws Exception {
// Hard-delete the domain to simulate a DB failure
deleteTestDomain(domain, clock.nowUtc());
@ -222,7 +220,7 @@ public class RelockDomainActionTest {
assertThat(response.getPayload()).startsWith("Re-lock failed: VKey");
}
@TestOfyAndSql
@Test
void testSuccess_afterSufficientFailures_sendsEmail() throws Exception {
action = createAction(oldLock.getRevisionId(), RelockDomainAction.FAILURES_BEFORE_EMAIL + 1);
action.run();
@ -230,7 +228,7 @@ public class RelockDomainActionTest {
assertSuccessEmailSent();
}
@TestOfyAndSql
@Test
void testFailure_relockAlreadySet() {
RegistryLock newLock =
domainLockUtils.administrativelyApplyLock(DOMAIN_NAME, CLIENT_ID, null, true);
@ -244,7 +242,7 @@ public class RelockDomainActionTest {
cloudTasksHelper.assertNoTasksEnqueued(QUEUE_ASYNC_ACTIONS);
}
@TestOfyAndSql
@Test
void testFailure_slowsDown() throws Exception {
deleteTestDomain(domain, clock.nowUtc());
action = createAction(oldLock.getRevisionId(), RelockDomainAction.ATTEMPTS_BEFORE_SLOWDOWN);

3
core/src/test/java/google/registry/batch/ResaveAllEppResourcesPipelineActionTest.java
View file

@ -35,8 +35,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
public class ResaveAllEppResourcesPipelineActionTest extends BeamActionTestBase {
@RegisterExtension
final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().build();
final AppEngineExtension appEngine = AppEngineExtension.builder().withCloudSql().build();
private final FakeClock fakeClock = new FakeClock();

10
core/src/test/java/google/registry/batch/ResaveEntityActionTest.java
View file

@ -40,13 +40,12 @@ import google.registry.request.Response;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.CloudTasksHelper;
import google.registry.testing.CloudTasksHelper.TaskMatcher;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.InjectExtension;
import google.registry.testing.TestOfyAndSql;
import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.Mock;
@ -56,12 +55,11 @@ import org.mockito.quality.Strictness;
/** Unit tests for {@link ResaveEntityAction}. */
@ExtendWith(MockitoExtension.class)
@DualDatabaseTest
public class ResaveEntityActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();
@ -88,7 +86,7 @@ public class ResaveEntityActionTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void test_domainPendingTransfer_isResavedAndTransferCompleted() {
DomainBase domain =
persistDomainWithPendingTransfer(
@ -113,7 +111,7 @@ public class ResaveEntityActionTest {
verify(response).setPayload("Entity re-saved.");
}
@TestOfyAndSql
@Test
void test_domainPendingDeletion_isResavedAndReenqueued() {
DomainBase newDomain = newDomainBase("domain.tld");
DomainBase domain =

60
core/src/test/java/google/registry/batch/SendExpiringCertificateNotificationEmailActionTest.java
View file

@ -38,11 +38,9 @@ import google.registry.model.registrar.RegistrarAddress;
import google.registry.model.registrar.RegistrarPoc;
import google.registry.model.registrar.RegistrarPoc.Type;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeResponse;
import google.registry.testing.InjectExtension;
import google.registry.testing.TestOfyAndSql;
import google.registry.util.SelfSignedCaCertificate;
import google.registry.util.SendEmailService;
import java.security.cert.X509Certificate;
@ -51,10 +49,10 @@ import javax.annotation.Nullable;
import javax.mail.internet.InternetAddress;
import org.joda.time.DateTime;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link SendExpiringCertificateNotificationEmailAction}. */
@DualDatabaseTest
class SendExpiringCertificateNotificationEmailActionTest {
private static final String EXPIRATION_WARNING_EMAIL_BODY_TEXT =
@ -75,7 +73,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();
private final FakeClock clock = new FakeClock(DateTime.parse("2021-05-24T20:21:22Z"));
@ -111,7 +109,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
persistResource(createRegistrar("clientId", "sampleRegistrar", null, null).build());
}
@TestOfyAndSql
@Test
void sendNotificationEmail_techEMailAsRecipient_returnsTrue() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -133,7 +131,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(true);
}
@TestOfyAndSql
@Test
void sendNotificationEmail_adminEMailAsRecipient_returnsTrue() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -155,7 +153,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(true);
}
@TestOfyAndSql
@Test
void sendNotificationEmail_returnsFalse_unsupportedEmailType() throws Exception {
Registrar registrar =
persistResource(
@ -185,7 +183,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(false);
}
@TestOfyAndSql
@Test
void sendNotificationEmail_returnsFalse_noEmailRecipients() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -201,7 +199,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(false);
}
@TestOfyAndSql
@Test
void sendNotificationEmail_throwsRunTimeException() throws Exception {
doThrow(new RuntimeException("this is a runtime exception"))
.when(sendEmailService)
@ -247,7 +245,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
registrar.getRegistrarName()));
}
@TestOfyAndSql
@Test
void sendNotificationEmail_returnsFalse_noCertificate() {
assertThat(
action.sendNotificationEmail(
@ -255,7 +253,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(false);
}
@TestOfyAndSql
@Test
void sendNotificationEmails_allEmailsBeingSent_onlyMainCertificates() throws Exception {
for (int i = 1; i <= 10; i++) {
Registrar registrar =
@ -275,7 +273,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(action.sendNotificationEmails()).isEqualTo(10);
}
@TestOfyAndSql
@Test
void sendNotificationEmails_allEmailsBeingSent_onlyFailOverCertificates() throws Exception {
for (int i = 1; i <= 10; i++) {
Registrar registrar =
@ -295,7 +293,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(action.sendNotificationEmails()).isEqualTo(10);
}
@TestOfyAndSql
@Test
void sendNotificationEmails_allEmailsBeingSent_mixedOfCertificates() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -334,7 +332,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(action.sendNotificationEmails()).isEqualTo(16);
}
@TestOfyAndSql
@Test
void updateLastNotificationSentDate_updatedSuccessfully_primaryCertificate() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -350,7 +348,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(clock.nowUtc());
}
@TestOfyAndSql
@Test
void updateLastNotificationSentDate_updatedSuccessfully_failOverCertificate() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -366,7 +364,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(clock.nowUtc());
}
@TestOfyAndSql
@Test
void updateLastNotificationSentDate_noUpdates_noLastNotificationSentDate() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -386,7 +384,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.contains("Failed to update the last notification sent date to Registrar");
}
@TestOfyAndSql
@Test
void updateLastNotificationSentDate_noUpdates_invalidCertificateType() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -406,7 +404,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(thrown).hasMessageThat().contains("No enum constant");
}
@TestOfyAndSql
@Test
void getRegistrarsWithExpiringCertificates_returnsPartOfRegistrars() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -434,7 +432,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(results).hasSize(numOfRegistrarsWithExpiringCertificates);
}
@TestOfyAndSql
@Test
void getRegistrarsWithExpiringCertificates_returnsPartOfRegistrars_failOverCertificateBranch()
throws Exception {
X509Certificate expiringCertificate =
@ -463,7 +461,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(numOfRegistrarsWithExpiringCertificates);
}
@TestOfyAndSql
@Test
void getRegistrarsWithExpiringCertificates_returnsAllRegistrars() throws Exception {
X509Certificate expiringCertificate =
SelfSignedCaCertificate.create(
@ -481,7 +479,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo(numOfRegistrarsWithExpiringCertificates);
}
@TestOfyAndSql
@Test
void getRegistrarsWithExpiringCertificates_returnsNoRegistrars() throws Exception {
X509Certificate certificate =
SelfSignedCaCertificate.create(
@ -496,18 +494,18 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(action.getRegistrarsWithExpiringCertificates()).isEmpty();
}
@TestOfyAndSql
@Test
void getRegistrarsWithExpiringCertificates_noRegistrarsInDatabase() {
assertThat(action.getRegistrarsWithExpiringCertificates()).isEmpty();
}
@TestOfyAndSql
@Test
void getEmailAddresses_success_returnsAnEmptyList() {
assertThat(action.getEmailAddresses(sampleRegistrar, Type.TECH)).isEmpty();
assertThat(action.getEmailAddresses(sampleRegistrar, Type.ADMIN)).isEmpty();
}
@TestOfyAndSql
@Test
void getEmailAddresses_success_returnsAListOfEmails() throws Exception {
Registrar registrar = persistResource(makeRegistrar1());
ImmutableList<RegistrarPoc> contacts =
@ -570,7 +568,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
new InternetAddress("john@example-registrar.tld"));
}
@TestOfyAndSql
@Test
void getEmailAddresses_failure_returnsPartialListOfEmails_skipInvalidEmails() {
// when building a new RegistrarContact object, there's already an email validation process.
// if the registrarContact is created successful, the email address of the contact object
@ -578,7 +576,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
// a new InternetAddress using the email address string of the contact object.
}
@TestOfyAndSql
@Test
void getEmailBody_returnsEmailBodyText() {
String registrarName = "good registrar";
String certExpirationDateStr = "2021-06-15";
@ -600,7 +598,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(emailBody).doesNotContain("%4$s");
}
@TestOfyAndSql
@Test
void getEmailBody_throwsIllegalArgumentException_noExpirationDate() {
IllegalArgumentException thrown =
assertThrows(
@ -611,7 +609,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(thrown).hasMessageThat().contains("Expiration date cannot be null");
}
@TestOfyAndSql
@Test
void getEmailBody_throwsIllegalArgumentException_noCertificateType() {
IllegalArgumentException thrown =
assertThrows(
@ -622,7 +620,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(thrown).hasMessageThat().contains("Certificate type cannot be null");
}
@TestOfyAndSql
@Test
void getEmailBody_throwsIllegalArgumentException_noRegistrarId() {
IllegalArgumentException thrown =
assertThrows(
@ -636,7 +634,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
assertThat(thrown).hasMessageThat().contains("Registrar Id cannot be null");
}
@TestOfyAndSql
@Test
void run_sentZeroEmail_responseStatusIs200() {
action.run();
assertThat(response.getStatus()).isEqualTo(SC_OK);
@ -644,7 +642,7 @@ class SendExpiringCertificateNotificationEmailActionTest {
.isEqualTo("Done. Sent 0 expiring certificate notification emails in total.");
}
@TestOfyAndSql
@Test
void run_sentEmails_responseStatusIs200() throws Exception {
for (int i = 1; i <= 5; i++) {
Registrar registrar =

22
core/src/test/java/google/registry/batch/WipeOutContactHistoryPiiActionTest.java
View file

@ -36,17 +36,15 @@ import google.registry.model.eppcommon.PresenceMarker;
import google.registry.model.eppcommon.StatusValue;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DatabaseHelper;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeResponse;
import google.registry.testing.InjectExtension;
import google.registry.testing.TestSqlOnly;
import org.joda.time.DateTime;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link WipeOutContactHistoryPiiAction}. */
@DualDatabaseTest
class WipeOutContactHistoryPiiActionTest {
private static final int TEST_BATCH_SIZE = 20;
@ -102,7 +100,7 @@ class WipeOutContactHistoryPiiActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();
private final FakeClock clock = new FakeClock(DateTime.parse("2021-08-26T20:21:22Z"));
@ -118,7 +116,7 @@ class WipeOutContactHistoryPiiActionTest {
clock, MIN_MONTHS_BEFORE_WIPE_OUT, TEST_BATCH_SIZE, response);
}
@TestSqlOnly
@Test
void getAllHistoryEntitiesOlderThan_returnsAllPersistedEntities() {
ImmutableList<ContactHistory> expectedToBeWipedOut =
persistLotsOfContactHistoryEntities(
@ -132,7 +130,7 @@ class WipeOutContactHistoryPiiActionTest {
.containsExactlyElementsIn(expectedToBeWipedOut));
}
@TestSqlOnly
@Test
void getAllHistoryEntitiesOlderThan_returnsOnlyOldEnoughPersistedEntities() {
ImmutableList<ContactHistory> expectedToBeWipedOut =
persistLotsOfContactHistoryEntities(
@ -151,7 +149,7 @@ class WipeOutContactHistoryPiiActionTest {
.containsExactlyElementsIn(expectedToBeWipedOut));
}
@TestSqlOnly
@Test
void run_withNoEntitiesToWipeOut_success() {
assertThat(
jpaTm()
@ -179,7 +177,7 @@ class WipeOutContactHistoryPiiActionTest {
.isEqualTo("Done. Wiped out PII of 0 ContactHistory entities in total.");
}
@TestSqlOnly
@Test
void run_withOneBatchOfEntities_success() {
int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
ImmutableList<ContactHistory> expectedToBeWipedOut =
@ -216,7 +214,7 @@ class WipeOutContactHistoryPiiActionTest {
assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
}
@TestSqlOnly
@Test
void run_withMultipleBatches_numOfEntitiesAsNonMultipleOfBatchSize_success() {
int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
ImmutableList<ContactHistory> expectedToBeWipedOut =
@ -252,7 +250,7 @@ class WipeOutContactHistoryPiiActionTest {
assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
}
@TestSqlOnly
@Test
void run_withMultipleBatches_numOfEntitiesAsMultiplesOfBatchSize_success() {
int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 2;
ImmutableList<ContactHistory> expectedToBeWipedOut =
@ -289,7 +287,7 @@ class WipeOutContactHistoryPiiActionTest {
assertAllPiiFieldsAreWipedOut(DatabaseHelper.loadByEntitiesIfPresent(expectedToBeWipedOut));
}
@TestSqlOnly
@Test
void wipeOutContactHistoryData_wipesOutNoEntity() {
jpaTm()
.transact(
@ -302,7 +300,7 @@ class WipeOutContactHistoryPiiActionTest {
});
}
@TestSqlOnly
@Test
void wipeOutContactHistoryData_wipesOutMultipleEntities() {
int numOfMonthsFromNow = MIN_MONTHS_BEFORE_WIPE_OUT + 3;
ImmutableList<ContactHistory> expectedToBeWipedOut =

79
core/src/test/java/google/registry/batch/WipeOutDatastoreActionTest.java
View file

@ -1,79 +0,0 @@
// Copyright 2021 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.batch;
import static com.google.common.truth.Truth.assertThat;
import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;
import static org.apache.http.HttpStatus.SC_INTERNAL_SERVER_ERROR;
import static org.apache.http.HttpStatus.SC_OK;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import google.registry.beam.BeamActionTestBase;
import google.registry.config.RegistryEnvironment;
import google.registry.testing.FakeClock;
import org.junit.jupiter.api.Test;
/** Unit tests for {@link WipeoutDatastoreAction}. */
class WipeOutDatastoreActionTest extends BeamActionTestBase {
private final FakeClock clock = new FakeClock();
@Test
void run_projectNotAllowed() {
try {
RegistryEnvironment.SANDBOX.setup();
WipeoutDatastoreAction action =
new WipeoutDatastoreAction(
"domain-registry-sandbox",
"us-central1",
"gs://some-bucket",
clock,
response,
dataflow);
action.run();
assertThat(response.getStatus()).isEqualTo(SC_FORBIDDEN);
verifyNoInteractions(dataflow);
} finally {
RegistryEnvironment.UNITTEST.setup();
}
}
@Test
void run_projectAllowed() throws Exception {
WipeoutDatastoreAction action =
new WipeoutDatastoreAction(
"domain-registry-qa", "us-central1", "gs://some-bucket", clock, response, dataflow);
action.run();
assertThat(response.getStatus()).isEqualTo(SC_OK);
verify(launch, times(1)).execute();
verifyNoMoreInteractions(launch);
}
@Test
void run_failure() throws Exception {
when(launch.execute()).thenThrow(new RuntimeException());
WipeoutDatastoreAction action =
new WipeoutDatastoreAction(
"domain-registry-qa", "us-central1", "gs://some-bucket", clock, response, dataflow);
action.run();
assertThat(response.getStatus()).isEqualTo(SC_INTERNAL_SERVER_ERROR);
verify(launch, times(1)).execute();
verifyNoMoreInteractions(launch);
}
}

152
core/src/test/java/google/registry/beam/datastore/BulkDeleteDatastorePipelineTest.java
View file

@ -1,152 +0,0 @@
// 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.datastore;
import static google.registry.beam.datastore.BulkDeleteDatastorePipeline.discoverEntityKinds;
import static google.registry.beam.datastore.BulkDeleteDatastorePipeline.getDeletionTags;
import static google.registry.beam.datastore.BulkDeleteDatastorePipeline.getOneDeletionTag;
import com.google.common.base.Verify;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.datastore.v1.Entity;
import com.google.datastore.v1.Key;
import com.google.datastore.v1.Key.PathElement;
import google.registry.beam.TestPipelineExtension;
import google.registry.beam.datastore.BulkDeleteDatastorePipeline.GenerateQueries;
import google.registry.beam.datastore.BulkDeleteDatastorePipeline.SplitEntities;
import java.io.Serializable;
import java.util.Map;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfSystemProperty;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link BulkDeleteDatastorePipeline}. */
class BulkDeleteDatastorePipelineTest implements Serializable {
@RegisterExtension
final transient TestPipelineExtension testPipeline =
TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
@Test
void generateQueries() {
PCollection<String> queries =
testPipeline
.apply("InjectKinds", Create.of("A", "B"))
.apply("GenerateQueries", ParDo.of(new GenerateQueries()));
PAssert.that(queries).containsInAnyOrder("select __key__ from `A`", "select __key__ from `B`");
testPipeline.run();
}
@Test
void mapKindsToTags() {
TupleTagList tags = getDeletionTags(2);
PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
PCollection<KV<String, TupleTag<Entity>>> kindToTagMapping =
BulkDeleteDatastorePipeline.mapKindsToDeletionTags(kinds, tags);
PAssert.thatMap(kindToTagMapping)
.isEqualTo(
ImmutableMap.of(
"A", new TupleTag<Entity>("0"),
"B", new TupleTag<Entity>("1")));
testPipeline.run();
}
@Test
void mapKindsToTags_fewerKindsThanTags() {
TupleTagList tags = getDeletionTags(3);
PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
PCollection<KV<String, TupleTag<Entity>>> kindToTagMapping =
BulkDeleteDatastorePipeline.mapKindsToDeletionTags(kinds, tags);
PAssert.thatMap(kindToTagMapping)
.isEqualTo(
ImmutableMap.of(
"A", new TupleTag<Entity>("0"),
"B", new TupleTag<Entity>("1")));
testPipeline.run();
}
@Test
void mapKindsToTags_moreKindsThanTags() {
TupleTagList tags = getDeletionTags(2);
PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B", "C"));
PCollection<KV<String, TupleTag<Entity>>> kindToTagMapping =
BulkDeleteDatastorePipeline.mapKindsToDeletionTags(kinds, tags);
PAssert.thatMap(kindToTagMapping)
.isEqualTo(
ImmutableMap.of(
"A", new TupleTag<Entity>("0"),
"B", new TupleTag<Entity>("1"),
"C", new TupleTag<Entity>("0")));
testPipeline.run();
}
@Test
void splitEntitiesByKind() {
TupleTagList tags = getDeletionTags(2);
PCollection<String> kinds = testPipeline.apply("InjectKinds", Create.of("A", "B"));
PCollectionView<Map<String, TupleTag<Entity>>> kindToTagMapping =
BulkDeleteDatastorePipeline.mapKindsToDeletionTags(kinds, tags).apply(View.asMap());
Entity entityA = createTestEntity("A", 1);
Entity entityB = createTestEntity("B", 2);
PCollection<Entity> entities =
testPipeline.apply("InjectEntities", Create.of(entityA, entityB));
PCollectionTuple allCollections =
entities.apply(
"SplitByKind",
ParDo.of(new SplitEntities(kindToTagMapping))
.withSideInputs(kindToTagMapping)
.withOutputTags(getOneDeletionTag("placeholder"), tags));
PAssert.that(allCollections.get((TupleTag<Entity>) tags.get(0))).containsInAnyOrder(entityA);
PAssert.that(allCollections.get((TupleTag<Entity>) tags.get(1))).containsInAnyOrder(entityB);
testPipeline.run();
}
private static Entity createTestEntity(String kind, long id) {
return Entity.newBuilder()
.setKey(Key.newBuilder().addPath(PathElement.newBuilder().setId(id).setKind(kind)))
.build();
}
@Test
@EnabledIfSystemProperty(named = "test.gcp_integration.env", matches = "\\S+")
void discoverKindsFromDatastore() {
String environmentName = System.getProperty("test.gcp_integration.env");
String project = "domain-registry-" + environmentName;
PCollection<String> kinds =
testPipeline.apply("DiscoverEntityKinds", discoverEntityKinds(project));
PAssert.that(kinds.apply(Count.globally()))
.satisfies(
longs -> {
Verify.verify(Iterables.size(longs) == 1 && Iterables.getFirst(longs, -1L) > 0);
return null;
});
testPipeline.run();
}
}

3
core/src/test/java/google/registry/cron/TldFanoutActionTest.java
View file

@ -48,8 +48,7 @@ class TldFanoutActionTest {
private final CloudTasksHelper cloudTasksHelper = new CloudTasksHelper(new FakeClock());
@RegisterExtension
final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().build();
final AppEngineExtension appEngine = AppEngineExtension.builder().withCloudSql().build();
private static ImmutableListMultimap<String, String> getParamsMap(String... keysAndValues) {
ImmutableListMultimap.Builder<String, String> params = new ImmutableListMultimap.Builder<>();

2
core/src/test/java/google/registry/dns/DnsInjectionTest.java
View file

@ -45,7 +45,7 @@ public final class DnsInjectionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();

2
core/src/test/java/google/registry/dns/DnsQueueTest.java
View file

@ -33,7 +33,7 @@ public class DnsQueueTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
private DnsQueue dnsQueue;
private final FakeClock clock = new FakeClock(DateTime.parse("2010-01-01T10:00:00Z"));

2
core/src/test/java/google/registry/dns/PublishDnsUpdatesActionTest.java
View file

@ -53,7 +53,7 @@ public class PublishDnsUpdatesActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();
private final FakeClock clock = new FakeClock(DateTime.parse("1971-01-01TZ"));

2
core/src/test/java/google/registry/dns/ReadDnsQueueActionTest.java
View file

@ -72,7 +72,7 @@ public class ReadDnsQueueActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder()
.withDatastoreAndCloudSql()
.withCloudSql()
.withTaskQueue(
Joiner.on('\n')
.join(

2
core/src/test/java/google/registry/dns/RefreshDnsActionTest.java
View file

@ -39,7 +39,7 @@ public class RefreshDnsActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
private final DnsQueue dnsQueue = mock(DnsQueue.class);
private final FakeClock clock = new FakeClock();

41
core/src/test/java/google/registry/dns/writer/clouddns/CloudDnsWriterTest.java
View file

@ -45,8 +45,6 @@ import google.registry.model.eppcommon.StatusValue;
import google.registry.model.host.HostResource;
import google.registry.persistence.VKey;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.TestOfyAndSql;
import google.registry.util.Retrier;
import google.registry.util.SystemClock;
import google.registry.util.SystemSleeper;
@ -56,6 +54,7 @@ import java.net.Inet6Address;
import java.net.InetAddress;
import org.joda.time.Duration;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.ArgumentCaptor;
@ -68,12 +67,10 @@ import org.mockito.quality.Strictness;
/** Test case for {@link CloudDnsWriter}. */
@ExtendWith(MockitoExtension.class)
@DualDatabaseTest
public class CloudDnsWriterTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().build();
public final AppEngineExtension appEngine = AppEngineExtension.builder().withCloudSql().build();
private static final Inet4Address IPv4 = (Inet4Address) InetAddresses.forString("127.0.0.1");
private static final Inet6Address IPv6 = (Inet6Address) InetAddresses.forString("::1");
@ -314,7 +311,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testLoadDomain_nonExistentDomain() {
writer.publishDomain("example.tld");
@ -322,7 +319,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testLoadDomain_noDsDataOrNameservers() {
persistResource(fakeDomain("example.tld", ImmutableSet.of(), 0));
writer.publishDomain("example.tld");
@ -330,7 +327,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 0, 0, 0, 0));
}
@TestOfyAndSql
@Test
void testLoadDomain_deleteOldData() {
stubZone = fakeDomainRecords("example.tld", 2, 2, 2, 2);
persistResource(fakeDomain("example.tld", ImmutableSet.of(), 0));
@ -339,7 +336,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 0, 0, 0, 0));
}
@TestOfyAndSql
@Test
void testLoadDomain_withExternalNs() {
persistResource(
fakeDomain("example.tld", ImmutableSet.of(persistResource(fakeHost("0.external"))), 0));
@ -348,7 +345,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 0, 0, 1, 0));
}
@TestOfyAndSql
@Test
void testLoadDomain_withDsData() {
persistResource(
fakeDomain("example.tld", ImmutableSet.of(persistResource(fakeHost("0.external"))), 1));
@ -357,7 +354,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 0, 0, 1, 1));
}
@TestOfyAndSql
@Test
void testLoadDomain_withInBailiwickNs_IPv4() {
persistResource(
fakeDomain(
@ -372,7 +369,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 1, 0, 0, 0));
}
@TestOfyAndSql
@Test
void testLoadDomain_withInBailiwickNs_IPv6() {
persistResource(
fakeDomain(
@ -387,7 +384,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 0, 1, 0, 0));
}
@TestOfyAndSql
@Test
void testLoadDomain_withNameserveThatEndsWithDomainName() {
persistResource(
fakeDomain(
@ -400,7 +397,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testLoadHost_externalHost() {
writer.publishHost("ns1.example.com");
@ -408,7 +405,7 @@ public class CloudDnsWriterTest {
verifyZone(ImmutableSet.of());
}
@TestOfyAndSql
@Test
void testLoadHost_removeStaleNsRecords() {
// Initialize the zone with both NS records
stubZone = fakeDomainRecords("example.tld", 2, 0, 0, 0);
@ -431,7 +428,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void retryMutateZoneOnError() {
CloudDnsWriter spyWriter = spy(writer);
// First call - throw. Second call - do nothing.
@ -445,7 +442,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testLoadDomain_withClientHold() {
persistResource(
fakeDomain(
@ -461,7 +458,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testLoadDomain_withServerHold() {
persistResource(
fakeDomain(
@ -478,7 +475,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testLoadDomain_withPendingDelete() {
persistResource(
fakeDomain(
@ -493,7 +490,7 @@ public class CloudDnsWriterTest {
verifyZone(ImmutableSet.of());
}
@TestOfyAndSql
@Test
void testDuplicateRecords() {
// In publishing DNS records, we can end up publishing information on the same host twice
// (through a domain change and a host change), so this scenario needs to work.
@ -511,7 +508,7 @@ public class CloudDnsWriterTest {
verifyZone(fakeDomainRecords("example.tld", 1, 0, 0, 0));
}
@TestOfyAndSql
@Test
void testInvalidZoneNames() {
createTld("triple.secret.tld");
persistResource(
@ -527,7 +524,7 @@ public class CloudDnsWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testEmptyCommit() {
writer.commit();
verify(dnsConnection, times(0)).changes();

32
core/src/test/java/google/registry/dns/writer/dnsupdate/DnsUpdateWriterTest.java
View file

@ -42,15 +42,14 @@ import google.registry.model.eppcommon.StatusValue;
import google.registry.model.host.HostResource;
import google.registry.model.ofy.Ofy;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.DualDatabaseTest;
import google.registry.testing.FakeClock;
import google.registry.testing.InjectExtension;
import google.registry.testing.TestOfyAndSql;
import java.util.ArrayList;
import java.util.Collections;
import org.joda.time.DateTime;
import org.joda.time.Duration;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.ArgumentCaptor;
@ -71,12 +70,11 @@ import org.xbill.DNS.Update;
/** Unit tests for {@link DnsUpdateWriter}. */
@ExtendWith(MockitoExtension.class)
@DualDatabaseTest
public class DnsUpdateWriterTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
AppEngineExtension.builder().withCloudSql().withTaskQueue().build();
@RegisterExtension public final InjectExtension inject = new InjectExtension();
@ -98,7 +96,7 @@ public class DnsUpdateWriterTest {
"tld", Duration.ZERO, Duration.ZERO, Duration.ZERO, mockResolver, clock);
}
@TestOfyAndSql
@Test
void testPublishDomainCreate_publishesNameServers() throws Exception {
HostResource host1 = persistActiveHost("ns1.example.tld");
HostResource host2 = persistActiveHost("ns2.example.tld");
@ -121,7 +119,7 @@ public class DnsUpdateWriterTest {
}
@MockitoSettings(strictness = Strictness.LENIENT)
@TestOfyAndSql
@Test
void testPublishAtomic_noCommit() {
HostResource host1 = persistActiveHost("ns.example1.tld");
DomainBase domain1 =
@ -145,7 +143,7 @@ public class DnsUpdateWriterTest {
verifyNoInteractions(mockResolver);
}
@TestOfyAndSql
@Test
void testPublishAtomic_oneUpdate() throws Exception {
HostResource host1 = persistActiveHost("ns.example1.tld");
DomainBase domain1 =
@ -177,7 +175,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 4); // The delete and NS sets for each TLD
}
@TestOfyAndSql
@Test
void testPublishDomainCreate_publishesDelegationSigner() throws Exception {
DomainBase domain =
persistActiveDomain("example.tld")
@ -201,7 +199,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 3); // The delete, the NS, and DS sets
}
@TestOfyAndSql
@Test
void testPublishDomainWhenNotActive_removesDnsRecords() throws Exception {
DomainBase domain =
persistActiveDomain("example.tld")
@ -221,7 +219,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 1); // Just the delete set
}
@TestOfyAndSql
@Test
void testPublishDomainDelete_removesDnsRecords() throws Exception {
persistDeletedDomain("example.tld", clock.nowUtc().minusDays(1));
@ -235,7 +233,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 1); // Just the delete set
}
@TestOfyAndSql
@Test
void testPublishHostCreate_publishesAddressRecords() throws Exception {
HostResource host =
persistResource(
@ -268,7 +266,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 5);
}
@TestOfyAndSql
@Test
void testPublishHostDelete_removesDnsRecords() throws Exception {
persistDeletedHost("ns1.example.tld", clock.nowUtc().minusDays(1));
persistActiveDomain("example.tld");
@ -284,7 +282,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 2); // Just the delete set
}
@TestOfyAndSql
@Test
void testPublishHostDelete_removesGlueRecords() throws Exception {
persistDeletedHost("ns1.example.tld", clock.nowUtc().minusDays(1));
persistResource(
@ -305,7 +303,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 3);
}
@TestOfyAndSql
@Test
void testPublishDomainExternalAndInBailiwickNameServer() throws Exception {
HostResource externalNameserver = persistResource(newHostResource("ns1.example.com"));
HostResource inBailiwickNameserver =
@ -342,7 +340,7 @@ public class DnsUpdateWriterTest {
assertThatTotalUpdateSetsIs(update, 5);
}
@TestOfyAndSql
@Test
void testPublishDomainDeleteOrphanGlues() throws Exception {
HostResource inBailiwickNameserver =
persistResource(
@ -380,7 +378,7 @@ public class DnsUpdateWriterTest {
@MockitoSettings(strictness = Strictness.LENIENT)
@SuppressWarnings("AssertThrowsMultipleStatements")
@TestOfyAndSql
@Test
void testPublishDomainFails_whenDnsUpdateReturnsError() throws Exception {
DomainBase domain =
persistActiveDomain("example.tld")
@ -401,7 +399,7 @@ public class DnsUpdateWriterTest {
@MockitoSettings(strictness = Strictness.LENIENT)
@SuppressWarnings("AssertThrowsMultipleStatements")
@TestOfyAndSql
@Test
void testPublishHostFails_whenDnsUpdateReturnsError() throws Exception {
HostResource host =
persistActiveSubordinateHost("ns1.example.tld", persistActiveDomain("example.tld"))

102
core/src/test/java/google/registry/export/AnnotatedEntitiesTest.java
View file

@ -1,102 +0,0 @@
// 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.export;
import static com.google.common.base.Strings.repeat;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.io.Resources.getResource;
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
import static google.registry.export.AnnotatedEntities.getBackupKinds;
import static google.registry.export.AnnotatedEntities.getCrossTldKinds;
import static google.registry.export.AnnotatedEntities.getReportingKinds;
import static google.registry.util.ResourceUtils.readResourceUtf8;
import com.google.common.base.Joiner;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Streams;
import com.google.re2j.Pattern;
import org.junit.jupiter.api.Test;
/** Unit tests for {@link AnnotatedEntities}. */
class AnnotatedEntitiesTest {
private static final String GOLDEN_BACKUP_KINDS_FILENAME = "backup_kinds.txt";
private static final String GOLDEN_REPORTING_KINDS_FILENAME = "reporting_kinds.txt";
private static final String GOLDEN_CROSSTLD_KINDS_FILENAME = "crosstld_kinds.txt";
private static final String UPDATE_INSTRUCTIONS_TEMPLATE = Joiner.on('\n').join(
"",
repeat("-", 80),
"Your changes affect the list of %s kinds in the golden file:",
" %s",
"If these changes are desired, update the golden file with the following contents:",
repeat("=", 80),
"%s",
repeat("=", 80),
"");
@Test
void testBackupKinds_matchGoldenFile() {
checkKindsMatchGoldenFile("backed-up", GOLDEN_BACKUP_KINDS_FILENAME, getBackupKinds());
}
@Test
void testReportingKinds_matchGoldenFile() {
checkKindsMatchGoldenFile("reporting", GOLDEN_REPORTING_KINDS_FILENAME, getReportingKinds());
}
@Test
void testCrossTldKinds_matchGoldenFile() {
checkKindsMatchGoldenFile("crosstld", GOLDEN_CROSSTLD_KINDS_FILENAME, getCrossTldKinds());
}
@Test
void testReportingKinds_areSubsetOfBackupKinds() {
assertThat(getBackupKinds()).containsAtLeastElementsIn(getReportingKinds());
}
private static void checkKindsMatchGoldenFile(
String kindsName, String goldenFilename, ImmutableSet<String> actualKinds) {
String updateInstructions =
String.format(
UPDATE_INSTRUCTIONS_TEMPLATE,
kindsName,
getResource(AnnotatedEntitiesTest.class, goldenFilename).toString(),
Joiner.on('\n').join(actualKinds));
assertWithMessage(updateInstructions)
.that(actualKinds)
.containsExactlyElementsIn(extractListFromFile(goldenFilename))
.inOrder();
}
/**
* Helper method to extract list from file
*
* @param filename
* @return ImmutableList<String>
*/
private static ImmutableList<String> extractListFromFile(String filename) {
String fileContents = readResourceUtf8(AnnotatedEntitiesTest.class, filename);
final Pattern stripComments = Pattern.compile("\\s*#.*$");
return Streams.stream(Splitter.on('\n').split(fileContents.trim()))
.map(line -> stripComments.matcher(line).replaceFirst(""))
.collect(toImmutableList());
}
}

92
core/src/test/java/google/registry/export/BackupDatastoreActionTest.java
View file

@ -1,92 +0,0 @@
// Copyright 2018 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.export;
import static com.google.common.truth.Truth.assertThat;
import static google.registry.export.CheckBackupAction.CHECK_BACKUP_KINDS_TO_LOAD_PARAM;
import static google.registry.export.CheckBackupAction.CHECK_BACKUP_NAME_PARAM;
import static org.mockito.Mockito.when;
import com.google.cloud.tasks.v2.HttpMethod;
import com.google.common.base.Joiner;
import com.google.protobuf.util.Timestamps;
import google.registry.export.datastore.DatastoreAdmin;
import google.registry.export.datastore.DatastoreAdmin.Export;
import google.registry.export.datastore.Operation;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.CloudTasksHelper;
import google.registry.testing.CloudTasksHelper.TaskMatcher;
import google.registry.testing.FakeClock;
import google.registry.testing.FakeResponse;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
/** Unit tests for {@link BackupDatastoreAction}. */
@ExtendWith(MockitoExtension.class)
public class BackupDatastoreActionTest {
@RegisterExtension
public final AppEngineExtension appEngine = AppEngineExtension.builder().withTaskQueue().build();
@Mock private DatastoreAdmin datastoreAdmin;
@Mock private Export exportRequest;
@Mock private Operation backupOperation;
private final FakeResponse response = new FakeResponse();
private CloudTasksHelper cloudTasksHelper = new CloudTasksHelper();
private final BackupDatastoreAction action = new BackupDatastoreAction();
@BeforeEach
void beforeEach() throws Exception {
action.datastoreAdmin = datastoreAdmin;
action.response = response;
action.cloudTasksUtils = cloudTasksHelper.getTestCloudTasksUtils();
action.clock = new FakeClock();
when(datastoreAdmin.export(
"gs://registry-project-id-datastore-backups", AnnotatedEntities.getBackupKinds()))
.thenReturn(exportRequest);
when(exportRequest.execute()).thenReturn(backupOperation);
when(backupOperation.getName())
.thenReturn("projects/registry-project-id/operations/ASA1ODYwNjc");
when(backupOperation.getExportFolderUrl())
.thenReturn("gs://registry-project-id-datastore-backups/some-id");
}
@Test
void testBackup_enqueuesPollTask() {
action.run();
cloudTasksHelper.assertTasksEnqueued(
CheckBackupAction.QUEUE,
new TaskMatcher()
.url(CheckBackupAction.PATH)
.param(CHECK_BACKUP_NAME_PARAM, "projects/registry-project-id/operations/ASA1ODYwNjc")
.param(
CHECK_BACKUP_KINDS_TO_LOAD_PARAM,
Joiner.on(",").join(AnnotatedEntities.getReportingKinds()))
.method(HttpMethod.POST)
.scheduleTime(
Timestamps.fromMillis(
action.clock.nowUtc().plus(CheckBackupAction.POLL_COUNTDOWN).getMillis())));
assertThat(response.getPayload())
.isEqualTo(
"Datastore backup started with name: "
+ "projects/registry-project-id/operations/ASA1ODYwNjc\n"
+ "Saving to gs://registry-project-id-datastore-backups/some-id");
}
}

161
core/src/test/java/google/registry/export/BigqueryPollJobActionTest.java
View file

@ -1,161 +0,0 @@
// 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.export;
import static com.google.common.truth.Truth.assertThat;
import static google.registry.testing.TestLogHandlerUtils.assertLogMessage;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.logging.Level.INFO;
import static java.util.logging.Level.SEVERE;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.ErrorProto;
import com.google.api.services.bigquery.model.Job;
import com.google.api.services.bigquery.model.JobStatus;
import com.google.cloud.tasks.v2.AppEngineHttpRequest;
import com.google.cloud.tasks.v2.HttpMethod;
import com.google.cloud.tasks.v2.Task;
import com.google.common.net.HttpHeaders;
import com.google.common.net.MediaType;
import com.google.protobuf.ByteString;
import google.registry.request.HttpException.BadRequestException;
import google.registry.request.HttpException.NotModifiedException;
import google.registry.testing.AppEngineExtension;
import google.registry.testing.CloudTasksHelper;
import google.registry.testing.CloudTasksHelper.TaskMatcher;
import google.registry.util.CapturingLogHandler;
import google.registry.util.JdkLoggerConfig;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
/** Unit tests for {@link BigqueryPollJobAction}. */
public class BigqueryPollJobActionTest {
@RegisterExtension
public final AppEngineExtension appEngine =
AppEngineExtension.builder().withDatastoreAndCloudSql().withTaskQueue().build();
private static final String PROJECT_ID = "project_id";
private static final String JOB_ID = "job_id";
private static final String CHAINED_QUEUE_NAME = UpdateSnapshotViewAction.QUEUE;
private final Bigquery bigquery = mock(Bigquery.class);
private final Bigquery.Jobs bigqueryJobs = mock(Bigquery.Jobs.class);
private final Bigquery.Jobs.Get bigqueryJobsGet = mock(Bigquery.Jobs.Get.class);
private final CapturingLogHandler logHandler = new CapturingLogHandler();
private BigqueryPollJobAction action = new BigqueryPollJobAction();
private CloudTasksHelper cloudTasksHelper = new CloudTasksHelper();
@BeforeEach
void beforeEach() throws Exception {
action.bigquery = bigquery;
when(bigquery.jobs()).thenReturn(bigqueryJobs);
when(bigqueryJobs.get(PROJECT_ID, JOB_ID)).thenReturn(bigqueryJobsGet);
action.cloudTasksUtils = cloudTasksHelper.getTestCloudTasksUtils();
action.projectId = PROJECT_ID;
action.jobId = JOB_ID;
action.chainedQueueName = () -> CHAINED_QUEUE_NAME;
JdkLoggerConfig.getConfig(BigqueryPollJobAction.class).addHandler(logHandler);
}
@Test
void testSuccess_jobCompletedSuccessfully() throws Exception {
when(bigqueryJobsGet.execute()).thenReturn(
new Job().setStatus(new JobStatus().setState("DONE")));
action.run();
assertLogMessage(
logHandler, INFO, String.format("Bigquery job succeeded - %s:%s", PROJECT_ID, JOB_ID));
}
@Test
void testSuccess_chainedPayloadAndJobSucceeded_enqueuesChainedTask() throws Exception {
when(bigqueryJobsGet.execute()).thenReturn(
new Job().setStatus(new JobStatus().setState("DONE")));
Task chainedTask =
Task.newBuilder()
.setName("my_task_name")
.setAppEngineHttpRequest(
AppEngineHttpRequest.newBuilder()
.setHttpMethod(HttpMethod.POST)
.setRelativeUri("/_dr/something")
.putHeaders("X-Test", "foo")
.putHeaders(HttpHeaders.CONTENT_TYPE, MediaType.FORM_DATA.toString())
.setBody(ByteString.copyFromUtf8("testing=bar")))
.build();
ByteArrayOutputStream bytes = new ByteArrayOutputStream();
new ObjectOutputStream(bytes).writeObject(chainedTask);
action.payload = ByteString.copyFrom(bytes.toByteArray());
action.run();
assertLogMessage(
logHandler, INFO, String.format("Bigquery job succeeded - %s:%s", PROJECT_ID, JOB_ID));
assertLogMessage(
logHandler,
INFO,
"Added chained task my_task_name for /_dr/something to queue " + CHAINED_QUEUE_NAME);
cloudTasksHelper.assertTasksEnqueued(
CHAINED_QUEUE_NAME,
new TaskMatcher()
.url("/_dr/something")
.header("X-Test", "foo")
.header(HttpHeaders.CONTENT_TYPE, MediaType.FORM_DATA.toString())
.param("testing", "bar")
.taskName("my_task_name")
.method(HttpMethod.POST));
}
@Test
void testJobFailed() throws Exception {
when(bigqueryJobsGet.execute()).thenReturn(new Job().setStatus(
new JobStatus()
.setState("DONE")
.setErrorResult(new ErrorProto().setMessage("Job failed"))));
action.run();
assertLogMessage(
logHandler, SEVERE, String.format("Bigquery job failed - %s:%s", PROJECT_ID, JOB_ID));
cloudTasksHelper.assertNoTasksEnqueued(CHAINED_QUEUE_NAME);
}
@Test
void testJobPending() throws Exception {
when(bigqueryJobsGet.execute()).thenReturn(
new Job().setStatus(new JobStatus().setState("PENDING")));
assertThrows(NotModifiedException.class, action::run);
}
@Test
void testJobStatusUnreadable() throws Exception {
when(bigqueryJobsGet.execute()).thenThrow(IOException.class);
assertThrows(NotModifiedException.class, action::run);
}
@Test
void testFailure_badChainedTaskPayload() throws Exception {
when(bigqueryJobsGet.execute()).thenReturn(
new Job().setStatus(new JobStatus().setState("DONE")));
action.payload = ByteString.copyFrom("payload".getBytes(UTF_8));
BadRequestException thrown = assertThrows(BadRequestException.class, action::run);
assertThat(thrown).hasMessageThat().contains("Cannot deserialize task from payload");
}
}

Some files were not shown because too many files have changed in this diff Show more