diff --git a/scripts/codemods/data-masking/unmask.ts b/scripts/codemods/data-masking/unmask.ts
index 1651b16494..77135c75a6 100644
--- a/scripts/codemods/data-masking/unmask.ts
+++ b/scripts/codemods/data-masking/unmask.ts
@@ -1,6 +1,7 @@
 import type { Collection, Transform, TemplateLiteral } from "jscodeshift";
 import type { DirectiveNode, DocumentNode } from "graphql";
 import { Kind, parse, visit, print } from "graphql";
+import path from "node:path";
 
 const LEADING_WHITESPACE = /^[\s\t]*(?=[\S\n])/;
 const TRAILING_WHITESPACE = /(?<=[\S\n])[\s\t]*$/;
@@ -8,9 +9,6 @@ const TRAILING_WHITESPACE = /(?<=[\S\n])[\s\t]*$/;
 const DEFAULT_TAGS = ["gql", "graphql"];
 
 const transform: Transform = function transform(file, api, options) {
-  const j = api.jscodeshift;
-  const source = j(file.source);
-
   const { tag = DEFAULT_TAGS, mode } = options;
 
   if (mode && mode !== "migrate") {
@@ -19,6 +17,15 @@ const transform: Transform = function transform(file, api, options) {
     );
   }
 
+  const extname = path.extname(file.path);
+
+  if (extname === ".graphql" || extname === ".gql") {
+    return transformGraphQLFile(file.source, mode);
+  }
+
+  const j = api.jscodeshift;
+  const source = j(file.source);
+
   const tagNames = Array.isArray(tag) ? tag : [tag];
 
   tagNames.forEach((tagName) => {
@@ -148,4 +155,8 @@ function getMatch(str: string, match: RegExp) {
   return str.match(match)?.at(0) ?? "";
 }
 
+function transformGraphQLFile(source: string, mode: string) {
+  return print(addUnmaskDirective(parse(source), mode));
+}
+
 export default transform;