Decorate app instead of pages with Apollo.

app/pages/_app.js

@@ -0,0 +1,18 @@
+import { ApolloProvider } from 'react-apollo'
+import App, { Container } from 'next/app'
+import withApollo from '../withApollo'
+
+class CustomApp extends App {
+  render() {
+    const { Component, pageProps, apolloClient } = this.props
+    return (
+      <Container>
+        <ApolloProvider client={apolloClient}>
+          <Component {...pageProps} />
+        </ApolloProvider>
+      </Container>
+    )
+  }
+}
+
+export default withApollo(CustomApp)

app/pages/index.js

@@ -4,7 +4,6 @@ import UploadFile from '../components/UploadFile'
 import UploadFileList from '../components/UploadFileList'
 import UploadBlob from '../components/UploadBlob'
 import Uploads from '../components/Uploads'
-import withData from '../providers/with-data'
 
 const IndexPage = () => (
   <Page title="Apollo upload examples">
@@ -29,4 +28,4 @@ const IndexPage = () => (
   </Page>
 )
 
-export default withData(IndexPage)
+export default IndexPage

app/providers/with-data.js

@@ -1,78 +0,0 @@
-import 'cross-fetch/polyfill'
-import React, { Component } from 'react'
-import getDisplayName from 'react-display-name'
-import { ApolloClient } from 'apollo-client'
-import { InMemoryCache } from 'apollo-cache-inmemory'
-import { createUploadLink } from 'apollo-upload-client'
-import { getDataFromTree, ApolloProvider } from 'react-apollo'
-import { Router } from 'next/router'
-import App from 'next/app'
-
-// Shares the browser ApolloClient instance between pages.
-let browserApolloClient
-
-const createApolloClient = (cache = {}) =>
-  new ApolloClient({
-    ssrMode: typeof window !== 'undefined',
-    cache: new InMemoryCache().restore(cache),
-    link: createUploadLink({ uri: process.env.API_URI })
-  })
-
-export default Composed =>
-  class WithData extends Component {
-    constructor(props) {
-      super(props)
-
-      // Set the ApolloClient instance used in render().
-      this.apolloClient =
-        typeof window !== 'undefined'
-          ? // Client: Shared instance, created at first render after SSR.
-            (browserApolloClient =
-              browserApolloClient || createApolloClient(props.cache))
-          : // Server: Private instance for SSR.
-            createApolloClient(props.cache)
-    }
-
-    static displayName = `WithApollo(${getDisplayName(Composed)})`
-
-    static async getInitialProps(ctx) {
-      // Use initial props from nested page decorators.
-      const props = Composed.getInitialProps
-        ? await Composed.getInitialProps(ctx)
-        : {}
-
-      // If server environment, preload the page.
-      if (ctx.req) {
-        const apolloClient = createApolloClient()
-
-        try {
-          await getDataFromTree(
-            <App
-              router={new Router(ctx.pathname, ctx.query, ctx.asPath)}
-              pageProps={{ apolloClient, pageProps: props }}
-              Component={this.renderPage}
-            />
-          )
-        } catch (error) {
-          // Prevent crash from GraphQL errors.
-        }
-
-        props.cache = apolloClient.cache.extract()
-      }
-
-      return props
-    }
-
-    static renderPage = ({ apolloClient, pageProps }) => (
-      <ApolloProvider client={apolloClient}>
-        <Composed {...pageProps} />
-      </ApolloProvider>
-    )
-
-    render() {
-      return this.constructor.renderPage({
-        apolloClient: this.apolloClient,
-        pageProps: this.props
-      })
-    }
-  }

app/withApollo.js

@@ -0,0 +1,77 @@
+import 'cross-fetch/polyfill'
+import { InMemoryCache } from 'apollo-cache-inmemory'
+import { ApolloClient } from 'apollo-client'
+import { createUploadLink } from 'apollo-upload-client'
+import Head from 'next/head'
+import { Component } from 'react'
+import { getDataFromTree } from 'react-apollo'
+import getDisplayName from 'react-display-name'
+
+let apolloClient
+
+const createApolloClient = (cache = {}) =>
+  new ApolloClient({
+    ssrMode: typeof window !== 'undefined',
+    cache: new InMemoryCache().restore(cache),
+    link: createUploadLink({ uri: process.env.API_URI })
+  })
+
+export default App =>
+  class WithApollo extends Component {
+    static displayName = `WithApollo(${getDisplayName(App)})`
+
+    static async getInitialProps(context) {
+      const props = App.getInitialProps
+        ? await App.getInitialProps(context)
+        : {}
+
+      // If server environment, preload the page.
+      if (context.ctx.req) {
+        const apolloClient = createApolloClient()
+
+        try {
+          await getDataFromTree(
+            <App
+              {...props}
+              Component={context.Component}
+              router={context.router}
+              apolloClient={apolloClient}
+            />
+          )
+        } catch (error) {
+          // Prevent crash from GraphQL errors.
+          // eslint-disable-next-line no-console
+          console.error(error)
+        }
+
+        props.apolloCache = apolloClient.cache.extract()
+
+        Head.rewind()
+      }
+
+      return props
+    }
+
+    constructor(props) {
+      super(props)
+
+      // Set the ApolloClient instance used in render().
+      this.apolloClient =
+        typeof window !== 'undefined'
+          ? // Client: Shared instance, created at first render after SSR.
+            (apolloClient =
+              apolloClient || createApolloClient(props.apolloCache))
+          : // Server: Private instance for SSR.
+            createApolloClient(props.apolloCache)
+    }
+
+    render() {
+      const {
+        // eslint-disable-next-line no-unused-vars
+        apolloCache,
+        ...appProps
+      } = this.props
+
+      return <App {...appProps} apolloClient={this.apolloClient} />
+    }
+  }