Skip to content

Commit

Permalink
Allow for particular flows to log all SQL statements executed
Browse files Browse the repository at this point in the history
We use this now for the DomainDeleteFlow in an attempt to figure out
what statements it's running (cross-referencing that with PSQL's own
statement logging to find slow statements).
  • Loading branch information
gbrodman committed Jan 28, 2025
1 parent 8987fd3 commit 75e3b82
Show file tree
Hide file tree
Showing 8 changed files with 99 additions and 5 deletions.
12 changes: 12 additions & 0 deletions core/src/main/java/google/registry/flows/FlowModule.java
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,13 @@ Optional<TransactionIsolationLevel> provideIsolationLevelOverride(
.map(IsolationLevel::value);
}

@Provides
@FlowScope
@LogSqlStatements
boolean provideShouldLogSqlStatements(Class<? extends Flow> flowClass) {
return SqlStatementLoggingFlow.class.isAssignableFrom(flowClass);
}

@Provides
@FlowScope
@Superuser
Expand Down Expand Up @@ -370,4 +377,9 @@ static class EppExceptionInProviderException extends RuntimeException {
@Qualifier
@Documented
public @interface Transactional {}

/** Dagger qualifier for if we should log all SQL statements in a flow. */
@Qualifier
@Documented
public @interface LogSqlStatements {}
}
5 changes: 4 additions & 1 deletion core/src/main/java/google/registry/flows/FlowRunner.java
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.google.common.flogger.FluentLogger;
import google.registry.flows.FlowModule.DryRun;
import google.registry.flows.FlowModule.InputXml;
import google.registry.flows.FlowModule.LogSqlStatements;
import google.registry.flows.FlowModule.RegistrarId;
import google.registry.flows.FlowModule.Superuser;
import google.registry.flows.FlowModule.Transactional;
Expand Down Expand Up @@ -49,6 +50,7 @@ public class FlowRunner {
@Inject @DryRun boolean isDryRun;
@Inject @Superuser boolean isSuperuser;
@Inject @Transactional boolean isTransactional;
@Inject @LogSqlStatements boolean logSqlStatements;
@Inject SessionMetadata sessionMetadata;
@Inject Trid trid;
@Inject FlowReporter flowReporter;
Expand Down Expand Up @@ -97,7 +99,8 @@ public EppOutput run(final EppMetric.Builder eppMetricBuilder) throws EppExcepti
} catch (EppException e) {
throw new EppRuntimeException(e);
}
});
},
logSqlStatements);
} catch (DryRunException e) {
return e.output;
} catch (EppRuntimeException e) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
// Copyright 2025 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.flows;

/**
* Interface for a {@link Flow} that logs its SQL statements when running transactionally.
*
* <p>We don't wish to log all SQL statements ever executed (that'll create too much log bloat) but
* for some flows and some occasions we may wish to know precisely what SQL statements are being
* run.
*/
public interface SqlStatementLoggingFlow extends TransactionalFlow {}
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@
import google.registry.flows.FlowModule.TargetId;
import google.registry.flows.MutatingFlow;
import google.registry.flows.SessionMetadata;
import google.registry.flows.SqlStatementLoggingFlow;
import google.registry.flows.annotations.ReportingSpec;
import google.registry.flows.custom.DomainDeleteFlowCustomLogic;
import google.registry.flows.custom.DomainDeleteFlowCustomLogic.AfterValidationParameters;
Expand Down Expand Up @@ -117,7 +118,7 @@
* @error {@link DomainFlowUtils.NotAuthorizedForTldException}
*/
@ReportingSpec(ActivityReportField.DOMAIN_DELETE)
public final class DomainDeleteFlow implements MutatingFlow {
public final class DomainDeleteFlow implements MutatingFlow, SqlStatementLoggingFlow {

private static final FluentLogger logger = FluentLogger.forEnclosingClass();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import jakarta.persistence.TypedQuery;
import jakarta.persistence.criteria.CriteriaQuery;
import jakarta.persistence.metamodel.Metamodel;
import java.util.concurrent.Callable;

/** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
public interface JpaTransactionManager extends TransactionManager {
Expand Down Expand Up @@ -92,4 +93,15 @@ static Query setQueryFetchSize(Query query, int fetchSize) {

/** Return the {@link TransactionIsolationLevel} used in the current transaction. */
TransactionIsolationLevel getCurrentTransactionIsolationLevel();

/** Executes the work with the given isolation level, possibly logging all SQL statements used. */
<T> T transact(
TransactionIsolationLevel isolationLevel, Callable<T> work, boolean logSqlStatements);

/**
* Executes the work with the given isolation level without retry, possibly logging all SQL
* statements used.
*/
<T> T transactNoRetry(
TransactionIsolationLevel isolationLevel, Callable<T> work, boolean logSqlStatements);
}
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@
import java.util.stream.StreamSupport;
import javax.annotation.Nullable;
import org.hibernate.Session;
import org.hibernate.SessionFactory;
import org.hibernate.cfg.Environment;
import org.joda.time.DateTime;

Expand Down Expand Up @@ -174,6 +175,12 @@ public <T> T reTransact(Callable<T> work) {

@Override
public <T> T transact(TransactionIsolationLevel isolationLevel, Callable<T> work) {
return transact(isolationLevel, work, false);
}

@Override
public <T> T transact(
TransactionIsolationLevel isolationLevel, Callable<T> work, boolean logSqlStatements) {
if (inTransaction()) {
if (!getHibernateAllowNestedTransactions()) {
throw new IllegalStateException(NESTED_TRANSACTION_MESSAGE);
Expand All @@ -186,7 +193,8 @@ public <T> T transact(TransactionIsolationLevel isolationLevel, Callable<T> work
return transactNoRetry(isolationLevel, work);
}
return retrier.callWithRetry(
() -> transactNoRetry(isolationLevel, work), JpaRetries::isFailedTxnRetriable);
() -> transactNoRetry(isolationLevel, work, logSqlStatements),
JpaRetries::isFailedTxnRetriable);
}

@Override
Expand All @@ -202,6 +210,14 @@ public <T> T transactNoRetry(Callable<T> work) {
@Override
public <T> T transactNoRetry(
@Nullable TransactionIsolationLevel isolationLevel, Callable<T> work) {
return transactNoRetry(isolationLevel, work, false);
}

@Override
public <T> T transactNoRetry(
@Nullable TransactionIsolationLevel isolationLevel,
Callable<T> work,
boolean logSqlStatements) {
if (inTransaction()) {
// This check will no longer be necessary when the transact() method always throws
// inside a nested transaction, as the only way to pass a non-null isolation level
Expand All @@ -224,7 +240,18 @@ public <T> T transactNoRetry(
}
}
TransactionInfo txnInfo = transactionInfo.get();
txnInfo.entityManager = emf.createEntityManager();

txnInfo.entityManager =
logSqlStatements
? emf.unwrap(SessionFactory.class)
.withOptions()
.statementInspector(
(s) -> {
logger.atInfo().log(s);
return s;
})
.openSession()
: emf.createEntityManager();
if (readOnly) {
// Disable Hibernate's dirty object check on flushing, it has become more aggressive in v6.
txnInfo.entityManager.unwrap(Session.class).setDefaultReadOnly(true);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@
public abstract class ResourceFlowTestCase<F extends Flow, R extends EppResource>
extends FlowTestCase<F> {

private final TestLogHandler logHandler = new TestLogHandler();
protected final TestLogHandler logHandler = new TestLogHandler();

@RegisterExtension
public final TestCacheExtension testCacheExtension =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -102,8 +102,10 @@
import google.registry.model.transfer.TransferStatus;
import google.registry.testing.CloudTasksHelper.TaskMatcher;
import google.registry.testing.DatabaseHelper;
import google.registry.testing.LogsSubject;
import java.util.Map;
import java.util.Optional;
import java.util.logging.Level;
import org.joda.money.Money;
import org.joda.time.DateTime;
import org.joda.time.Duration;
Expand Down Expand Up @@ -1258,4 +1260,17 @@ void testSuccess_skipsPollMessage_whenConfigured() throws Exception {
runFlowAssertResponse(loadFile("domain_delete_response_pending.xml"));
assertPollMessages();
}

@Test
void testSuccess_logsSqlStatements() throws Exception {
setUpSuccessfulTest();
runFlowAssertResponse(loadFile("domain_delete_response_pending.xml"));
LogsSubject.assertAboutLogs()
.that(logHandler)
.hasLogAtLevelWithMessage(
Level.INFO,
"insert into \"DomainHistory\" (history_by_superuser,history_registrar_id,"
+ "history_modification_time,history_other_registrar_id,history_period_unit,"
+ "history_period_value,history_reason,history");
}
}

0 comments on commit 75e3b82

Please sign in to comment.