feat(develop): add query on demand behind feature flag

packages/gatsby/cache-dir/dev-loader.js

@@ -36,6 +36,8 @@ class DevLoader extends BaseLoader {
           this.handleStaticQueryResultHotUpdate(msg)
         } else if (msg.type === `pageQueryResult`) {
           this.handlePageQueryResultHotUpdate(msg)
+        } else if (msg.type === `dirtyQueries`) {
+          this.handleDirtyPageQueryMessage(msg)
         }
       })
     } else {
@@ -75,6 +77,9 @@ class DevLoader extends BaseLoader {
   }
 
   doPrefetch(pagePath) {
+    if (process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND) {
+      return Promise.resolve()
+    }
     return super.doPrefetch(pagePath).then(result => result.payload)
   }
 
@@ -138,6 +143,41 @@ class DevLoader extends BaseLoader {
       ___emitter.emit(`pageQueryResult`, newPageData)
     }
   }
+
+  handleDirtyPageQueryMessage(msg) {
+    msg.payload.dirtyQueries.forEach(dirtyQueryId => {
+      if (dirtyQueryId === `/dev-404-page/` || dirtyQueryId === `/404.html`) {
+        // those pages are not on demand so skipping
+        return
+      }
+
+      const normalizedId = normalizePagePath(dirtyQueryId)
+
+      // We can't just delete items in caches, because then
+      // using history.back() would show dev-404 page
+      // due to our special handling of it in root.js (loader.isPageNotFound check)
+      // so instead we mark it as stale and instruct loader's async methods
+      // to refetch resources if they are marked as stale
+
+      const cachedPageData = this.pageDataDb.get(normalizedId)
+      if (cachedPageData) {
+        // if we have page data in cache, mark it as stale
+        this.pageDataDb.set(normalizedId, {
+          ...cachedPageData,
+          stale: true,
+        })
+      }
+
+      const cachedPage = this.pageDb.get(normalizedId)
+      if (cachedPage) {
+        // if we have page data in cache, mark it as stale
+        this.pageDb.set(normalizedId, {
+          ...cachedPage,
+          payload: { ...cachedPage.payload, stale: true },
+        })
+      }
+    })
+  }
 }
 
 export default DevLoader

packages/gatsby/cache-dir/ensure-resources.js

@@ -47,6 +47,14 @@ class EnsureResources extends React.Component {
       return false
     }
 
+    if (
+      process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND &&
+      nextState.pageResources.stale
+    ) {
+      this.loadResources(nextProps.location.pathname)
+      return false
+    }
+
     // Check if the component or json have changed.
     if (this.state.pageResources !== nextState.pageResources) {
       return true

packages/gatsby/cache-dir/loader.js

@@ -190,7 +190,10 @@ export class BaseLoader {
   loadPageDataJson(rawPath) {
     const pagePath = findPath(rawPath)
     if (this.pageDataDb.has(pagePath)) {
-      return Promise.resolve(this.pageDataDb.get(pagePath))
+      const pageData = this.pageDataDb.get(pagePath)
+      if (!process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND || !pageData.stale) {
+        return Promise.resolve(pageData)
+      }
     }
 
     return this.fetchPageDataJson({ pagePath }).then(pageData => {
@@ -209,7 +212,12 @@ export class BaseLoader {
     const pagePath = findPath(rawPath)
     if (this.pageDb.has(pagePath)) {
       const page = this.pageDb.get(pagePath)
-      return Promise.resolve(page.payload)
+      if (
+        !process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND ||
+        !page.payload.stale
+      ) {
+        return Promise.resolve(page.payload)
+      }
     }
 
     if (this.inFlightDb.has(pagePath)) {

packages/gatsby/src/redux/__tests__/__snapshots__/index.js.snap

@@ -57,6 +57,7 @@ Object {
     "byConnection": Map {},
     "byNode": Map {},
     "deletedQueries": Set {},
+    "dirtyQueriesListToEmitViaWebsocket": Array [],
     "queryNodes": Map {},
     "trackedComponents": Map {
       "/Users/username/dev/site/src/templates/my-sweet-new-page.js" => Object {

packages/gatsby/src/redux/actions/internal.ts

@@ -20,6 +20,7 @@ import {
   ISetGraphQLDefinitionsAction,
   IQueryStartAction,
   IApiFinishedAction,
+  IQueryClearDirtyQueriesListToEmitViaWebsocket,
 } from "../types"
 
 import { gatsbyConfigSchema } from "../../joi-schemas/joi"
@@ -249,6 +250,12 @@ export const queryStart = (
   }
 }
 
+export const clearDirtyQueriesListToEmitViaWebsocket = (): IQueryClearDirtyQueriesListToEmitViaWebsocket => {
+  return {
+    type: `QUERY_CLEAR_DIRTY_QUERIES_LIST_TO_EMIT_VIA_WEBSOCKET`,
+  }
+}
+
 /**
  * Remove jobs which are marked as stale (inputPath doesn't exists)
  * @private

packages/gatsby/src/redux/reducers/__tests__/__snapshots__/queries.ts.snap

@@ -13,6 +13,7 @@ Object {
     },
   },
   "deletedQueries": Set {},
+  "dirtyQueriesListToEmitViaWebsocket": Array [],
   "queryNodes": Map {
     "/hi/" => Set {
       "SuperCoolNode",

packages/gatsby/src/redux/reducers/__tests__/queries.ts

@@ -88,6 +88,7 @@ it(`has expected initial state`, () => {
       "byConnection": Map {},
       "byNode": Map {},
       "deletedQueries": Set {},
+      "dirtyQueriesListToEmitViaWebsocket": Array [],
       "queryNodes": Map {},
       "trackedComponents": Map {},
       "trackedQueries": Map {},

packages/gatsby/src/redux/reducers/queries.ts

@@ -26,6 +26,7 @@ const initialState = (): IGatsbyState["queries"] => {
     trackedQueries: new Map<QueryId, IQueryState>(),
     trackedComponents: new Map<ComponentPath, IComponentState>(),
     deletedQueries: new Set<QueryId>(),
+    dirtyQueriesListToEmitViaWebsocket: [],
   }
 }
 
@@ -69,6 +70,7 @@ export function queriesReducer(
       if (!query || action.contextModified) {
         query = registerQuery(state, path)
         query.dirty = setFlag(query.dirty, FLAG_DIRTY_PAGE)
+        state = trackDirtyQuery(state, path)
       }
       registerComponent(state, componentPath).pages.add(path)
       state.deletedQueries.delete(path)
@@ -114,6 +116,7 @@ export function queriesReducer(
           const query = state.trackedQueries.get(queryId)
           if (query) {
             query.dirty = setFlag(query.dirty, FLAG_DIRTY_TEXT)
+            state = trackDirtyQuery(state, queryId)
           }
         }
         component.query = query
@@ -134,6 +137,9 @@ export function queriesReducer(
       // TODO: unify the behavior?
       const query = registerQuery(state, action.payload.id)
       query.dirty = setFlag(query.dirty, FLAG_DIRTY_TEXT)
+      // static queries are not on demand, so skipping tracking which
+      // queries were marked dirty recently
+      // state = trackDirtyQuery(state, action.payload.id)
       state.deletedQueries.delete(action.payload.id)
       return state
     }
@@ -176,12 +182,14 @@ export function queriesReducer(
         const query = state.trackedQueries.get(queryId)
         if (query) {
           query.dirty = setFlag(query.dirty, FLAG_DIRTY_DATA)
+          state = trackDirtyQuery(state, queryId)
         }
       }
       for (const queryId of queriesByConnection) {
         const query = state.trackedQueries.get(queryId)
         if (query) {
           query.dirty = setFlag(query.dirty, FLAG_DIRTY_DATA)
+          state = trackDirtyQuery(state, queryId)
         }
       }
       return state
@@ -199,9 +207,15 @@ export function queriesReducer(
         for (const [, query] of state.trackedQueries) {
           query.running = 0
         }
+        // Reset list of dirty queries (this is used only to notify runtime and it could've been persisted)
+        state.dirtyQueriesListToEmitViaWebsocket = []
       }
       return state
     }
+    case `QUERY_CLEAR_DIRTY_QUERIES_LIST_TO_EMIT_VIA_WEBSOCKET`: {
+      state.dirtyQueriesListToEmitViaWebsocket = []
+      return state
+    }
     default:
       return state
   }
@@ -305,3 +319,14 @@ function registerComponent(
   }
   return component
 }
+
+function trackDirtyQuery(
+  state: IGatsbyState["queries"],
+  queryId: QueryId
+): IGatsbyState["queries"] {
+  if (process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND) {
+    state.dirtyQueriesListToEmitViaWebsocket.push(queryId)
+  }
+
+  return state
+}

packages/gatsby/src/redux/types.ts

@@ -215,6 +215,7 @@ export interface IGatsbyState {
     trackedQueries: Map<Identifier, IQueryState>
     trackedComponents: Map<string, IComponentState>
     deletedQueries: Set<Identifier>
+    dirtyQueriesListToEmitViaWebsocket: Array<string>
   }
   components: Map<
     SystemPath,
@@ -307,6 +308,7 @@ export type ActionsUnion =
   | IDeletePageAction
   | IPageQueryRunAction
   | IPrintTypeDefinitions
+  | IQueryClearDirtyQueriesListToEmitViaWebsocket
   | IQueryExtractedAction
   | IQueryExtractedBabelSuccessAction
   | IQueryExtractionBabelErrorAction
@@ -473,6 +475,10 @@ export interface IReplaceStaticQueryAction {
   }
 }
 
+export interface IQueryClearDirtyQueriesListToEmitViaWebsocket {
+  type: `QUERY_CLEAR_DIRTY_QUERIES_LIST_TO_EMIT_VIA_WEBSOCKET`
+}
+
 export interface IQueryExtractedAction {
   type: `QUERY_EXTRACTED`
   plugin: IGatsbyPlugin

packages/gatsby/src/services/calculate-dirty-queries.ts

@@ -5,11 +5,48 @@ import { assertStore } from "../utils/assert-store"
 
 export async function calculateDirtyQueries({
   store,
+  websocketManager,
+  currentlyHandledPendingQueryRuns,
 }: Partial<IQueryRunningContext>): Promise<{
   queryIds: IGroupedQueryIds
 }> {
   assertStore(store)
   const state = store.getState()
   const queryIds = calcDirtyQueryIds(state)
-  return { queryIds: groupQueryIds(queryIds) }
+
+  let queriesToRun: Array<string> = queryIds
+
+  if (
+    process.env.gatsby_executing_command === `develop` &&
+    process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND
+  ) {
+    // 404 are special cases in our runtime that ideally use
+    // generic things to work, but for now they have special handling
+    const pagePathFilter = new Set([`/404.html`, `/dev-404-page/`])
+
+    // we want to make sure we run queries for pages that user currently
+    // view in the browser
+    if (websocketManager?.activePaths) {
+      for (const activePath of websocketManager.activePaths) {
+        pagePathFilter.add(activePath)
+      }
+    }
+
+    // we also want to make sure we include pages that were requested from
+    // via `page-data` fetches or websocket requests
+    if (currentlyHandledPendingQueryRuns) {
+      for (const pendingQuery of currentlyHandledPendingQueryRuns) {
+        pagePathFilter.add(pendingQuery)
+      }
+    }
+
+    // static queries are also not on demand
+    queriesToRun = queryIds.filter(
+      queryId => queryId.startsWith(`sq--`) || pagePathFilter.has(queryId)
+    )
+  }
+
+  return {
+    queryIds: groupQueryIds(queriesToRun),
+  }
 }

packages/gatsby/src/services/initialize.ts

@@ -1,5 +1,5 @@
 import _ from "lodash"
-import { slash } from "gatsby-core-utils"
+import { slash, isCI } from "gatsby-core-utils"
 import fs from "fs-extra"
 import md5File from "md5-file"
 import crypto from "crypto"
@@ -162,6 +162,22 @@ export async function initialize({
 
   activity.end()
 
+  if (process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND) {
+    if (process.env.gatsby_executing_command !== `develop`) {
+      // we don't want to ever have this flag enabled for anything than develop
+      // in case someone have this env var globally set
+      delete process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND
+    } else if (isCI()) {
+      delete process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND
+      reporter.warn(
+        `Experimental Query on Demand feature is not available in CI environment. Continuing with regular mode.`
+      )
+    } else {
+      reporter.info(`Using experimental Query on Demand feature`)
+      telemetry.trackFeatureIsUsed(`QueryOnDemand`)
+    }
+  }
+
   // run stale jobs
   store.dispatch(removeStaleJobs(store.getState()))
 

packages/gatsby/src/utils/page-data.ts

@@ -124,6 +124,7 @@ export async function flush(): Promise<void> {
     pages,
     program,
     staticQueriesByTemplate,
+    queries,
   } = store.getState()
 
   const { pagePaths } = pendingPageDataWrites
@@ -140,6 +141,28 @@ export async function flush(): Promise<void> {
     // them, a page might not exist anymore щ（ﾟДﾟщ）
     // This is why we need this check
     if (page) {
+      if (
+        program?._?.[0] === `develop` &&
+        process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND
+      ) {
+        // check if already did run query for this page
+        // with query-on-demand we might have pending page-data write due to
+        // changes in static queries assigned to page template, but we might not
+        // have query result for it
+        const query = queries.trackedQueries.get(page.path)
+        if (!query) {
+          // this should not happen ever
+          throw new Error(
+            `We have a page, but we don't have registered query for it (???)`
+          )
+        }
+
+        if (query.dirty !== 0) {
+          // query results are not up to date, it's not safe to write page-data and emit results
+          continue
+        }
+      }
+
       const staticQueryHashes =
         staticQueriesByTemplate.get(page.componentPath) || []
 

packages/gatsby/src/utils/webpack.config.js

@@ -85,6 +85,9 @@ module.exports = async (
     envObject.PUBLIC_DIR = JSON.stringify(`${process.cwd()}/public`)
     envObject.BUILD_STAGE = JSON.stringify(stage)
     envObject.CYPRESS_SUPPORT = JSON.stringify(process.env.CYPRESS_SUPPORT)
+    envObject.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND = JSON.stringify(
+      !!process.env.GATSBY_EXPERIMENTAL_QUERY_ON_DEMAND
+    )
 
     if (stage === `develop`) {
       envObject.GATSBY_SOCKET_IO_DEFAULT_TRANSPORT = JSON.stringify(