diff --git a/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java b/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java index a86d4ae6cbe18..a14b51afb54df 100644 --- a/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java +++ b/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java @@ -338,6 +338,8 @@ public final class SystemSessionProperties public static final String NATIVE_MIN_COLUMNAR_ENCODING_CHANNELS_TO_PREFER_ROW_WISE_ENCODING = "native_min_columnar_encoding_channels_to_prefer_row_wise_encoding"; public static final String NATIVE_ENFORCE_JOIN_BUILD_INPUT_PARTITION = "native_enforce_join_build_input_partition"; public static final String NATIVE_EXECUTION_SCALE_WRITER_THREADS_ENABLED = "native_execution_scale_writer_threads_enabled"; + public static final String IS_QUERY_REWRITER_PLUGIN_ENABLED = "is_query_rewriter_plugin_enabled"; + public static final String IS_QUERY_REWRITER_PLUGIN_SUCCEEDED = "is_query_rewriter_plugin_succeeded"; private final List> sessionProperties; @@ -1858,7 +1860,27 @@ public SystemSessionProperties( EXPRESSION_OPTIMIZER_NAME, "Configure which expression optimizer to use", featuresConfig.getExpressionOptimizerName(), - false)); + false), + booleanProperty( + IS_QUERY_REWRITER_PLUGIN_ENABLED, + "Use queries rewriter plugin", + false, + true), + booleanProperty( + IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, + "Query rewrite success", + false, + true)); + } + + public static boolean isQueryRewriterPluginSucceeded(Session session) + { + return session.getSystemProperty(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, Boolean.class); + } + + public static boolean isQueryRewriterPluginEnabled(Session session) + { + return session.getSystemProperty(IS_QUERY_REWRITER_PLUGIN_ENABLED, Boolean.class); } public static boolean isSpoolingOutputBufferEnabled(Session session) diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java index 1b57af2b25880..9db5a0c5fd061 100644 --- a/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java @@ -15,6 +15,7 @@ import com.facebook.airlift.concurrent.BoundedExecutor; import com.facebook.presto.Session; +import com.facebook.presto.SystemSessionProperties; import com.facebook.presto.common.analyzer.PreparedQuery; import com.facebook.presto.common.resourceGroups.QueryType; import com.facebook.presto.execution.QueryIdGenerator; @@ -37,6 +38,7 @@ import com.facebook.presto.spi.resourceGroups.SelectionContext; import com.facebook.presto.spi.resourceGroups.SelectionCriteria; import com.facebook.presto.spi.security.AccessControl; +import com.facebook.presto.sql.analyzer.AnalyzerProviderManager; import com.facebook.presto.sql.analyzer.QueryPreparerProviderManager; import com.facebook.presto.transaction.TransactionManager; import com.google.common.util.concurrent.AbstractFuture; @@ -92,9 +94,13 @@ public class DispatchManager private final QueryPreparerProviderManager queryPreparerProviderManager; + private final AnalyzerProviderManager analyzerProviderManager; + + private final QueryRewriterManager queryRewriterManager; + /** * Dispatch Manager is used for the pre-queuing part of queries prior to the query execution phase. - * + *

* Dispatch Manager object is instantiated when the presto server is launched by server bootstrap time. It is a critical component in resource management section of the query. * * @param queryIdGenerator query ID generator for generating a new query ID when a query is created @@ -126,7 +132,9 @@ public DispatchManager( QueryManagerConfig queryManagerConfig, DispatchExecutor dispatchExecutor, ClusterStatusSender clusterStatusSender, - Optional clusterQueryTrackerService) + Optional clusterQueryTrackerService, + AnalyzerProviderManager analyzerProviderManager, + QueryRewriterManager queryRewriterManager) { this.queryIdGenerator = requireNonNull(queryIdGenerator, "queryIdGenerator is null"); this.queryPreparerProviderManager = requireNonNull(queryPreparerProviderManager, "queryPreparerProviderManager is null"); @@ -147,6 +155,10 @@ public DispatchManager( this.clusterStatusSender = requireNonNull(clusterStatusSender, "clusterStatusSender is null"); this.queryTracker = new QueryTracker<>(queryManagerConfig, dispatchExecutor.getScheduledExecutor(), clusterQueryTrackerService); + + this.analyzerProviderManager = requireNonNull(analyzerProviderManager, "analyzerProviderManager is null"); + + this.queryRewriterManager = requireNonNull(queryRewriterManager, "queryRewriterManager is null"); } /** @@ -181,7 +193,7 @@ public QueryManagerStats getStats() /** * Create a query id - * + *

* This method is called when a {@code Query} object is created * * @return {@link QueryId} @@ -279,6 +291,17 @@ private void createQueryInternal(QueryId queryId, String slug, int retryCoun preparedQuery = queryPreparerProvider.getQueryPreparer().prepareQuery(analyzerOptions, query, sessionBuilder.getPreparedStatements(), sessionBuilder.getWarningCollector()); query = preparedQuery.getFormattedQuery().orElse(query); + // Rewrite the query + if (SystemSessionProperties.isQueryRewriterPluginEnabled(session)) { + QueryAndSessionProperties queryAndSessionProperties = queryRewriterManager.rewriteQueryAndSession(query, session, analyzerOptions, + analyzerProviderManager.getAnalyzerProvider(getAnalyzerType(session)), queryPreparerProvider); + if (queryAndSessionProperties.getPreparedQuery().isPresent()) { + queryAndSessionProperties.getSystemSessionProperties().forEach(sessionBuilder::setSystemProperty); + preparedQuery = queryAndSessionProperties.getPreparedQuery().get(); + query = queryAndSessionProperties.getQuery(); + } + } + // select resource group Optional queryType = preparedQuery.getQueryType(); sessionBuilder.setQueryType(queryType); @@ -377,7 +400,6 @@ public ListenableFuture waitForDispatched(QueryId queryId) /** * Return a list of {@link BasicQueryInfo}. - * */ public List getQueries() { diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/QueryAndSessionProperties.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/QueryAndSessionProperties.java new file mode 100644 index 0000000000000..25784c1204422 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/QueryAndSessionProperties.java @@ -0,0 +1,49 @@ +/* + * 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.dispatcher; + +import com.facebook.presto.common.analyzer.PreparedQuery; + +import java.util.Map; +import java.util.Optional; + +public class QueryAndSessionProperties +{ + private final String query; + private final Map systemSessionProperties; + private final Optional preparedQuery; + + public QueryAndSessionProperties(String query, Map systemSessionProperties, Optional preparedQuery) + { + this.query = query; + this.systemSessionProperties = systemSessionProperties; + this.preparedQuery = preparedQuery; + } + + public String getQuery() + { + return query; + } + + public Map getSystemSessionProperties() + { + return systemSessionProperties; + } + + public Optional getPreparedQuery() + { + return preparedQuery; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/dispatcher/QueryRewriterManager.java b/presto-main/src/main/java/com/facebook/presto/dispatcher/QueryRewriterManager.java new file mode 100644 index 0000000000000..399454a6092ee --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/dispatcher/QueryRewriterManager.java @@ -0,0 +1,139 @@ +/* + * 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.dispatcher; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.Session; +import com.facebook.presto.common.analyzer.PreparedQuery; +import com.facebook.presto.eventlistener.EventListenerManager; +import com.facebook.presto.server.SessionContext; +import com.facebook.presto.spi.QueryId; +import com.facebook.presto.spi.analyzer.AnalyzerOptions; +import com.facebook.presto.spi.analyzer.AnalyzerProvider; +import com.facebook.presto.spi.analyzer.QueryPreparerProvider; +import com.facebook.presto.spi.eventlistener.EventListener; +import com.facebook.presto.spi.rewriter.QueryRewriterInput; +import com.facebook.presto.spi.rewriter.QueryRewriterOutput; +import com.facebook.presto.spi.rewriter.QueryRewriterProvider; +import com.facebook.presto.spi.rewriter.QueryRewriterProviderFactory; +import com.google.common.collect.ImmutableList; +import com.google.inject.Inject; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import static com.facebook.presto.SystemSessionProperties.IS_QUERY_REWRITER_PLUGIN_ENABLED; +import static com.google.common.base.Preconditions.checkState; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +/** + * To provide a query rewriter plugin, i.e. a plugin that inputs a set of session properties and query and returns the + * updated session properties and rewritten query. + * 1) Provide implementation for QueryRewriterProviderFactory + * 2) Provide implementation for QueryRewriterProvider + * 3) Implement com.facebook.presto.spi.Plugin and provide implementation for Iterable getQueryRewriterProviderFactory() + * 4) Finally provide implementation of QueryRewriter, which does actual query rewrite. + * For example: + * {@link presto-tests/com.facebook.presto.plugin.rewriter.UpperCasingQueryRewriterPlugin} + */ +public class QueryRewriterManager +{ + private static final Logger log = Logger.get(QueryRewriterManager.class); + + private final AtomicReference> providerFactory = new AtomicReference<>(Optional.empty()); + private final AtomicReference> provider = new AtomicReference<>(Optional.empty()); + private final EventListenerManager eventListenerManager; + + @Inject + public QueryRewriterManager(EventListenerManager eventListenerManager) + { + this.eventListenerManager = requireNonNull(eventListenerManager, "eventListenerManager is null"); + } + + public static Boolean isQueryRewriterPluginEnabled(SessionContext sessionContext) + { + return Boolean.parseBoolean(sessionContext.getSystemProperties().getOrDefault(IS_QUERY_REWRITER_PLUGIN_ENABLED, "false")); + } + + public void addQueryRewriterProviderFactory(QueryRewriterProviderFactory queryRewriterProviderFactory) + { + requireNonNull(queryRewriterProviderFactory, "queryRewriterProviderFactory is null"); + checkState(providerFactory.compareAndSet(Optional.empty(), Optional.of(queryRewriterProviderFactory)), + format("A query rewriter factory is already registered with name %s", queryRewriterProviderFactory.getName())); + } + + public void loadQueryRewriterProvider() + { + List configuredEventListener = ImmutableList.of(); + if (eventListenerManager.getConfiguredEventListener().isPresent()) { + configuredEventListener = ImmutableList.of(eventListenerManager.getConfiguredEventListener().get()); + } + Optional queryRewriterProviderFactory = providerFactory.get(); + if (queryRewriterProviderFactory.isPresent()) { + QueryRewriterProvider queryRewriterProvider = queryRewriterProviderFactory.get().create(configuredEventListener); + checkState(provider.compareAndSet(Optional.empty(), Optional.of(queryRewriterProvider)), + format("A query rewriter provider is already registered %s", queryRewriterProvider)); + } + } + + public Optional getQueryRewriterProvider() + { + return provider.get(); + } + + public QueryAndSessionProperties rewriteQueryAndSession( + String query, + Session session, + AnalyzerOptions analyzerOptions, + AnalyzerProvider analyzerProvider, + QueryPreparerProvider queryPreparerProvider) + { + QueryId queryId = session.getQueryId(); + QueryAndSessionProperties rewrittenQueryAndSessionProperties = new QueryAndSessionProperties(query, session.getSystemProperties(), Optional.empty()); + if (getQueryRewriterProvider().isPresent()) { + QueryRewriterInput queryRewriterInput = new QueryRewriterInput.Builder() + .setQuery(query) + .setQueryId(queryId.getId()) + .setCatalog(session.getCatalog()) + .setSchema(session.getSchema()) + .setPreparedStatements(session.getPreparedStatements()) + .setWarningCollector(session.getWarningCollector()) + .setSessionProperties(session.getSystemProperties()) + .setAnalyzerOptions(analyzerOptions) + .setAnalyzerProvider(analyzerProvider) + .setQueryPreparer(queryPreparerProvider.getQueryPreparer()) + .build(); + try { + QueryRewriterProvider provider = getQueryRewriterProvider().get(); + QueryRewriterOutput queryRewriterOutput = provider.getQueryRewriter().rewriteSQL(queryRewriterInput); + String rewrittenQuery = queryRewriterOutput.getRewrittenQuery(); + // Checking if the rewritten query is parseable. + PreparedQuery preparedQuery = queryPreparerProvider.getQueryPreparer() + .prepareQuery(analyzerOptions, rewrittenQuery, session.getPreparedStatements(), session.getWarningCollector()); + // apply updated session properties. + Map systemPropertyOverrides = queryRewriterOutput.getSessionProperties(); + rewrittenQueryAndSessionProperties = new QueryAndSessionProperties(rewrittenQuery, systemPropertyOverrides, Optional.of(preparedQuery)); + log.info("createQueryInternal :: QueryId [%s] - Replacing with optimized query", queryId.getId()); + } + catch (Exception e) { + // TODO : Implement a better way to test if rewritten query is parseable. + log.warn(format("rewritten query for query with id %s is discarded", session.getQueryId()), e); + } + } + return rewrittenQueryAndSessionProperties; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/eventlistener/EventListenerManager.java b/presto-main/src/main/java/com/facebook/presto/eventlistener/EventListenerManager.java index 24a1d04b54bb1..1f098cfbcdbdc 100644 --- a/presto-main/src/main/java/com/facebook/presto/eventlistener/EventListenerManager.java +++ b/presto-main/src/main/java/com/facebook/presto/eventlistener/EventListenerManager.java @@ -57,6 +57,11 @@ public void addEventListenerFactory(EventListenerFactory eventListenerFactory) } } + public Optional getConfiguredEventListener() + { + return configuredEventListener.get(); + } + public void loadConfiguredEventListener() throws Exception { diff --git a/presto-main/src/main/java/com/facebook/presto/server/PluginManager.java b/presto-main/src/main/java/com/facebook/presto/server/PluginManager.java index 274f90bed1204..648bebf47e8e5 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/PluginManager.java +++ b/presto-main/src/main/java/com/facebook/presto/server/PluginManager.java @@ -23,6 +23,7 @@ import com.facebook.presto.connector.ConnectorManager; import com.facebook.presto.cost.HistoryBasedPlanStatisticsManager; import com.facebook.presto.dispatcher.QueryPrerequisitesManager; +import com.facebook.presto.dispatcher.QueryRewriterManager; import com.facebook.presto.eventlistener.EventListenerManager; import com.facebook.presto.execution.resourceGroups.ResourceGroupManager; import com.facebook.presto.metadata.Metadata; @@ -42,6 +43,7 @@ import com.facebook.presto.spi.plan.PlanCheckerProviderFactory; import com.facebook.presto.spi.prerequisites.QueryPrerequisitesFactory; import com.facebook.presto.spi.resourceGroups.ResourceGroupConfigurationManagerFactory; +import com.facebook.presto.spi.rewriter.QueryRewriterProviderFactory; import com.facebook.presto.spi.security.PasswordAuthenticatorFactory; import com.facebook.presto.spi.security.PrestoAuthenticatorFactory; import com.facebook.presto.spi.security.SystemAccessControlFactory; @@ -144,6 +146,7 @@ public class PluginManager private final ClientRequestFilterManager clientRequestFilterManager; private final PlanCheckerProviderManager planCheckerProviderManager; private final ExpressionOptimizerManager expressionOptimizerManager; + private final QueryRewriterManager queryRewriterManager; @Inject public PluginManager( @@ -169,7 +172,8 @@ public PluginManager( NodeStatusNotificationManager nodeStatusNotificationManager, ClientRequestFilterManager clientRequestFilterManager, PlanCheckerProviderManager planCheckerProviderManager, - ExpressionOptimizerManager expressionOptimizerManager) + ExpressionOptimizerManager expressionOptimizerManager, + QueryRewriterManager queryRewriterManager) { requireNonNull(nodeInfo, "nodeInfo is null"); requireNonNull(config, "config is null"); @@ -205,6 +209,7 @@ public PluginManager( this.clientRequestFilterManager = requireNonNull(clientRequestFilterManager, "clientRequestFilterManager is null"); this.planCheckerProviderManager = requireNonNull(planCheckerProviderManager, "planCheckerProviderManager is null"); this.expressionOptimizerManager = requireNonNull(expressionOptimizerManager, "expressionManager is null"); + this.queryRewriterManager = requireNonNull(queryRewriterManager, "queryRewriterManager is null"); } public void loadPlugins() @@ -379,6 +384,11 @@ public void installPlugin(Plugin plugin) log.info("Registering client request filter factory"); clientRequestFilterManager.registerClientRequestFilterFactory(clientRequestFilterFactory); } + + for (QueryRewriterProviderFactory queryRewriterProviderFactory : plugin.getQueryRewriterProviderFactories()) { + log.info("Registering query rewriter provider factory %s", queryRewriterProviderFactory.getName()); + queryRewriterManager.addQueryRewriterProviderFactory(queryRewriterProviderFactory); + } } public void installCoordinatorPlugin(CoordinatorPlugin plugin) diff --git a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java index 19250dd963eb7..466e83ce5929b 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java @@ -35,6 +35,7 @@ import com.facebook.presto.ClientRequestFilterModule; import com.facebook.presto.dispatcher.QueryPrerequisitesManager; import com.facebook.presto.dispatcher.QueryPrerequisitesManagerModule; +import com.facebook.presto.dispatcher.QueryRewriterManager; import com.facebook.presto.eventlistener.EventListenerManager; import com.facebook.presto.eventlistener.EventListenerModule; import com.facebook.presto.execution.resourceGroups.ResourceGroupManager; @@ -190,6 +191,7 @@ public void run() injector.getInstance(NodeStatusNotificationManager.class).loadNodeStatusNotificationProvider(); injector.getInstance(GracefulShutdownHandler.class).loadNodeStatusNotification(); injector.getInstance(SessionPropertyManager.class).loadSessionPropertyProviders(); + injector.getInstance(QueryRewriterManager.class).loadQueryRewriterProvider(); PlanCheckerProviderManager planCheckerProviderManager = injector.getInstance(PlanCheckerProviderManager.class); InternalNodeManager nodeManager = injector.getInstance(DiscoveryNodeManager.class); NodeInfo nodeInfo = injector.getInstance(NodeInfo.class); diff --git a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java index 881588658003e..6983d33c59c9d 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java +++ b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java @@ -49,6 +49,7 @@ import com.facebook.presto.cost.HistoryBasedPlanStatisticsManager; import com.facebook.presto.cost.ScalarStatsCalculator; import com.facebook.presto.cost.StatsNormalizer; +import com.facebook.presto.dispatcher.QueryRewriterManager; import com.facebook.presto.event.SplitMonitor; import com.facebook.presto.execution.ExecutionFailureInfo; import com.facebook.presto.execution.ExplainAnalyzeContext; @@ -797,6 +798,9 @@ public ListeningExecutorService createResourceManagerExecutor(ResourceManagerCon // cleanup binder.bind(ExecutorCleanup.class).in(Scopes.SINGLETON); + // Pluggable query rewriter + binder.bind(QueryRewriterManager.class).in(Scopes.SINGLETON); + // Distributed tracing configBinder(binder).bindConfig(TracingConfig.class); binder.bind(TracerProviderManager.class).in(Scopes.SINGLETON); diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java index a34b46003ad89..106cfa0a7648f 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java @@ -38,6 +38,7 @@ import com.facebook.presto.cost.StatsCalculator; import com.facebook.presto.dispatcher.DispatchManager; import com.facebook.presto.dispatcher.QueryPrerequisitesManagerModule; +import com.facebook.presto.dispatcher.QueryRewriterManager; import com.facebook.presto.eventlistener.EventListenerManager; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryManager; @@ -184,6 +185,7 @@ public class TestingPrestoServer private final PlanCheckerProviderManager planCheckerProviderManager; private final NodeManager pluginNodeManager; private final ClientRequestFilterManager clientRequestFilterManager; + private final QueryRewriterManager queryRewriterManager; public static class TestShutdownAction implements ShutdownAction @@ -330,6 +332,7 @@ public TestingPrestoServer( binder.bind(GracefulShutdownHandler.class).in(Scopes.SINGLETON); binder.bind(ProcedureTester.class).in(Scopes.SINGLETON); binder.bind(RequestBlocker.class).in(Scopes.SINGLETON); + binder.bind(QueryRewriterManager.class).in(Scopes.SINGLETON); newSetBinder(binder, Filter.class, TheServlet.class).addBinding() .to(RequestBlocker.class).in(Scopes.SINGLETON); }); @@ -365,6 +368,7 @@ public TestingPrestoServer( pluginManager = injector.getInstance(PluginManager.class); connectorManager = injector.getInstance(ConnectorManager.class); + queryRewriterManager = injector.getInstance(QueryRewriterManager.class); server = injector.getInstance(TestingHttpServer.class); catalogManager = injector.getInstance(CatalogManager.class); @@ -465,6 +469,11 @@ else if (catalogServer) { refreshNodes(); } + public QueryRewriterManager getQueryRewriterManager() + { + return queryRewriterManager; + } + private Map getServerProperties( boolean resourceManagerEnabled, boolean catalogServerEnabled, diff --git a/presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java b/presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java index 8d2a9ac22e5bb..7f0282fcc5ab2 100644 --- a/presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java +++ b/presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java @@ -51,6 +51,7 @@ import com.facebook.presto.cost.TaskCountEstimator; import com.facebook.presto.dispatcher.NoOpQueryManager; import com.facebook.presto.dispatcher.QueryPrerequisitesManager; +import com.facebook.presto.dispatcher.QueryRewriterManager; import com.facebook.presto.eventlistener.EventListenerManager; import com.facebook.presto.execution.AlterFunctionTask; import com.facebook.presto.execution.CommitTask; @@ -523,7 +524,7 @@ private LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig, BuiltInAnalyzerProvider analyzerProvider = new BuiltInAnalyzerProvider(queryAnalyzer); BuiltInQueryPreparer queryPreparer = new BuiltInQueryPreparer(sqlParser); BuiltInQueryPreparerProvider queryPreparerProvider = new BuiltInQueryPreparerProvider(queryPreparer); - + EventListenerManager eventListenerManager = new EventListenerManager(); this.pluginManager = new PluginManager( nodeInfo, new PluginManagerConfig(), @@ -547,7 +548,8 @@ private LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig, new NodeStatusNotificationManager(), new ClientRequestFilterManager(), planCheckerProviderManager, - expressionOptimizerManager); + expressionOptimizerManager, + new QueryRewriterManager(eventListenerManager)); connectorManager.addConnectorFactory(globalSystemConnectorFactory); connectorManager.createConnection(GlobalSystemConnector.NAME, GlobalSystemConnector.NAME, ImmutableMap.of()); diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/Plugin.java b/presto-spi/src/main/java/com/facebook/presto/spi/Plugin.java index 2da8ad1970eec..5f7ab4b42850e 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/Plugin.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/Plugin.java @@ -24,6 +24,7 @@ import com.facebook.presto.spi.nodestatus.NodeStatusNotificationProviderFactory; import com.facebook.presto.spi.prerequisites.QueryPrerequisitesFactory; import com.facebook.presto.spi.resourceGroups.ResourceGroupConfigurationManagerFactory; +import com.facebook.presto.spi.rewriter.QueryRewriterProviderFactory; import com.facebook.presto.spi.security.PasswordAuthenticatorFactory; import com.facebook.presto.spi.security.PrestoAuthenticatorFactory; import com.facebook.presto.spi.security.SystemAccessControlFactory; @@ -144,6 +145,11 @@ default Iterable getQueryPreparerProviders() return emptyList(); } + default Iterable getQueryRewriterProviderFactories() + { + return emptyList(); + } + default Iterable getNodeStatusNotificationProviderFactory() { return emptyList(); diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriter.java b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriter.java new file mode 100644 index 0000000000000..0913424850ee6 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriter.java @@ -0,0 +1,20 @@ +/* + * 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.spi.rewriter; + +public interface QueryRewriter +{ + QueryRewriterOutput rewriteSQL(QueryRewriterInput queryRewriterInput); +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterInput.java b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterInput.java new file mode 100644 index 0000000000000..3ae470910634e --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterInput.java @@ -0,0 +1,207 @@ +/* + * 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.spi.rewriter; + +import com.facebook.presto.spi.WarningCollector; +import com.facebook.presto.spi.analyzer.AnalyzerOptions; +import com.facebook.presto.spi.analyzer.AnalyzerProvider; +import com.facebook.presto.spi.analyzer.QueryPreparer; + +import java.util.Map; +import java.util.Optional; + +import static com.facebook.presto.common.Utils.checkArgument; +import static java.util.Collections.emptyMap; +import static java.util.Objects.requireNonNull; + +public class QueryRewriterInput +{ + private final String queryId; + private final Optional catalog; + private final Optional schema; + private final String query; + private final Map sessionProperties; + private final WarningCollector warningCollector; + private final Map preparedStatements; + private final AnalyzerProvider analyzerProvider; + private final QueryPreparer queryPreparer; + private final AnalyzerOptions analyzerOptions; + + private QueryRewriterInput( + String queryId, + Optional catalog, + Optional schema, + String query, + Map sessionProperties, + WarningCollector warningCollector, + Map preparedStatements, + AnalyzerProvider analyzerProvider, + AnalyzerOptions analyzerOptions, + QueryPreparer queryPreparer) + { + this.queryId = queryId; + this.catalog = catalog; + this.schema = schema; + this.query = query; + this.sessionProperties = sessionProperties; + this.warningCollector = warningCollector; + this.preparedStatements = preparedStatements; + this.analyzerOptions = analyzerOptions; + this.analyzerProvider = analyzerProvider; + this.queryPreparer = queryPreparer; + } + + public Optional getCatalog() + { + return catalog; + } + + public Optional getSchema() + { + return schema; + } + + public WarningCollector getWarningCollector() + { + return warningCollector; + } + + public Map getPreparedStatements() + { + return preparedStatements; + } + + public String getQueryId() + { + return queryId; + } + + public String getQuery() + { + return query; + } + + public Map getSessionProperties() + { + return sessionProperties; + } + + public AnalyzerProvider getAnalyzerProvider() + { + return analyzerProvider; + } + + public AnalyzerOptions getAnalyzerOptions() + { + return analyzerOptions; + } + + public QueryPreparer getQueryPreparer() + { + return queryPreparer; + } + + public static class Builder + { + private String queryId; + private Optional catalog = Optional.empty(); + private Optional schema = Optional.empty(); + private String query; + private Map sessionProperties = emptyMap(); + private WarningCollector warningCollector = WarningCollector.NOOP; + private Map preparedStatements = emptyMap(); + private AnalyzerProvider analyzerProvider; + private AnalyzerOptions analyzerOptions; + private QueryPreparer queryPreparer; + + public Builder setQueryId(String queryId) + { + this.queryId = requireNonNull(queryId, "queryId is null"); + return this; + } + + public Builder setCatalog(Optional catalog) + { + this.catalog = catalog; + return this; + } + + public Builder setSchema(Optional schema) + { + this.schema = schema; + return this; + } + + public Builder setQuery(String query) + { + this.query = requireNonNull(query, "query is null"); + return this; + } + + public Builder setSessionProperties(Map sessionProperties) + { + this.sessionProperties = sessionProperties; + return this; + } + + public Builder setWarningCollector(WarningCollector warningCollector) + { + this.warningCollector = warningCollector; + return this; + } + + public Builder setPreparedStatements(Map preparedStatements) + { + this.preparedStatements = preparedStatements; + return this; + } + + public Builder setAnalyzerProvider(AnalyzerProvider analyzerProvider) + { + this.analyzerProvider = analyzerProvider; + return this; + } + + public Builder setAnalyzerOptions(AnalyzerOptions analyzerOptions) + { + this.analyzerOptions = analyzerOptions; + return this; + } + + public Builder setQueryPreparer(QueryPreparer queryPreparer) + { + this.queryPreparer = queryPreparer; + return this; + } + + public QueryRewriterInput build() + { + requireNonNull(queryId, "queryId is null"); + requireNonNull(query, "query is null"); + checkArgument(!query.isEmpty() && !queryId.isEmpty()); + return new QueryRewriterInput( + queryId, + catalog, + schema, + query, + sessionProperties, + warningCollector, + preparedStatements, + analyzerProvider, + analyzerOptions, + queryPreparer); + } + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterOutput.java b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterOutput.java new file mode 100644 index 0000000000000..cdd9310fe162a --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterOutput.java @@ -0,0 +1,98 @@ +/* + * 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.spi.rewriter; + +import java.util.Map; + +import static com.facebook.presto.common.Utils.checkArgument; +import static java.util.Collections.emptyMap; +import static java.util.Objects.requireNonNull; + +public class QueryRewriterOutput +{ + private final String queryId; + private final String originalQuery; + private final String rewrittenQuery; + private final Map sessionProperties; + + private QueryRewriterOutput(String queryId, String originalQuery, String rewrittenQuery, Map sessionProperties) + { + this.queryId = queryId; + this.originalQuery = originalQuery; + this.rewrittenQuery = rewrittenQuery; + this.sessionProperties = sessionProperties; + } + + public String getOriginalQuery() + { + return originalQuery; + } + + public String getRewrittenQuery() + { + return rewrittenQuery; + } + + public Map getSessionProperties() + { + return sessionProperties; + } + + public String getQueryId() + { + return queryId; + } + + public static class Builder + { + private String queryId; + private String originalQuery; + private String rewrittenQuery; + private Map sessionProperties = emptyMap(); + + public Builder setQueryId(String queryId) + { + this.queryId = requireNonNull(queryId, "queryId is null"); + return this; + } + + public Builder setOriginalQuery(String originalQuery) + { + this.originalQuery = requireNonNull(originalQuery, "originalQuery is null"); + return this; + } + + public Builder setRewrittenQuery(String rewrittenQuery) + { + this.rewrittenQuery = requireNonNull(rewrittenQuery, "rewrittenQuery is null"); + return this; + } + + public Builder setSessionProperties(Map sessionProperties) + { + this.sessionProperties = sessionProperties; + return this; + } + + public QueryRewriterOutput build() + { + requireNonNull(rewrittenQuery, "rewrittenQuery is null"); + requireNonNull(originalQuery, "originalQuery is null"); + requireNonNull(queryId, "queryId is null"); + checkArgument(!rewrittenQuery.isEmpty() && !originalQuery.isEmpty() && !queryId.isEmpty()); + return new QueryRewriterOutput(queryId, originalQuery, rewrittenQuery, sessionProperties); + } + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterProvider.java b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterProvider.java new file mode 100644 index 0000000000000..4be013f2c738a --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterProvider.java @@ -0,0 +1,19 @@ +/* + * 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.spi.rewriter; + +public interface QueryRewriterProvider +{ + QueryRewriter getQueryRewriter(); +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterProviderFactory.java b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterProviderFactory.java new file mode 100644 index 0000000000000..dcc5bd9bdbff9 --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/rewriter/QueryRewriterProviderFactory.java @@ -0,0 +1,25 @@ +/* + * 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.spi.rewriter; + +import com.facebook.presto.spi.eventlistener.EventListener; + +import java.util.List; + +public interface QueryRewriterProviderFactory +{ + String getName(); + + QueryRewriterProvider create(List eventListener); +} diff --git a/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/FailingQueryRewriterPlugin.java b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/FailingQueryRewriterPlugin.java new file mode 100644 index 0000000000000..c35f214df3788 --- /dev/null +++ b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/FailingQueryRewriterPlugin.java @@ -0,0 +1,64 @@ +/* + * 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.plugin.rewriter; + +import com.facebook.presto.spi.Plugin; +import com.facebook.presto.spi.eventlistener.EventListener; +import com.facebook.presto.spi.rewriter.QueryRewriter; +import com.facebook.presto.spi.rewriter.QueryRewriterInput; +import com.facebook.presto.spi.rewriter.QueryRewriterOutput; +import com.facebook.presto.spi.rewriter.QueryRewriterProvider; +import com.facebook.presto.spi.rewriter.QueryRewriterProviderFactory; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +public class FailingQueryRewriterPlugin + implements QueryRewriter, Plugin, QueryRewriterProvider, QueryRewriterProviderFactory +{ + public FailingQueryRewriterPlugin() + { + } + + @Override + public QueryRewriterOutput rewriteSQL(QueryRewriterInput queryRewriterInput) + { + throw new RuntimeException("This exception is thrown on purpose to verify plugin behaviour, during a failure."); + } + + @Override + public String getName() + { + return "FailingQueryRewriterPlugin"; + } + + @Override + public QueryRewriterProvider create(List eventListener) + { + return this; + } + + @Override + public QueryRewriter getQueryRewriter() + { + return this; + } + + @Override + public Iterable getQueryRewriterProviderFactories() + { + return ImmutableList.of(this); + } +} diff --git a/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/RewriteToFixedQueryRewriterPlugin.java b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/RewriteToFixedQueryRewriterPlugin.java new file mode 100644 index 0000000000000..6d95fd6785c2f --- /dev/null +++ b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/RewriteToFixedQueryRewriterPlugin.java @@ -0,0 +1,82 @@ +/* + * 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.plugin.rewriter; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.spi.Plugin; +import com.facebook.presto.spi.eventlistener.EventListener; +import com.facebook.presto.spi.rewriter.QueryRewriter; +import com.facebook.presto.spi.rewriter.QueryRewriterInput; +import com.facebook.presto.spi.rewriter.QueryRewriterOutput; +import com.facebook.presto.spi.rewriter.QueryRewriterProvider; +import com.facebook.presto.spi.rewriter.QueryRewriterProviderFactory; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; + +import java.util.List; + +import static com.facebook.presto.SystemSessionProperties.IS_QUERY_REWRITER_PLUGIN_SUCCEEDED; + +public class RewriteToFixedQueryRewriterPlugin + implements QueryRewriter, Plugin, QueryRewriterProvider, QueryRewriterProviderFactory +{ + private static final Logger log = Logger.get(RewriteToFixedQueryRewriterPlugin.class); + private List eventListener; + private String targetSql; + + public RewriteToFixedQueryRewriterPlugin(String targetSql) + { + this.targetSql = targetSql; + } + + @Override + public QueryRewriterOutput rewriteSQL(QueryRewriterInput queryRewriterInput) + { + log.debug("rewrite input sql: " + queryRewriterInput.getQuery() + " to " + targetSql); + QueryRewriterOutput.Builder outputBuilder = new QueryRewriterOutput.Builder(); + outputBuilder.setQueryId(queryRewriterInput.getQueryId()); + outputBuilder.setOriginalQuery(queryRewriterInput.getQuery()); + outputBuilder.setRewrittenQuery(targetSql); + if (!targetSql.equals("invalid")) { + outputBuilder.setSessionProperties(ImmutableMap.of(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, "true")); + } + return outputBuilder.build(); + } + + @Override + public String getName() + { + return "RewriteToFixedQueryRewriterPlugin: " + targetSql; + } + + @Override + public QueryRewriterProvider create(List eventListener) + { + this.eventListener = eventListener; + return this; + } + + @Override + public QueryRewriter getQueryRewriter() + { + return this; + } + + @Override + public Iterable getQueryRewriterProviderFactories() + { + return ImmutableList.of(this); + } +} diff --git a/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/TestQueryRewriterPlugin.java b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/TestQueryRewriterPlugin.java new file mode 100644 index 0000000000000..cdd5e67c84bcd --- /dev/null +++ b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/TestQueryRewriterPlugin.java @@ -0,0 +1,121 @@ +/* + * 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.plugin.rewriter; + +import com.facebook.presto.Session; +import com.facebook.presto.SystemSessionProperties; +import com.facebook.presto.execution.QueryInfo; +import com.facebook.presto.spi.QueryId; +import com.facebook.presto.tests.DistributedQueryRunner; +import com.google.common.collect.ImmutableSet; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static com.facebook.airlift.testing.Closeables.closeQuietly; +import static com.facebook.presto.SystemSessionProperties.IS_QUERY_REWRITER_PLUGIN_SUCCEEDED; +import static com.facebook.presto.execution.QueryState.FINISHED; +import static com.facebook.presto.execution.QueryState.RUNNING; +import static com.facebook.presto.execution.TestQueryRunnerUtil.createQuery; +import static com.facebook.presto.execution.TestQueryRunnerUtil.createQueryRunner; +import static com.facebook.presto.execution.TestQueryRunnerUtil.waitForQueryState; +import static com.facebook.presto.testing.TestingSession.testSessionBuilder; +import static com.facebook.presto.testing.assertions.Assert.assertEquals; + +@Test(singleThreaded = true) +public class TestQueryRewriterPlugin +{ + public static final String TARGET_SQL = "SELECT avg(l.extendedprice) AS avg_price, count(*) AS count_order FROM lineitem AS l"; + private DistributedQueryRunner queryRunner; + private Session session; + + @BeforeMethod + public void setup() + throws Exception + { + session = testSessionBuilder() + .setCatalog("tpch") + .setSchema("tiny") + .setSystemProperty(SystemSessionProperties.IS_QUERY_REWRITER_PLUGIN_ENABLED, "true") + .build(); + queryRunner = createQueryRunner(); + } + + @AfterMethod(alwaysRun = true) + public void tearDown() + { + closeQuietly(queryRunner); + queryRunner = null; + } + + @Test(timeOut = 240_000) + public void testRewriteToFixedQueryRewriter() + throws InterruptedException + { + queryRunner.installPlugin(new RewriteToFixedQueryRewriterPlugin(TARGET_SQL)); + queryRunner.getCoordinator().getQueryRewriterManager().loadQueryRewriterProvider(); + queryRunner.waitForClusterToGetReady(); + String sql = "select * from customer"; + QueryId queryId = createQuery(queryRunner, session, sql); + waitForQueryState(queryRunner, 0, queryId, ImmutableSet.of(RUNNING, FINISHED)); + QueryInfo queryInfo = queryRunner.getQueryInfo(queryId); + assertEquals(queryInfo.getQuery(), TARGET_SQL); + assertEquals(queryInfo.getSession().getSystemProperties().getOrDefault(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, "false"), "true"); + } + + @Test(timeOut = 240_000) + public void testUpperCasingQueryRewriter() + throws InterruptedException + { + queryRunner.installPlugin(new UpperCasingQueryRewriterPlugin()); + queryRunner.getCoordinator().getQueryRewriterManager().loadQueryRewriterProvider(); + queryRunner.waitForClusterToGetReady(); + QueryId queryId = createQuery(queryRunner, session, TARGET_SQL); + waitForQueryState(queryRunner, 0, queryId, ImmutableSet.of(RUNNING, FINISHED)); + QueryInfo queryInfo = queryRunner.getQueryInfo(queryId); + assertEquals(queryInfo.getQuery(), TARGET_SQL.toUpperCase()); + assertEquals(queryInfo.getSession().getSystemProperties().getOrDefault(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, "false"), "true"); + } + + @Test(timeOut = 240_000) + public void testRewriteToFixedInvalidQueryRewriter() + throws InterruptedException + { + String invalidQuery = "invalid"; + queryRunner.installPlugin(new RewriteToFixedQueryRewriterPlugin(invalidQuery)); + queryRunner.getCoordinator().getQueryRewriterManager().loadQueryRewriterProvider(); + queryRunner.waitForClusterToGetReady(); + QueryId queryId = createQuery(queryRunner, session, TARGET_SQL); + // Actual user given query is executed, an invalid query is generated by query rewriter plugin (and is rejected). + waitForQueryState(queryRunner, 0, queryId, ImmutableSet.of(RUNNING, FINISHED)); + QueryInfo queryInfo = queryRunner.getQueryInfo(queryId); + assertEquals(queryInfo.getQuery(), TARGET_SQL); + assertEquals(queryInfo.getSession().getSystemProperties().getOrDefault(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, "false"), "false"); + } + + @Test(timeOut = 240_000) + public void testExceptionThrowingQueryRewriterPlugin() + throws InterruptedException + { + queryRunner.installPlugin(new FailingQueryRewriterPlugin()); + queryRunner.getCoordinator().getQueryRewriterManager().loadQueryRewriterProvider(); + queryRunner.waitForClusterToGetReady(); + QueryId queryId = createQuery(queryRunner, session, TARGET_SQL); + waitForQueryState(queryRunner, 0, queryId, ImmutableSet.of(RUNNING, FINISHED)); + QueryInfo queryInfo = queryRunner.getQueryInfo(queryId); + assertEquals(queryInfo.getQuery(), TARGET_SQL); + assertEquals(queryInfo.getSession().getSystemProperties().getOrDefault(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, "false"), "false"); + } +} diff --git a/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/UpperCasingQueryRewriterPlugin.java b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/UpperCasingQueryRewriterPlugin.java new file mode 100644 index 0000000000000..32adbc16ac2b9 --- /dev/null +++ b/presto-tests/src/test/java/com/facebook/presto/plugin/rewriter/UpperCasingQueryRewriterPlugin.java @@ -0,0 +1,77 @@ +/* + * 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.plugin.rewriter; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.spi.Plugin; +import com.facebook.presto.spi.eventlistener.EventListener; +import com.facebook.presto.spi.rewriter.QueryRewriter; +import com.facebook.presto.spi.rewriter.QueryRewriterInput; +import com.facebook.presto.spi.rewriter.QueryRewriterOutput; +import com.facebook.presto.spi.rewriter.QueryRewriterProvider; +import com.facebook.presto.spi.rewriter.QueryRewriterProviderFactory; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; + +import java.util.List; + +import static com.facebook.presto.SystemSessionProperties.IS_QUERY_REWRITER_PLUGIN_SUCCEEDED; + +/** + * Rewrites each query to upper case + */ +public class UpperCasingQueryRewriterPlugin + implements QueryRewriter, Plugin, QueryRewriterProvider, QueryRewriterProviderFactory +{ + private static final Logger log = Logger.get(UpperCasingQueryRewriterPlugin.class); + private List eventListener; + + @Override + public QueryRewriterOutput rewriteSQL(QueryRewriterInput queryRewriterInput) + { + log.debug("Upper casing query writer"); + QueryRewriterOutput.Builder outputBuilder = new QueryRewriterOutput.Builder(); + outputBuilder.setQueryId(queryRewriterInput.getQueryId()); + outputBuilder.setOriginalQuery(queryRewriterInput.getQuery()); + outputBuilder.setRewrittenQuery(queryRewriterInput.getQuery().toUpperCase()); + outputBuilder.setSessionProperties(ImmutableMap.of(IS_QUERY_REWRITER_PLUGIN_SUCCEEDED, "true")); + return outputBuilder.build(); + } + + @Override + public QueryRewriter getQueryRewriter() + { + return this; + } + + @Override + public String getName() + { + return "UpperCasingQueryRewriterPlugin"; + } + + @Override + public QueryRewriterProvider create(List eventListener) + { + this.eventListener = eventListener; + return this; + } + + @Override + public Iterable getQueryRewriterProviderFactories() + { + return ImmutableList.of(this); + } +}