diff --git a/Changelog.md b/Changelog.md
index 048b725d22..619fc105f3 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -21,6 +21,35 @@
   necessary anymore.
   [PR #2533](https://github.com/apollographql/react-apollo/pull/2533)
 
+- Restore original `getDataFromTree(tree, context)` API, and introduce a
+  new function called `getMarkupFromTree` to enable custom rendering
+  functions:
+  ```typescript
+  export default function getDataFromTree(
+    tree: React.ReactNode,
+    context: { [key: string]: any } = {},
+  ) {
+    return getMarkupFromTree({
+      tree,
+      context,
+      renderFunction: renderToStaticMarkup,
+    });
+  }
+
+  export type GetMarkupFromTreeOptions = {
+    tree: React.ReactNode;
+    context?: { [key: string]: any };
+    renderFunction?: typeof renderToStaticMarkup;
+  };
+
+  export function getMarkupFromTree({
+    tree,
+    context = {},
+    renderFunction = renderToStaticMarkup,
+  }: GetMarkupFromTreeOptions): Promise<string> {...}
+  ```
+  [PR #2586](https://github.com/apollographql/react-apollo/pull/2586)
+
 ## 2.2.4 (October 2, 2018)
 
 ### Bug Fixes
diff --git a/examples/ssr/server/main.js b/examples/ssr/server/main.js
index 485cebc056..433f8c63a2 100644
--- a/examples/ssr/server/main.js
+++ b/examples/ssr/server/main.js
@@ -1,7 +1,7 @@
 import { renderToString } from 'react-dom/server';
 import { onPageLoad } from 'meteor/server-render';
 import { ApolloClient } from 'apollo-client';
-import { getDataFromTree, ApolloProvider } from 'react-apollo';
+import { getMarkupFromTree, ApolloProvider } from 'react-apollo';
 import { InMemoryCache } from 'apollo-cache-inmemory';
 import { HttpLink } from 'apollo-link-http';
 import { WebApp } from 'meteor/webapp';
@@ -27,12 +27,14 @@ export const render = async sink => {
     </ApolloProvider>
   );
 
-  const start = +new Date;
+  const start = Date.now();
   // Load all data from local server
-  await getDataFromTree(WrappedApp);
-  const body = renderToString(WrappedApp);
-  console.log("server rendering took", new Date - start, "ms");
-  sink.renderIntoElementById('app', body);
+  const markup = await getMarkupFromTree({
+    tree: WrappedApp,
+    renderFunction: renderToString,
+  });
+  console.log("server rendering took", Date.now() - start, "ms");
+  sink.renderIntoElementById('app', markup);
   sink.appendToBody(`
     <script>
       window.__APOLLO_STATE__=${JSON.stringify(client.extract())};
diff --git a/src/getDataFromTree.ts b/src/getDataFromTree.ts
index ed3954e630..6745fe8b66 100755
--- a/src/getDataFromTree.ts
+++ b/src/getDataFromTree.ts
@@ -84,46 +84,60 @@ export class RenderPromises {
   }
 }
 
-class RenderPromisesProvider extends React.Component<{
-  renderPromises: RenderPromises;
-}> {
-  static childContextTypes = {
-    renderPromises: PropTypes.object,
-  };
-
-  getChildContext() {
-    return {
-      renderPromises: this.props.renderPromises,
-    };
-  }
-
-  render() {
-    return this.props.children;
-  }
+export default function getDataFromTree(
+  tree: React.ReactNode,
+  context: { [key: string]: any } = {},
+) {
+  return getMarkupFromTree({
+    tree,
+    context,
+    // If you need to configure this renderFunction, call getMarkupFromTree
+    // directly instead of getDataFromTree.
+    renderFunction: renderToStaticMarkup,
+  });
 }
 
-export default function getDataFromTree(
-  rootElement: React.ReactNode,
+export type GetMarkupFromTreeOptions = {
+  tree: React.ReactNode;
+  context?: { [key: string]: any };
+  renderFunction?: typeof renderToStaticMarkup;
+};
+
+export function getMarkupFromTree({
+  tree,
+  context = {},
   // The rendering function is configurable! We use renderToStaticMarkup as
   // the default, because it's a little less expensive than renderToString,
   // and legacy usage of getDataFromTree ignores the return value anyway.
   renderFunction = renderToStaticMarkup,
-): Promise<string> {
+}: GetMarkupFromTreeOptions): Promise<string> {
   const renderPromises = new RenderPromises();
 
-  function process(): Promise<string> | string {
-    const html = renderFunction(
-      React.createElement(RenderPromisesProvider, {
-        renderPromises,
-        // Always re-render from the rootElement, even though it might seem
-        // better to render the children of the component responsible for the
-        // promise, because it is not possible to reconstruct the full context
-        // of the original rendering (including all unknown context provider
-        // elements) for a subtree of the orginal component tree.
-        children: rootElement,
-      })
-    );
+  class RenderPromisesProvider extends React.Component {
+    static childContextTypes: { [key: string]: any } = {
+      renderPromises: PropTypes.object,
+    };
+
+    getChildContext() {
+      return { ...context, renderPromises };
+    }
+
+    render() {
+      // Always re-render from the rootElement, even though it might seem
+      // better to render the children of the component responsible for the
+      // promise, because it is not possible to reconstruct the full context
+      // of the original rendering (including all unknown context provider
+      // elements) for a subtree of the orginal component tree.
+      return tree;
+    }
+  }
 
+  Object.keys(context).forEach(key => {
+    RenderPromisesProvider.childContextTypes[key] = PropTypes.any;
+  });
+
+  function process(): Promise<string> | string {
+    const html = renderFunction(React.createElement(RenderPromisesProvider));
     return renderPromises.hasPromises()
       ? renderPromises.consumeAndAwaitPromises().then(process)
       : html;
diff --git a/test/client/getDataFromTree.test.tsx b/test/client/getDataFromTree.test.tsx
index 9fde7c67c6..fe4a6eaf3e 100644
--- a/test/client/getDataFromTree.test.tsx
+++ b/test/client/getDataFromTree.test.tsx
@@ -8,6 +8,7 @@ import {
   ApolloProvider,
   walkTree,
   getDataFromTree,
+  getMarkupFromTree,
   DataValue,
   ChildProps,
 } from '../../src';
@@ -511,13 +512,34 @@ describe('SSR', () => {
         expect(markup).toMatch(/James/);
       });
 
-      await getDataFromTree(app, ReactDOM.renderToString).then(html => {
+      await getMarkupFromTree({
+        tree: app,
+        renderFunction: ReactDOM.renderToString,
+      }).then(html => {
         const markup = ReactDOM.renderToString(app);
         expect(markup).toEqual(html);
         expect(markup).toMatch(/James/);
       });
     });
 
+    it('should support passing a root context', () => {
+      class Consumer extends React.Component {
+        static contextTypes = {
+          text: PropTypes.string.isRequired,
+        };
+
+        render() {
+          return <div>{this.context.text}</div>;
+        }
+      }
+
+      return getDataFromTree(<Consumer/>, {
+        text: "oyez"
+      }).then(html => {
+        expect(html).toEqual('<div>oyez</div>');
+      });
+    });
+
     it('should run through all of the queries (also defined via Query component) that want SSR', () => {
       const query = gql`
         {