From f01e058c037059cdff68e833681e71b49516bd22 Mon Sep 17 00:00:00 2001
From: louiszawadzki <louis.zawadzki@datadoghq.com>
Date: Tue, 12 Sep 2023 11:53:51 +0200
Subject: [PATCH] Add variables mapper

---
 .../src/DatadogLink.ts                        | 16 +++-
 .../src/__tests__/mappers.test.ts             | 77 +++++++++++++++++++
 .../react-native-apollo-client/src/mappers.ts | 22 ++++++
 3 files changed, 111 insertions(+), 4 deletions(-)
 create mode 100644 packages/react-native-apollo-client/src/__tests__/mappers.test.ts
 create mode 100644 packages/react-native-apollo-client/src/mappers.ts

diff --git a/packages/react-native-apollo-client/src/DatadogLink.ts b/packages/react-native-apollo-client/src/DatadogLink.ts
index 414afd29b..a28b50683 100644
--- a/packages/react-native-apollo-client/src/DatadogLink.ts
+++ b/packages/react-native-apollo-client/src/DatadogLink.ts
@@ -12,12 +12,20 @@ import {
 } from '@datadog/mobile-react-native';
 
 import { getOperationName, getVariables, getOperationType } from './helpers';
+import type { VariablesMapper } from './mappers';
+import { mapVariables } from './mappers';
 
 export class DatadogLink extends ApolloLink {
-    constructor() {
+    private variablesMapper: undefined | VariablesMapper;
+
+    constructor({ variablesMapper }: { variablesMapper?: VariablesMapper }) {
         super((operation, forward) => {
             const operationName = getOperationName(operation);
             const formattedVariables = getVariables(operation);
+            const mappedVariables = mapVariables(
+                formattedVariables,
+                this.variablesMapper
+            );
             const operationType = getOperationType(operation);
 
             operation.setContext(({ headers = {} }) => {
@@ -31,9 +39,7 @@ export class DatadogLink extends ApolloLink {
                 newHeaders[
                     DATADOG_GRAPH_QL_OPERATION_NAME_HEADER
                 ] = operationName;
-                newHeaders[
-                    DATADOG_GRAPH_QL_VARIABLES_HEADER
-                ] = formattedVariables;
+                newHeaders[DATADOG_GRAPH_QL_VARIABLES_HEADER] = mappedVariables;
 
                 return {
                     headers: newHeaders
@@ -42,5 +48,7 @@ export class DatadogLink extends ApolloLink {
 
             return forward(operation);
         });
+
+        this.variablesMapper = variablesMapper;
     }
 }
diff --git a/packages/react-native-apollo-client/src/__tests__/mappers.test.ts b/packages/react-native-apollo-client/src/__tests__/mappers.test.ts
new file mode 100644
index 000000000..021a8d065
--- /dev/null
+++ b/packages/react-native-apollo-client/src/__tests__/mappers.test.ts
@@ -0,0 +1,77 @@
+import { mapVariables } from '../mappers';
+
+describe('mappers', () => {
+    describe('mapVariables', () => {
+        it('executes the provided mapper on the variables as JSON', () => {
+            const originalVariables = JSON.stringify({
+                user: {
+                    email: 'user@email.com',
+                    password: 'pass',
+                    type: 'admin'
+                }
+            });
+            const mapper = (variables: Record<string, any> | null) => {
+                if (variables && variables.user?.email) {
+                    delete variables.user.email;
+                }
+                if (variables && variables.user?.password) {
+                    delete variables.user.password;
+                }
+                return variables;
+            };
+
+            expect(mapVariables(originalVariables, mapper)).toBe(
+                '{"user":{"type":"admin"}}'
+            );
+        });
+        it('returns the original variables if they are not well formatted as JSON', () => {
+            const originalVariables = '{user: [Object object]}';
+            const mapper = (variables: Record<string, any> | null) => {
+                if (variables && variables.user?.email) {
+                    delete variables.user.email;
+                }
+                if (variables && variables.user?.password) {
+                    delete variables.user.password;
+                }
+                return variables;
+            };
+
+            expect(mapVariables(originalVariables, mapper)).toBe(
+                originalVariables
+            );
+        });
+        it('returns the original variables if the mapper errors', () => {
+            const originalVariables = JSON.stringify({
+                user: {
+                    email: 'user@email.com',
+                    password: 'pass',
+                    type: 'admin'
+                }
+            });
+            const mapper = (variables: Record<string, any> | null) => {
+                if (variables && variables.user?.email) {
+                    delete variables.user.email;
+                }
+                if (variables && variables.user?.password) {
+                    delete variables.user.password;
+                }
+                throw new Error('mapper error');
+            };
+
+            expect(mapVariables(originalVariables, mapper)).toBe(
+                originalVariables
+            );
+        });
+        it('returns the original variables if no mapper was provided', () => {
+            const originalVariables = JSON.stringify({
+                user: {
+                    email: 'user@email.com',
+                    password: 'pass',
+                    type: 'admin'
+                }
+            });
+
+            expect(mapVariables(originalVariables)).toBe(originalVariables);
+        });
+    });
+});
diff --git a/packages/react-native-apollo-client/src/mappers.ts b/packages/react-native-apollo-client/src/mappers.ts
new file mode 100644
index 000000000..ed9ac49b8
--- /dev/null
+++ b/packages/react-native-apollo-client/src/mappers.ts
@@ -0,0 +1,22 @@
+export type VariablesMapper = (
+    variables: Record<string, unknown> | null
+) => Record<string, unknown> | null;
+
+export const mapVariables = (
+    originalVariables: string | null,
+    mapper?: VariablesMapper
+): string | null => {
+    if (!mapper) {
+        return originalVariables;
+    }
+    try {
+        const mappedVariables = mapper(
+            originalVariables === null ? null : JSON.parse(originalVariables)
+        );
+
+        return JSON.stringify(mappedVariables);
+    } catch (e) {
+        // TODO: telemetry
+        return originalVariables;
+    }
+};