Support logging formatted prepared query

Logging prepared query is useful especially for EXECUTE queries as it could make debugging easier.
prestodb · Apr 22, 2021 · 24243ad · 24243ad
1 parent 60edac5
commit 24243ad
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 1 deletion.
diff --git a/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java b/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java
@@ -179,6 +179,7 @@ public final class SystemSessionProperties
     public static final String OPTIMIZE_JOINS_WITH_EMPTY_SOURCES = "optimize_joins_with_empty_sources";
     public static final String SPOOLING_OUTPUT_BUFFER_ENABLED = "spooling_output_buffer_enabled";
     public static final String SPARK_ASSIGN_BUCKET_TO_PARTITION_FOR_PARTITIONED_TABLE_WRITE_ENABLED = "spark_assign_bucket_to_partition_for_partitioned_table_write_enabled";
+    public static final String LOG_FORMATTED_QUERY_ENABLED = "log_formatted_query_enabled";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -938,7 +939,12 @@ public SystemSessionProperties(
                         SPARK_ASSIGN_BUCKET_TO_PARTITION_FOR_PARTITIONED_TABLE_WRITE_ENABLED,
                         "Assign bucket to partition map for partitioned table write when adding an exchange",
                         featuresConfig.isPrestoSparkAssignBucketToPartitionForPartitionedTableWriteEnabled(),
-                        true));
+                        true),
+                booleanProperty(
+                        LOG_FORMATTED_QUERY_ENABLED,
+                        "Log formatted prepared query instead of raw query when enabled",
+                        featuresConfig.isLogFormattedQueryEnabled(),
+                        false));
     }
 
     public static boolean isEmptyJoinOptimization(Session session)
@@ -1589,4 +1595,9 @@ public static boolean isPrestoSparkAssignBucketToPartitionForPartitionedTableWri
     {
         return session.getSystemProperty(SPARK_ASSIGN_BUCKET_TO_PARTITION_FOR_PARTITIONED_TABLE_WRITE_ENABLED, Boolean.class);
     }
+
+    public static boolean isLogFormattedQueryEnabled(Session session)
+    {
+        return session.getSystemProperty(LOG_FORMATTED_QUERY_ENABLED, Boolean.class);
+    }
 }
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
@@ -51,7 +51,9 @@
 import java.util.concurrent.Executor;
 
 import static com.facebook.presto.SystemSessionProperties.getWarningHandlingLevel;
+import static com.facebook.presto.SystemSessionProperties.isLogFormattedQueryEnabled;
 import static com.facebook.presto.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE;
+import static com.facebook.presto.sql.SqlFormatter.formatSql;
 import static com.facebook.presto.util.StatementUtils.getQueryType;
 import static com.facebook.presto.util.StatementUtils.isTransactionControlStatement;
 import static com.google.common.base.Preconditions.checkArgument;
@@ -181,6 +183,9 @@ private <C> void createQueryInternal(QueryId queryId, String slug, SessionContex
             // prepare query
             WarningCollector warningCollector = warningCollectorFactory.create(getWarningHandlingLevel(session));
             preparedQuery = queryPreparer.prepareQuery(session, query, warningCollector);
+            if (isLogFormattedQueryEnabled(session)) {
+                query = formatSql(preparedQuery.getStatement(), Optional.of(preparedQuery.getParameters()));
+            }
 
             // select resource group
             Optional<QueryType> queryType = getQueryType(preparedQuery.getStatement().getClass());

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java b/presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java
@@ -111,6 +111,7 @@ public class FeaturesConfig
     private boolean pagesIndexEagerCompactionEnabled;
     private boolean distributedSort = true;
     private boolean optimizeJoinsWithEmptySources;
+    private boolean logFormattedQueryEnabled;
 
     private boolean dictionaryAggregation;
 
@@ -1549,6 +1550,19 @@ public FeaturesConfig setEmptyJoinOptimization(boolean value)
         return this;
     }
 
+    public boolean isLogFormattedQueryEnabled()
+    {
+        return logFormattedQueryEnabled;
+    }
+
+    @Config("log-formatted-query-enabled")
+    @ConfigDescription("Log formatted prepared query instead of raw query when enabled")
+    public FeaturesConfig setLogFormattedQueryEnabled(boolean logFormattedQueryEnabled)
+    {
+        this.logFormattedQueryEnabled = logFormattedQueryEnabled;
+        return this;
+    }
+
     public boolean isSpoolingOutputBufferEnabled()
     {
         return spoolingOutputBufferEnabled;

diff --git a/presto-main/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java b/presto-main/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
@@ -157,6 +157,7 @@ public void testDefaults()
                 .setAllowWindowOrderByLiterals(true)
                 .setEnforceFixedDistributionForOutputOperator(false)
                 .setEmptyJoinOptimization(false)
+                .setLogFormattedQueryEnabled(false)
                 .setSpoolingOutputBufferEnabled(false)
                 .setSpoolingOutputBufferThreshold(new DataSize(8, MEGABYTE))
                 .setSpoolingOutputBufferTempStorage("local")
@@ -270,6 +271,7 @@ public void testExplicitPropertyMappings()
                 .put("is-allow-window-order-by-literals", "false")
                 .put("enforce-fixed-distribution-for-output-operator", "true")
                 .put("optimizer.optimize-joins-with-empty-sources", "true")
+                .put("log-formatted-query-enabled", "true")
                 .put("spooling-output-buffer-enabled", "true")
                 .put("spooling-output-buffer-threshold", "16MB")
                 .put("spooling-output-buffer-temp-storage", "tempfs")
@@ -381,6 +383,7 @@ public void testExplicitPropertyMappings()
                 .setAllowWindowOrderByLiterals(false)
                 .setEnforceFixedDistributionForOutputOperator(true)
                 .setEmptyJoinOptimization(true)
+                .setLogFormattedQueryEnabled(true)
                 .setSpoolingOutputBufferEnabled(true)
                 .setSpoolingOutputBufferThreshold(new DataSize(16, MEGABYTE))
                 .setSpoolingOutputBufferTempStorage("tempfs")