Add Cache-Control header to statement endpoints

xieandrew · xieandrew · commit d2eb154d5276 · 2025-07-10T10:36:57.000-07:00
diff --git a/presto-main-base/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java b/presto-main-base/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java
@@ -412,6 +412,11 @@ public Optional<DispatchInfo> getDispatchInfo(QueryId queryId)
                 });
     }
 
+    public long getDurationUntilExpirationInMillis(QueryId queryId)
+    {
+        return queryTracker.getQuery(queryId).getDurationUntilExpirationInMillis();
+    }
+
     /**
      * Check if a given queryId exists in query tracker
      *
diff --git a/presto-main-base/src/main/java/com/facebook/presto/dispatcher/NoOpQueryManager.java b/presto-main-base/src/main/java/com/facebook/presto/dispatcher/NoOpQueryManager.java
@@ -75,6 +75,12 @@ public QueryInfo getFullQueryInfo(QueryId queryId)
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public long getDurationUntilExpirationInMillis(QueryId queryId)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public Session getQuerySession(QueryId queryId)
     {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/QueryManager.java b/presto-main-base/src/main/java/com/facebook/presto/execution/QueryManager.java
@@ -65,6 +65,12 @@ BasicQueryInfo getQueryInfo(QueryId queryId)
     QueryInfo getFullQueryInfo(QueryId queryId)
             throws NoSuchElementException;
 
+    /**
+     * @throws NoSuchElementException if query does not exist
+     */
+    long getDurationUntilExpirationInMillis(QueryId queryId)
+            throws NoSuchElementException;
+
     /**
      * @throws NoSuchElementException if query does not exist
      */
diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/QueryTracker.java b/presto-main-base/src/main/java/com/facebook/presto/execution/QueryTracker.java
@@ -405,6 +405,13 @@ public interface TrackedQuery
 
         long getEndTimeInMillis();
 
+        default long getDurationUntilExpirationInMillis()
+        {
+            Duration queryClientTimeout = getQueryClientTimeout(getSession());
+            long expireTime = getLastHeartbeatInMillis() + queryClientTimeout.toMillis();
+            return Math.max(0, expireTime - currentTimeMillis());
+        }
+
         Optional<ResourceGroupQueryLimits> getResourceGroupQueryLimits();
 
         void fail(Throwable cause);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/protocol/ExecutingQueryResponseProvider.java b/presto-main-base/src/main/java/com/facebook/presto/server/protocol/ExecutingQueryResponseProvider.java
@@ -60,5 +60,6 @@ Optional<ListenableFuture<Response>> waitForExecutingResponse(
             DataSize targetResultSize,
             boolean compressionEnabled,
             boolean nestedDataSerializationEnabled,
-            boolean binaryResults);
+            boolean binaryResults,
+            long durationUntilExpirationMs);
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/protocol/QueryResourceUtil.java b/presto-main-base/src/main/java/com/facebook/presto/server/protocol/QueryResourceUtil.java
@@ -39,6 +39,7 @@
 import io.airlift.units.Duration;
 
 import javax.ws.rs.WebApplicationException;
+import javax.ws.rs.core.CacheControl;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.UriBuilder;
@@ -100,7 +101,7 @@ public final class QueryResourceUtil
 
     private QueryResourceUtil() {}
 
-    public static Response toResponse(Query query, QueryResults queryResults, boolean compressionEnabled)
+    public static Response toResponse(Query query, QueryResults queryResults, boolean compressionEnabled, long durationUntilExpirationMs)
     {
         Response.ResponseBuilder response = Response.ok(queryResults);
 
@@ -158,10 +159,18 @@ public static Response toResponse(Query query, QueryResults queryResults, boolea
             response.header(PRESTO_REMOVED_SESSION_FUNCTION, urlEncode(SQL_FUNCTION_ID_JSON_CODEC.toJson(signature)));
         }
 
+        response.cacheControl(getCacheControlMaxAge(durationUntilExpirationMs));
+
         return response.build();
     }
 
-    public static Response toResponse(Query query, QueryResults queryResults, String xPrestoPrefixUri, boolean compressionEnabled, boolean nestedDataSerializationEnabled)
+    public static Response toResponse(
+            Query query,
+            QueryResults queryResults,
+            String xPrestoPrefixUri,
+            boolean compressionEnabled,
+            boolean nestedDataSerializationEnabled,
+            long durationUntilExpirationMs)
     {
         Iterable<List<Object>> queryResultsData = queryResults.getData();
         if (nestedDataSerializationEnabled) {
@@ -181,7 +190,12 @@ public static Response toResponse(Query query, QueryResults queryResults, String
                 queryResults.getUpdateType(),
                 queryResults.getUpdateCount());
 
-        return toResponse(query, resultsClone, compressionEnabled);
+        return toResponse(query, resultsClone, compressionEnabled, durationUntilExpirationMs);
+    }
+
+    public static CacheControl getCacheControlMaxAge(long durationUntilExpirationMs)
+    {
+        return CacheControl.valueOf("max-age=" + MILLISECONDS.toSeconds(durationUntilExpirationMs));
     }
 
     public static void abortIfPrefixUrlInvalid(String xPrestoPrefixUrl)
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/TestQueryManager.java b/presto-main-base/src/test/java/com/facebook/presto/execution/TestQueryManager.java
@@ -53,6 +53,12 @@ public QueryInfo getFullQueryInfo(QueryId queryId)
         return null;
     }
 
+    @Override
+    public long getDurationUntilExpirationInMillis(QueryId queryId)
+    {
+        return 0;
+    }
+
     public Session getQuerySession(QueryId queryId)
     {
         return null;
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java
@@ -252,6 +252,13 @@ public QueryInfo getFullQueryInfo(QueryId queryId)
         return queryTracker.getQuery(queryId).getQueryInfo();
     }
 
+    @Override
+    public long getDurationUntilExpirationInMillis(QueryId queryId)
+            throws NoSuchElementException
+    {
+        return queryTracker.getQuery(queryId).getDurationUntilExpirationInMillis();
+    }
+
     @Override
     public Session getQuerySession(QueryId queryId)
             throws NoSuchElementException
diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/ExecutingStatementResource.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/ExecutingStatementResource.java
@@ -16,6 +16,7 @@
 import com.facebook.airlift.concurrent.BoundedExecutor;
 import com.facebook.airlift.stats.TimeStat;
 import com.facebook.presto.client.QueryResults;
+import com.facebook.presto.execution.QueryManager;
 import com.facebook.presto.server.ForStatementResource;
 import com.facebook.presto.server.ServerConfig;
 import com.facebook.presto.spi.QueryId;
@@ -68,6 +69,7 @@ public class ExecutingStatementResource
 
     private final BoundedExecutor responseExecutor;
     private final LocalQueryProvider queryProvider;
+    private final QueryManager queryManager;
     private final boolean compressionEnabled;
     private final boolean nestedDataSerializationEnabled;
     private final QueryBlockingRateLimiter queryRateLimiter;
@@ -76,11 +78,13 @@ public class ExecutingStatementResource
     public ExecutingStatementResource(
             @ForStatementResource BoundedExecutor responseExecutor,
             LocalQueryProvider queryProvider,
+            QueryManager queryManager,
             ServerConfig serverConfig,
             QueryBlockingRateLimiter queryRateLimiter)
     {
         this.responseExecutor = requireNonNull(responseExecutor, "responseExecutor is null");
         this.queryProvider = requireNonNull(queryProvider, "queryProvider is null");
+        this.queryManager = requireNonNull(queryManager, "queryManager is null");
         this.compressionEnabled = requireNonNull(serverConfig, "serverConfig is null").isQueryResultsCompressionEnabled();
         this.nestedDataSerializationEnabled = requireNonNull(serverConfig, "serverConfig is null").isNestedDataSerializationEnabled();
         this.queryRateLimiter = requireNonNull(queryRateLimiter, "queryRateLimiter is null");
@@ -132,9 +136,10 @@ public void getQueryResults(
                     return query.waitForResults(token, uriInfo, effectiveFinalProto, wait, effectiveFinalTargetResultSize, binaryResults);
                 },
                 responseExecutor);
+        long durationUntilExpirationMs = queryManager.getDurationUntilExpirationInMillis(queryId);
         ListenableFuture<Response> queryResultsFuture = transform(
                 waitForResultsAsync,
-                results -> toResponse(query, results, xPrestoPrefixUrl, compressionEnabled, nestedDataSerializationEnabled),
+                results -> toResponse(query, results, xPrestoPrefixUrl, compressionEnabled, nestedDataSerializationEnabled, durationUntilExpirationMs),
                 directExecutor());
         bindAsyncResponse(asyncResponse, queryResultsFuture, responseExecutor);
     }
diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/LocalExecutingQueryResponseProvider.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/LocalExecutingQueryResponseProvider.java
@@ -26,6 +26,7 @@
 
 import java.util.Optional;
 
+import static com.facebook.presto.server.protocol.QueryResourceUtil.toResponse;
 import static com.google.common.util.concurrent.Futures.transform;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static java.util.Objects.requireNonNull;
@@ -53,7 +54,8 @@ public Optional<ListenableFuture<Response>> waitForExecutingResponse(
             DataSize targetResultSize,
             boolean compressionEnabled,
             boolean nestedDataSerializationEnabled,
-            boolean binaryResults)
+            boolean binaryResults,
+            long durationUntilExpirationMs)
     {
         Query query;
         try {
@@ -64,7 +66,7 @@ public Optional<ListenableFuture<Response>> waitForExecutingResponse(
         }
         return Optional.of(transform(
                 query.waitForResults(0, uriInfo, scheme, maxWait, targetResultSize, binaryResults),
-                results -> QueryResourceUtil.toResponse(query, results, xPrestoPrefixUrl, compressionEnabled, nestedDataSerializationEnabled),
+                results -> toResponse(query, results, xPrestoPrefixUrl, compressionEnabled, nestedDataSerializationEnabled, durationUntilExpirationMs),
                 directExecutor()));
     }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java
@@ -56,6 +56,7 @@
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.container.AsyncResponse;
 import javax.ws.rs.container.Suspended;
+import javax.ws.rs.core.CacheControl;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.Response.Status;
@@ -77,6 +78,7 @@
 import static com.facebook.presto.server.protocol.QueryResourceUtil.NO_DURATION;
 import static com.facebook.presto.server.protocol.QueryResourceUtil.abortIfPrefixUrlInvalid;
 import static com.facebook.presto.server.protocol.QueryResourceUtil.createQueuedQueryResults;
+import static com.facebook.presto.server.protocol.QueryResourceUtil.getCacheControlMaxAge;
 import static com.facebook.presto.server.protocol.QueryResourceUtil.getQueuedUri;
 import static com.facebook.presto.server.protocol.QueryResourceUtil.getScheme;
 import static com.facebook.presto.server.security.RoleType.USER;
@@ -91,6 +93,7 @@
 import static com.google.common.util.concurrent.Futures.transformAsync;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
+import static java.lang.System.currentTimeMillis;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
 import static java.util.UUID.randomUUID;
@@ -221,7 +224,9 @@ public Response postStatement(
         Query query = new Query(statement, sessionContext, dispatchManager, executingQueryResponseProvider, 0);
         queries.put(query.getQueryId(), query);
 
-        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled).build();
+        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled)
+                .cacheControl(query.getDefaultCacheControl())
+                .build();
     }
 
     /**
@@ -273,7 +278,9 @@ public Response putStatement(
             throw badRequest(CONFLICT, "Query already exists");
         }
 
-        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled).build();
+        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled)
+                .cacheControl(query.getDefaultCacheControl())
+                .build();
     }
 
     /**
@@ -322,7 +329,9 @@ public Response retryFailedQuery(
             }
         }
 
-        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled).build();
+        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled)
+                .cacheControl(query.getDefaultCacheControl())
+                .build();
     }
 
     /**
@@ -435,6 +444,7 @@ private static Response.ResponseBuilder withCompressionConfiguration(Response.Re
 
     private static final class Query
     {
+        private static final int CACHE_CONTROL_MAX_AGE_SEC = 60;
         private final String query;
         private final SessionContext sessionContext;
         private final DispatchManager dispatchManager;
@@ -443,6 +453,7 @@ private static final class Query
         private final String slug;
         private final AtomicLong lastToken = new AtomicLong();
         private final int retryCount;
+        private final long expirationTime;
 
         @GuardedBy("this")
         private ListenableFuture<?> querySubmissionFuture;
@@ -468,6 +479,7 @@ public Query(
             this.retryCount = retryCount;
             this.queryId = requireNonNull(queryId, "queryId is null");
             this.slug = requireNonNull(slug, "slug is null");
+            this.expirationTime = currentTimeMillis() + SECONDS.toMillis(CACHE_CONTROL_MAX_AGE_SEC);
         }
 
         /**
@@ -518,6 +530,15 @@ public int getRetryCount()
             return retryCount;
         }
 
+        /**
+         * Returns a cache control with the default max age value
+         */
+        public CacheControl getDefaultCacheControl()
+        {
+            long maxAgeMillis = Math.max(0, expirationTime - currentTimeMillis());
+            return CacheControl.valueOf("max-age=" + MILLISECONDS.toSeconds(maxAgeMillis));
+        }
+
         /**
          * Checks whether the query has been processed by the dispatchManager
          */
@@ -591,7 +612,10 @@ public ListenableFuture<Response> toResponse(
                             xPrestoPrefixUrl,
                             DispatchInfo.waitingForPrerequisites(NO_DURATION, NO_DURATION),
                             binaryResults);
-                    return immediateFuture(withCompressionConfiguration(Response.ok(queryResults), compressionEnabled).build());
+
+                    return immediateFuture(withCompressionConfiguration(Response.ok(queryResults), compressionEnabled)
+                            .cacheControl(getDefaultCacheControl())
+                            .build());
                 }
             }
 
@@ -602,6 +626,7 @@ public ListenableFuture<Response> toResponse(
                         .status(NOT_FOUND)
                         .build()));
             }
+            long durationUntilExpirationMs = dispatchManager.getDurationUntilExpirationInMillis(queryId);
 
             if (waitForDispatched().isDone()) {
                 Optional<ListenableFuture<Response>> executingQueryResponse = executingQueryResponseProvider.waitForExecutingResponse(
@@ -615,7 +640,8 @@ public ListenableFuture<Response> toResponse(
                         TARGET_RESULT_SIZE,
                         compressionEnabled,
                         nestedDataSerializationEnabled,
-                        binaryResults);
+                        binaryResults,
+                        durationUntilExpirationMs);
 
                 if (executingQueryResponse.isPresent()) {
                     return executingQueryResponse.get();
@@ -624,6 +650,7 @@ public ListenableFuture<Response> toResponse(
 
             return immediateFuture(withCompressionConfiguration(Response.ok(
                     createQueryResults(token + 1, uriInfo, xForwardedProto, xPrestoPrefixUrl, dispatchInfo.get(), binaryResults)), compressionEnabled)
+                    .cacheControl(getCacheControlMaxAge(durationUntilExpirationMs))
                     .build());
         }
 
diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestServer.java b/presto-main/src/test/java/com/facebook/presto/server/TestServer.java
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkQueryManager.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkQueryManager.java

Original file line number	Diff line number	Diff line change
`@@ -412,6 +412,11 @@ public Optional<DispatchInfo> getDispatchInfo(QueryId queryId)`
`412`	`412`	`});`
`413`	`413`	`}`
`414`	`414`
	`415`	`+ public long getDurationUntilExpirationInMillis(QueryId queryId)`
	`416`	`+ {`
	`417`	`+ return queryTracker.getQuery(queryId).getDurationUntilExpirationInMillis();`
	`418`	`+ }`
	`419`	`+`
`415`	`420`	`/**`
`416`	`421`	`* Check if a given queryId exists in query tracker`
`417`	`422`	`*`
Original file line number	Diff line number	Diff line change
`@@ -75,6 +75,12 @@ public QueryInfo getFullQueryInfo(QueryId queryId)`
`75`	`75`	`throw new UnsupportedOperationException();`
`76`	`76`	`}`
`77`	`77`
	`78`	`+ @Override`
	`79`	`+ public long getDurationUntilExpirationInMillis(QueryId queryId)`
	`80`	`+ {`
	`81`	`+ throw new UnsupportedOperationException();`
	`82`	`+ }`
	`83`	`+`
`78`	`84`	`@Override`
`79`	`85`	`public Session getQuerySession(QueryId queryId)`
`80`	`86`	`{`
Original file line number	Diff line number	Diff line change
`@@ -60,5 +60,6 @@ Optional<ListenableFuture<Response>> waitForExecutingResponse(`
`60`	`60`	`DataSize targetResultSize,`
`61`	`61`	`boolean compressionEnabled,`
`62`	`62`	`boolean nestedDataSerializationEnabled,`
`63`		`- boolean binaryResults);`
	`63`	`+ boolean binaryResults,`
	`64`	`+ long durationUntilExpirationMs);`
`64`	`65`	`}`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,12 @@ public QueryInfo getFullQueryInfo(QueryId queryId)`
`53`	`53`	`return null;`
`54`	`54`	`}`
`55`	`55`
	`56`	`+ @Override`
	`57`	`+ public long getDurationUntilExpirationInMillis(QueryId queryId)`
	`58`	`+ {`
	`59`	`+ return 0;`
	`60`	`+ }`
	`61`	`+`
`56`	`62`	`public Session getQuerySession(QueryId queryId)`
`57`	`63`	`{`
`58`	`64`	`return null;`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@`
`26`	`26`
`27`	`27`	`import java.util.Optional;`
`28`	`28`
	`29`	`+import static com.facebook.presto.server.protocol.QueryResourceUtil.toResponse;`
`29`	`30`	`import static com.google.common.util.concurrent.Futures.transform;`
`30`	`31`	`import static com.google.common.util.concurrent.MoreExecutors.directExecutor;`
`31`	`32`	`import static java.util.Objects.requireNonNull;`
`@@ -53,7 +54,8 @@ public Optional<ListenableFuture<Response>> waitForExecutingResponse(`
`53`	`54`	`DataSize targetResultSize,`
`54`	`55`	`boolean compressionEnabled,`
`55`	`56`	`boolean nestedDataSerializationEnabled,`
`56`		`- boolean binaryResults)`
	`57`	`+ boolean binaryResults,`
	`58`	`+ long durationUntilExpirationMs)`
`57`	`59`	`{`
`58`	`60`	`Query query;`
`59`	`61`	`try {`
`@@ -64,7 +66,7 @@ public Optional<ListenableFuture<Response>> waitForExecutingResponse(`
`64`	`66`	`}`
`65`	`67`	`return Optional.of(transform(`
`66`	`68`	`query.waitForResults(0, uriInfo, scheme, maxWait, targetResultSize, binaryResults),`
`67`		`- results -> QueryResourceUtil.toResponse(query, results, xPrestoPrefixUrl, compressionEnabled, nestedDataSerializationEnabled),`
	`69`	`+ results -> toResponse(query, results, xPrestoPrefixUrl, compressionEnabled, nestedDataSerializationEnabled, durationUntilExpirationMs),`
`68`	`70`	`directExecutor()));`
`69`	`71`	`}`
`70`	`72`	`}`