refactor proxy check code into helper function

graphprotocol · Jan 9, 2025 · 3f1b940 · 3f1b940
1 parent 6b4b00b
commit 3f1b940
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 59 deletions.
diff --git a/packages/cli/src/command-helpers/proxy.ts b/packages/cli/src/command-helpers/proxy.ts
@@ -0,0 +1,55 @@
+import { prompt } from 'gluegun';
+import EthereumABI from '../protocols/ethereum/abi.js';
+import { ContractService } from './contracts.js';
+import { retryWithPrompt } from './retry.js';
+import { withSpinner } from './spinner.js';
+
+export interface CheckForProxyResult {
+  implementationAbi: EthereumABI | null;
+  implementationAddress: string | null;
+}
+
+export async function checkForProxy(
+  contractService: ContractService,
+  network: string,
+  address: string,
+  abi: EthereumABI,
+): Promise<CheckForProxyResult> {
+  let implementationAddress = null;
+  let implementationAbi = null;
+
+  const maybeProxy = abi.callFunctionSignatures()?.includes('upgradeTo(address)');
+  if (maybeProxy) {
+    const impl = await retryWithPrompt(() =>
+      withSpinner(
+        'Fetching proxy implementation address...',
+        'Failed to fetch proxy implementation address',
+        'Warning fetching proxy implementation address',
+        () => contractService.getProxyImplementation(network, address),
+      ),
+    );
+
+    if (impl) {
+      const useImplementation = await prompt
+        .confirm(`Proxy contract detected. Use implementation contract ABI at ${impl}?`, true)
+        .catch(() => false);
+
+      if (useImplementation) {
+        implementationAddress = impl;
+        implementationAbi = await retryWithPrompt(() =>
+          withSpinner(
+            'Fetching implementation contract ABI...',
+            'Failed to fetch implementation ABI',
+            'Warning fetching implementation ABI',
+            () => contractService.getABI(EthereumABI, network, implementationAddress!),
+          ),
+        );
+      }
+    }
+  }
+
+  return {
+    implementationAbi,
+    implementationAddress,
+  };
+}
diff --git a/packages/cli/src/commands/add.ts b/packages/cli/src/commands/add.ts
@@ -4,6 +4,7 @@ import { Args, Command, Errors, Flags } from '@oclif/core';
 import { ContractService } from '../command-helpers/contracts.js';
 import * as DataSourcesExtractor from '../command-helpers/data-sources.js';
 import { updateNetworksFile } from '../command-helpers/network.js';
+import { checkForProxy } from '../command-helpers/proxy.js';
 import { loadRegistry } from '../command-helpers/registry.js';
 import { retryWithPrompt } from '../command-helpers/retry.js';
 import {
@@ -101,35 +102,18 @@ export default class AddCommand extends Command {
           ),
         );
         if (!ethabi) throw Error;
-        const isProxy = ethabi.callFunctionSignatures()?.includes('upgradeTo(address)');
-        if (isProxy) {
-          const impl = await retryWithPrompt(() =>
-            withSpinner(
-              'Fetching proxy implementation address...',
-              'Failed to fetch proxy implementation address',
-              'Warning fetching proxy implementation address',
-              () => contractService.getProxyImplementation(network, address),
-            ),
-          );
-          if (impl) {
-            const useImplementation = await prompt
-              .confirm(`Proxy contract detected. Use implementation contract ABI at ${impl}?`, true)
-              .catch(() => false);
-
-            if (useImplementation) {
-              implAddress = impl;
-              ethabi = await retryWithPrompt(() =>
-                withSpinner(
-                  'Fetching implementation contract ABI...',
-                  'Failed to fetch implementation ABI',
-                  'Warning fetching implementation ABI',
-                  () => contractService.getABI(EthereumABI, network, implAddress!),
-                ),
-              );
-            }
-            if (!ethabi) throw Error;
-          }
+
+        const { implementationAbi, implementationAddress } = await checkForProxy(
+          contractService,
+          network,
+          address,
+          ethabi,
+        );
+        if (implementationAddress) {
+          implAddress = implementationAddress;
+          ethabi = implementationAbi!;
         }
+        if (!ethabi) throw Error;
       } catch (error) {
         // we cannot ask user to do prompt in test environment
         if (process.env.NODE_ENV !== 'test') {

diff --git a/packages/cli/src/commands/init.ts b/packages/cli/src/commands/init.ts
@@ -11,6 +11,7 @@ import { DEFAULT_IPFS_URL } from '../command-helpers/ipfs.js';
 import { initNetworksConfig } from '../command-helpers/network.js';
 import { chooseNodeUrl } from '../command-helpers/node.js';
 import { PromptManager } from '../command-helpers/prompt-manager.js';
+import { checkForProxy } from '../command-helpers/proxy.js';
 import { loadRegistry } from '../command-helpers/registry.js';
 import { retryWithPrompt } from '../command-helpers/retry.js';
 import { generateScaffold, writeScaffold } from '../command-helpers/scaffold.js';
@@ -622,38 +623,22 @@ async function processInitForm(
               () => contractService.getABI(protocolInstance.getABI(), network.id, address),
             ),
           );
-          initDebugger.extend('processInitForm')("abiFromEtherscan len: '%s'", abiFromApi?.name);
-          const isProxy = abiFromApi?.callFunctionSignatures().includes('upgradeTo(address)');
-          if (isProxy) {
-            const impl = await retryWithPrompt(() =>
-              withSpinner(
-                'Fetching proxy implementation address...',
-                'Failed to fetch proxy implementation address',
-                'Warning fetching proxy implementation address',
-                () => contractService.getProxyImplementation(network.id, address),
-              ),
+          initDebugger.extend('processInitForm')("ABI: '%s'", abiFromApi?.name);
+          if (abiFromApi) {
+            const { implementationAbi, implementationAddress } = await checkForProxy(
+              contractService,
+              network.id,
+              address,
+              abiFromApi,
             );
-            initDebugger.extend('processInitForm')("proxyImplementation: '%s'", impl);
-            if (impl) {
-              const useImplementation = await prompt
-                .confirm(
-                  `Proxy contract detected. Use implementation contract ABI at ${impl}?`,
-                  true,
-                )
-                .catch(() => false);
-
-              if (useImplementation) {
-                implAddress = impl;
-                abiFromApi = await retryWithPrompt(() =>
-                  withSpinner(
-                    'Fetching implementation contract ABI...',
-                    'Failed to fetch implementation ABI',
-                    'Warning fetching implementation ABI',
-                    () =>
-                      contractService.getABI(protocolInstance.getABI(), network.id, implAddress!),
-                  ),
-                );
-              }
+            if (implementationAddress) {
+              implAddress = implementationAddress;
+              abiFromApi = implementationAbi!;
+              initDebugger.extend('processInitForm')(
+                "Impl ABI: '%s', Impl Address: '%s'",
+                abiFromApi?.name,
+                implAddress,
+              );
             }
           }
         }