Use RpcResponseTransformer in PubSubSubscriptionPlan

Author: lorisleiva

packages/rpc-subscriptions-spec/src/rpc-subscriptions-pubsub-plan.ts

@@ -5,7 +5,7 @@ import {
     SolanaError,
 } from '@solana/errors';
 import { safeRace } from '@solana/promises';
-import { createRpcMessage, RpcRequest, RpcResponseData } from '@solana/rpc-spec-types';
+import { createRpcMessage, RpcRequest, RpcResponseData, RpcResponseTransformer } from '@solana/rpc-spec-types';
 import { DataPublisher } from '@solana/subscribable';
 import { demultiplexDataPublisher } from '@solana/subscribable';
 
@@ -14,7 +14,7 @@ import { RpcSubscriptionsChannel } from './rpc-subscriptions-channel';
 
 type Config<TNotification> = Readonly<{
     channel: RpcSubscriptionsChannel<unknown, RpcNotification<TNotification> | RpcResponseData<RpcSubscriptionId>>;
-    responseTransformer?: <T>(response: unknown, notificationName: string) => T;
+    responseTransformer?: RpcResponseTransformer;
     signal: AbortSignal;
     subscribeRequest: RpcRequest;
     unsubscribeMethodName: string;
@@ -65,7 +65,8 @@ function augmentSubscriberCountAndReturnNewCount(
 const cache = new WeakMap();
 function getMemoizedDemultiplexedNotificationPublisherFromChannelAndResponseTransformer<TNotification>(
     channel: RpcSubscriptionsChannel<unknown, RpcNotification<TNotification>>,
-    responseTransformer?: <T>(response: unknown, notificationName: string) => T,
+    subscribeRequest: RpcRequest,
+    responseTransformer?: RpcResponseTransformer,
 ): DataPublisher<{
     [channelName: `notification:${number}`]: TNotification;
 }> {
@@ -84,7 +85,7 @@ function getMemoizedDemultiplexedNotificationPublisherFromChannelAndResponseTran
                     return;
                 }
                 const transformedNotification = responseTransformer
-                    ? responseTransformer(message.params.result, message.method)
+                    ? responseTransformer(message.params.result, subscribeRequest)
                     : message.params.result;
                 return [`notification:${message.params.subscription}`, transformedNotification];
             })),
@@ -189,6 +190,7 @@ export async function executeRpcPubSubSubscriptionPlan<TNotification>({
      */
     const notificationPublisher = getMemoizedDemultiplexedNotificationPublisherFromChannelAndResponseTransformer(
         channel,
+        subscribeRequest,
         responseTransformer,
     );
     const notificationKey = `notification:${subscriptionId}` as const;

packages/rpc-transformers/src/response-transformer.ts

@@ -2,10 +2,9 @@ import { pipe } from '@solana/functional';
 import { RpcRequest, RpcResponse, RpcResponseTransformer } from '@solana/rpc-spec-types';
 
 import { AllowedNumericKeypaths } from './response-transformer-allowed-numeric-values';
-import { getBigIntUpcastResponseTransformer, getBigIntUpcastVisitor } from './response-transformer-bigint-upcast';
+import { getBigIntUpcastResponseTransformer } from './response-transformer-bigint-upcast';
 import { getResultResponseTransformer } from './response-transformer-result';
 import { getThrowSolanaErrorResponseTransformer } from './response-transformer-throw-solana-error';
-import { getTreeWalker } from './tree-traversal';
 
 export type ResponseTransformerConfig<TApi> = Readonly<{
     allowedNumericKeyPaths?: AllowedNumericKeypaths<TApi>;
@@ -29,16 +28,11 @@ export function getDefaultResponseTransformerForSolanaRpc<TApi>(
 
 export function getDefaultResponseTransformerForSolanaRpcSubscriptions<TApi>(
     config?: ResponseTransformerConfig<TApi>,
-): <T>(notification: unknown, notificationName: string) => T {
-    return <T>(notification: unknown, notificationName: string): T => {
+): RpcResponseTransformer {
+    return (response: RpcResponse, request: RpcRequest): RpcResponse => {
+        const methodName = request.methodName as keyof TApi;
         const keyPaths =
-            config?.allowedNumericKeyPaths && notificationName
-                ? config.allowedNumericKeyPaths[notificationName as keyof TApi]
-                : undefined;
-        const traverse = getTreeWalker([getBigIntUpcastVisitor(keyPaths ?? [])]);
-        const initialState = {
-            keyPath: [],
-        };
-        return traverse(notification, initialState) as T;
+            config?.allowedNumericKeyPaths && methodName ? config.allowedNumericKeyPaths[methodName] : undefined;
+        return pipe(response, r => getBigIntUpcastResponseTransformer(keyPaths ?? [])(r, request));
     };
 }