From 7c6f43da559f47e071c4adf216c76d69aa2cdec6 Mon Sep 17 00:00:00 2001
From: Adam Petro <adam.petro@shopify.com>
Date: Tue, 20 Aug 2024 10:44:56 -0400
Subject: [PATCH] Allow variable definition directives in validation

---
 .../directives_are_in_valid_locations.rb      |  1 +
 .../directives_are_in_valid_locations_spec.rb | 50 ++++++++++++-------
 spec/support/dummy/schema.rb                  |  8 +++
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/lib/graphql/static_validation/rules/directives_are_in_valid_locations.rb b/lib/graphql/static_validation/rules/directives_are_in_valid_locations.rb
index 6e7ae3f3e3..b9aa0c7147 100644
--- a/lib/graphql/static_validation/rules/directives_are_in_valid_locations.rb
+++ b/lib/graphql/static_validation/rules/directives_are_in_valid_locations.rb
@@ -26,6 +26,7 @@ def on_directive(node, parent)
         Nodes::InlineFragment =>      GraphQL::Schema::Directive::INLINE_FRAGMENT,
         Nodes::FragmentSpread =>      GraphQL::Schema::Directive::FRAGMENT_SPREAD,
         Nodes::FragmentDefinition =>  GraphQL::Schema::Directive::FRAGMENT_DEFINITION,
+        Nodes::VariableDefinition =>  GraphQL::Schema::Directive::VARIABLE_DEFINITION,
       }
 
       SIMPLE_LOCATION_NODES = SIMPLE_LOCATIONS.keys
diff --git a/spec/graphql/static_validation/rules/directives_are_in_valid_locations_spec.rb b/spec/graphql/static_validation/rules/directives_are_in_valid_locations_spec.rb
index 50a01970f9..a9566c96c2 100644
--- a/spec/graphql/static_validation/rules/directives_are_in_valid_locations_spec.rb
+++ b/spec/graphql/static_validation/rules/directives_are_in_valid_locations_spec.rb
@@ -3,35 +3,36 @@
 
 describe GraphQL::StaticValidation::DirectivesAreInValidLocations do
   include StaticValidationHelpers
-  let(:query_string) {"
-    query getCheese @skip(if: true) {
-      okCheese: cheese(id: 1) {
-        id @skip(if: true),
-        source
-        ... on Cheese @skip(if: true) {
-          flavor
-          ... whatever
+
+  describe "invalid directive locations" do
+    let(:query_string) {"
+      query getCheese @skip(if: true) {
+        okCheese: cheese(id: 1) {
+          id @skip(if: true),
+          source
+          ... on Cheese @skip(if: true) {
+            flavor
+            ... whatever
+          }
         }
       }
-    }
-
-    fragment whatever on Cheese @skip(if: true) {
-      id
-    }
-  "}
+  
+      fragment whatever on Cheese @skip(if: true) {
+        id
+      }
+    "}
 
-  describe "invalid directive locations" do
     it "makes errors for them" do
       expected = [
         {
           "message"=> "'@skip' can't be applied to queries (allowed: fields, fragment spreads, inline fragments)",
-          "locations"=>[{"line"=>2, "column"=>21}],
+          "locations"=>[{"line"=>2, "column"=>23}],
           "path"=>["query getCheese"],
           "extensions"=>{"code"=>"directiveCannotBeApplied", "targetName"=>"queries", "name"=>"skip"}
         },
         {
           "message"=>"'@skip' can't be applied to fragment definitions (allowed: fields, fragment spreads, inline fragments)",
-          "locations"=>[{"line"=>13, "column"=>33}],
+          "locations"=>[{"line"=>13, "column"=>35}],
            "path"=>["fragment whatever"],
            "extensions"=>{"code"=>"directiveCannotBeApplied", "targetName"=>"fragment definitions", "name"=>"skip"}
         },
@@ -39,4 +40,19 @@
       assert_equal(expected, errors)
     end
   end
+
+  describe "valid directive locations" do
+    let(:query_string) {"
+      query getCheese($id: Int! @directiveForVariableDefinition) {
+        cheese(id: $id) {
+          id
+        }
+      }
+    "}
+
+    it "does not make errors for them" do
+      expected = []
+      assert_equal(expected, errors)
+    end
+  end
 end
diff --git a/spec/support/dummy/schema.rb b/spec/support/dummy/schema.rb
index 195ab6fb1a..0fbcb6c7f1 100644
--- a/spec/support/dummy/schema.rb
+++ b/spec/support/dummy/schema.rb
@@ -33,6 +33,9 @@ class BaseInputObject < GraphQL::Schema::InputObject
   class BaseScalar < GraphQL::Schema::Scalar
   end
 
+  class BaseDirective < GraphQL::Schema::Directive
+  end
+
   module LocalProduct
     include BaseInterface
     description "Something that comes from somewhere"
@@ -530,6 +533,10 @@ def replace_values(input:)
     end
   end
 
+  class DirectiveForVariableDefinition < BaseDirective
+    locations(VARIABLE_DEFINITION)
+  end
+
   class Subscription < BaseObject
     field :test, String
     def test; "Test"; end
@@ -542,6 +549,7 @@ class Schema < GraphQL::Schema
     max_depth 5
     orphan_types Honey
     trace_with GraphQL::Tracing::CallLegacyTracers
+    directives(DirectiveForVariableDefinition)
 
     rescue_from(NoSuchDairyError) { |err| raise GraphQL::ExecutionError, err.message  }