From 8afe9561d9b31f5caa3ee980112c6034a191d314 Mon Sep 17 00:00:00 2001
From: Marco polo <salazarm@users.noreply.github.com>
Date: Wed, 28 Aug 2024 10:44:17 -0400
Subject: [PATCH] [1/4] Wrap useQuery/useLazyQuery + lint rule (#23860)

## Summary & Motivation

Wrapping `useQuery` and `useLazyQuery` to automatically integrate with
our performance tracking so that we can get rid of
`useBlockTraceOnQueryResult`

## How I Tested These Changes

I ran yarn lint --fix and saw it auto-fix all of the @apollo/client
imports in ui-core. I will put these changes in the next PR to keep this
one small.

## Changelog [New | Bug | Docs]
NOCHANGELOG
---
 .../packages/eslint-config/rules/index.js     |  2 +
 .../eslint-config/rules/no-apollo-client.js   | 60 +++++++++++++++++++
 .../ui-core/src/apollo-client/index.tsx       | 46 ++++++++++++++
 3 files changed, 108 insertions(+)
 create mode 100644 js_modules/dagster-ui/packages/eslint-config/rules/no-apollo-client.js
 create mode 100644 js_modules/dagster-ui/packages/ui-core/src/apollo-client/index.tsx

diff --git a/js_modules/dagster-ui/packages/eslint-config/rules/index.js b/js_modules/dagster-ui/packages/eslint-config/rules/index.js
index 2e5eb2f1527ed..6ad2643f0d846 100644
--- a/js_modules/dagster-ui/packages/eslint-config/rules/index.js
+++ b/js_modules/dagster-ui/packages/eslint-config/rules/index.js
@@ -4,6 +4,7 @@ const projectName = 'dagster-rules';
 const rules = {
   'missing-graphql-variables-type': require('./missing-graphql-variables-type').rule,
   'no-oss-imports': require('./no-oss-imports'),
+  'no-apollo-client': require('./no-apollo-client'),
 };
 
 module.exports = {
@@ -14,6 +15,7 @@ module.exports = {
       rules: {
         [`${projectName}/missing-graphql-variables-type`]: 'error',
         [`${projectName}/no-oss-imports`]: 'error',
+        [`${projectName}/no-apollo-client`]: 'error',
       },
     },
   },
diff --git a/js_modules/dagster-ui/packages/eslint-config/rules/no-apollo-client.js b/js_modules/dagster-ui/packages/eslint-config/rules/no-apollo-client.js
new file mode 100644
index 0000000000000..236b600ef4faf
--- /dev/null
+++ b/js_modules/dagster-ui/packages/eslint-config/rules/no-apollo-client.js
@@ -0,0 +1,60 @@
+/* eslint-disable */
+
+const fs = require('fs');
+const path = require('path');
+
+module.exports = {
+  meta: {
+    type: 'suggestion',
+    fixable: 'code',
+    messages: {
+      useWrappedApolloClient:
+        'Please use our wrapped apollo-client module which includes performance instrumentation.',
+    },
+  },
+  create(context) {
+    return {
+      ImportDeclaration(node) {
+        if (context.getFilename().endsWith('/apollo-client')) {
+          return;
+        }
+        if (node.source.value === '@apollo/client') {
+          context.report({
+            node,
+            messageId: 'useWrappedApolloClient',
+            fix(fixer) {
+              const currentFilePath = context.getFilename();
+              const relativeImportPath = findRelativeImportPath(
+                currentFilePath,
+                'src/apollo-client',
+                'index.tsx',
+              );
+
+              return fixer.replaceText(node.source, `'${relativeImportPath}'`);
+            },
+          });
+        }
+      },
+    };
+  },
+};
+
+function findRelativeImportPath(currentFilePath, targetDirName, targetFileName) {
+  let currentDir = path.dirname(currentFilePath);
+
+  while (currentDir !== path.parse(currentDir).root) {
+    const srcPath = path.join(currentDir, targetDirName);
+    const targetFilePath = path.join(srcPath, targetFileName);
+
+    if (fs.existsSync(targetFilePath)) {
+      return path.relative(path.dirname(currentFilePath), srcPath);
+    }
+
+    currentDir = path.dirname(currentDir);
+  }
+
+  // If we can't find the file then it means we're not in the ui-core package (we might be in cloud) so
+  // grab the import from @dagster-io/ui-core.
+
+  return '@dagster-io/ui-core/apollo-client';
+}
diff --git a/js_modules/dagster-ui/packages/ui-core/src/apollo-client/index.tsx b/js_modules/dagster-ui/packages/ui-core/src/apollo-client/index.tsx
new file mode 100644
index 0000000000000..a2796a4a85ee4
--- /dev/null
+++ b/js_modules/dagster-ui/packages/ui-core/src/apollo-client/index.tsx
@@ -0,0 +1,46 @@
+/* eslint-disable no-restricted-imports */
+import {
+  LazyQueryHookOptions,
+  OperationVariables,
+  QueryHookOptions,
+  useLazyQuery as useLazyQueryApollo,
+  useQuery as useQueryApollo,
+} from '@apollo/client';
+
+import {useBlockTraceOnQueryResult} from '../performance/TraceContext';
+
+export * from '@apollo/client';
+
+interface ExtendedQueryOptions<TData = any, TVariables extends OperationVariables = any>
+  extends QueryHookOptions<TData, TVariables> {
+  blocking?: boolean;
+}
+
+interface ExtendedLazyQueryOptions<TData = any, TVariables extends OperationVariables = any>
+  extends LazyQueryHookOptions<TData, TVariables> {
+  blocking?: boolean;
+}
+
+export function useQuery<TData = any, TVariables extends OperationVariables = any>(
+  query: any,
+  options?: ExtendedQueryOptions<TData, TVariables>,
+) {
+  const {blocking = true, ...restOptions} = options || {};
+  const queryResult = useQueryApollo<TData, TVariables>(query, restOptions);
+  useBlockTraceOnQueryResult(queryResult, 'graphql', {
+    skip: !blocking,
+  });
+  return queryResult;
+}
+
+export function useLazyQuery<TData = any, TVariables extends OperationVariables = any>(
+  query: any,
+  options?: ExtendedLazyQueryOptions<TData, TVariables>,
+) {
+  const {blocking = true, ...restOptions} = options || {};
+  const result = useLazyQueryApollo<TData, TVariables>(query, restOptions);
+  useBlockTraceOnQueryResult(result[1], 'graphql', {
+    skip: !blocking,
+  });
+  return result;
+}