From e1e2605b30efc488b57f62ba43436606a38a3607 Mon Sep 17 00:00:00 2001
From: Chris Lenfest <clenfest@apollographql.com>
Date: Wed, 25 Sep 2024 10:41:09 -0500
Subject: [PATCH] Fixes issue where there can be naming collisions in
 contextual parameter names (#3155)

Subgraph number set for a contextual paramter was colliding if used in
multiple subgraphs
---
 .changeset/tame-paws-return.md                |  6 +++++
 query-graphs-js/src/querygraph.ts             | 27 +++++++++++--------
 .../src/__tests__/buildPlan.test.ts           |  4 +--
 3 files changed, 24 insertions(+), 13 deletions(-)
 create mode 100644 .changeset/tame-paws-return.md

diff --git a/.changeset/tame-paws-return.md b/.changeset/tame-paws-return.md
new file mode 100644
index 000000000..819bf0d1b
--- /dev/null
+++ b/.changeset/tame-paws-return.md
@@ -0,0 +1,6 @@
+---
+"@apollo/query-planner": patch
+"@apollo/query-graphs": patch
+---
+
+Fixes issue where contextual parameters can have naming collisions if used in multiple subgraphs
diff --git a/query-graphs-js/src/querygraph.ts b/query-graphs-js/src/querygraph.ts
index d2f738d40..e077771f5 100644
--- a/query-graphs-js/src/querygraph.ts
+++ b/query-graphs-js/src/querygraph.ts
@@ -934,17 +934,6 @@ function federateSubgraphs(supergraph: Schema, subgraphs: QueryGraph[]): QueryGr
       }
     }
     
-    for (const [subgraphName, args] of subgraphToArgs) {
-      args.sort();
-      const argToIndex = new Map();
-      for (let idx=0; idx < args.length; idx++) {
-        argToIndex.set(args[idx], `contextualArgument_${i}_${idx}`);
-      }
-      subgraphToArgIndices.set(subgraphName, argToIndex);
-    }
-    
-    builder.setContextMaps(subgraphToArgs, subgraphToArgIndices);
-    
     simpleTraversal(
       subgraph,
       _v => { return undefined; },
@@ -965,6 +954,22 @@ function federateSubgraphs(supergraph: Schema, subgraphs: QueryGraph[]): QueryGr
    
   }
 
+  // add contextual argument maps to builder
+  for (const [i, subgraph] of subgraphs.entries()) {
+    const subgraphName = subgraph.name;
+    const args = subgraphToArgs.get(subgraph.name);
+    if (args) {
+      args.sort();
+      const argToIndex = new Map();
+      for (let idx=0; idx < args.length; idx++) {
+        argToIndex.set(args[idx], `contextualArgument_${i+1}_${idx}`);
+      }
+      subgraphToArgIndices.set(subgraphName, argToIndex);
+    }
+  }
+  
+  builder.setContextMaps(subgraphToArgs, subgraphToArgIndices);
+      
   // Now we handle @provides
   let provideId = 0;
   for (const [i, subgraph] of subgraphs.entries()) {
diff --git a/query-planner-js/src/__tests__/buildPlan.test.ts b/query-planner-js/src/__tests__/buildPlan.test.ts
index af995c066..9822c0ccc 100644
--- a/query-planner-js/src/__tests__/buildPlan.test.ts
+++ b/query-planner-js/src/__tests__/buildPlan.test.ts
@@ -9164,7 +9164,7 @@ describe('@fromContext impacts on query planning', () => {
               } =>
               {
                 ... on U {
-                  field(a: $contextualArgument_1_0)
+                  field(a: $contextualArgument_2_0)
                 }
               }
             },
@@ -9176,7 +9176,7 @@ describe('@fromContext impacts on query planning', () => {
       {
         kind: 'KeyRenamer',
         path: ['..', '... on T', 'prop'],
-        renameKeyTo: 'contextualArgument_1_0',
+        renameKeyTo: 'contextualArgument_2_0',
       },
     ]);
   });