Remove DatabaseSnapshot (#2105)

It is no longer being used.
This commit is contained in:
Lai Jiang 2023-08-17 22:40:54 -04:00 committed by GitHub
parent b2d8eec4f2
commit 2c9d60f890
8 changed files with 14 additions and 342 deletions

View file

@ -109,13 +109,6 @@ PRESUBMITS = {
"System.(out|err).println is only allowed in tools/ packages. Please "
"use a logger instead.",
# PostgreSQLContainer instantiation must specify docker tag
# TODO(b/204572437): Fix the pattern to pass DatabaseSnapshotTest.java
PresubmitCheck(
r"[\s\S]*new\s+PostgreSQLContainer(<[\s\S]*>)?\(\s*\)[\s\S]*",
"java", {"DatabaseSnapshotTest.java"}):
"PostgreSQLContainer instantiation must specify docker tag.",
# Various Soy linting checks
PresubmitCheck(
r".* (/\*)?\* {?@param ",

View file

@ -1,88 +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.beam.common;
import static com.google.common.base.Preconditions.checkState;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import com.google.common.flogger.FluentLogger;
import java.util.List;
import javax.persistence.EntityManager;
import javax.persistence.EntityTransaction;
/**
* A database snapshot shareable by concurrent queries from multiple database clients. A snapshot is
* uniquely identified by its {@link #getSnapshotId snapshotId}, and must stay open until all
* concurrent queries to this snapshot have attached to it by calling {@link
* google.registry.persistence.transaction.JpaTransactionManager#setDatabaseSnapshot}. However, it
* can be closed before those queries complete.
*
* <p>This feature is <em>Postgresql-only</em>.
*
* <p>To support large queries, transaction isolation level is fixed at the REPEATABLE_READ to avoid
* exhausting predicate locks at the SERIALIZABLE level.
*/
// TODO(b/193662898): vendor-independent support for richer transaction semantics.
public class DatabaseSnapshot implements AutoCloseable {
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
private String snapshotId;
private EntityManager entityManager;
private EntityTransaction transaction;
private DatabaseSnapshot() {}
public String getSnapshotId() {
checkState(entityManager != null, "Snapshot not opened yet.");
checkState(entityManager.isOpen(), "Snapshot already closed.");
return snapshotId;
}
private DatabaseSnapshot open() {
entityManager = tm().getStandaloneEntityManager();
transaction = entityManager.getTransaction();
transaction.setRollbackOnly();
transaction.begin();
entityManager
.createNativeQuery("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ")
.executeUpdate();
List<?> snapshotIds =
entityManager.createNativeQuery("SELECT pg_export_snapshot();").getResultList();
checkState(snapshotIds.size() == 1, "Unexpected number of snapshots: %s", snapshotIds.size());
snapshotId = (String) snapshotIds.get(0);
return this;
}
@Override
public void close() {
if (transaction != null && transaction.isActive()) {
try {
transaction.rollback();
} catch (Exception e) {
logger.atWarning().withCause(e).log("Failed to close a Database Snapshot");
}
}
if (entityManager != null && entityManager.isOpen()) {
entityManager.close();
}
}
public static DatabaseSnapshot createSnapshot() {
return new DatabaseSnapshot().open();
}
}

View file

@ -122,6 +122,7 @@ public final class RegistryJpaIO {
@AutoValue
public abstract static class Read<R, T> extends PTransform<PBegin, PCollection<T>> {
private static final long serialVersionUID = 6906842877429561700L;
public static final String DEFAULT_NAME = "RegistryJpaIO.Read";
abstract String name();
@ -133,9 +134,6 @@ public final class RegistryJpaIO {
@Nullable
abstract Coder<T> coder();
@Nullable
abstract String snapshotId();
abstract Builder<R, T> toBuilder();
@Override
@ -145,8 +143,7 @@ public final class RegistryJpaIO {
input
.apply("Starting " + name(), Create.of((Void) null))
.apply(
"Run query for " + name(),
ParDo.of(new QueryRunner<>(query(), resultMapper(), snapshotId())));
"Run query for " + name(), ParDo.of(new QueryRunner<>(query(), resultMapper())));
if (coder() != null) {
output = output.setCoder(coder());
}
@ -165,18 +162,6 @@ public final class RegistryJpaIO {
return toBuilder().coder(coder).build();
}
/**
* Specifies the database snapshot to use for this query.
*
* <p>This feature is <em>Postgresql-only</em>. User is responsible for keeping the snapshot
* available until all JVM workers have started using it by calling {@link
* JpaTransactionManager#setDatabaseSnapshot}.
*/
// TODO(b/193662898): vendor-independent support for richer transaction semantics.
public Read<R, T> withSnapshot(@Nullable String snapshotId) {
return toBuilder().snapshotId(snapshotId).build();
}
static <R, T> Builder<R, T> builder() {
return new AutoValue_RegistryJpaIO_Read.Builder<R, T>().name(DEFAULT_NAME);
}
@ -192,8 +177,6 @@ public final class RegistryJpaIO {
abstract Builder<R, T> coder(Coder<T> coder);
abstract Builder<R, T> snapshotId(@Nullable String sharedSnapshotId);
abstract Read<R, T> build();
Builder<R, T> criteriaQuery(CriteriaQuerySupplier<R> criteriaQuery) {
@ -214,27 +197,20 @@ public final class RegistryJpaIO {
}
static class QueryRunner<R, T> extends DoFn<Void, T> {
private static final long serialVersionUID = 7293891513058653334L;
private final RegistryQuery<R> query;
private final SerializableFunction<R, T> resultMapper;
// java.util.Optional is not serializable. Use of Guava Optional is discouraged.
@Nullable private final String snapshotId;
QueryRunner(
RegistryQuery<R> query,
SerializableFunction<R, T> resultMapper,
@Nullable String snapshotId) {
QueryRunner(RegistryQuery<R> query, SerializableFunction<R, T> resultMapper) {
this.query = query;
this.resultMapper = resultMapper;
this.snapshotId = snapshotId;
}
@ProcessElement
public void processElement(OutputReceiver<T> outputReceiver) {
tm().transactNoRetry(
() -> {
if (snapshotId != null) {
tm().setDatabaseSnapshot(snapshotId);
}
query.stream().map(resultMapper::apply).forEach(outputReceiver::output);
});
}
@ -256,8 +232,8 @@ public final class RegistryJpaIO {
@AutoValue
public abstract static class Write<T> extends PTransform<PCollection<T>, PCollection<Void>> {
private static final long serialVersionUID = -4023583243078410323L;
public static final String DEFAULT_NAME = "RegistryJpaIO.Write";
public static final int DEFAULT_BATCH_SIZE = 1;
public abstract String name();
@ -321,6 +297,8 @@ public final class RegistryJpaIO {
/** Writes a batch of entities to a SQL database through a {@link JpaTransactionManager}. */
private static class SqlBatchWriter<T> extends DoFn<KV<ShardedKey<Integer>, Iterable<T>>, Void> {
private static final long serialVersionUID = -7519944406319472690L;
private final Counter counter;
private final SerializableFunction<T, Object> jpaConverter;
@ -337,7 +315,7 @@ public final class RegistryJpaIO {
private void actuallyProcessElement(@Element KV<ShardedKey<Integer>, Iterable<T>> kv) {
ImmutableList<Object> entities =
Streams.stream(kv.getValue())
.map(this.jpaConverter::apply)
.map(jpaConverter::apply)
// TODO(b/177340730): post migration delete the line below.
.filter(Objects::nonNull)
.collect(ImmutableList.toImmutableList());
@ -373,7 +351,7 @@ public final class RegistryJpaIO {
}
/** Returns this entity's primary key field(s) in a string. */
private String toEntityKeyString(Object entity) {
private static String toEntityKeyString(Object entity) {
try {
return tm().transact(
() ->

View file

@ -31,21 +31,6 @@ public interface JpaTransactionManager extends TransactionManager {
*/
EntityManager getStandaloneEntityManager();
/**
* Specifies a database snapshot exported by another transaction to use in the current
* transaction.
*
* <p>This is a Postgresql-specific feature. This method must be called before any other SQL
* commands in a transaction.
*
* <p>To support large queries, transaction isolation level is fixed at the REPEATABLE_READ to
* avoid exhausting predicate locks at the SERIALIZABLE level.
*
* @see google.registry.beam.common.DatabaseSnapshot
*/
// TODO(b/193662898): vendor-independent support for richer transaction semantics.
JpaTransactionManager setDatabaseSnapshot(String snapshotId);
/**
* Returns the {@link EntityManager} for the current request.
*
@ -56,8 +41,8 @@ public interface JpaTransactionManager extends TransactionManager {
/**
* Creates a JPA SQL query for the given query string and result class.
*
* <p>This is a convenience method for the longer <code>
* jpaTm().getEntityManager().createQuery(...)</code>.
* <p>This is a convenience method for the longer {@code
* jpaTm().getEntityManager().createQuery(...)}.
*/
<T> TypedQuery<T> query(String sqlString, Class<T> resultClass);
@ -67,8 +52,8 @@ public interface JpaTransactionManager extends TransactionManager {
/**
* Creates a JPA SQL query for the given query string.
*
* <p>This is a convenience method for the longer <code>
* jpaTm().getEntityManager().createQuery(...)</code>.
* <p>This is a convenience method for the longer {@code
* jpaTm().getEntityManager().createQuery(...)}.
*
* <p>Note that while this method can legally be used for queries that return results, <u>it
* should not be</u>, as it does not correctly detach entities as must be done for nomulus model

View file

@ -109,22 +109,6 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
return entityManager;
}
@Override
public JpaTransactionManager setDatabaseSnapshot(String snapshotId) {
// Postgresql-specific: 'set transaction' command must be called inside a transaction
assertInTransaction();
EntityManager entityManager = getEntityManager();
// Isolation is hardcoded to REPEATABLE READ, as specified by parent's Javadoc.
entityManager
.createNativeQuery("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ")
.executeUpdate();
entityManager
.createNativeQuery(String.format("SET TRANSACTION SNAPSHOT '%s'", snapshotId))
.executeUpdate();
return this;
}
@Override
public <T> TypedQuery<T> query(String sqlString, Class<T> resultClass) {
return new DetachingTypedQuery<>(getEntityManager().createQuery(sqlString, resultClass));

View file

@ -1,174 +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.beam.common;
import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
import com.google.common.collect.ImmutableMap;
import com.google.common.truth.Truth;
import google.registry.beam.TestPipelineExtension;
import google.registry.beam.common.RegistryJpaIO.Read;
import google.registry.model.tld.Tld;
import google.registry.persistence.NomulusPostgreSql;
import google.registry.persistence.PersistenceModule;
import google.registry.persistence.transaction.CriteriaQueryBuilder;
import google.registry.persistence.transaction.JpaTransactionManager;
import google.registry.persistence.transaction.JpaTransactionManagerImpl;
import google.registry.persistence.transaction.TransactionManagerFactory;
import google.registry.testing.DatabaseHelper;
import google.registry.testing.FakeClock;
import javax.persistence.Persistence;
import org.apache.beam.sdk.coders.SerializableCoder;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.values.PCollection;
import org.hibernate.cfg.Environment;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
/** Unit tests for {@link DatabaseSnapshot}. */
@Testcontainers
public class DatabaseSnapshotTest {
/**
* Directly start a PSQL database instead of going through the test extensions.
*
* <p>For reasons unknown, an EntityManagerFactory created by {@code JpaIntegrationTestExtension}
* or {@code JpaUnitTestExtension} enters a bad state after exporting the first snapshot. Starting
* with the second attempt, exports alternate between error ("cannot export a snapshot from a
* subtransaction") and success. The {@link #createSnapshot_twiceNoRead} test below fails with
* either extension. EntityManagerFactory created for production does not have this problem.
*/
@Container
private static PostgreSQLContainer sqlContainer =
new PostgreSQLContainer<>(NomulusPostgreSql.getDockerTag())
.withInitScript("sql/schema/nomulus.golden.sql");
@RegisterExtension
final transient TestPipelineExtension testPipeline =
TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
static JpaTransactionManager origJpa;
static JpaTransactionManager jpa;
static Tld registry;
@BeforeAll
static void setup() {
ImmutableMap<String, String> jpaProperties =
new ImmutableMap.Builder<String, String>()
.put(Environment.URL, sqlContainer.getJdbcUrl())
.put(Environment.USER, sqlContainer.getUsername())
.put(Environment.PASS, sqlContainer.getPassword())
.putAll(PersistenceModule.provideDefaultDatabaseConfigs())
.build();
jpa =
new JpaTransactionManagerImpl(
Persistence.createEntityManagerFactory("nomulus", jpaProperties), new FakeClock());
origJpa = tm();
TransactionManagerFactory.setJpaTm(() -> jpa);
Tld tld = DatabaseHelper.newTld("tld", "TLD");
tm().transact(() -> tm().put(tld));
registry = tm().transact(() -> tm().loadByEntity(tld));
}
@AfterAll
static void tearDown() {
TransactionManagerFactory.setJpaTm(() -> origJpa);
if (jpa != null) {
jpa.teardown();
}
}
@Test
void createSnapshot_onceNoRead() {
try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {}
}
@Test
void createSnapshot_twiceNoRead() {
try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {}
try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {}
}
@Test
void readSnapshot() {
try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {
Tld snapshotRegistry =
tm().transact(
() ->
tm().setDatabaseSnapshot(databaseSnapshot.getSnapshotId())
.loadByEntity(registry));
Truth.assertThat(snapshotRegistry).isEqualTo(registry);
}
}
@Test
void readSnapshot_withSubsequentChange() {
try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {
Tld updated =
registry
.asBuilder()
.setCreateBillingCost(registry.getCreateBillingCost().plus(1))
.build();
tm().transact(() -> tm().put(updated));
Tld persistedUpdate = tm().transact(() -> tm().loadByEntity(registry));
Truth.assertThat(persistedUpdate).isNotEqualTo(registry);
Tld snapshotRegistry =
tm().transact(
() ->
tm().setDatabaseSnapshot(databaseSnapshot.getSnapshotId())
.loadByEntity(registry));
Truth.assertThat(snapshotRegistry).isEqualTo(registry);
} finally {
// Revert change to registry in DB, which is shared by all test methods.
tm().transact(() -> tm().put(registry));
}
}
@Test
void readWithRegistryJpaIO() {
try (DatabaseSnapshot databaseSnapshot = DatabaseSnapshot.createSnapshot()) {
Tld updated =
registry
.asBuilder()
.setCreateBillingCost(registry.getCreateBillingCost().plus(1))
.build();
tm().transact(() -> tm().put(updated));
Read<Tld, Tld> read =
RegistryJpaIO.read(() -> CriteriaQueryBuilder.create(Tld.class).build(), x -> x)
.withCoder(SerializableCoder.of(Tld.class))
.withSnapshot(databaseSnapshot.getSnapshotId());
PCollection<Tld> registries = testPipeline.apply(read);
// This assertion depends on Registry being Serializable, which may change if the
// UnsafeSerializable interface is removed after migration.
PAssert.that(registries).containsInAnyOrder(registry);
testPipeline.run();
} finally {
// Revert change to registry in DB, which is shared by all test methods.
tm().transact(() -> tm().put(registry));
}
}
}

View file

@ -134,7 +134,6 @@ public final class TldTest extends EntityTestCase {
assertThat(yaml).isEqualTo(loadFile(getClass(), "tld.yaml"));
}
// TODO (sarahbot): re-enable this test after we figure out why it fails in presubmits.
@Test
void testYamlToTld() throws Exception {
fakeClock.setTo(START_OF_TIME);

View file

@ -59,11 +59,6 @@ public class ReplicaSimulatingJpaTransactionManager implements JpaTransactionMan
return delegate.getEntityManager();
}
@Override
public JpaTransactionManager setDatabaseSnapshot(String snapshotId) {
return delegate.setDatabaseSnapshot(snapshotId);
}
@Override
public <T> TypedQuery<T> query(String sqlString, Class<T> resultClass) {
return delegate.query(sqlString, resultClass);