From 1243b004cff484f71b84ee070f3d75cb22664572 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ti=E1=BA=BFn=20Nguy=E1=BB=85n=20Kh=E1=BA=AFc?=
 <tien.nguyenkhac@icloud.com>
Date: Wed, 12 Jun 2024 13:29:09 +1200
Subject: [PATCH] refactor: implement query preflight to use simple atom when
 possible

---
 packages/core/src/index.ts         |  2 +-
 packages/core/src/query.ts         | 23 +++++++++++++++++++++++
 packages/react/src/stores/query.ts | 30 +++++++++++++++++++-----------
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 76435564..07727538 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -6,5 +6,5 @@ export {
   type QueryInstruction,
 } from "./QueryBuilder.js";
 export * from "./errors.js";
-export { default as query } from "./query.js";
+export { preflight, default as query } from "./query.js";
 export * from "./symbols.js";
diff --git a/packages/core/src/query.ts b/packages/core/src/query.ts
index bcbc98c7..98f3fb03 100644
--- a/packages/core/src/query.ts
+++ b/packages/core/src/query.ts
@@ -5,6 +5,29 @@ import type {
 import type { ReDotDescriptor } from "@reactive-dot/types";
 import type { ChainDefinition, TypedApi } from "polkadot-api";
 
+export const preflight = <TInstruction extends QueryInstruction>(
+  instruction: TInstruction,
+) => {
+  type Return = TInstruction["instruction"] extends "fetch-constant"
+    ? "promise"
+    : TInstruction["instruction"] extends "call-api"
+      ? "promise"
+      : TInstruction["instruction"] extends "read-storage-entries"
+        ? "promise"
+        : TInstruction["instruction"] extends "read-storage"
+          ? "observable"
+          : "promise" | "observable";
+
+  switch (instruction.instruction) {
+    case "fetch-constant":
+    case "call-api":
+    case "read-storage-entries":
+      return "promise" as Return;
+    case "read-storage":
+      return "observable" as Return;
+  }
+};
+
 const query = <
   TInstruction extends QueryInstruction,
   TDescriptor extends ChainDefinition = ReDotDescriptor,
diff --git a/packages/react/src/stores/query.ts b/packages/react/src/stores/query.ts
index 61695c75..ea9e798c 100644
--- a/packages/react/src/stores/query.ts
+++ b/packages/react/src/stores/query.ts
@@ -4,6 +4,7 @@ import {
   MultiInstruction,
   QueryError,
   QueryInstruction,
+  preflight,
   query,
 } from "@reactive-dot/core";
 import type { ChainId } from "@reactive-dot/types";
@@ -20,17 +21,24 @@ const _queryAtomFamily = atomFamily(
       // eslint-disable-next-line @typescript-eslint/ban-types
       {}>
     >;
-  }) =>
-    atomWithObservable((get) =>
-      from(get(typedApiAtomFamily(param.chainId))).pipe(
-        switchMap(
-          (api) =>
-            query(api, param.instruction) as
-              | Promise<unknown>
-              | Observable<unknown>,
-        ),
-      ),
-    ),
+  }) => {
+    switch (preflight(param.instruction)) {
+      case "promise":
+        return atom(async (get, { signal }) => {
+          const api = await get(typedApiAtomFamily(param.chainId));
+
+          return query(api, param.instruction, { signal }) as Promise<unknown>;
+        });
+      case "observable":
+        return atomWithObservable((get) =>
+          from(get(typedApiAtomFamily(param.chainId))).pipe(
+            switchMap(
+              (api) => query(api, param.instruction) as Observable<unknown>,
+            ),
+          ),
+        );
+    }
+  },
   (a, b) => stringify(a) === stringify(b),
 );