mirror of
https://github.com/google/nomulus
synced 2026-09-02 22:27:07 +00:00
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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(
|
||||
() ->
|
||||
|
||||
+4
-19
@@ -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
|
||||
|
||||
-16
@@ -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));
|
||||
|
||||
Reference in New Issue
Block a user