Update Next.js hydration handling

src/consts.ts

@@ -8,6 +8,7 @@ export const TYPESCRIPT_PACKAGE = "typescript";
 export const BABEL_RELAY_PACKAGE = "babel-plugin-relay";
 export const BABEL_RELAY_MACRO = BABEL_RELAY_PACKAGE + "/macro";
 export const REACT_RELAY_PACKAGE = "react-relay";
+export const RELAY_RUNTIME_PACKAGE = "relay-runtime";
 export const GRAPHQL_WS_PACKAGE = "graphql-ws";
 export const VITE_RELAY_PACKAGE = "vite-plugin-relay";
 

src/tasks/GenerateRelayEnvironmentTask.ts

@@ -51,11 +51,6 @@ export class GenerateRelayEnvironmentTask extends TaskBase {
       if (this.context.args.subscriptions) {
         relayRuntimeImports.push("SubscribeFunction");
       }
-
-      if (this.context.is("next")) {
-        // prettier-ignore
-        b.addLine(`import type { RecordMap } from "relay-runtime/lib/store/RelayStoreTypes";`)
-      }
     }
 
     // prettier-ignore
@@ -184,15 +179,9 @@ export class GenerateRelayEnvironmentTask extends TaskBase {
     if (this.context.is("next")) {
       let initEnv = `let relayEnvironment: Environment | undefined;
 
-        export function initRelayEnvironment(initialRecords?: RecordMap) {
+        export function initRelayEnvironment() {
           const environment = relayEnvironment ?? createRelayEnvironment();
 
-          // If your page has Next.js data fetching methods that use Relay,
-          // the initial records will get hydrated here.
-          if (initialRecords) {
-            environment.getStore().publish(new RecordSource(initialRecords));
-          }
-
           // For SSG and SSR always create a new Relay environment.
           if (typeof window === "undefined") {
             return environment;
@@ -208,12 +197,6 @@ export class GenerateRelayEnvironmentTask extends TaskBase {
         }`;
 
       if (!this.context.args.typescript) {
-        // Remove Typescript type
-        initEnv = initEnv.replace(
-          "initialRecords?: RecordMap",
-          "initialRecords"
-        );
-
         initEnv = initEnv.replace(": Environment | undefined", "");
       }
 

src/tasks/cra/Cra_AddRelayEnvironmentProvider.ts

@@ -7,7 +7,12 @@ import {
 } from "../../consts.js";
 import { ProjectContext } from "../../misc/ProjectContext.js";
 import { RelativePath } from "../../misc/RelativePath.js";
-import { insertNamedImport, parseAst, printAst } from "../../utils/ast.js";
+import {
+  insertNamedImport,
+  insertNamedImports,
+  parseAst,
+  printAst,
+} from "../../utils/ast.js";
 import { h } from "../../utils/cli.js";
 import { TaskBase, TaskSkippedError } from "../TaskBase.js";
 import t from "@babel/types";

src/tasks/next/Next_AddRelayEnvironmentProvider.ts

@@ -1,9 +1,19 @@
 import traverse from "@babel/traverse";
 import path from "path";
-import { REACT_RELAY_PACKAGE, RELAY_ENV_PROVIDER } from "../../consts.js";
+import {
+  REACT_RELAY_PACKAGE,
+  RELAY_ENV_PROVIDER,
+  RELAY_RUNTIME_PACKAGE,
+} from "../../consts.js";
 import { ProjectContext } from "../../misc/ProjectContext.js";
 import { RelativePath } from "../../misc/RelativePath.js";
-import { insertNamedImport, parseAst, printAst } from "../../utils/ast.js";
+import {
+  astToString,
+  insertNamedImport,
+  insertNamedImports,
+  parseAst,
+  prettifyCode,
+} from "../../utils/ast.js";
 import { h } from "../../utils/cli.js";
 import { TaskBase, TaskSkippedError } from "../TaskBase.js";
 import t from "@babel/types";
@@ -14,11 +24,19 @@ import {
 } from "../cra/Cra_AddRelayEnvironmentProvider.js";
 import { Next_AddTypeHelpers } from "./Next_AddTypeHelpers.js";
 
-const envCreation = `
-const environment = useMemo(
-    () => initRelayEnvironment(pageProps.initialRecords),
-    [pageProps.initialRecords]
-);
+const envCreationAndHydration = `
+const environment = useMemo(initRelayEnvironment, []);
+
+useEffect(() => {
+  const store = environment.getStore();
+
+  // Hydrate the store.
+  store.publish(new RecordSource(pageProps.initialRecords));
+
+  // Notify any existing subscribers.
+  store.notify();
+}, [environment, pageProps.initialRecords])
+
 `;
 
 const APP_PROPS = "AppProps";
@@ -47,8 +65,6 @@ export class Next_AddRelayEnvironmentProvider extends TaskBase {
 
     const ast = parseAst(code);
 
-    const envCreationAst = parseAst(envCreation).program.body[0];
-
     let providerWrapped = false;
 
     traverse.default(ast, {
@@ -58,6 +74,8 @@ export class Next_AddRelayEnvironmentProvider extends TaskBase {
           return;
         }
 
+        const functionReturn = path.parentPath;
+
         const isProviderConfigured = hasRelayProvider(path);
 
         if (isProviderConfigured) {
@@ -79,7 +97,7 @@ export class Next_AddRelayEnvironmentProvider extends TaskBase {
           insertNamedImport(path, RELAY_PAGE_PROPS, relayTypesImportPath.rel);
 
           // Change argument of type AppProps to AppProps<RelayPageProps>.
-          const functionBodyPath = path.parentPath.parentPath;
+          const functionBodyPath = functionReturn.parentPath;
           if (!functionBodyPath.isBlockStatement()) {
             throw new Error("Expected parentPath to be a block statement.");
           }
@@ -110,7 +128,8 @@ export class Next_AddRelayEnvironmentProvider extends TaskBase {
           appPropsArg.typeAnnotation = t.typeAnnotation(genericAppProps);
         }
 
-        insertNamedImport(path, "useMemo", "react");
+        insertNamedImports(path, ["useMemo", "useEffect"], "react");
+        insertNamedImport(path, "RecordSource", RELAY_RUNTIME_PACKAGE);
 
         const relayEnvImportPath = new RelativePath(
           mainFile.parentDirectory,
@@ -119,8 +138,7 @@ export class Next_AddRelayEnvironmentProvider extends TaskBase {
 
         insertNamedImport(path, "initRelayEnvironment", relayEnvImportPath.rel);
 
-        // Insert the useMemo creating the environment in the function body.
-        path.parentPath.insertBefore(envCreationAst);
+        functionReturn.addComment("leading", "--MARKER", true);
 
         const envProviderId = t.jsxIdentifier(
           insertNamedImport(path, RELAY_ENV_PROVIDER, REACT_RELAY_PACKAGE).name
@@ -142,7 +160,10 @@ export class Next_AddRelayEnvironmentProvider extends TaskBase {
       throw new Error("Could not find JSX");
     }
 
-    const updatedCode = printAst(ast, code);
+    let updatedCode = astToString(ast, code);
+
+    updatedCode = updatedCode.replace("//--MARKER", envCreationAndHydration);
+    updatedCode = prettifyCode(updatedCode);
 
     await this.context.fs.writeToFile(mainFile.abs, updatedCode);
   }

src/tasks/next/Next_AddTypeHelpers.ts

@@ -44,6 +44,11 @@ export class Next_AddTypeHelpers extends TaskBase {
 
     this.updateMessage(this.message + " " + h(filepath.rel));
 
+    if (this.context.fs.exists(filepath.abs)) {
+      this.skip("File exists");
+      return;
+    }
+
     const prettifiedCode = prettifyCode(code);
 
     await this.context.fs.writeToFile(filepath.abs, prettifiedCode);

src/utils/ast.ts

@@ -11,8 +11,12 @@ export function parseAst(code: string): ParseResult<t.File> {
   });
 }
 
+export function astToString(ast: ParseResult<t.File>, oldCode: string): string {
+  return generate.default(ast, { retainLines: true }, oldCode).code;
+}
+
 export function printAst(ast: ParseResult<t.File>, oldCode: string): string {
-  const newCode = generate.default(ast, { retainLines: true }, oldCode).code;
+  const newCode = astToString(ast, oldCode);
 
   return prettifyCode(newCode);
 }
@@ -30,25 +34,64 @@ export function insertNamedImport(
   importName: string,
   packageName: string
 ): t.Identifier {
-  const importIdentifier = t.identifier(importName);
+  return insertNamedImports(path, [importName], packageName)[0];
+}
 
+export function insertNamedImports(
+  path: NodePath,
+  imports: string[],
+  packageName: string
+): t.Identifier[] {
   const program = path.findParent((p) => p.isProgram()) as NodePath<t.Program>;
 
-  const existingImport = getNamedImport(program, importName, packageName);
+  const identifiers: t.Identifier[] = [];
+  const missingImports: string[] = [];
 
-  if (!!existingImport) {
-    return importIdentifier;
+  for (const namedImport of imports) {
+    const importIdentifier = t.identifier(namedImport);
+
+    const existingImport = getNamedImport(program, namedImport, packageName);
+
+    if (!!existingImport) {
+      identifiers.push(importIdentifier);
+      continue;
+    }
+
+    missingImports.push(namedImport);
   }
 
-  const importDeclaration = t.importDeclaration(
-    [t.importSpecifier(t.cloneNode(importIdentifier), importIdentifier)],
-    t.stringLiteral(packageName)
-  );
+  console.log({ imports, missingImports, identifiers });
 
-  // Insert import at start of file.
-  program.node.body.unshift(importDeclaration);
+  let importDeclaration: t.ImportDeclaration;
+  const isFirstImportFromPackage = missingImports.length === imports.length;
 
-  return importIdentifier;
+  if (isFirstImportFromPackage) {
+    console.log("create new");
+    importDeclaration = t.importDeclaration([], t.stringLiteral(packageName));
+  } else {
+    console.log("get existing");
+    importDeclaration = getImportDeclaration(program, packageName)!;
+  }
+
+  for (const namedImport of missingImports) {
+    const importIdentifier = t.identifier(namedImport);
+
+    const newImport = t.importSpecifier(
+      t.cloneNode(importIdentifier),
+      importIdentifier
+    );
+
+    importDeclaration.specifiers.push(newImport);
+
+    identifiers.push(importIdentifier);
+  }
+
+  if (isFirstImportFromPackage) {
+    // Insert import at start of file.
+    program.node.body.unshift(importDeclaration);
+  }
+
+  return identifiers;
 }
 
 export function insertDefaultImport(
@@ -78,6 +121,15 @@ export function insertDefaultImport(
   return importIdentifier;
 }
 
+function getImportDeclaration(
+  path: NodePath<t.Program>,
+  packageName: string
+): t.ImportDeclaration | null {
+  return path.node.body.find(
+    (s) => t.isImportDeclaration(s) && s.source.value === packageName
+  ) as t.ImportDeclaration | null;
+}
+
 export function getNamedImport(
   path: NodePath<t.Program>,
   importName: string,
@@ -93,7 +145,7 @@ export function getNamedImport(
   ) as t.ImportDeclaration;
 }
 
-export function getDefaultImport(
+function getDefaultImport(
   path: NodePath<t.Program>,
   importName: string,
   packageName: string