From a4b02bd4e865ead7e7cc35507fb0ef7a8bd0cf42 Mon Sep 17 00:00:00 2001
From: handtrix <handtrix@gmail.com>
Date: Thu, 19 Sep 2019 16:22:14 +0200
Subject: [PATCH] Make withApollo work with _app.js components

---
 examples/with-apollo/lib/apollo.js | 32 +++++++++++++++++++-----------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/examples/with-apollo/lib/apollo.js b/examples/with-apollo/lib/apollo.js
index 303826b44c09be5..6f61f5e07c5f1e0 100644
--- a/examples/with-apollo/lib/apollo.js
+++ b/examples/with-apollo/lib/apollo.js
@@ -1,4 +1,5 @@
 import React, { useMemo } from 'react'
+import App from 'next/app'
 import Head from 'next/head'
 import { ApolloProvider } from '@apollo/react-hooks'
 import { ApolloClient } from 'apollo-client'
@@ -17,6 +18,9 @@ let apolloClient = null
  * @param {Boolean} [config.ssr=true]
  */
 export function withApollo (PageComponent, { ssr = true } = {}) {
+  const isAppHoc =
+    PageComponent === App || PageComponent.prototype instanceof App
+
   const WithApollo = ({ apolloClient, apolloState, ...pageProps }) => {
     const client = useMemo(
       () => apolloClient || initApolloClient(apolloState),
@@ -29,15 +33,19 @@ export function withApollo (PageComponent, { ssr = true } = {}) {
     )
   }
 
+  if (process.env.NODE_ENV !== 'production') {
+    if (isAppHoc && ssr) {
+      console.warn(
+        'You are using the "withApollo" HOC on "_app.js" level. Please note that this disables project wide automatic static optimization. Better wrap your PageComponents directly.'
+      )
+    }
+  }
+
   // Set the correct displayName in development
   if (process.env.NODE_ENV !== 'production') {
     const displayName =
       PageComponent.displayName || PageComponent.name || 'Component'
 
-    if (displayName === 'App') {
-      console.warn('This withApollo HOC only works with PageComponents.')
-    }
-
     WithApollo.displayName = `withApollo(${displayName})`
   }
 
@@ -68,14 +76,14 @@ export function withApollo (PageComponent, { ssr = true } = {}) {
           try {
             // Run all GraphQL queries
             const { getDataFromTree } = await import('@apollo/react-ssr')
-            await getDataFromTree(
-              <AppTree
-                pageProps={{
-                  ...pageProps,
-                  apolloClient
-                }}
-              />
-            )
+
+            // Since AppComponents and PageComponents have different context types
+            // we need to modify their props a little.
+            const props = isAppHoc
+              ? { ...pageProps, apolloClient }
+              : { pageProps: { ...pageProps, apolloClient } }
+
+            await getDataFromTree(<AppTree {...props} />)
           } catch (error) {
             // Prevent Apollo Client GraphQL errors from crashing SSR.
             // Handle them in components via the data.error prop: