Abstract client info for internal and external use

prestodb · May 30, 2024 · 33d94e8 · 33d94e8
1 parent ab51711
commit 33d94e8
Show file tree

Hide file tree

Showing 14 changed files with 104 additions and 30 deletions.
diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/PrestoVerifyCommand.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/PrestoVerifyCommand.java
@@ -16,6 +16,8 @@
 import com.facebook.presto.sql.parser.SqlParserOptions;
 import com.facebook.presto.verifier.framework.AbstractVerifyCommand;
 import com.facebook.presto.verifier.framework.SourceQuery;
+import com.facebook.presto.verifier.prestoaction.ClientInfoFactory;
+import com.facebook.presto.verifier.prestoaction.DefaultClientInfoFactory;
 import com.facebook.presto.verifier.prestoaction.PrestoExceptionClassifier;
 import com.facebook.presto.verifier.prestoaction.SqlExceptionClassifier;
 import com.google.common.collect.ImmutableList;
@@ -72,4 +74,10 @@ public SqlExceptionClassifier getSqlExceptionClassifier()
     {
         return PrestoExceptionClassifier.defaultBuilder().build();
     }
+
+    @Override
+    public Class<? extends ClientInfoFactory> getClientInfoFactory()
+    {
+        return DefaultClientInfoFactory.class;
+    }
 }
diff --git a/...-verifier/src/main/java/com/facebook/presto/verifier/framework/AbstractVerifyCommand.java b/...-verifier/src/main/java/com/facebook/presto/verifier/framework/AbstractVerifyCommand.java
@@ -44,7 +44,8 @@ public void run()
         Bootstrap app = new Bootstrap(ImmutableList.<Module>builder()
                 .add(new VerifierModule(
                         getSqlParserOptions(),
-                        getCustomQueryFilterClasses()))
+                        getCustomQueryFilterClasses(),
+                        getClientInfoFactory()))
                 .add(new SourceQueryModule(getCustomSourceQuerySupplierTypes()))
                 .add(new EventClientModule(getCustomEventClientTypes()))
                 .add(new QueryActionsModule(getSqlExceptionClassifier(), getCustomQueryActionTypes()))

diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/VerifierModule.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/VerifierModule.java
@@ -39,6 +39,7 @@
 import com.facebook.presto.verifier.checksum.RowColumnValidator;
 import com.facebook.presto.verifier.checksum.SimpleColumnValidator;
 import com.facebook.presto.verifier.framework.Column.Category;
+import com.facebook.presto.verifier.prestoaction.ClientInfoFactory;
 import com.facebook.presto.verifier.resolver.FailureResolverModule;
 import com.facebook.presto.verifier.rewrite.VerificationQueryRewriterModule;
 import com.google.common.collect.ImmutableList;
@@ -77,13 +78,16 @@ public class VerifierModule
 {
     private final SqlParserOptions sqlParserOptions;
     private final List<Class<? extends Predicate<SourceQuery>>> customQueryFilterClasses;
+    private final Class<? extends ClientInfoFactory> clientInfoFactory;
 
     public VerifierModule(
             SqlParserOptions sqlParserOptions,
-            List<Class<? extends Predicate<SourceQuery>>> customQueryFilterClasses)
+            List<Class<? extends Predicate<SourceQuery>>> customQueryFilterClasses,
+            Class<? extends ClientInfoFactory> clientInfoFactory)
     {
         this.sqlParserOptions = requireNonNull(sqlParserOptions, "sqlParserOptions is null");
         this.customQueryFilterClasses = ImmutableList.copyOf(customQueryFilterClasses);
+        this.clientInfoFactory = requireNonNull(clientInfoFactory, "clientInfoFactory is null");
     }
 
     protected final void setup(Binder binder)
@@ -114,6 +118,7 @@ protected final void setup(Binder binder)
 
         // parser
         binder.bind(SqlParserOptions.class).toInstance(sqlParserOptions);
+        binder.bind(ClientInfoFactory.class).to(clientInfoFactory).in(SINGLETON);
         binder.bind(SqlParser.class).in(SINGLETON);
 
         // transaction

diff --git a/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/VerifyCommand.java b/presto-verifier/src/main/java/com/facebook/presto/verifier/framework/VerifyCommand.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.verifier.framework;
 
 import com.facebook.presto.sql.parser.SqlParserOptions;
+import com.facebook.presto.verifier.prestoaction.ClientInfoFactory;
 import com.facebook.presto.verifier.prestoaction.SqlExceptionClassifier;
 import com.google.inject.Module;
 
@@ -37,4 +38,6 @@ public interface VerifyCommand
     List<Class<? extends Predicate<SourceQuery>>> getCustomQueryFilterClasses();
 
     SqlExceptionClassifier getSqlExceptionClassifier();
+
+    Class<? extends ClientInfoFactory> getClientInfoFactory();
 }
diff --git a/...o-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/ClientInfoFactory.java b/...o-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/ClientInfoFactory.java
@@ -0,0 +1,22 @@
+/*
+ * 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 com.facebook.presto.verifier.prestoaction;
+
+import com.facebook.presto.verifier.framework.QueryStage;
+import com.facebook.presto.verifier.framework.VerificationContext;
+
+public interface ClientInfoFactory
+{
+    public ClientInfo create(VerificationContext verificationContext, QueryStage queryStage);
+}
diff --git a/...ier/src/main/java/com/facebook/presto/verifier/prestoaction/DefaultClientInfoFactory.java b/...ier/src/main/java/com/facebook/presto/verifier/prestoaction/DefaultClientInfoFactory.java
@@ -0,0 +1,39 @@
+/*
+ * 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 com.facebook.presto.verifier.prestoaction;
+
+import com.facebook.presto.verifier.framework.QueryStage;
+import com.facebook.presto.verifier.framework.VerificationContext;
+import com.facebook.presto.verifier.framework.VerifierConfig;
+
+import javax.inject.Inject;
+
+import static java.util.Objects.requireNonNull;
+
+public class DefaultClientInfoFactory
+        implements ClientInfoFactory
+{
+    private final VerifierConfig verifierConfig;
+
+    @Inject
+    public DefaultClientInfoFactory(VerifierConfig verifierConfig)
+    {
+        this.verifierConfig = requireNonNull(verifierConfig, "verifierConfig is null");
+    }
+
+    public ClientInfo create(VerificationContext verificationContext, QueryStage queryStage)
+    {
+        return new ClientInfo(verifierConfig.getTestId(), verifierConfig.getTestName(), verificationContext.getSourceQueryName(), verificationContext.getSuite(), queryStage.name());
+    }
+}
diff --git a/...to-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/JdbcPrestoAction.java b/...to-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/JdbcPrestoAction.java
@@ -26,7 +26,6 @@
 import com.facebook.presto.verifier.framework.QueryStage;
 import com.facebook.presto.verifier.framework.ThrottlingException;
 import com.facebook.presto.verifier.framework.VerificationContext;
-import com.facebook.presto.verifier.framework.VerifierConfig;
 import com.facebook.presto.verifier.retry.ForClusterConnection;
 import com.facebook.presto.verifier.retry.ForPresto;
 import com.facebook.presto.verifier.retry.RetryConfig;
@@ -62,13 +61,12 @@ public class JdbcPrestoAction
     private final Duration queryTimeout;
     private final Duration metadataTimeout;
     private final Duration checksumTimeout;
-    private final String testId;
-    private final Optional<String> testName;
     private final String applicationName;
     private final boolean removeMemoryRelatedSessionProperties;
 
     private final RetryDriver<QueryException> networkRetry;
     private final RetryDriver<QueryException> prestoRetry;
+    private final ClientInfoFactory clientInfoFactory;
 
     public JdbcPrestoAction(
             SqlExceptionClassifier exceptionClassifier,
@@ -80,7 +78,7 @@ public JdbcPrestoAction(
             Duration checksumTimeout,
             @ForClusterConnection RetryConfig networkRetryConfig,
             @ForPresto RetryConfig prestoRetryConfig,
-            VerifierConfig verifierConfig)
+            ClientInfoFactory clientInfoFactory)
     {
         this.exceptionClassifier = requireNonNull(exceptionClassifier, "exceptionClassifier is null");
         this.queryConfiguration = requireNonNull(queryConfiguration, "queryConfiguration is null");
@@ -92,8 +90,7 @@ public JdbcPrestoAction(
         this.removeMemoryRelatedSessionProperties = prestoActionConfig.isRemoveMemoryRelatedSessionProperties();
         this.metadataTimeout = requireNonNull(metadataTimeout, "metadataTimeout is null");
         this.checksumTimeout = requireNonNull(checksumTimeout, "checksumTimeout is null");
-        this.testId = requireNonNull(verifierConfig.getTestId(), "testId is null");
-        this.testName = requireNonNull(verifierConfig.getTestName(), "testName is null");
+        this.clientInfoFactory = requireNonNull(clientInfoFactory, "clientInfoFactory is null");
 
         this.networkRetry = new RetryDriver<>(
                 networkRetryConfig,
@@ -131,14 +128,8 @@ private <T> T execute(Statement statement, QueryStage queryStage, StatementExecu
     private <T> T executeOnce(Statement statement, QueryStage queryStage, StatementExecutor<T> statementExecutor)
     {
         String query = formatSql(statement, Optional.empty());
-        String clientInfo = new ClientInfo(
-                testId,
-                testName,
-                verificationContext.getSourceQueryName(),
-                verificationContext.getSuite(),
-                queryStage.name()).serialize();
-
-        try (PrestoConnection connection = getConnection(queryStage, clientInfo)) {
+
+        try (PrestoConnection connection = getConnection(queryStage, clientInfoFactory.create(verificationContext, queryStage).serialize())) {
             try (java.sql.Statement jdbcStatement = connection.createStatement()) {
                 PrestoStatement prestoStatement = jdbcStatement.unwrap(PrestoStatement.class);
                 prestoStatement.setProgressMonitor(statementExecutor.getProgressMonitor());
@@ -154,9 +145,9 @@ private PrestoConnection getConnection(QueryStage queryStage, String clientInfo)
             throws SQLException
     {
         PrestoConnection connection = DriverManager.getConnection(
-                jdbcUrlSelector.next(),
-                queryConfiguration.getUsername().orElse(null),
-                queryConfiguration.getPassword().orElse(null))
+                        jdbcUrlSelector.next(),
+                        queryConfiguration.getUsername().orElse(null),
+                        queryConfiguration.getPassword().orElse(null))
                 .unwrap(PrestoConnection.class);
 
         try {

diff --git a/...fier/src/main/java/com/facebook/presto/verifier/prestoaction/JdbcPrestoActionFactory.java b/...fier/src/main/java/com/facebook/presto/verifier/prestoaction/JdbcPrestoActionFactory.java
@@ -15,7 +15,6 @@
 
 import com.facebook.presto.verifier.framework.QueryConfiguration;
 import com.facebook.presto.verifier.framework.VerificationContext;
-import com.facebook.presto.verifier.framework.VerifierConfig;
 import com.facebook.presto.verifier.retry.RetryConfig;
 import io.airlift.units.Duration;
 
@@ -30,7 +29,7 @@ public class JdbcPrestoActionFactory
     private final RetryConfig prestoRetryConfig;
     private final Duration metadataTimeout;
     private final Duration checksumTimeout;
-    private final VerifierConfig verifierConfig;
+    private final ClientInfoFactory clientInfoFactory;
     private final JdbcUrlSelector jdbcUrlSelector;
 
     public JdbcPrestoActionFactory(
@@ -40,16 +39,16 @@ public JdbcPrestoActionFactory(
             RetryConfig prestoRetryConfig,
             Duration metadataTimeout,
             Duration checksumTimeout,
-            VerifierConfig verifierConfig)
+            ClientInfoFactory clientInfoFactory)
     {
         this.exceptionClassifier = requireNonNull(exceptionClassifier, "exceptionClassifier is null");
         this.prestoActionConfig = requireNonNull(prestoActionConfig, "prestoClusterConfig is null");
         this.networkRetryConfig = requireNonNull(networkRetryConfig, "networkRetryConfig is null");
         this.prestoRetryConfig = requireNonNull(prestoRetryConfig, "prestoRetryConfig is null");
         this.metadataTimeout = requireNonNull(metadataTimeout, "metadataTimeout is null");
         this.checksumTimeout = requireNonNull(checksumTimeout, "checksumTimeout is null");
-        this.verifierConfig = requireNonNull(verifierConfig, "verifierConfig is null");
         this.jdbcUrlSelector = new JdbcUrlSelector(prestoActionConfig.getJdbcUrls());
+        this.clientInfoFactory = requireNonNull(clientInfoFactory, "clientInfoFactory is null");
     }
 
     @Override
@@ -65,6 +64,6 @@ public JdbcPrestoAction create(QueryConfiguration queryConfiguration, Verificati
                 checksumTimeout,
                 networkRetryConfig,
                 prestoRetryConfig,
-                verifierConfig);
+                clientInfoFactory);
     }
 }
diff --git a/...-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/QueryActionsModule.java b/...-verifier/src/main/java/com/facebook/presto/verifier/prestoaction/QueryActionsModule.java
@@ -17,7 +17,6 @@
 import com.facebook.presto.verifier.annotation.ForControl;
 import com.facebook.presto.verifier.annotation.ForHelper;
 import com.facebook.presto.verifier.annotation.ForTest;
-import com.facebook.presto.verifier.framework.VerifierConfig;
 import com.facebook.presto.verifier.retry.ForClusterConnection;
 import com.facebook.presto.verifier.retry.ForPresto;
 import com.facebook.presto.verifier.retry.RetryConfig;
@@ -115,7 +114,7 @@ public PrestoActionFactory get()
                     injector.getInstance(Key.get(RetryConfig.class, ForPresto.class)),
                     queryActionsConfig.getMetadataTimeout(),
                     queryActionsConfig.getChecksumTimeout(),
-                    injector.getInstance(VerifierConfig.class));
+                    injector.getInstance(ClientInfoFactory.class));
         }
     }
 }
diff --git a/...rifier/src/test/java/com/facebook/presto/verifier/framework/AbstractVerificationTest.java b/...rifier/src/test/java/com/facebook/presto/verifier/framework/AbstractVerificationTest.java
@@ -21,6 +21,7 @@
 import com.facebook.presto.sql.parser.SqlParserOptions;
 import com.facebook.presto.tests.StandaloneQueryRunner;
 import com.facebook.presto.verifier.event.VerifierQueryEvent;
+import com.facebook.presto.verifier.prestoaction.DefaultClientInfoFactory;
 import com.facebook.presto.verifier.prestoaction.JdbcPrestoAction;
 import com.facebook.presto.verifier.prestoaction.JdbcUrlSelector;
 import com.facebook.presto.verifier.prestoaction.PrestoAction;
@@ -173,7 +174,7 @@ protected PrestoAction getPrestoAction(Optional<QueryConfiguration> queryConfigu
                 queryActionsConfig.getChecksumTimeout(),
                 retryConfig,
                 retryConfig,
-                verifierConfig);
+                new DefaultClientInfoFactory(verifierConfig));
     }
 
     private Optional<VerifierQueryEvent> verify(

diff --git a/...erifier/src/test/java/com/facebook/presto/verifier/framework/TestDeterminismAnalyzer.java b/...erifier/src/test/java/com/facebook/presto/verifier/framework/TestDeterminismAnalyzer.java
@@ -19,6 +19,7 @@
 import com.facebook.presto.sql.parser.SqlParserOptions;
 import com.facebook.presto.sql.tree.QualifiedName;
 import com.facebook.presto.verifier.checksum.ChecksumValidator;
+import com.facebook.presto.verifier.prestoaction.DefaultClientInfoFactory;
 import com.facebook.presto.verifier.prestoaction.JdbcPrestoAction;
 import com.facebook.presto.verifier.prestoaction.JdbcUrlSelector;
 import com.facebook.presto.verifier.prestoaction.PrestoAction;
@@ -81,7 +82,7 @@ private static DeterminismAnalyzer createDeterminismAnalyzer(String mutableCatal
                 queryActionsConfig.getChecksumTimeout(),
                 retryConfig,
                 retryConfig,
-                verifierConfig);
+                new DefaultClientInfoFactory(verifierConfig));
         QueryRewriter queryRewriter = new QueryRewriter(
                 sqlParser,
                 typeManager,

diff --git a/...erifier/src/test/java/com/facebook/presto/verifier/prestoaction/TestJdbcPrestoAction.java b/...erifier/src/test/java/com/facebook/presto/verifier/prestoaction/TestJdbcPrestoAction.java
@@ -80,6 +80,7 @@ public void setup()
         PrestoActionConfig prestoActionConfig = new PrestoActionConfig()
                 .setHosts(queryRunner.getServer().getAddress().getHost())
                 .setJdbcPort(queryRunner.getServer().getAddress().getPort());
+        VerifierConfig verifierConfig = new VerifierConfig().setTestId("test");
         prestoAction = new JdbcPrestoAction(
                 PrestoExceptionClassifier.defaultBuilder().build(),
                 CONFIGURATION,
@@ -90,7 +91,7 @@ public void setup()
                 queryActionsConfig.getChecksumTimeout(),
                 new RetryConfig(),
                 new RetryConfig(),
-                new VerifierConfig().setTestId("test"));
+                new DefaultClientInfoFactory(new VerifierConfig().setTestId("test")));
     }
 
     @Test

diff --git a/presto-verifier/src/test/java/com/facebook/presto/verifier/rewrite/TestQueryRewriter.java b/presto-verifier/src/test/java/com/facebook/presto/verifier/rewrite/TestQueryRewriter.java
@@ -26,6 +26,7 @@
 import com.facebook.presto.verifier.framework.QueryObjectBundle;
 import com.facebook.presto.verifier.framework.VerificationContext;
 import com.facebook.presto.verifier.framework.VerifierConfig;
+import com.facebook.presto.verifier.prestoaction.DefaultClientInfoFactory;
 import com.facebook.presto.verifier.prestoaction.JdbcPrestoAction;
 import com.facebook.presto.verifier.prestoaction.JdbcUrlSelector;
 import com.facebook.presto.verifier.prestoaction.PrestoAction;
@@ -96,7 +97,7 @@ public void setup()
                 new QueryActionsConfig().getChecksumTimeout(),
                 new RetryConfig(),
                 new RetryConfig(),
-                new VerifierConfig().setTestId("test"));
+                new DefaultClientInfoFactory(new VerifierConfig().setTestId("test")));
     }
 
     @AfterClass

diff --git a/...src/test/java/com/facebook/presto/verifier/source/TestPrestoQuerySourceQuerySupplier.java b/...src/test/java/com/facebook/presto/verifier/source/TestPrestoQuerySourceQuerySupplier.java
@@ -22,6 +22,8 @@
 import com.facebook.presto.verifier.framework.QueryConfiguration;
 import com.facebook.presto.verifier.framework.SourceQuery;
 import com.facebook.presto.verifier.framework.VerifierConfig;
+import com.facebook.presto.verifier.prestoaction.ClientInfoFactory;
+import com.facebook.presto.verifier.prestoaction.DefaultClientInfoFactory;
 import com.facebook.presto.verifier.prestoaction.PrestoExceptionClassifier;
 import com.facebook.presto.verifier.prestoaction.QueryActionsModule;
 import com.google.common.collect.ImmutableList;
@@ -97,6 +99,7 @@ public void setup()
                             configBinder(binder).bindConfig(VerifierConfig.class);
                             binder.bind(SqlParserOptions.class).toInstance(new SqlParserOptions());
                             binder.bind(SqlParser.class).in(SINGLETON);
+                            binder.bind(ClientInfoFactory.class).to(DefaultClientInfoFactory.class);
                         })
                         .build());
         injector = app